Resize the object feature embeddings with torch to optimize feature computation time

micro_sam/object_classification.py

@@ -13,6 +13,9 @@
 except ImportError:
     from tqdm import tqdm
 
+import torch
+
+from bioimage_cpp.transformation import resample
 from bioimage_cpp.utils import Blocking, take_dict
 
 from . import util
@@ -24,7 +27,6 @@ def _anyup_object_resize(embeds_chw, image_region, target_hw, upsampler):
     # Upsample a (C, h, w) embedding to (target_h, target_w, C) with AnyUp, using the image region
     # as guidance. The encoders pad the image to a square, so we square-pad the image here to match
     # how the segmentation is padded.
-    import torch
     import torch.nn.functional as F
     from .pixel_classification import _to_anyup_image
 
@@ -76,11 +78,22 @@ def _compute_object_features_impl(
     if upsampler is not None:
         # AnyUp upsamples the embedding using the image region instead of plain interpolation.
         embeddings = _anyup_object_resize(embeddings_chw, image_region, resize_hw, upsampler)
+    elif all(t >= s for t, s in zip(resize_hw, embeddings_chw.shape[1:])):
+        # Upsampling (the usual case): torch is much faster than skimage and matches it, except at the edge.
+        embeddings = torch.nn.functional.interpolate(
+            torch.from_numpy(np.ascontiguousarray(embeddings_chw, dtype="float32"))[None], size=resize_hw,
+            mode="bilinear", align_corners=False,
+        )[0].permute(1, 2, 0).numpy().astype(embeddings_chw.dtype, copy=False)  # put the channel axis last
     else:
-        embeddings = embeddings_chw.transpose(1, 2, 0)  # put the channel axis last
-        embeddings = resize(
-            embeddings, resize_hw + (embeddings.shape[-1],), preserve_range=True
-        ).astype(embeddings.dtype)
+        # Downsampling (tiny images or corner tiles): bioimage-cpp matches the skimage anti-aliasing, torch does not.
+        scale = [s / t for s, t in zip(embeddings_chw.shape[1:], resize_hw)]
+        matrix = np.array([
+            [1.0, 0.0, 0.0, 0.0], [0.0, scale[0], 0.0, (scale[0] - 1) / 2], [0.0, 0.0, scale[1], (scale[1] - 1) / 2]
+        ])
+        bounding_box = (slice(0, embeddings_chw.shape[0]),) + tuple(slice(0, t) for t in resize_hw)
+        embeddings = resample(
+            np.ascontiguousarray(embeddings_chw, dtype="float32"), matrix, bounding_box=bounding_box
+        ).transpose(1, 2, 0).astype(embeddings_chw.dtype, copy=False)  # put the channel axis last
 
     segmentation_rescaled = resize(
         segmentation_rescaled, embeddings.shape[:2], order=0, anti_aliasing=False, preserve_range=True

test/test_classification_geometry.py

@@ -6,6 +6,7 @@
 content, which silently misplaces every feature for non-square inputs.
 """
 
+import pytest
 import numpy as np
 
 from micro_sam.object_classification import compute_object_features
@@ -50,3 +51,24 @@ def test_object_features_exclude_padding():
     assert seg_ids.tolist() == [1]
     # Feature layout is [area, per-channel mean], so the embedding means follow the area column.
     assert np.allclose(features[0, 1:], 1.0), f"padded region leaked into the object features: {features[0, 1:]}"
+
+
+@pytest.mark.parametrize("shape, objects", [
+    ((64, 64), [np.s_[8:20, 8:20], np.s_[30:40, 40:56]]),  # upsampling of the embedding, with torch
+    ((8, 8), [np.s_[2:4, 2:4], np.s_[4:6, 3:6]]),  # downsampling of the embedding, with bioimage-cpp
+])
+def test_object_features_match_skimage_interpolation(shape, objects):
+    # The objects stay off the 2 px edge, where torch repeats the edge pixel and skimage mirrors it.
+    from skimage.transform import resize
+
+    features = np.random.default_rng(0).normal(size=(1, N_CHANNELS, EMBEDDING_SIZE, EMBEDDING_SIZE))
+    embeddings = {**_padded_embeddings(), "features": features.astype("float32"), "original_size": shape}
+    segmentation = np.zeros(shape, dtype="uint32")
+    for seg_id, bb in enumerate(objects, start=1):
+        segmentation[bb] = seg_id
+
+    _, object_features = compute_object_features(embeddings, segmentation, verbose=False)
+
+    resized = resize(features[0].transpose(1, 2, 0), shape + (N_CHANNELS,), preserve_range=True)
+    expected = [resized[segmentation == seg_id].mean(axis=0) for seg_id in range(1, len(objects) + 1)]
+    assert np.allclose(object_features[:, 1:], expected, atol=1e-5)