From ef2cfee924cfef8a62e6881c336e65f40c235268 Mon Sep 17 00:00:00 2001 From: titusgriebel Date: Fri, 16 Oct 2026 14:47:38 +0200 Subject: [PATCH 1/2] Resize the object feature embeddings with torch 'skimage.transform.resize' took 94% of the object feature time: about 0.5 s per tile, so 8 s for a 2048 px field with the default tiling. 'torch.nn.functional.interpolate' takes about 0.01 s per tile. The values match, except in the 2 px edge of the resized embedding, where torch repeats the edge pixel and skimage mirrors it. Co-Authored-By: Claude Opus 5.5 --- micro_sam/object_classification.py | 11 +++++++---- test/test_classification_geometry.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/micro_sam/object_classification.py b/micro_sam/object_classification.py index 579445847..c766ddcf8 100644 --- a/micro_sam/object_classification.py +++ b/micro_sam/object_classification.py @@ -4,6 +4,7 @@ import numpy as np import pandas as pd +import torch from skimage.transform import resize from skimage.measure import regionprops_table @@ -77,10 +78,12 @@ def _compute_object_features_impl( # AnyUp upsamples the embedding using the image region instead of plain interpolation. embeddings = _anyup_object_resize(embeddings_chw, image_region, resize_hw, upsampler) 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) + # Torch is about 50x faster than 'skimage.transform.resize' here and matches it, except in the + # 2 px edge: torch repeats the edge pixel where skimage mirrors it. + embeddings = torch.nn.functional.interpolate( + torch.from_numpy(np.ascontiguousarray(embeddings_chw, dtype="float32"))[None], size=resize_hw, + mode="bilinear", align_corners=False, antialias=True, + )[0].permute(1, 2, 0).numpy().astype(embeddings_chw.dtype) # put the channel axis last segmentation_rescaled = resize( segmentation_rescaled, embeddings.shape[:2], order=0, anti_aliasing=False, preserve_range=True diff --git a/test/test_classification_geometry.py b/test/test_classification_geometry.py index e5f6b8f44..e5362439b 100644 --- a/test/test_classification_geometry.py +++ b/test/test_classification_geometry.py @@ -50,3 +50,22 @@ 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:]}" + + +def test_object_features_match_skimage_interpolation(): + # The embedding resize must keep the pixel-center alignment of 'skimage.transform.resize'. The objects + # are away from the 2 px edge, where torch repeats the edge pixel and skimage mirrors it. + from skimage.transform import resize + + shape = (64, 64) + 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") + segmentation[8:20, 8:20] = 1 + segmentation[30:40, 40:56] = 2 + + _, object_features = compute_object_features(embeddings, segmentation, verbose=False) + + upsampled = resize(features[0].transpose(1, 2, 0), shape + (N_CHANNELS,), preserve_range=True) + expected = [upsampled[segmentation == seg_id].mean(axis=0) for seg_id in (1, 2)] + assert np.allclose(object_features[:, 1:], expected, atol=1e-5) From bc5fc3a88a6d5f10840b4a5db95211ea69b24039 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Mon, 19 Oct 2026 18:03:50 +0200 Subject: [PATCH 2/2] Downsample the object feature embeddings with bioimage-cpp to match skimage --- micro_sam/object_classification.py | 24 +++++++++++++++++------- test/test_classification_geometry.py | 19 +++++++++++-------- 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/micro_sam/object_classification.py b/micro_sam/object_classification.py index c766ddcf8..44ddb3a12 100644 --- a/micro_sam/object_classification.py +++ b/micro_sam/object_classification.py @@ -4,7 +4,6 @@ import numpy as np import pandas as pd -import torch from skimage.transform import resize from skimage.measure import regionprops_table @@ -14,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 @@ -25,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 @@ -77,13 +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) - else: - # Torch is about 50x faster than 'skimage.transform.resize' here and matches it, except in the - # 2 px edge: torch repeats the edge pixel where skimage mirrors it. + 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, antialias=True, - )[0].permute(1, 2, 0).numpy().astype(embeddings_chw.dtype) # put the channel axis last + mode="bilinear", align_corners=False, + )[0].permute(1, 2, 0).numpy().astype(embeddings_chw.dtype, copy=False) # put the channel axis last + else: + # 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 diff --git a/test/test_classification_geometry.py b/test/test_classification_geometry.py index e5362439b..c5cbc3c49 100644 --- a/test/test_classification_geometry.py +++ b/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 @@ -52,20 +53,22 @@ def test_object_features_exclude_padding(): assert np.allclose(features[0, 1:], 1.0), f"padded region leaked into the object features: {features[0, 1:]}" -def test_object_features_match_skimage_interpolation(): - # The embedding resize must keep the pixel-center alignment of 'skimage.transform.resize'. The objects - # are away from the 2 px edge, where torch repeats the edge pixel and skimage mirrors it. +@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 - shape = (64, 64) 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") - segmentation[8:20, 8:20] = 1 - segmentation[30:40, 40:56] = 2 + for seg_id, bb in enumerate(objects, start=1): + segmentation[bb] = seg_id _, object_features = compute_object_features(embeddings, segmentation, verbose=False) - upsampled = resize(features[0].transpose(1, 2, 0), shape + (N_CHANNELS,), preserve_range=True) - expected = [upsampled[segmentation == seg_id].mean(axis=0) for seg_id in (1, 2)] + 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)