diff --git a/micro_sam/object_classification.py b/micro_sam/object_classification.py index 579445847..44ddb3a12 100644 --- a/micro_sam/object_classification.py +++ b/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 diff --git a/test/test_classification_geometry.py b/test/test_classification_geometry.py index e5f6b8f44..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 @@ -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)