From 6d632e2fcb9fbe500e2331e5bff26121f1486013 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 10 Jun 2026 13:13:16 -0700 Subject: [PATCH] feat(cytoland): non-square rotation TTA + reuse helpers (#450) Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> (cherry picked from commit 674dbeef1561cb45659b515f6fea1b0192b75271) --- .../cytoland/src/cytoland/__init__.py | 2 + applications/cytoland/src/cytoland/engine.py | 73 ++++++++++++++++++- applications/cytoland/tests/test_engine.py | 43 +++++++++++ .../viscy-data/src/viscy_data/__init__.py | 4 + packages/viscy-data/src/viscy_data/_utils.py | 10 ++- 5 files changed, 129 insertions(+), 3 deletions(-) diff --git a/applications/cytoland/src/cytoland/__init__.py b/applications/cytoland/src/cytoland/__init__.py index 6df466c9e..d399e9802 100644 --- a/applications/cytoland/src/cytoland/__init__.py +++ b/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", ] diff --git a/applications/cytoland/src/cytoland/engine.py b/applications/cytoland/src/cytoland/engine.py index 03b272c08..83ca7ea66 100644 --- a/applications/cytoland/src/cytoland/engine.py +++ b/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] diff --git a/applications/cytoland/tests/test_engine.py b/applications/cytoland/tests/test_engine.py index de1ed5138..ac2a03759 100644 --- a/applications/cytoland/tests/test_engine.py +++ b/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() diff --git a/packages/viscy-data/src/viscy_data/__init__.py b/packages/viscy-data/src/viscy_data/__init__.py index edbec8ed6..4d2d54214 100644 --- a/packages/viscy-data/src/viscy_data/__init__.py +++ b/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 diff --git a/packages/viscy-data/src/viscy_data/_utils.py b/packages/viscy-data/src/viscy_data/_utils.py index e6a6523c7..cad79cb7d 100644 --- a/packages/viscy-data/src/viscy_data/_utils.py +++ b/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.