refactor(cytoland): unify VSUNet rotation TTA on rotation_tta_transforms

[ieivanov]

Stacked on #450.

Summary

VSUNet.perform_test_time_augmentations (the viscy predict TTA path) had its own rotation implementation — _rotate_volume (square-pad + monai Rotate90) + _crop_to_original — separate from AugmentedPredictionVSUNet's transform-list path. This unifies both on the rotation_tta_transforms() helper added in #450, so rotation TTA has a single source of truth.

• perform_test_time_augmentations now builds its rotations from rotation_tta_transforms() and reuses _pad_forward_crop, which crops to the augmented shape before the inverse rotation — so non-square FOVs work with bare torch.rot90, no square padding.
• Removes _rotate_volume, _crop_to_original, the _original_shape_yx attribute, and the now-unused Rotate90 import.
• mean/median/product reductions preserved.

Behavior change (intended)

For non-square FOVs, viscy predict TTA output changes: the old path zero-padded to a square before rotating (introducing border regions that were cropped back out); the new path rotates without padding. Square FOVs are unchanged. This is a correctness improvement for non-square inputs.

Test plan

• New test_vsunet_perform_test_time_augmentations_nonsquare (64×64, 64×48, 48×64 × mean/median).
• pytest test_engine.py -k "tta or rotation or augment or sliding_windows" → 15 passed.

🤖 Generated with Claude Code

[ieivanov]

Ed, Alex, this is optional, but I think nice to have. It unifies how AugmentedPredictionVSUNet and VSUNet do TTA transforms. As a result, now viscy predict gives better results (at least according to Claude) for non-square inputs. Let me know what you think, happy to refactor or scrap it altogether if you think it's not needed.

[Copilot]

Pull request overview

This PR refactors Cytoland’s VSUNet.perform_test_time_augmentations (the viscy predict TTA path) to reuse the shared rotation_tta_transforms() helper, aligning it with the transform-list TTA implementation and fixing non-square FOV rotation behavior by removing square-padding rotation logic.

Changes:

• Switched VSUNet.perform_test_time_augmentations to use rotation_tta_transforms() + _pad_forward_crop, and removed the bespoke square-pad rotation helpers.
• Removed now-unused rotation-related state (_original_shape_yx) and Rotate90 import.
• Added a regression test to ensure non-square FOVs preserve YX shape under VSUNet rotation TTA for mean/median reductions.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
applications/cytoland/src/cytoland/engine.py	Unifies VSUNet rotation TTA implementation on rotation_tta_transforms() and removes square-padding rotation helpers.
applications/cytoland/tests/test_engine.py	Adds regression coverage for VSUNet rotation TTA on non-square YX shapes.

Comments suppressed due to low confidence (1)

applications/cytoland/src/cytoland/engine.py:491

• perform_test_time_augmentations will raise UnboundLocalError if self.tta_type is anything other than "mean"/"median"/"product" because prediction is never assigned in that case. Consider failing fast with a clear ValueError so misconfiguration is easier to debug.

        if self.tta_type == "mean":
            prediction = torch.stack(predictions).mean(dim=0)
        elif self.tta_type == "median":
            prediction = torch.stack(predictions).median(dim=0).values
        elif self.tta_type == "product":

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

[ieivanov]

⚠️ Known failing test: product TTA produces NaN

test_vsunet_perform_test_time_augmentations_nonsquare[product-*] fails for all 3 FOV shapes (yx0/yx1/yx2). mean and median pass.

Cause: tta_type="product" aggregates predictions in log space (torch.log(p + 1e-9)). The test uses an untrained model with random weights, whose outputs are frequently negative → log() of a negative is NaN, so the torch.isfinite(output).all() assertion fails.

Not introduced by the main merge (a31eaa45) — confirmed identical failures on the pre-merge branch tip cf1a4a71 ("add 'product' to tta tests"), where the product parametrization was first added.

Suggested fix (separate from this PR): either

• clamp/relu predictions to be non-negative before the log in perform_test_time_augmentations, or
• give the product case a fixture with positive outputs (the product method only makes sense for non-negative predictions anyway).