Rescale and LOT correction#424
Open
Soorya19Pradeep wants to merge 13 commits into
Open
Conversation
edyoshikun
force-pushed
the
rescale-lot-correction
branch
from
a93d071 to
37b81ff
Compare
…staging Scale-aware patch extraction (packages/viscy-data): - Add _read_pixel_size() helper to read X pixel size from OME-Zarr metadata - Add reference_pixel_size parameter to TripletDataModule: when set, computes initial_yx_patch_size from the pixel-size ratio so the same physical area is covered at inference time - Replace BatchedRescaleYXd (removed per review) with existing BatchedZoomd using scale_factor=(1.0, final_y/initial_y, final_x/initial_x) and mode="bilinear" with antialias; import is lazy to avoid a hard dep LOT batch correction (applications/dynaclr): - Add dynaclr.evaluation.lot_correction submodule with core logic (fit, apply, save, load), Pydantic configs, and Click CLI entry points - Register fit-lot-correction and apply-lot-correction in cli.py via LazyCommand - Add pot and joblib to optional-dependencies.eval in pyproject.toml Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Plots n cell patches side-by-side: - Left: raw patch at initial_yx_patch_size (larger physical area) - Right: same patch bilinearly downscaled to final_yx_patch_size (model input) Both columns share the same percentile contrast window so differences are spatial. Physical scale (µm × µm) is shown in each panel title. Usage: python visualize_triplet_rescaling.py --data-path /path/to/data.zarr --tracks-path /path/to/tracks --source-channel Phase3D --z-range 0 5 --final-yx-patch-size 224 224 --reference-pixel-size 0.325 --output rescaling_comparison.png Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…meta
_collate_norm_meta iterated every normalization level and called
torch.stack on each stat, but timepoint_statistics is nested
{timepoint: {stat: tensor}} rather than flat {stat: tensor}. Any zarr
carrying timepoint_statistics (alongside fov/dataset stats) crashed
batch collation in TripletDataModule with "expected Tensor as element 0
... but got dict". Stack within each timepoint sub-dict instead.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…etDataModule
Lets a 3D OME-Zarr feed a 2D model without materializing a separate MIP
dataset, and centers the extracted Z window on each FOV's focus plane.
- z_reduction ("mip"/"center"): collapse the extracted z_range to one
slice via BatchedChannelWiseZReductiond. Label-free channels (resolved
by parse_channel_name) take the center slice; others are max-projected.
on_after_batch_transfer expects Z=1 when reduction is on.
- z_extraction_window/z_focus_offset/focus_channel: resolve a per-FOV
focus-centered window from each position's
focus_slice[ch].fov_statistics.z_focus_mean (fallback z_total//2), all
windows the same width. z_range stays as an explicit override; exactly
one of z_range / z_extraction_window must be given. Per-FOV windows are
resolved at setup() and looked up per patch in the dataset.
Tests cover both reduction strategies (discriminating center vs MIP),
the normalize-then-reduce order, per-FOV focus resolution, and the
z_range/z_extraction_window XOR guard.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Documents the TripletDataModule predict path (zarr + tracking, not parquet) and adds a runnable sample config demonstrating z_reduction + reference_pixel_size to feed a 3D dataset to a 2D model. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Notebook-style script (no CLI) that loads a 3D OME-Zarr + tracking, extracts a per-FOV focus-centered Z window, collapses it via z_reduction, applies a random affine so anchor/positive diverge, and visualizes a couple of batches to a PNG. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
edyoshikun
force-pushed
the
rescale-lot-correction
branch
from
901d5b7 to
dc39073
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
This PR expands the Triplet-based data loading path to better support cross-dataset inference (pixel-size aware patch extraction + on-the-fly Z reduction + focus-centered per-FOV Z windows), and adds a LOT (Linear Optimal Transport) batch-correction pipeline with CLI entrypoints for embedding-space correction.
Changes:
- Extend
TripletDataModule/TripletDatasetto support per-FOVz_rangeresolution viaz_extraction_window+focus_slicemetadata, plus optionalz_reductionand pixel-size aware rescaling. - Add LOT correction core functions + Pydantic config models +
dynaclrCLI commands for fitting/applying correction pipelines. - Add tests for the new Triplet Z-window/Z-reduction behavior and update docs/config examples for inference.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| packages/viscy-data/src/viscy_data/triplet.py | Adds pixel-size rescaling, per-FOV focus-centered Z windows, and optional Z-reduction in the Triplet datamodule/dataset. |
| packages/viscy-data/tests/test_triplet.py | Adds tests for z_extraction_window XOR validation, per-FOV focus-centered windows, and Z-reduction behavior/order. |
| packages/viscy-data/src/viscy_data/_utils.py | Extends norm-meta collation to correctly stack nested timepoint_statistics. |
| applications/dynaclr/src/dynaclr/evaluation/lot_correction/lot_correction.py | Implements LOT correction fit/apply + pipeline save/load helpers. |
| applications/dynaclr/src/dynaclr/evaluation/lot_correction/config.py | Adds Pydantic models for fit/apply configs and filter specs. |
| applications/dynaclr/src/dynaclr/evaluation/lot_correction/fit_lot_correction.py | CLI wrapper for fitting and saving a LOT pipeline from YAML config. |
| applications/dynaclr/src/dynaclr/evaluation/lot_correction/apply_lot_correction.py | CLI wrapper for applying a saved LOT pipeline to an embedding zarr. |
| applications/dynaclr/src/dynaclr/evaluation/lot_correction/init.py | Exposes LOT correction functions at the package level. |
| applications/dynaclr/src/dynaclr/cli.py | Registers new fit-lot-correction and apply-lot-correction CLI subcommands. |
| applications/dynaclr/pyproject.toml | Adds joblib + pot to eval extras for LOT correction dependencies. |
| applications/dynaclr/scripts/dataloader_inspection/visualize_triplet_rescaling.py | Adds a visualization utility for the new pixel-size rescaling behavior. |
| applications/dynaclr/scripts/dataloader_inspection/triplet_dataloader_zprojection.py | Adds an inspection notebook/script for focus-centered Z windows + Z-reduction. |
| applications/dynaclr/docs/DAGs/inference_triplet.md | Documents Triplet inference path and new options (needs a small wording update for z-window semantics). |
| applications/dynaclr/configs/prediction/predict_triplet_2d_from_3d.yml | Provides a sample prediction config demonstrating reference_pixel_size + z_reduction. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
- Declare viscy-transforms in viscy-data's triplet extra and hoist the BatchedZoomd/BatchedChannelWiseZReductiond imports to module top; the inline imports were masking a missing runtime dependency that would ImportError for anyone using z_reduction or reference_pixel_size. - Raise in _resolve_per_fov_z_ranges when a FOV's Z is smaller than z_extraction_window, instead of silently emitting a narrower window that breaks cross-FOV batch stacking. - Remove a dead duplicate break in test_focus_centered_z_range. - Correct the inference DAG doc: embeddings live in .X (the embedding_key array), mirrored to obsm["X_backbone"]/["X_projections"]. uv.lock intentionally omitted: the workspace glob currently entangles the untracked applications/eet package, so a clean regen of the single viscy-transforms edge is not possible until eet is committed or removed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The following have been implemented in this branch:
TripletDataModulethat rescales the patches to a reference pixel and a uses the Z-MIP projection transform