attr-eomt 1.1.4__tar.gz → 1.2.0__tar.gz

{attr_eomt-1.1.4 → attr_eomt-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: attr-eomt
-Version: 1.1.4
+Version: 1.2.0
 Summary: Standalone EoMT (Encoder-only Mask Transformer) for instance segmentation, with DINOv2 init, COCO training/validation and inference.
 Project-URL: Homepage, https://github.com/imagra93/attr-eomt
 Project-URL: Documentation, https://imagra93.github.io/attr-eomt

{attr_eomt-1.1.4 → attr_eomt-1.2.0}/eomt/preprocess.py

@@ -65,3 +65,35 @@ def preprocess_numpy(
     canvas[:nh, :nw] = arr
     chw = np.ascontiguousarray(canvas.transpose(2, 0, 1))
     return chw, make_preprocess_meta(True, (nh, nw), input_size)
+
+
+def preprocess_torch(
+    img_rgb_hwc: np.ndarray,
+    input_size: int,
+    *,
+    letterbox: bool = True,
+    mean=None,
+    std=None,
+    device="cpu",
+):
+    """:func:`preprocess_numpy` run with torch on ``device`` (e.g. the model's GPU).
+
+    The resize is an antialiased bilinear one rounded to uint8, as PIL's, so the result matches
+    :func:`preprocess_numpy` within one grey level. Returns ``(chw_float32_tensor, meta)``.
+    """
+    import torch
+    import torch.nn.functional as F  # noqa: N812
+
+    h, w = img_rgb_hwc.shape[:2]
+    if letterbox:
+        scale = input_size / max(h, w)
+        nh, nw = max(1, round(h * scale)), max(1, round(w * scale))
+    else:
+        nh, nw = input_size, input_size
+    x = torch.from_numpy(np.ascontiguousarray(img_rgb_hwc)).to(device).permute(2, 0, 1)[None].float()
+    x = F.interpolate(x, size=(nh, nw), mode="bilinear", antialias=True, align_corners=False).round_()
+    mean = torch.as_tensor(IMAGENET_MEAN if mean is None else mean, dtype=torch.float32, device=device).view(3, 1, 1)
+    std = torch.as_tensor(IMAGENET_STD if std is None else std, dtype=torch.float32, device=device).view(3, 1, 1)
+    chw = x.new_zeros((3, input_size, input_size))  # 0 == mean post-normalize
+    chw[:, :nh, :nw] = (x[0] / 255.0 - mean) / std
+    return chw, make_preprocess_meta(letterbox, (nh, nw), input_size)

{attr_eomt-1.1.4 → attr_eomt-1.2.0}/eomt/tta.py

@@ -38,6 +38,7 @@ repair / replace rules). Attribute probabilities and embeddings are score-weight
 from __future__ import annotations
 
 import contextlib
+import functools
 import itertools
 import math
 from dataclasses import dataclass
@@ -47,7 +48,7 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 
 from .postprocess import _masks_to_original, aux_result, aux_value_result, boxes_from_masks, query_scores
-from .preprocess import preprocess_numpy
+from .preprocess import preprocess_numpy, preprocess_torch
 
 _AUGMENT_KEYS = {"flip", "scales"}
 _TILES_KEYS = {"size", "overlap", "full"}
@@ -288,12 +289,16 @@ def predict_views(
     amp_dtype: torch.dtype | None = None,
     aux_scopes: dict | None = None,
     embed: bool = False,
+    device_preprocess: bool = False,
 ) -> dict:
     """Instances of the RGB ``image`` ``(H, W, 3)`` uint8, merged over the views of ``cfg``.
 
     Returns the :func:`~eomt.postprocess.postprocess_instance` dict (``boxes`` / ``scores`` / ``classes`` / ``masks``
     and, with attribute heads, ``aux``), plus ``embed`` ``(N, hidden)`` with ``embed=True``. There is no ``query_idx``:
     a merged instance comes from several queries.
+
+    With ``device_preprocess=True`` the views are preprocessed on the model's device
+    (:func:`~eomt.preprocess.preprocess_torch`) instead of with PIL on the CPU; the inputs match within one grey level.
     """
     if getattr(model, "family", "instance") != "instance":
         raise NotImplementedError("TTA / tiled inference supports the instance family only.")
@@ -311,15 +316,14 @@ def predict_views(
     view_cells = [_grid_cells(v.rect, g, hg, wg) for v in views]
 
     flip_perms = {s.name: s.flip_perm() for s in getattr(model, "aux_specs", []) if s.flip_perm() is not None}
+    prep = functools.partial(preprocess_torch, device=dev) if device_preprocess else preprocess_numpy
     found: list[dict] = []
     for size in dict.fromkeys(v.size for v in views):
         idx = [i for i, v in enumerate(views) if v.size == size]
         for start in range(0, len(idx), cfg.batch):
             chunk = idx[start : start + cfg.batch]
-            pre = [
-                preprocess_numpy(_crop(image, views[i]), size, letterbox=letterbox, mean=mean, std=std) for i in chunk
-            ]
-            x = torch.from_numpy(np.stack([chw for chw, _ in pre])).to(dev)
+            pre = [prep(_crop(image, views[i]), size, letterbox=letterbox, mean=mean, std=std) for i in chunk]
+            x = torch.stack([torch.as_tensor(chw) for chw, _ in pre]).to(dev)
             ctx = torch.autocast(dev.type, dtype=amp_dtype) if amp_dtype is not None else contextlib.nullcontext()
             with ctx:
                 out = model(x)

{attr_eomt-1.1.4 → attr_eomt-1.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "attr-eomt"
-version = "1.1.4"
+version = "1.2.0"
 description = "Standalone EoMT (Encoder-only Mask Transformer) for instance segmentation, with DINOv2 init, COCO training/validation and inference."
 readme = "README.md"
 requires-python = ">=3.10"

{attr_eomt-1.1.4 → attr_eomt-1.2.0}/tests/test_smoke.py

@@ -514,6 +514,20 @@ def test_preprocess_letterbox_meta_and_padding():
     assert not meta2["letterbox"] and meta2["content_hw"] == (20, 20)
 
 
+@pytest.mark.parametrize("letterbox", [True, False])
+def test_preprocess_torch_matches_preprocess_numpy_within_one_grey_level(letterbox):
+    import numpy as np
+
+    from eomt.preprocess import IMAGENET_STD, preprocess_numpy, preprocess_torch
+
+    img = np.random.default_rng(0).integers(0, 255, (97, 151, 3), dtype=np.uint8)
+    chw, meta = preprocess_torch(img, 64, letterbox=letterbox)
+    expected, expected_meta = preprocess_numpy(img, 64, letterbox=letterbox)
+    assert chw.shape == (3, 64, 64) and chw.dtype == torch.float32 and meta == expected_meta
+    diff = np.abs(chw.numpy() - expected) * IMAGENET_STD.reshape(3, 1, 1) * 255  # in grey levels
+    assert diff.max() <= 1.01 and diff.mean() < 0.25
+
+
 def test_letterbox_inverse_crops_content_not_padding():
     from eomt.postprocess import _masks_to_original
     from eomt.preprocess import make_preprocess_meta

{attr_eomt-1.1.4 → attr_eomt-1.2.0}/tests/test_tta.py

@@ -172,6 +172,15 @@ def test_flip_tta_merges_the_two_views_and_scores_their_agreement():
     assert _iou(mean["masks"][0], 14, 35, 56, 91) > 0.95
 
 
+def test_device_preprocessed_views_find_the_same_instances():
+    img = _image(280, 420, 70, 42, 315, 98)
+    cfg = TTAConfig(flip=True, scales=(1.0, 1.5), tiles=True, tile_size=140)
+    cpu = predict_views(_Blobs(), img, cfg, imgsz=IMGSZ, conf_thres=0.3)
+    dev = predict_views(_Blobs(), img, cfg, imgsz=IMGSZ, conf_thres=0.3, device_preprocess=True)
+    assert dev["num_detections"] == cpu["num_detections"] == 1
+    assert torch.allclose(dev["scores"], cpu["scores"]) and _same(dev["masks"], cpu["masks"])
+
+
 def test_predict_image_with_tta_keeps_the_result_contract():
     from eomt.engine.predict import predict_image