Restore input normalization back to min-max (#1296)

Author: anwai98

micro_sam/sam_annotator/_widgets.py

@@ -2046,6 +2046,14 @@ def _validate_existing_embeddings(self, state):
             "info", "Embeddings have already been precomputed. Press OK to recompute the embeddings."
         )
 
+    @staticmethod
+    def _clear_autosegment_cache(state):
+        """Discard predictions derived from the embeddings that were just replaced."""
+        widget = state.widgets.get("autosegment")
+        if widget is not None:
+            widget._segmenter = None
+            widget._segmenter_key = None
+
     def _n_tiles(self, tile_shape, shape):
         """The number of tiles the embedding computation is split into (1 without tiling)."""
         if tile_shape is None:
@@ -2173,6 +2181,7 @@ def pbar_update(update):
             pbar_signals.pbar_stop.emit()
 
         compute_image_embedding()
+        self._clear_autosegment_cache(state)
         self._update_model(state)
         # worker = compute_image_embedding()
         # worker.returned.connect(self._update_model)

micro_sam/v2/normalization.py

@@ -5,7 +5,9 @@
 import numpy as np
 from torch_em.transform.raw import normalize_percentile
 
-RAW_NORMALIZATION = "percentile_2_98"
+# The released UniSAM2 microscopy checkpoint and the legacy 2D embedding path use min-max normalization.
+# Persist the policy in embedding caches so percentile-normalized features are not silently reused.
+RAW_NORMALIZATION = "minmax_per_channel"
 
 
 def normalize_raw(
@@ -60,13 +62,13 @@ def normalize_raw(
 
 
 def to_image(image: np.ndarray) -> np.ndarray:
-    """Map a 2D or channel-last image to percentile-normalized, channel-last uint8 RGB.
+    """Map a 2D or channel-last image to min-max-normalized, channel-last uint8 RGB.
 
     Args:
         image: The input image. Either 2D or channel-last with up to three channels.
 
     Returns:
         The channel-last uint8 RGB image, with each channel normalized independently.
     """
-    from micro_sam.util import _ensure_rgb
-    return normalize_raw(_ensure_rgb(image), axis=(0, 1), output_dtype="uint8")
+    from micro_sam.util import _to_image
+    return _to_image(image)

micro_sam/v2/util.py

@@ -294,11 +294,12 @@ def _check_saved_embeddings(input_, predictor, f, save_path, tile_shape, halo):
     data (data signature mismatch).
     """
     # We may have an empty zarr file that was already created to save the embeddings in. A
-    # feature-bearing file without the completion metadata is a partial cache: only resume it when
-    # it already records the current normalization policy. Otherwise it may contain legacy min-max
-    # features, which must not be mixed with newly computed percentile-normalized features.
+    # feature-bearing file without the completion metadata is a partial cache: only reject it when
+    # it explicitly records a different normalization policy. Caches without this metadata predate
+    # percentile normalization and therefore use the restored min-max policy.
     if "input_size" not in f.attrs:
-        return "features" in f and f.attrs.get("normalization") != RAW_NORMALIZATION
+        normalization = f.attrs.get("normalization")
+        return "features" in f and normalization is not None and normalization != RAW_NORMALIZATION
 
     # Creates all the metadta that is stored along with the embeddings.
     # TODO: This is currently paired with `micro_sam`-level metadata. Should we get separate for `micro_sam.v2`?
@@ -308,9 +309,9 @@ def _check_saved_embeddings(input_, predictor, f, save_path, tile_shape, halo):
 
     stale = False
     for key, val in signature.items():
-        # Embeddings without normalization metadata used the former min-max policy.
+        # Embeddings without normalization metadata predate percentile normalization and used the
+        # same min-max policy as the released checkpoints, so they remain compatible.
         if key not in f.attrs:
-            stale = stale or key == "normalization"
             continue
         if f.attrs[key] == val:
             continue

test/test_sam_annotator/test_annotator.py

@@ -512,3 +512,14 @@ def _add_flow_integration_params(self, settings, n_iter, dt=0.5, sigma=1.0):
         assert autoseg.sigma == defaults["sigma"]
         assert autoseg.n_iter == defaults["n_iter"]
         assert autoseg.dt == defaults["dt"]
+
+    def test_embedding_recompute_clears_cached_prediction(self):
+        from types import SimpleNamespace
+
+        from micro_sam.sam_annotator._widgets import EmbeddingWidget
+
+        autosegment = SimpleNamespace(_segmenter=object(), _segmenter_key=object())
+        state = SimpleNamespace(widgets={"autosegment": autosegment})
+        EmbeddingWidget._clear_autosegment_cache(state)
+        assert autosegment._segmenter is None
+        assert autosegment._segmenter_key is None

test/test_util.py

@@ -177,6 +177,7 @@ def test_normalize_raw_percentile_params(self):
 
     def test_normalize_raw_per_channel(self):
         from micro_sam.v2.normalization import normalize_raw, to_image
+        from micro_sam.util import _to_image
 
         channel = np.arange(100, dtype="float32").reshape(10, 10)
         image = np.stack([channel, channel * 100 + 42], axis=-1)
@@ -188,8 +189,9 @@ def test_normalize_raw_per_channel(self):
         self.assertEqual(rgb.dtype, np.uint8)
         self.assertTrue(np.array_equal(rgb[..., 0], rgb[..., 1]))
         self.assertFalse(np.any(rgb[..., 2]))
+        self.assertTrue(np.array_equal(rgb, _to_image(image)))
 
-    def test_normalization_invalidates_old_embeddings(self):
+    def test_normalization_invalidates_incompatible_embeddings(self):
         from types import SimpleNamespace
 
         from micro_sam.util import _get_embedding_signature
@@ -200,22 +202,29 @@ def test_normalization_invalidates_old_embeddings(self):
         raw = np.arange(100).reshape(10, 10)
         signature = _get_embedding_signature(raw, predictor, tile_shape=None, halo=None)
         signature["normalization"] = RAW_NORMALIZATION
-        self.assertEqual(signature["normalization"], "percentile_2_98")
+        self.assertEqual(signature["normalization"], "minmax_per_channel")
 
         attrs = {"input_size": [10, 10], **signature}
         embeddings = SimpleNamespace(attrs=attrs)
         self.assertFalse(_check_saved_embeddings(raw, predictor, embeddings, "cache.zarr", None, None))
 
-        del attrs["normalization"]
+        attrs["normalization"] = "percentile_2_98"
         self.assertTrue(_check_saved_embeddings(raw, predictor, embeddings, "cache.zarr", None, None))
 
+        # Untagged caches predate percentile normalization and used the same min-max policy.
+        del attrs["normalization"]
+        self.assertFalse(_check_saved_embeddings(raw, predictor, embeddings, "cache.zarr", None, None))
+
         class PartialEmbeddings(dict):
             def __init__(self, normalization=None):
                 super().__init__(features=object())
                 self.attrs = {} if normalization is None else {"normalization": normalization}
 
         legacy_partial = PartialEmbeddings()
-        self.assertTrue(_check_saved_embeddings(raw, predictor, legacy_partial, "cache.zarr", None, None))
+        self.assertFalse(_check_saved_embeddings(raw, predictor, legacy_partial, "cache.zarr", None, None))
+
+        percentile_partial = PartialEmbeddings("percentile_2_98")
+        self.assertTrue(_check_saved_embeddings(raw, predictor, percentile_partial, "cache.zarr", None, None))
 
         current_partial = PartialEmbeddings(RAW_NORMALIZATION)
         self.assertFalse(_check_saved_embeddings(raw, predictor, current_partial, "cache.zarr", None, None))