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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: attr-eomt
-Version: 1.1.4
+Version: 1.2.1
 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.2.1/eomt/preprocess.py

@@ -0,0 +1,144 @@
+"""Preprocessing helpers for EoMT.
+
+EoMT / DINOv2 expect a square, ImageNet-normalized RGB tensor. Two modes:
+
+* **letterbox** (default): resize the long side to ``input_size`` preserving
+  aspect ratio, then pad bottom/right to a square — object shapes are undistorted.
+* **stretch**: resize straight to a square (DETR-style, legacy) — faster but
+  distorts aspect ratio.
+
+Both return ``(chw_float32, meta)`` where ``meta`` records what postprocessing
+needs to map masks back to the original image (see :func:`make_preprocess_meta`).
+"""
+
+from __future__ import annotations
+
+import functools
+
+import numpy as np
+
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+#: Fixed-point precision of PIL's 8-bit resampling (Pillow's Resample.c).
+_PIL_PRECISION_BITS = 32 - 8 - 2
+
+
+def make_preprocess_meta(letterbox: bool, content_hw: tuple[int, int], input_size: int) -> dict:
+    """Pack the info postprocessing needs to invert preprocessing.
+
+    ``content_hw`` is the ``(height, width)`` the real image occupies inside the
+    square canvas (equal to ``(input_size, input_size)`` for stretch).
+    """
+    return {"letterbox": bool(letterbox), "content_hw": tuple(content_hw), "input_size": int(input_size)}
+
+
+def preprocess_numpy(
+    img_rgb_hwc: np.ndarray,
+    input_size: int,
+    *,
+    letterbox: bool = True,
+    mean=None,
+    std=None,
+):
+    """Resize an RGB HWC uint8 image to ``(C, input_size, input_size)``, normalized.
+
+    ``mean`` / ``std`` default to :data:`IMAGENET_MEAN` / :data:`IMAGENET_STD`; pass
+    a model's stored ``pixel_mean`` / ``pixel_std`` to reproduce its exact training
+    normalization. Returns ``(chw_float32, meta)``; ``meta`` is from
+    :func:`make_preprocess_meta`. The padded border (letterbox) is filled with 0 in
+    normalized space, i.e. the mean colour.
+    """
+    from PIL import Image
+
+    mean = IMAGENET_MEAN if mean is None else np.asarray(mean, dtype=np.float32)
+    std = IMAGENET_STD if std is None else np.asarray(std, dtype=np.float32)
+
+    h, w = img_rgb_hwc.shape[:2]
+    if not letterbox:
+        img = Image.fromarray(img_rgb_hwc).resize((input_size, input_size), Image.BILINEAR)
+        arr = np.asarray(img, dtype=np.float32) / 255.0
+        arr = (arr - mean) / std
+        chw = np.ascontiguousarray(arr.transpose(2, 0, 1))
+        return chw, make_preprocess_meta(False, (input_size, input_size), input_size)
+
+    scale = input_size / max(h, w)
+    nh, nw = max(1, round(h * scale)), max(1, round(w * scale))
+    img = Image.fromarray(img_rgb_hwc).resize((nw, nh), Image.BILINEAR)
+    arr = (np.asarray(img, dtype=np.float32) / 255.0 - mean) / std  # (nh, nw, 3)
+    canvas = np.zeros((input_size, input_size, 3), dtype=np.float32)  # 0 == mean post-normalize
+    canvas[:nh, :nw] = arr
+    chw = np.ascontiguousarray(canvas.transpose(2, 0, 1))
+    return chw, make_preprocess_meta(True, (nh, nw), input_size)
+
+
+@functools.lru_cache(maxsize=64)
+def _pil_bilinear_taps(in_size: int, out_size: int) -> tuple[np.ndarray, np.ndarray]:
+    """PIL's fixed-point bilinear coefficients for resizing ``in_size`` px to ``out_size``.
+
+    Mirrors ``precompute_coeffs`` and ``normalize_coeffs_8bpc`` of Pillow's Resample.c. Returns the input index and
+    the integer weight of every tap, both ``(out_size, taps)``; padding taps have weight 0.
+    """
+    scale = in_size / out_size
+    filterscale = max(scale, 1.0)
+    support = filterscale  # the bilinear filter's support is 1
+    center = (np.arange(out_size) + 0.5) * scale
+    xmin = np.maximum((center - support + 0.5).astype(np.int64), 0)  # C's (int) truncation
+    xmax = np.minimum((center + support + 0.5).astype(np.int64), in_size)
+    x = xmin[:, None] + np.arange(int(np.ceil(support)) * 2 + 1)
+    w = np.maximum(1.0 - np.abs((x - center[:, None] + 0.5) * (1.0 / filterscale)), 0.0) * (x < xmax[:, None])
+    total = np.cumsum(w, axis=1)[:, -1:]  # summed in order, as PIL does
+    w = np.divide(w, total, out=np.zeros_like(w), where=total != 0)
+    return np.minimum(x, in_size - 1), np.floor(0.5 + w * (1 << _PIL_PRECISION_BITS)).astype(np.int32)
+
+
+def _pil_resize(x, nh: int, nw: int):
+    """Resize a ``(C, H, W)`` int32 image as PIL's ``Image.resize(..., BILINEAR)`` does an 8-bit one, bit for bit:
+    a horizontal then a vertical pass, each rounded to uint8."""
+    import torch
+
+    half = 1 << (_PIL_PRECISION_BITS - 1)
+    if nw != x.shape[2]:
+        idx, k = (torch.from_numpy(a).to(x.device) for a in _pil_bilinear_taps(x.shape[2], nw))
+        acc = torch.full((x.shape[0], x.shape[1], nw), half, dtype=torch.int32, device=x.device)
+        for t in range(k.shape[1]):
+            acc += x[:, :, idx[:, t]] * k[:, t]
+        x = (acc >> _PIL_PRECISION_BITS).clamp_(0, 255)
+    if nh != x.shape[1]:
+        idx, k = (torch.from_numpy(a).to(x.device) for a in _pil_bilinear_taps(x.shape[1], nh))
+        acc = torch.full((x.shape[0], nh, x.shape[2]), half, dtype=torch.int32, device=x.device)
+        for t in range(k.shape[1]):
+            acc += x[:, idx[:, t], :] * k[:, t, None]
+        x = (acc >> _PIL_PRECISION_BITS).clamp_(0, 255)
+    return x
+
+
+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), with the same result.
+
+    The resize reproduces PIL's fixed-point bilinear one, so the output equals :func:`preprocess_numpy`'s.
+    Returns ``(chw_float32_tensor, meta)``.
+    """
+    import torch
+
+    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).to(torch.int32)
+    x = _pil_resize(x, nh, nw).float()
+    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)
+    levels = torch.full((3, 1, 1), 255.0, device=device)  # not a scalar, which torch would divide by as x * (1 / 255)
+    chw = x.new_zeros((3, input_size, input_size))  # 0 == mean post-normalize
+    chw[:, :nh, :nw] = (x / levels - mean) / std
+    return chw, make_preprocess_meta(letterbox, (nh, nw), input_size)

{attr_eomt-1.1.4 → attr_eomt-1.2.1}/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.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "attr-eomt"
-version = "1.1.4"
+version = "1.2.1"
 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.1}/tests/test_smoke.py

@@ -514,6 +514,25 @@ def test_preprocess_letterbox_meta_and_padding():
     assert not meta2["letterbox"] and meta2["content_hw"] == (20, 20)
 
 
+@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="no GPU"))])
+@pytest.mark.parametrize("letterbox", [True, False])
+@pytest.mark.parametrize("shape, size", [((97, 151), 64), ((720, 1280), 644), ((300, 200), 644), ((5, 900), 140)])
+def test_preprocess_torch_equals_preprocess_numpy(shape, size, letterbox, device):
+    """Downscaled, upscaled and extreme aspect ratios: the same input as PIL's, value for value."""
+    import numpy as np
+
+    from eomt.preprocess import preprocess_numpy, preprocess_torch
+
+    img = np.random.default_rng(0).integers(0, 256, (*shape, 3), dtype=np.uint8)
+    chw, meta = preprocess_torch(img, size, letterbox=letterbox, mean=(0.5, 0.4, 0.3), std=(0.2, 0.3, 0.25),
+                                 device=device)
+    expected, expected_meta = preprocess_numpy(img, size, letterbox=letterbox, mean=(0.5, 0.4, 0.3),
+                                               std=(0.2, 0.3, 0.25))
+    assert chw.shape == (3, size, size) and chw.dtype == torch.float32 and meta == expected_meta
+    assert np.array_equal(chw.cpu().numpy(), expected)
+
+
 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.1}/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.equal(dev["scores"], cpu["scores"]) and torch.equal(dev["masks"], cpu["masks"])
+
+
 def test_predict_image_with_tta_keeps_the_result_contract():
     from eomt.engine.predict import predict_image
 

attr_eomt-1.1.4/eomt/preprocess.py

@@ -1,67 +0,0 @@
-"""Preprocessing helpers for EoMT.
-
-EoMT / DINOv2 expect a square, ImageNet-normalized RGB tensor. Two modes:
-
-* **letterbox** (default): resize the long side to ``input_size`` preserving
-  aspect ratio, then pad bottom/right to a square — object shapes are undistorted.
-* **stretch**: resize straight to a square (DETR-style, legacy) — faster but
-  distorts aspect ratio.
-
-Both return ``(chw_float32, meta)`` where ``meta`` records what postprocessing
-needs to map masks back to the original image (see :func:`make_preprocess_meta`).
-"""
-
-from __future__ import annotations
-
-import numpy as np
-
-IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
-IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
-
-
-def make_preprocess_meta(letterbox: bool, content_hw: tuple[int, int], input_size: int) -> dict:
-    """Pack the info postprocessing needs to invert preprocessing.
-
-    ``content_hw`` is the ``(height, width)`` the real image occupies inside the
-    square canvas (equal to ``(input_size, input_size)`` for stretch).
-    """
-    return {"letterbox": bool(letterbox), "content_hw": tuple(content_hw), "input_size": int(input_size)}
-
-
-def preprocess_numpy(
-    img_rgb_hwc: np.ndarray,
-    input_size: int,
-    *,
-    letterbox: bool = True,
-    mean=None,
-    std=None,
-):
-    """Resize an RGB HWC uint8 image to ``(C, input_size, input_size)``, normalized.
-
-    ``mean`` / ``std`` default to :data:`IMAGENET_MEAN` / :data:`IMAGENET_STD`; pass
-    a model's stored ``pixel_mean`` / ``pixel_std`` to reproduce its exact training
-    normalization. Returns ``(chw_float32, meta)``; ``meta`` is from
-    :func:`make_preprocess_meta`. The padded border (letterbox) is filled with 0 in
-    normalized space, i.e. the mean colour.
-    """
-    from PIL import Image
-
-    mean = IMAGENET_MEAN if mean is None else np.asarray(mean, dtype=np.float32)
-    std = IMAGENET_STD if std is None else np.asarray(std, dtype=np.float32)
-
-    h, w = img_rgb_hwc.shape[:2]
-    if not letterbox:
-        img = Image.fromarray(img_rgb_hwc).resize((input_size, input_size), Image.BILINEAR)
-        arr = np.asarray(img, dtype=np.float32) / 255.0
-        arr = (arr - mean) / std
-        chw = np.ascontiguousarray(arr.transpose(2, 0, 1))
-        return chw, make_preprocess_meta(False, (input_size, input_size), input_size)
-
-    scale = input_size / max(h, w)
-    nh, nw = max(1, round(h * scale)), max(1, round(w * scale))
-    img = Image.fromarray(img_rgb_hwc).resize((nw, nh), Image.BILINEAR)
-    arr = (np.asarray(img, dtype=np.float32) / 255.0 - mean) / std  # (nh, nw, 3)
-    canvas = np.zeros((input_size, input_size, 3), dtype=np.float32)  # 0 == mean post-normalize
-    canvas[:nh, :nw] = arr
-    chw = np.ascontiguousarray(canvas.transpose(2, 0, 1))
-    return chw, make_preprocess_meta(True, (nh, nw), input_size)