feat(cytoland): non-square rotation TTA + reuse helpers (#450) (#465)

(cherry picked from commit 674dbee)

Co-authored-by: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: 4 people

applications/cytoland/src/cytoland/__init__.py

@@ -5,6 +5,7 @@
     FcmaeUNet,
     MaskedMSELoss,
     VSUNet,
+    rotation_tta_transforms,
 )
 from cytoland.evaluation import SegmentationMetrics2D
 
@@ -14,4 +15,5 @@
     "MaskedMSELoss",
     "SegmentationMetrics2D",
     "VSUNet",
+    "rotation_tta_transforms",
 ]

applications/cytoland/src/cytoland/engine.py

@@ -3,6 +3,7 @@
 import inspect
 import logging
 import os
+from functools import partial
 from typing import Callable, Literal, Sequence
 
 import numpy as np
@@ -70,6 +71,36 @@ def _center_crop_to_shape(tensor: Tensor, spatial_shape: tuple[int, ...]) -> Ten
     return tensor[tuple(slices)]
 
 
+def rotation_tta_transforms(
+    n: int = 4,
+) -> tuple[list[Callable[[Tensor], Tensor]], list[Callable[[Tensor], Tensor]]]:
+    """Build forward/inverse 90-degree rotation transforms for test-time augmentation.
+
+    Returns ``(forward_transforms, inverse_transforms)`` suitable for
+    :class:`AugmentedPredictionVSUNet`. Each forward transform rotates the YX
+    plane by ``k * 90`` degrees (``k = 0 .. n-1``) and the matching inverse
+    rotates back. Combined with ``reduction="median"`` this reproduces the
+    rotation TTA used by :meth:`VSUNet.perform_test_time_augmentations`, and
+    (unlike passing rotations through other code paths) works for non-square
+    fields of view.
+
+    Parameters
+    ----------
+    n : int, optional
+        Number of 90-degree rotations, by default 4 (0, 90, 180, 270 degrees).
+
+    Returns
+    -------
+    tuple[list[Callable], list[Callable]]
+        The forward and inverse transform lists.
+    """
+    if n < 1:
+        raise ValueError(f"n must be >= 1, got {n}")
+    forward = [partial(torch.rot90, k=k, dims=(-2, -1)) for k in range(n)]
+    inverse = [partial(torch.rot90, k=-k, dims=(-2, -1)) for k in range(n)]
+    return forward, inverse
+
+
 class MaskedMSELoss(nn.Module):
     """Masked MSE loss for FCMAE pre-training."""
 
@@ -603,6 +634,40 @@ def __init__(
         self._inverse_transforms = inverse_transforms or [_identity]
         self._reduction = reduction
 
+    @classmethod
+    def with_rotation_tta(
+        cls,
+        model: nn.Module,
+        n_rotations: int = 4,
+        reduction: Literal["mean", "median"] = "median",
+    ) -> "AugmentedPredictionVSUNet":
+        """Build a predictor that applies 90-degree rotation test-time augmentation.
+
+        Convenience constructor that wires :func:`rotation_tta_transforms` into
+        the forward/inverse transform lists, so callers do not have to build
+        them by hand. Works for non-square fields of view.
+
+        Parameters
+        ----------
+        model : nn.Module
+            The model to wrap.
+        n_rotations : int, optional
+            Number of 90-degree rotations, by default 4.
+        reduction : {"mean", "median"}, optional
+            How to aggregate the rotated predictions, by default "median".
+
+        Returns
+        -------
+        AugmentedPredictionVSUNet
+        """
+        forward_transforms, inverse_transforms = rotation_tta_transforms(n_rotations)
+        return cls(
+            model=model,
+            forward_transforms=forward_transforms,
+            inverse_transforms=inverse_transforms,
+            reduction=reduction,
+        )
+
     def forward(self, x: Tensor) -> Tensor:
         """Run forward pass through the model.
 
@@ -659,9 +724,15 @@ def _predict_with_tta(self, source: Tensor) -> Tensor:
         preds = []
         for fwd_t, inv_t in zip(self._forward_transforms, self._inverse_transforms):
             aug_source = fwd_t(source)
+            # Crop back to the augmented (post-forward-transform) spatial shape,
+            # not the original one: a shape-changing transform such as a 90/270
+            # degree rotation swaps Y and X, so the prediction lives in the
+            # augmented frame until ``inv_t`` undoes the transform. Cropping to
+            # ``source.shape[2:]`` here would be wrong for non-square inputs.
+            aug_shape = aug_source.shape[2:]
             aug_source = self._predict_pad(aug_source)
             pred = self.forward(aug_source)
-            pred = _center_crop_to_shape(pred, source.shape[2:])
+            pred = _center_crop_to_shape(pred, aug_shape)
             preds.append(inv_t(pred))
         if len(preds) == 1:
             return preds[0]

applications/cytoland/tests/test_engine.py

@@ -206,3 +206,46 @@ def test_predict_sliding_windows_missing_out_stack_depth():
     vs = AugmentedPredictionVSUNet(model=model)
     with pytest.raises(ValueError, match="out_stack_depth"):
         vs.predict_sliding_windows(torch.randn(1, 1, 10, 4, 4))
+
+
+def test_rotation_tta_transforms():
+    """Verify the rotation TTA factory builds matched forward/inverse rotations."""
+    from cytoland.engine import rotation_tta_transforms
+
+    forward, inverse = rotation_tta_transforms()
+    assert len(forward) == len(inverse) == 4
+    x = torch.randn(1, 1, 5, 6, 8)  # non-square YX
+    for fwd_t, inv_t in zip(forward, inverse):
+        # inverse(forward(x)) is the identity and restores the original shape
+        assert torch.allclose(inv_t(fwd_t(x)), x)
+
+
+@pytest.mark.parametrize("yx", [(64, 64), (64, 48), (48, 64)])
+def test_predict_sliding_windows_rotation_tta_nonsquare(yx):
+    """Verify rotation TTA + sliding windows works for non-square FOVs.
+
+    Regression test: ``_predict_with_tta`` must crop to the augmented (rotated)
+    shape, otherwise 90/270-degree rotations on non-square inputs produce
+    mismatched shapes and fail to reduce.
+    """
+    z_window, depth, out_channels = 5, 8, 2
+    height, width = yx
+    model = VSUNet(
+        architecture="fcmae",
+        model_config={
+            "in_channels": 1,
+            "out_channels": out_channels,
+            "encoder_blocks": [2, 2, 2, 2],
+            "dims": [4, 8, 16, 32],
+            "decoder_conv_blocks": 1,
+            "stem_kernel_size": [z_window, 4, 4],
+            "in_stack_depth": z_window,
+            "pretraining": False,
+        },
+    )
+    vs = AugmentedPredictionVSUNet.with_rotation_tta(model.model, reduction="median").eval()
+    x = torch.randn(1, 1, depth, height, width)
+    with torch.inference_mode():
+        output = vs.predict_sliding_windows(x, out_channel=out_channels, step=1)
+    assert output.shape == (1, out_channels, depth, height, width)
+    assert torch.isfinite(output).all()

packages/viscy-data/src/viscy_data/__init__.py

@@ -72,6 +72,9 @@
 except ImportError:
     pass
 
+# Normalization metadata reader (from _utils.py)
+from viscy_data._utils import read_norm_meta
+
 # Channel dropout augmentation (from channel_dropout.py)
 from viscy_data.channel_dropout import ChannelDropout
 
@@ -150,6 +153,7 @@
     "ChannelDropout",
     # Utilities
     "FlexibleBatchSampler",
+    "read_norm_meta",
     "SelectWell",
     "ShardedDistributedSampler",
     # Core

packages/viscy-data/src/viscy_data/_utils.py

@@ -2,7 +2,7 @@
 
 This module centralizes helper functions that are used by multiple data modules:
 - From ``hcs.py``: ``_ensure_channel_list``, ``_search_int_in_str``,
-  ``_collate_samples``, ``_read_norm_meta``
+  ``_collate_samples``, ``read_norm_meta``
 - From ``triplet.py``: ``_scatter_channels``, ``_gather_channels``,
   ``_transform_channel_wise``
 """
@@ -24,7 +24,7 @@
     "_collate_samples",
     "_ensure_channel_list",
     "_gather_channels",
-    "_read_norm_meta",
+    "read_norm_meta",
     "_scatter_channels",
     "_search_int_in_str",
     "_transform_channel_wise",
@@ -165,6 +165,12 @@ def _read_norm_meta(fov: Position) -> NormMeta | None:
     return norm_meta
 
 
+# Public alias re-exported as ``viscy_data.read_norm_meta``; the canonical
+# definition keeps its private name since this is a private ``_utils`` module
+# and all internal callers import ``_read_norm_meta``.
+read_norm_meta = _read_norm_meta
+
+
 def _collate_norm_meta(norm_metas: list[NormMeta]) -> NormMeta:
     """Stack per-sample norm_meta dicts into batched tensors.