Add random percentile normalization augmentation (#1299)

* Add random percentile normalization augmentation

* Round sampled percentiles to one decimal place

* Shim down percentile upper and lower bound ranges

* Move percentile normalization transform to torch-em

* Switch validation normalization to 2-98

Author: anwai98

micro_sam/v2/datasets/generalist_loader.py

@@ -17,14 +17,14 @@
 from .sampler import UniBatchSampler, _build_group_map
 from ..transforms.raw import (
     _identity, _cellpose_raw_trafo, _to_8bit, _normalize_percentile, _resize_raw_to_512, _resize_to_512,
+    get_random_percentile_normalization,
 )
 from ..transforms.labels import (
     _em_cell_label_trafo, _joint_em_cell_label_trafo,
     _plantseg_label_trafo, _axondeepseg_pre_label_transform, _instance_labels,
     _JointLabelTransform,
 )
 
-
 # Cap on validation samples drawn per dataset, to keep the per-epoch validation pass cheap.
 # Each access is a random crop (see UniDataWrapper.max_samples), so this is N random samples.
 N_SAMPLES_VAL = 50
@@ -34,6 +34,11 @@
 # Sam2Trainer._validate_impl, so the validation metric is comparable across epochs.
 VALIDATION_SEED = 42
 
+# Train with uniformly sampled symmetric percentiles; validate deterministically with 2nd/98th percentiles
+# to match the inference-time normalization in normalize_raw.
+TRAIN_LOWER_PERCENTILE_BOUNDS = (0.0, 5.0)
+VALIDATION_LOWER_PERCENTILE_BOUNDS = (2.0, 2.0)
+
 
 def seed_worker(worker_id):
     """DataLoader worker_init_fn that pins per-worker RNG for deterministic validation crops.
@@ -54,6 +59,41 @@ def _ensure_native_byte_order(y):
     return y.byteswap().view(y.dtype.newbyteorder()) if not y.dtype.isnative else y
 
 
+def _set_percentile_normalization(dataset, lower_percentile_bounds):
+    """Replace fixed normalization in all torch-em leaves of a dataset tree."""
+    if isinstance(dataset, (list, tuple)):
+        for ds in dataset:
+            _set_percentile_normalization(ds, lower_percentile_bounds)
+        return
+
+    if isinstance(dataset, UniDataWrapper):
+        _set_percentile_normalization(dataset.ds, lower_percentile_bounds)
+        return
+
+    children = getattr(dataset, "datasets", None)
+    if children is not None:
+        for ds in children:
+            _set_percentile_normalization(ds, lower_percentile_bounds)
+        return
+
+    if not hasattr(dataset, "raw_transform"):
+        raise TypeError(f"Cannot configure raw normalization for dataset of type {type(dataset).__name__}.")
+
+    dataset.raw_transform = get_random_percentile_normalization(
+        dataset.raw_transform, lower_percentile_bounds=lower_percentile_bounds
+    )
+
+
+def _configure_training_normalization(train_datasets, val_datasets):
+    """Enable random percentile augmentation for training and deterministic 2nd/98th validation."""
+    _set_percentile_normalization(
+        train_datasets, lower_percentile_bounds=TRAIN_LOWER_PERCENTILE_BOUNDS,
+    )
+    _set_percentile_normalization(
+        val_datasets, lower_percentile_bounds=VALIDATION_LOWER_PERCENTILE_BOUNDS,
+    )
+
+
 def _prepare_data_loader(dataset, batch_size, shuffle, batch_size_per_group=None, num_workers=32, deterministic=False):
     # For deterministic validation, re-seed workers every epoch via worker_init_fn.
     # This requires non-persistent workers, since persistent workers run worker_init_fn only once.
@@ -163,7 +203,7 @@ def _get_embedseg_datasets(split_choice, z):
                 "Mouse-Organoid-Cells-CBG", "Mouse-Skull-Nuclei-CBG",
                 "Platynereis-ISH-Nuclei-CBG", "Platynereis-Nuclei-CBG",
             ]
-        else:   # Only two datasets have the test split.
+        else:  # Only two datasets have the test split.
             names = ["Mouse-Skull-Nuclei-CBG", "Platynereis-ISH-Nuclei-CBG"]
 
         all_embedseg_datasets = [
@@ -397,8 +437,8 @@ def _get_em_datasets(input_path, patch_shape, z_slices, kwargs, label_trafo, _em
     """Get all electron microscopy (EM) datasets for generalist training.
 
     Args:
-        _em_label_trafo: EM cell label transform function to use.  Defaults to
-            :func:`_em_cell_label_trafo`.  Pass :func:`_joint_em_cell_label_trafo`
+        _em_label_trafo: EM cell label transform function to use. Defaults to
+            :func:`_em_cell_label_trafo`. Pass :func:`_joint_em_cell_label_trafo`
             when building joint interactive+automatic datasets.
 
     Returns:
@@ -705,6 +745,8 @@ def get_dataloaders(
         train_ds.extend(em_train)
         val_ds.extend(em_val)
 
+    _configure_training_normalization(train_ds, val_ds)
+
     # Finally, we prepare a 'ConcatDataset' for all the available datasets.
     train_ds = ConcatDataset(*train_ds)
     val_ds = ConcatDataset(*val_ds)
@@ -739,7 +781,7 @@ def get_interactive_dataloaders(
 
     Identical dataset composition to :func:`get_dataloaders` but returns raw
     integer instance labels (``label_dtype=torch.int64``) instead of distance
-    transforms.  Used with :class:`micro_sam.v2.training.ConvertToSam2VideoBatch`.
+    transforms. Used with :class:`micro_sam.v2.training.ConvertToSam2VideoBatch`.
 
     Args:
         input_path: Root path to the generalist training data.
@@ -815,6 +857,7 @@ def _build_automatic_datasets(input_path, z_slices, dataset_choice):
         train_ds.extend(em_train)
         val_ds.extend(em_val)
 
+    _configure_training_normalization(train_ds, val_ds)
     return ConcatDataset(*train_ds), ConcatDataset(*val_ds)
 
 
@@ -849,6 +892,8 @@ def _build_interactive_datasets(input_path, z_slices, dataset_choice):
         train_ds.extend(em_train)
         val_ds.extend(em_val)
 
+    _configure_training_normalization(train_ds, val_ds)
+
     # Cap each validation dataset to N_SAMPLES_VAL random samples so the per-epoch
     # validation pass stays cheap (train datasets are left at full size).
     for w in val_ds:
@@ -897,6 +942,8 @@ def _build_joint_datasets(input_path, z_slices, dataset_choice):
         train_ds.extend(em_train)
         val_ds.extend(em_val)
 
+    _configure_training_normalization(train_ds, val_ds)
+
     # Cap each validation dataset to N_SAMPLES_VAL random samples so the per-epoch
     # validation pass stays cheap (matches the interactive builder; train datasets are full size).
     for w in val_ds:

micro_sam/v2/transforms/raw.py

@@ -1,9 +1,12 @@
 import random
+from functools import partial
+from typing import Callable, Dict, Optional, Tuple
 
 import numpy as np
 import torch
 import torchvision.transforms.functional as TF
 from torchvision.transforms import ColorJitter
+from torch_em.transform.raw import RandomPercentileNormalization, RawTransform
 
 from micro_sam.v2.normalization import normalize_raw
 
@@ -29,25 +32,35 @@ def to_rgb(image):
     return image
 
 
-def _to_8bit(raw):
-    "Ensures three channels for inputs and percentile-normalizes them to [0, 1]."
+def _prepare_to_8bit(raw):
+    """Prepare raw inputs for ``_to_8bit`` without changing their dynamic range."""
     if raw.ndim == 3 and raw.shape[0] == 1:  # If the inputs have 1 channel, we triplicate it.
         raw = np.concatenate([raw] * 3, axis=0)
 
-    raw = to_rgb(raw)  # Ensure all images are in 3-channels: triplicate one channel to three channels.
+    return to_rgb(raw)  # Ensure channel-first RGB without rescaling the original intensities.
+
+
+def _to_8bit(raw):
+    "Ensures three channels for inputs and percentile-normalizes them to [0, 1]."
+    raw = _prepare_to_8bit(raw)
     raw = normalize_raw(raw, axis=(1, 2))
     return raw
 
 
+def _prepare_identity(x):
+    """Prepare raw inputs for ``_identity`` without changing their dynamic range."""
+    return to_rgb(x)
+
+
 def _identity(x):
     "Ensures three channels for inputs and percentile-normalizes them to [0, 1]."
-    x = to_rgb(x)
+    x = _prepare_identity(x)
     x = normalize_raw(x, axis=(1, 2))
     return x
 
 
-def _cellpose_raw_trafo(x):
-    """Transforms input images to desired format.
+def _prepare_cellpose_raw(x):
+    """Prepare CellPose inputs without changing their dynamic range.
 
     NOTE: The input channel logic is arranged a bit strangely in `cyto` dataset.
     This function takes care of it here.
@@ -64,7 +77,12 @@ def _cellpose_raw_trafo(x):
         # The image is 2 channels and we sort the channels such that - 0: cell, 1: nucleus
         x = np.stack([g, r, np.zeros_like(b)], axis=0)
 
-    x = to_rgb(x)  # Ensures three channels for inputs and avoids rescaling inputs.
+    return to_rgb(x)  # Ensures three channels for inputs and avoids rescaling inputs.
+
+
+def _cellpose_raw_trafo(x):
+    """Prepare and percentile-normalize CellPose inputs."""
+    x = _prepare_cellpose_raw(x)
     x = normalize_raw(x, axis=(1, 2))
     return x
 
@@ -77,8 +95,13 @@ def _resize_to_512(x, is_label=False):
 
 def _resize_raw_to_512(x):
     """Resize small raw volume patch to 512×512 and normalize."""
-    x = _resize_to_512(x, is_label=False)
-    return _identity(x)
+    x = _prepare_resize_raw_to_512(x)
+    return normalize_raw(x, axis=(1, 2))
+
+
+def _prepare_resize_raw_to_512(x):
+    """Resize a raw patch without normalizing it."""
+    return _prepare_identity(_resize_to_512(x, is_label=False))
 
 
 class VideoAugment:
@@ -354,3 +377,52 @@ def _normalize_percentile(x, axis=None):
     'rgb' format of TissueNet image data for the expected axes.
     """
     return normalize_raw(x, axis=axis)
+
+
+def get_random_percentile_normalization(
+    raw_transform: Callable,
+    lower_percentile_bounds: Tuple[float, float] = (0.0, 5.0),
+    distribution: str = "uniform",
+    distribution_kwargs: Optional[Dict[str, float]] = None,
+) -> RawTransform:
+    """Convert a generalist raw transform to random percentile normalization.
+
+    The data-specific shape/channel preparation is retained as ``RawTransform.augmentation1``, so the
+    normalizer receives data in its original intensity range. An existing ``augmentation2`` is preserved.
+    """
+    augmentation2 = None
+    if isinstance(raw_transform, RawTransform):
+        if not isinstance(raw_transform.normalizer, RandomPercentileNormalization):
+            raise ValueError(f"Unsupported generalist raw transform: {raw_transform!r}")
+        augmentation1, augmentation2 = raw_transform.augmentation1, raw_transform.augmentation2
+        axis = raw_transform.normalizer.axis
+    elif isinstance(raw_transform, RandomPercentileNormalization):
+        augmentation1, axis = None, raw_transform.axis
+    elif raw_transform is _identity:
+        augmentation1, axis = _prepare_identity, (1, 2)
+    elif raw_transform is _to_8bit:
+        augmentation1, axis = _prepare_to_8bit, (1, 2)
+    elif raw_transform is _cellpose_raw_trafo:
+        augmentation1, axis = _prepare_cellpose_raw, (1, 2)
+    elif raw_transform is _resize_raw_to_512:
+        augmentation1, axis = _prepare_resize_raw_to_512, (1, 2)
+    elif raw_transform is _normalize_percentile:
+        augmentation1, axis = None, None
+    elif isinstance(raw_transform, partial) and raw_transform.func is _normalize_percentile:
+        if raw_transform.args:
+            raise ValueError("Positional arguments are not supported for _normalize_percentile transforms.")
+        unexpected_kwargs = set(raw_transform.keywords) - {"axis"}
+        if unexpected_kwargs:
+            raise ValueError(f"Unsupported _normalize_percentile arguments: {sorted(unexpected_kwargs)}")
+        augmentation1, axis = None, raw_transform.keywords.get("axis")
+    else:
+        raise ValueError(f"Unsupported generalist raw transform: {raw_transform!r}")
+
+    normalizer = RandomPercentileNormalization(
+        lower_percentile_bounds=lower_percentile_bounds,
+        distribution=distribution,
+        distribution_kwargs=distribution_kwargs,
+        rounding_decimals=1,
+        axis=axis,
+    )
+    return RawTransform(normalizer=normalizer, augmentation1=augmentation1, augmentation2=augmentation2)

test/test_raw_transforms.py

@@ -0,0 +1,72 @@
+import unittest
+from functools import partial
+
+import numpy as np
+
+
+class TestRandomPercentileNormalization(unittest.TestCase):
+    def test_factory_preserves_data_specific_preprocessing(self):
+        from torch_em.transform.raw import RandomPercentileNormalization, RawTransform
+
+        from micro_sam.v2.transforms.raw import _normalize_percentile, _to_8bit, get_random_percentile_normalization
+
+        raw = np.arange(16, dtype="uint16").reshape(1, 4, 4)
+        transform = get_random_percentile_normalization(_to_8bit)
+        self.assertIsInstance(transform, RawTransform)
+        self.assertIsInstance(transform.normalizer, RandomPercentileNormalization)
+        self.assertEqual(transform.normalizer.lower_percentile_bounds, (0.0, 5.0))
+        self.assertEqual(transform.normalizer.distribution, "uniform")
+        self.assertIsNone(transform.normalizer.distribution_kwargs)
+        self.assertEqual(transform.normalizer.rounding_decimals, 1)
+        self.assertEqual(transform(raw).shape, (3, 4, 4))
+
+        transform = get_random_percentile_normalization(
+            partial(_normalize_percentile, axis=(1, 2)), lower_percentile_bounds=(0.0, 0.0),
+        )
+        self.assertEqual(transform.normalizer.axis, (1, 2))
+
+        def postprocessing(x):
+            return x
+
+        transform.augmentation2 = postprocessing
+        reconfigured = get_random_percentile_normalization(transform, lower_percentile_bounds=(0.0, 1.0))
+        self.assertIs(reconfigured.augmentation1, transform.augmentation1)
+        self.assertIs(reconfigured.augmentation2, postprocessing)
+
+
+class TestGeneralistNormalizationConfiguration(unittest.TestCase):
+    def test_training_is_random_and_validation_is_deterministic(self):
+        from torch_em.transform.raw import RandomPercentileNormalization, RawTransform
+
+        from micro_sam.v2.datasets.generalist_loader import (
+            TRAIN_LOWER_PERCENTILE_BOUNDS, VALIDATION_LOWER_PERCENTILE_BOUNDS,
+            _configure_training_normalization,
+        )
+        from micro_sam.v2.datasets.wrapper import UniDataWrapper
+        from micro_sam.v2.transforms.raw import _identity
+
+        class Dataset:
+            def __init__(self):
+                self.raw_transform = _identity
+
+        train_leaf, val_leaf = Dataset(), Dataset()
+        train_datasets = [UniDataWrapper(train_leaf)]
+        val_datasets = [UniDataWrapper(val_leaf)]
+        _configure_training_normalization(train_datasets, val_datasets)
+
+        self.assertIsInstance(train_leaf.raw_transform, RawTransform)
+        self.assertIsInstance(train_leaf.raw_transform.normalizer, RandomPercentileNormalization)
+        self.assertEqual(train_leaf.raw_transform.normalizer.distribution, "uniform")
+        self.assertEqual(train_leaf.raw_transform.normalizer.lower_percentile_bounds, TRAIN_LOWER_PERCENTILE_BOUNDS)
+
+        self.assertIsInstance(val_leaf.raw_transform, RawTransform)
+        self.assertIsInstance(val_leaf.raw_transform.normalizer, RandomPercentileNormalization)
+        self.assertEqual(val_leaf.raw_transform.normalizer.distribution, "uniform")
+        self.assertEqual(
+            val_leaf.raw_transform.normalizer.lower_percentile_bounds, VALIDATION_LOWER_PERCENTILE_BOUNDS
+        )
+        self.assertEqual(val_leaf.raw_transform.normalizer.sample_percentiles(), (2.0, 98.0))
+
+
+if __name__ == "__main__":
+    unittest.main()