feat(cytoland): non-square rotation TTA + reuse helpers by ieivanov · Pull Request #450 · mehta-lab/VisCy

ieivanov · 2026-06-03T23:23:18Z

Summary

Make cytoland's prediction-time augmentation usable from external pipelines (e.g. biahub virtual staining) without re-implementing normalization or TTA.

Fix non-square rotation TTA. AugmentedPredictionVSUNet._predict_with_tta cropped the prediction to the original spatial shape before applying the inverse transform. Shape-changing transforms (90°/270° rotations) swap Y and X, so for non-square FOVs the crop happened in the wrong frame and the per-rotation predictions could not be stacked/reduced. Now it crops to the augmented shape (fwd_t(source).shape[2:]) before inv_t. No square padding needed — works for any aspect ratio.
rotation_tta_transforms(n=4) factory returning the standard 90° rotation (forward, inverse) transform lists.
AugmentedPredictionVSUNet.with_rotation_tta(model, reduction="median") convenience constructor that wires the factory in, so callers don't hand-build transforms. (The library now uses its own helper.)
viscy_data.read_norm_meta(fov) — public wrapper over _read_norm_meta for reading precomputed normalization stats from a FOV's zattrs.

Test plan

New regression tests: test_predict_sliding_windows_rotation_tta_nonsquare (64×64, 64×48, 48×64) and test_rotation_tta_transforms.
pytest test_engine.py -k "rotation_tta or nonsquare or sliding_windows" → 8 passed.

🤖 Generated with Claude Code

Delegate normalization and test-time augmentation to VisCy instead of re-implementing them, eliminating the divergence risk (the previous manual normalization omitted VisCy's divisor epsilon): - Normalization: apply the config's instantiated NormalizeSampled transforms (via Compose) to a per-timepoint sample dict, reading precomputed stats with viscy_data.read_norm_meta. Drops _normalization_params and _read_fov_statistics. - TTA: use AugmentedPredictionVSUNet.with_rotation_tta, which handles non-square FOVs natively. Drops the local _predict_volume square-pad workaround. Depends on the cytoland feat/vs-nonsquare-tta-and-helpers branch (mehta-lab/VisCy#450). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Picks up mehta-lab/VisCy#450's rename of _read_norm_meta to the public read_norm_meta (commit eb1ffa1). No biahub source change; the import already used the public name. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

- Fix AugmentedPredictionVSUNet._predict_with_tta to crop to the augmented (post-transform) spatial shape instead of the original. Shape-changing transforms such as 90/270-degree rotations swap Y and X, so cropping to the original shape was wrong for non-square FOVs (mismatched shapes failed to reduce). Rotation TTA now works for non-square inputs with no square padding. - Add rotation_tta_transforms() factory for the standard 90-degree rotation forward/inverse transform set, and AugmentedPredictionVSUNet.with_rotation_tta() convenience constructor that uses it, so callers don't hand-build transforms. - Rename viscy_data._read_norm_meta to public read_norm_meta (updating all callers) and export it, so external pipelines can read precomputed normalization statistics from a FOV's zattrs. - Add regression tests for non-square rotation TTA and the helper. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Repoints to mehta-lab/VisCy#450 head after its lint fixes (commit 8f42e30). No functional change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

ieivanov · 2026-06-03T23:56:29Z

Refactored biahub virtual-stain implementation: czbiohub-sf/biahub#267

Copilot

Pull request overview

This PR improves Cytoland’s prediction-time augmentation so rotation-based test-time augmentation (TTA) works on non-square fields of view, and it exposes reusable helpers (rotation TTA factory + normalization-metadata reader) for external pipelines.

Changes:

Fix AugmentedPredictionVSUNet._predict_with_tta to crop predictions in the augmented (rotated) frame before applying the inverse transform, enabling 90°/270° TTA on non-square inputs.
Add rotation_tta_transforms() factory and AugmentedPredictionVSUNet.with_rotation_tta() convenience constructor, plus regression tests for non-square rotation TTA.
Promote normalization metadata reading to a public helper read_norm_meta and update call sites across datasets/data modules.

Reviewed changes

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

Show a summary per file

File	Description
packages/viscy-data/src/viscy_data/triplet.py	Switch norm-meta reads to `read_norm_meta`.
packages/viscy-data/src/viscy_data/sliding_window.py	Use `read_norm_meta` when building window metadata.
packages/viscy-data/src/viscy_data/mmap_cache.py	Use `read_norm_meta` when populating cache metadata.
packages/viscy-data/src/viscy_data/gpu_aug.py	Use `read_norm_meta` when populating cache metadata.
packages/viscy-data/src/viscy_data/cell_classification.py	Use `read_norm_meta` for per-FOV normalization statistics.
packages/viscy-data/src/viscy_data/_utils.py	Rename/expose normalization metadata reader as `read_norm_meta`.
packages/viscy-data/src/viscy_data/init.py	Export `read_norm_meta` at package level.
applications/dynaclr/src/dynaclr/data/dataset.py	Update to `read_norm_meta` for cached norm-meta reads.
applications/cytoland/tests/test_engine.py	Add regression tests for rotation TTA factory + non-square sliding-window inference.
applications/cytoland/src/cytoland/engine.py	Add rotation TTA factory + convenience ctor; fix TTA crop logic for non-square rotations.
applications/cytoland/src/cytoland/init.py	Export `rotation_tta_transforms`.

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

ieivanov · 2026-06-05T16:52:52Z

@edyoshikun @alxndrkalinin it would be great if we can squeeze #450 and #451 in the upcoming VisCy release if you agree with the changes I've suggested. If you think you need more time for review, I'm happy if we merge these later also.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Instead of renaming every caller, keep the private name working: define the public read_norm_meta and add `_read_norm_meta = read_norm_meta`. Revert the call sites (triplet, sliding_window, gpu_aug, mmap_cache, cell_classification, dynaclr) back to the private name so existing code is untouched. __all__ still exports the public read_norm_meta. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

edyoshikun

I think there is only one switch for the read_norm_meta . other than that this is good.

edyoshikun · 2026-06-05T19:47:28Z

+    :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


do we want just rotations or do you want to add flips in x and y?

they are really fast operations but you add 4 more transforms. I'm guessing you dont need them right now.

Claude said the current codebase (viscy predict) does only rotations, you here I implemented the same. Let's bundle flips in a separate PR? We could package as with_rotation_tta, with_reflection_tta, and with_tta - which allows you to pick between rotation, reflection, or both.

Opened #462 to track this feature

I mean the ideal case in stead of bloating the code is to just pass the composable list of transforms. The gain from TTA is minimal so we can just do the rotations for now.

…the alias Flip the alias direction so the canonical definition keeps its private name in the private _utils module (where all internal callers import _read_norm_meta) and read_norm_meta is the public re-exported alias. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

ieivanov · 2026-06-09T17:27:25Z

@edyoshikun please take another look, I think I've addressed all your comments here

edyoshikun

Thanks for the changes. LGTM. I think this unblocks you now? @ieivanov

(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>

…ows (#267) * feat(virtual-stain): rewrite CLI on cytoland predict_sliding_windows Replace the viscy-CLI subprocess pipeline (preprocess/predict/combine via conda) with in-process cytoland inference: - Add cytoland (modular VisCy) under the `stain` extra, sourced from the modular-viscy-staging branch. - Validate the predict config against VisCy's own VSUNet/HCSDataModule signatures via jsonargparse, so biahub keeps no copy of the model schema. - Run predict_sliding_windows per position, looping over timepoints on a single GPU; read precomputed normalization stats from the store. - Rotation TTA (default on) is applied at the volume level with square padding so it is correct for non-square FOVs. - Add minimal example config pointing at the vscyto3d finetune checkpoint. Adopt the deskew --cluster/--init pattern: add cluster() and init_only() parsing decorators, extend get_submitit_cluster() to take an explicit cluster string, emit RESOURCES:, support --init, and run --cluster debug in-process in the foreground. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor(virtual-stain): reuse VisCy normalization and rotation TTA Delegate normalization and test-time augmentation to VisCy instead of re-implementing them, eliminating the divergence risk (the previous manual normalization omitted VisCy's divisor epsilon): - Normalization: apply the config's instantiated NormalizeSampled transforms (via Compose) to a per-timepoint sample dict, reading precomputed stats with viscy_data.read_norm_meta. Drops _normalization_params and _read_fov_statistics. - TTA: use AugmentedPredictionVSUNet.with_rotation_tta, which handles non-square FOVs natively. Drops the local _predict_volume square-pad workaround. Depends on the cytoland feat/vs-nonsquare-tta-and-helpers branch (mehta-lab/VisCy#450). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * chore: bump cytoland pin to VisCy read_norm_meta rename Picks up mehta-lab/VisCy#450's rename of _read_norm_meta to the public read_norm_meta (commit eb1ffa1). No biahub source change; the import already used the public name. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * chore: bump cytoland pin to lint-fixed VisCy commit Repoints to mehta-lab/VisCy#450 head after its lint fixes (commit 8f42e30). No functional change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * style: ruff-format virtual_stain.py Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(virtual-stain): expose sliding_window_step, device, output_ome_zarr_version Add the three optional orchestration knobs to the example config (defaults shown). output_ome_zarr_version (default None, mirroring DeskewSettings) now dictates the output store's OME-Zarr version, threaded through _init_output_plate to resolve_ome_zarr_version; None preserves the input store's version. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(virtual-stain): hardcode per-position resources Timepoints are processed sequentially on a single GPU, so resource needs no longer depend on dataset size. Use num_cpus=16, gb_ram=4 instead of estimate_resources. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(virtual-stain): lower slurm_array_parallelism to 20 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(virtual-stain): generous TTA-aware slurm_time budget Base the wall-clock budget on T * Z * 5s (per sliding window) with a 60-minute floor, assuming median TTA. Measured ~1.4 s/window with TTA, so this gives a ~3-4x margin and accounts for stack depth (number of z-windows), which the previous T*3 formula ignored. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs(virtual-stain): add batch_size to example config batch_size has no effect in biahub (predict_sliding_windows runs one z-window at a time); set to 1 with a clarifying comment for parity with viscy predict. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test: register virtual-stain example config (jsonargparse-validated) example_virtual_stain_settings.yml is a viscy predict config validated against VisCy's classes, not a biahub pydantic Settings model, so it broke test_all_example_settings_tested's file count. Track it in a separate list, include it in the expected count, and add a dedicated validation test that skips when cytoland (the `stain` extra) is not installed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(virtual-stain): add click.echo progress to per-position worker Per-position SLURM logs previously showed only submitit start/complete lines. Emit position-prefixed progress: device, checkpoint, TTA status, geometry (T / Z,Y,X / windows-per-timepoint), per-timepoint timing, and a completion summary with total elapsed time. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * build(deps): pin cytoland to VisCy PR #465 commit The feat/vs-nonsquare-tta-and-helpers branch was removed; pin to the underlying commit 7fb3bc1 (VisCy PR #465) instead. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor(virtual-stain): address PR review comments - Remove module-level torch.multiprocessing.set_start_method("spawn", force=True): it was a global import-time side effect and is unused. The inference path has no DataLoader workers or torch multiprocessing, and submitit launches jobs as fresh interpreters (or in-process for debug), so the start method never applied. - Delegate checkpoint loading to VisCy: route the top-level ckpt_path into the model's init args so VSUNet.__init__ extracts the Lightning checkpoint["state_dict"] itself, instead of biahub assuming that layout. - Rename the misleading in_stack_depth log variable to z_window_depth and document that it mirrors VisCy's predict_sliding_windows (which reads model.out_stack_depth as the input Z-window depth); guard the window count with max(0, ...). - Reject cluster='local' with >1 position: local jobs run concurrently without a parallelism cap and all target cuda:0, oversubscribing a single GPU. Point users to slurm (multi-GPU) or debug (sequential). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * update vs resources for nextflow * refactor(virtual-stain): align CLI pattern with deskew/flat-field - virtual_stain_cli: default monitor to False (matching deskew_cli / flat_field_correction_cli; the library functions keep monitor=True) and document the SLURM fan-out / --init / --cluster debug usage examples. - Fix flat_field_correction_cli docstring: it used a raw string, which turned click's `\b` no-wrap markers into literal text and rewrapped the usage examples. Use a normal docstring with `# noqa: D301`, matching deskew. virtual-stain follows the same (correct) pattern. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * update vs resource allocation --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

ieivanov force-pushed the feat/vs-nonsquare-tta-and-helpers branch from a629fa2 to eb1ffa1 Compare June 3, 2026 23:28

ieivanov mentioned this pull request Jun 3, 2026

refactor(cytoland): unify VSUNet rotation TTA on rotation_tta_transforms #451

Open

2 tasks

ieivanov force-pushed the feat/vs-nonsquare-tta-and-helpers branch from eb1ffa1 to 8f42e30 Compare June 3, 2026 23:45

ieivanov requested review from alxndrkalinin and edyoshikun June 3, 2026 23:46

ieivanov mentioned this pull request Jun 3, 2026

refactor(virtual-stain): rewrite CLI on cytoland predict_sliding_windows czbiohub-sf/biahub#267

Merged

3 tasks

edyoshikun requested a review from Copilot June 4, 2026 02:37

Copilot started reviewing on behalf of edyoshikun June 4, 2026 02:37 View session

Copilot AI reviewed Jun 4, 2026

View reviewed changes

Comment thread applications/cytoland/src/cytoland/engine.py

Comment thread packages/viscy-data/src/viscy_data/_utils.py Outdated

ieivanov and others added 2 commits June 5, 2026 09:53

Limit n>=1 in rotation_tta_transforms

e0f2fc0

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

edyoshikun reviewed Jun 5, 2026

View reviewed changes

Comment thread applications/cytoland/src/cytoland/engine.py

ieivanov mentioned this pull request Jun 9, 2026

Add option for reflection TTA #462

Open

edyoshikun approved these changes Jun 10, 2026

View reviewed changes

edyoshikun merged commit 674dbee into modular-viscy-staging Jun 10, 2026
26 checks passed

edyoshikun deleted the feat/vs-nonsquare-tta-and-helpers branch June 10, 2026 20:13

edyoshikun restored the feat/vs-nonsquare-tta-and-helpers branch June 10, 2026 20:18

This was referenced Jun 10, 2026

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

Merged

Revert "feat(cytoland): non-square rotation TTA + reuse helpers (#450)" #466

Open

ieivanov deleted the feat/vs-nonsquare-tta-and-helpers branch June 10, 2026 22:18

Uh oh!

Conversation

ieivanov commented Jun 3, 2026

Summary

Test plan

Uh oh!

ieivanov commented Jun 3, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

ieivanov commented Jun 5, 2026

Uh oh!

edyoshikun left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

edyoshikun Jun 5, 2026

Choose a reason for hiding this comment

Uh oh!

edyoshikun Jun 5, 2026

Choose a reason for hiding this comment

Uh oh!

ieivanov Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

ieivanov Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

edyoshikun Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ieivanov commented Jun 9, 2026

Uh oh!

edyoshikun left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants