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

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,34 @@ 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.
+    """
+    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 +632,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 +722,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()

applications/dynaclr/src/dynaclr/data/dataset.py

@@ -36,7 +36,7 @@
 from dynaclr.data.index import MultiExperimentIndex
 from dynaclr.data.tau_sampling import sample_tau
 from viscy_data._typing import ULTRACK_INDEX_COLUMNS, NormMeta, SampleMeta
-from viscy_data._utils import _read_norm_meta
+from viscy_data._utils import read_norm_meta
 
 
 def _pick_temporal_candidate(
@@ -742,7 +742,7 @@ def _build_norm_meta(
         cache_key = (store_path, fov_name)
         if cache_key not in self._norm_meta_cache:
             position = self._get_position(store_path, fov_name)
-            self._norm_meta_cache[cache_key] = _read_norm_meta(position)
+            self._norm_meta_cache[cache_key] = read_norm_meta(position)
         cached = self._norm_meta_cache[cache_key]
         if cached is None:
             return None

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",
@@ -136,7 +136,7 @@ def _collate_samples(batch: Sequence[Sample]) -> Sample:
     return collated
 
 
-def _read_norm_meta(fov: Position) -> NormMeta | None:
+def read_norm_meta(fov: Position) -> NormMeta | None:
     """Read normalization metadata from the FOV.
 
     Convert to float32 tensors to avoid automatic casting to float64.

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

@@ -21,7 +21,7 @@
 from torch.utils.data import DataLoader, Dataset
 
 from viscy_data._typing import ULTRACK_INDEX_COLUMNS, AnnotationColumns
-from viscy_data._utils import _read_norm_meta
+from viscy_data._utils import read_norm_meta
 
 
 class ClassificationDataset(Dataset):
@@ -100,7 +100,7 @@ def __getitem__(self, idx) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, dict
                 slice(x - x_half, x + x_half),
             ]
         ).float()[None]
-        norm_meta = _read_norm_meta(fov)
+        norm_meta = read_norm_meta(fov)
         if norm_meta is None:
             raise ValueError(f"Normalization metadata not found for FOV '{fov_name}'.")
         norm_meta = norm_meta[self.channel_name]["fov_statistics"]

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

@@ -19,7 +19,7 @@
 from torch.utils.data import DataLoader, Dataset
 
 from viscy_data._typing import DictTransform, NormMeta
-from viscy_data._utils import _ensure_channel_list, _read_norm_meta
+from viscy_data._utils import _ensure_channel_list, read_norm_meta
 from viscy_data.distributed import ShardedDistributedSampler
 from viscy_data.select import SelectWell
 
@@ -163,7 +163,7 @@ def __init__(
         self._metadata_map: dict[int, _CacheMetadata] = {}
         for position in positions:
             img = position[array_key]
-            norm_meta = _read_norm_meta(position)
+            norm_meta = read_norm_meta(position)
             for time_idx in range(img.frames):
                 cache_map[key] = None
                 self._metadata_map[key] = (position, time_idx, norm_meta)

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

@@ -23,7 +23,7 @@
     MemoryMappedTensor = None
 
 from viscy_data._typing import DictTransform, NormMeta
-from viscy_data._utils import _ensure_channel_list, _read_norm_meta
+from viscy_data._utils import _ensure_channel_list, read_norm_meta
 from viscy_data.gpu_aug import GPUTransformDataModule
 from viscy_data.select import SelectWell
 
@@ -75,7 +75,7 @@ def __init__(
         self._metadata_map: dict[int, _CacheMetadata] = {}
         for position in positions:
             img = position[array_key]
-            norm_meta = _read_norm_meta(position)
+            norm_meta = read_norm_meta(position)
             for time_idx in range(img.frames):
                 cache_map[key] = None
                 self._metadata_map[key] = (position, time_idx, norm_meta)

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

@@ -12,7 +12,7 @@
 from torch.utils.data import Dataset
 
 from viscy_data._typing import ChannelMap, DictTransform, HCSStackIndex, NormMeta, Sample
-from viscy_data._utils import _ensure_channel_list, _read_norm_meta, _search_int_in_str
+from viscy_data._utils import _ensure_channel_list, _search_int_in_str, read_norm_meta
 from viscy_data.foreground_masks import ForegroundMaskSupport
 
 _logger = logging.getLogger("lightning.pytorch")
@@ -134,7 +134,7 @@ def _get_windows(self) -> None:
             w += ts * zs
             self.window_keys.append(w)
             self.window_arrays.append(img_arr)
-            self.window_norm_meta.append(_read_norm_meta(fov))
+            self.window_norm_meta.append(read_norm_meta(fov))
             if self.fg_mask_support is not None:
                 self.fg_mask_support.validate_and_store(fov, img_arr, self.target_ch_idx)
         self._max_window = w

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

@@ -31,8 +31,8 @@
 
 from viscy_data._typing import ULTRACK_INDEX_COLUMNS, NormMeta
 from viscy_data._utils import (
-    _read_norm_meta,
     _transform_channel_wise,
+    read_norm_meta,
 )
 from viscy_data.hcs import HCSDataModule
 from viscy_data.select import _filter_fovs, _filter_wells
@@ -239,7 +239,7 @@ def _slice_patch(self, track_row: "pd.Series") -> "tuple[ts.TensorStore, NormMet
             slice(y_center - y_half, y_center + y_half),
             slice(x_center - x_half, x_center + x_half),
         ]
-        return patch, _read_norm_meta(position)
+        return patch, read_norm_meta(position)
 
     def _slice_patches(self, track_rows: "pd.DataFrame"):
         """Slice and stack patches for multiple track rows."""