From 540ac523dbcf50b282066ec9b28f48f067a9fe9a Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 7 Oct 2026 18:25:13 -0700 Subject: [PATCH 01/54] feat(track): write labels and tracks.geff into the assembled plate Tracking no longer creates a separate 5-track/.zarr plate holding the labels as an image channel. It writes into each -i position of the assembled plate (#360, #361): - labels/ as an OME-NGFF 0.5 label image: uint32, TZYX with the image's full Z, chunks (16,1,128,128) (1 MB, the DCA SHOULD), shards of 16 timepoints by the YX plane. 2D masks go in the middle plane of the tracked z-window (middle of the volume for a full-stack projection, the plane itself for a one-plane window, middle of a projected section); 3D masks fill the window. The track settings and z placement are recorded under `biahub-track` on the label image. - tracks.geff via ultrack's Tracker.to_geff (zarr v3, seg_id = track_id, related_objects pointing at the labels, axes with units and scale). - tracks___.csv, kept for downstream readers for now. `track --init` creates the empty label images with create_empty_plate(label_names=...). -o is now a working directory for the Ultrack databases and SLURM logs. TrackingSettings drops output_ome_zarr_version, since the labels take the input store's version. Nextflow: --qc_track_config and qc_track.yaml are removed. a549/qc.yaml counts instances in the label image (imaging-qc#243, `label:`), so QC is planned after track --init and computed after tracking. Pins: iohub dev/dca-labels (iohub#485, #486), the ieivanov/ultrack fork (royerlab/ultrack#288), imaging-qc f5ae656, geff/geff-spec >= 1.3. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../skills/reconstruct-with-nextflow/SKILL.md | 68 ++-- .../references/caveats.md | 41 ++- .../references/datasets.md | 12 +- .../references/monitoring.md | 2 +- .../references/recovery.md | 7 +- .../templates/check_qc_channels.py | 64 ++-- .../templates/run_mantis_v2.sh | 14 +- biahub/settings.py | 7 +- biahub/track.py | 340 +++++++++++++----- nextflow/configs/a549/qc.yaml | 27 +- nextflow/configs/a549/qc_track.yaml | 51 --- nextflow/configs/zebrafish/qc.yaml | 1 - nextflow/mantis-v2.nf | 80 ++--- nextflow/modules/tracking.nf | 47 +-- nextflow/nextflow.config | 6 +- nextflow/qc-standalone.nf | 2 +- pyproject.toml | 48 ++- tests/test_cli/test_track_cli.py | 213 ++++++++--- uv.lock | 79 ++-- 19 files changed, 705 insertions(+), 404 deletions(-) delete mode 100644 nextflow/configs/a549/qc_track.yaml diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index 1c3ce329..a2e01e84 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -77,7 +77,7 @@ git ls-remote https://github.com/czbiohub-sf/imaging-qc-pipeline HEAD >/dev/null Still refused after logging in means their account lacks access to that repo, which no local setup can fix — **tell them to ask a biahub developer (Ivan, Taylla) to be added**, and offer to continue without QC by dropping `--qc_config` -and `--qc_track_config` (§7). The rest of the pipeline needs no GitHub +(§7). The rest of the pipeline needs no GitHub credential at all, so this never blocks a reconstruction. A user who already uses SSH keys for GitHub needs nothing here; if they would @@ -173,10 +173,10 @@ say which you think it is and let the user decide. A fresh reprocess goes to a ## 5. Start from the configs on `main` The templates in `/nextflow/configs/{a549,zebrafish}/` are the source -of truth — they version with the pipeline and are reviewed. QC configs sit with them, -one self-contained file per store: `qc.yaml` for the assembled store and -`qc_track.yaml` for the tracking store. They carry `.yaml` rather than `.yml` -because imaging-qc loads a stage config through Hydra, which strips the +of truth — they version with the pipeline and are reviewed. The QC config sits with +them, one self-contained `qc.yaml` for the assembled store; A549's also counts +cells in the label image tracking writes into that store. It carries `.yaml` +rather than `.yml` because imaging-qc loads a stage config through Hydra, which strips the extension and looks for `.yaml`. A previous run's `configs/` directory is a fallback only: an unreviewed snapshot that drifts with the schema and carries dataset-specific edits. @@ -191,7 +191,7 @@ the next dataset in the family would also need), fix it in a PR against `nextflow/configs//`, stating which dataset exposed it. Genuinely per-dataset values stay in the run directory. -### 5a. Check the QC channel filters against the stores they will QC +### 5a. Check the QC channel filters against the store they will QC Once the configs are edited, run this from the biahub checkout: @@ -202,17 +202,20 @@ python .claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py \ It exits 1 and names every `channels:` entry that the store will not have. -The stores QC runs over do not exist yet — the pipeline's init phase creates -them — so the script predicts their channels from the configs that produce +The assembled store QC runs over does not exist yet — the pipeline's init phase +creates it — so the script predicts its channels from the configs that produce them: raw channel names for the deskewed data, waveorder's -`output_channel_names` for the phase channel, virtual-stain's `target_channel` -for the predictions, and `_labels` for the tracking -store. - -**That last one is the coupling that actually breaks.** `qc_track.yaml` names -`nuclei_prediction_labels`, which only resolves if `track.yml` says -`target_channel: nuclei_prediction` — two files with no shared source of truth, -and changing one is the ordinary way to desynchronise them. +`output_channel_names` for the phase channel, and virtual-stain's +`target_channel` for the predictions. When `qc.yaml` has an `instance_count` +block and `track.yml` exists, it also checks that block against tracking. + +**That last one is the coupling that actually breaks.** Tracking writes its +labels into each position at `labels/`, and +`qc.yaml`'s `instance_count` names them twice: `label: nuclei_prediction` and +its anchor `channels: [nuclei_prediction]`, the channel they were segmented +from. Both only resolve if `track.yml` says `target_channel: nuclei_prediction` +— two files with no shared source of truth, and changing one is the ordinary +way to desynchronise them. Since imaging-qc#226 the pipeline refuses an unknown channel name itself, at `plan-stage`, so a mistake fails during the init phase rather than silently @@ -342,7 +345,8 @@ Do not run anything yet. Show the user: **State the directory numbers this run will produce.** The number is the step's position among the steps performed, not a fixed label, so a neuromast run writes `4-assemble` as its last directory and an A549 run writes - `4-assemble` then `5-track`. Older A549 runs on disk say `5-assemble` / + `4-assemble` then `5-track` (tracking's work directory — its labels and + tracks go into the assembled plate). Older A549 runs on disk say `5-assemble` / `4-track`, from when the numbers were fixed — say so if the user is comparing against one. 7. Known caveats that apply to this dataset. @@ -359,7 +363,8 @@ Do not run anything yet. Show the user: pipeline deletes the flat-field, deskew, reconstruct and virtual-stain directories, the `slurm_output/` placeholders and `.iohub-progress/` resume markers beside the final stores, and the Nextflow work directory. Only the - assembled store, the tracking store, the QC report and the logs under + assembled store (with tracking's labels and tracks inside it), tracking's + per-FOV Ultrack databases in `5-track/`, the QC report and the logs under `nextflow/` remain, and the run is final: any rerun recomputes from raw. The template's `CLEANUP_INTERMEDIATES="auto"` lets the pipeline decide at launch (`cleanup_decision` in @@ -396,13 +401,12 @@ copied file's group bits follow the *source* file's mode — `cp` from the repo (644) yields `rw-r--r--`. Hence the `chmod` after every copy. Plain `cp` only: `cp -p`/`-a` would also carry over the checkout's group. -The run script passes all four optional configs — `--concatenate_config`, -`--track_config`, `--qc_config` (`qc.yaml`), `--qc_track_config` -(`qc_track.yaml`) — so an A549 run needs no edit. **For a neuromast/zebrafish run, or any step the user asked to skip, +The run script passes all three optional configs — `--concatenate_config`, +`--track_config`, `--qc_config` (`qc.yaml`) — so an A549 run needs no edit. **For a neuromast/zebrafish run, or any step the user asked to skip, DELETE that flag's line** from the `nextflow run` call and note the skip in a comment above it, so the script still records what this run did. A neuromast run -deletes `--track_config` and `--qc_track_config`; neither `track.yml` nor -`qc_track.yaml` need exist for that family, and neither does for zebrafish. +deletes `--track_config`; `track.yml` need not exist for that family, and does +not for zebrafish. `--qc_track_config` is gone: passing it errors at launch. Delete rather than comment: a `#` inside a backslash-continued command does not start a comment line — the continuation swallows it, every flag below is dropped @@ -521,14 +525,18 @@ directory (`_rerun`) over relaunching in place. whether a `rename-channels` CLI has landed on main, making this obsolete). 2. Verify the assembled store (`-assemble/.zarr`, `4-assemble` unless earlier steps were skipped) opens with iohub; report shape, channels, - size on disk. This is the deliverable for every family; tracking, when it - ran, is reported beside it rather than as a by-product. -3. If QC ran, report its verdict: the `QC_SUMMARY` line per store from the + size on disk. This is the deliverable for every family. When tracking ran, + its output is inside it: check a position has `labels/nuclei_prediction` + (the track config's `target_channel`; uint32 TZYX at the image's full Z), + `tracks.geff` and `tracks___.csv`, and report them beside + the image data. +3. If QC ran, report its verdict: the `QC_SUMMARY` line from the pipeline log (`pass=`/`fail=`/`gates_fail=`), and point at the report at - `/qc/report/index.html` — one page, one tab per QC'd store. A gate - failure does NOT fail the run: `imaging-qc gate` exits 0 either way, so a - failing verdict is only visible in the summary line, the report, and the - `tables/qc/` parquet inside each store. + `/qc/report/index.html` — one page, one tab: the assembled store + (A549's includes the cell counts). A gate failure does NOT fail the run: + `imaging-qc gate` exits 0 either way, so a failing verdict is only visible + in the summary line, the report, and the `tables/qc/` parquet inside the + store. 4. Report per-step task counts, failures, retries, and wall time from `/nextflow/trace.txt`; point at `report.html` and `timeline.html`. If cleanup was on, report what `/nextflow/intermediates_cleaned.txt` diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 81dbd3f0..2be94790 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -136,19 +136,27 @@ deliberately (§2). **`track.yml` — schema.** `TrackingSettings` forbids extra keys. Current fields: `output_mode`, `z_slicing`, `target_channel`, `fov`, `input_images`, -`segmentation_method`, `cellpose_config`, `tracking_config`. Old configs under -`/hpc/projects/tlg2_mantis/` still carry rejected fields — migrate: +`blank_frames_path`, `segmentation_method`, `cellpose_config`, +`tracking_config`. Old configs under `/hpc/projects/tlg2_mantis/` still carry +rejected fields — migrate: ``` mode -> output_mode z_range + focus_config.z_window -> z_slicing.{method, window_size, focus_channel} +output_ome_zarr_version -> (delete: labels take the assembled store's version) ``` (`focus_config`'s `NA_det`/`lambda_ill`/`pixel_size` have no equivalent.) `target_channel` and `cellpose_config.input_channel` must match **pre-rename** -names in the **assembled** plate (§2, §5). `qc_track.yaml`'s `channels:` must -then match `_labels`, which is what the tracking store actually -holds — SKILL.md §5a checks both couplings before launch. +names in the **assembled** plate (§2, §5). `target_channel` also names the +label image tracking writes into each position, `labels/` (no +`_labels` suffix), so `qc.yaml`'s `instance_count.label` and its anchor +`channels:` entry must equal it — SKILL.md §5a checks both couplings before +launch. + +**QC configs reject `upstream_qc_dir`.** The imaging-qc pin (f5ae656) dropped +the key; a `qc.yaml` copied from an older run that still carries +`upstream_qc_dir: null` fails at `plan-stage`. Delete the line. **`z_slicing.focus_channel` resolves against `input_images`, not the store.** `apply_focus_slicing` raises if the focus channel is not among the loaded @@ -172,8 +180,8 @@ assembled store, and track's parameters (cellpose `diameter`, `min_area`/ `max_area`, linking `max_distance`) are tuned for A549 cells and do not transfer. So simply do not run it: -- Omit `--track_config` — delete that line, and `--qc_track_config` with it, - from the `nextflow run` call in the run script. A step runs only if its config +- Omit `--track_config` — delete that line from the `nextflow run` call in + the run script. A step runs only if its config is passed ([biahub#306](https://github.com/czbiohub-sf/biahub/issues/306), fixed). Delete rather than comment: `#` inside a backslash-continued command drops every flag below it, `-resume` included. @@ -183,13 +191,24 @@ transfer. So simply do not run it: - Do not report tracking results or tune the track config for these datasets unless the user asks for neuromast tracking explicitly. - QC still runs: a neuromast run does image QC of the assembled store - (`--qc_config`). Only the tracking tab is absent. + (`--qc_config`). `zebrafish/qc.yaml` has no `instance_count` block, since + there are no labels to count; do not copy A549's `qc.yaml` over it — its + `instance_count` names a label image that only tracking creates, and + `plan-stage` refuses it. ## 5. Track reads the assembled plate, not the intermediates -`track_wf` takes the assembled `.zarr` for *both* inputs. So any Z/Y/X -crop or `time_indices` subset in `concatenate.yml` is what tracking sees, and -tracking starts only after the whole plate assembles. +`track_wf` takes the assembled `.zarr` for *both* inputs, and writes +its output back INTO it: per position, `labels/` (uint32, TZYX +at the image's full Z; 2D-tracking masks sit in the middle plane of the tracked +z-window), `tracks.geff`, and `tracks___.csv`. `5-track/` is only +the work directory (`-o`): per-FOV Ultrack databases at `__/` and +`slurm_output/`. So any Z/Y/X crop or `time_indices` subset in +`concatenate.yml` is what tracking sees, and tracking starts only after the +whole plate assembles. `track --init` creates the empty label image in every +position during the init phase, and with tracking on, QC is planned after that +init and computes after tracking finishes — not concurrently with it, so QC +now sits on the critical path behind tracking. ## 6. Assemble is a single job on one reserved node diff --git a/.claude/skills/reconstruct-with-nextflow/references/datasets.md b/.claude/skills/reconstruct-with-nextflow/references/datasets.md index 6ea22e0e..4c090817 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/datasets.md +++ b/.claude/skills/reconstruct-with-nextflow/references/datasets.md @@ -70,10 +70,14 @@ layout is defined once in `nextflow/mantis-v2.nf` (`directory_layout()`): ├── 1-deskew/.zarr ├── 2-reconstruct/.zarr ├── 3-virtual-stain/.zarr -├── 4-assemble/.zarr # the deliverable -├── 5-track/.zarr # A549 only — not run for neuromast +├── 4-assemble/.zarr # the deliverable; with tracking, each position +│ # also holds labels/, tracks.geff, +│ # tracks___.csv +├── 5-track/ # A549 only — tracking's work dir, not a store +│ ├── __/ # per-FOV Ultrack database +│ └── slurm_output/ ├── qc/ -│ ├── report/index.html # one report, one tab per QC'd store +│ ├── report/index.html # one report, one tab: the assembled store │ └── report_spec.yaml # written at launch, before any QC task ├── configs/ # copied + edited per dataset │ ├── flat_field.yml @@ -82,7 +86,7 @@ layout is defined once in `nextflow/mantis-v2.nf` (`directory_layout()`): │ ├── virtual_stain.yml # older runs name this predict.yml │ ├── concatenate.yml │ ├── track.yml # A549 only -│ └── qc/{assemble,track}/ # QC configs, one dir per store kind +│ └── qc.yaml # QC config for the assembled store ├── nextflow/ │ ├── work/ # Nextflow work dir (default: /nextflow/work) │ ├── slurm_output//%x_%j.{out,err} diff --git a/.claude/skills/reconstruct-with-nextflow/references/monitoring.md b/.claude/skills/reconstruct-with-nextflow/references/monitoring.md index fe3d257a..af5382d0 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/monitoring.md +++ b/.claude/skills/reconstruct-with-nextflow/references/monitoring.md @@ -152,7 +152,7 @@ Steps run strictly in sequence; within a step, positions fan out with | `compute_transfer_function` + `run_apply_inv_tf` | one-shot + fan-out, CPU | TF is quick; apply is the bulk | | `run_virtual_stain_preprocess` + `run_virtual_stain` | one-shot + fan-out, **GPU** | `preempted` partition, any GPU node (A40 workstations included); a reclaimed task resumes at the timepoint it lost; longest per-position step | | `run_concatenate` | **single job**, one reserved node | longest single job; whole-plate I/O | -| `run_track` | fan-out, GPU (cellpose) | reads the assembled plate | +| `run_track` | fan-out, GPU (cellpose) | reads the assembled plate and writes `labels/`, `tracks.geff` and a tracks CSV into each position; `5-track/` holds only the Ultrack databases | Per-step wall-time is derived from data volume (`estimate_time_minutes()` in `biahub/utils/cluster.py`) and passed through the init step's `RESOURCES:` payload, diff --git a/.claude/skills/reconstruct-with-nextflow/references/recovery.md b/.claude/skills/reconstruct-with-nextflow/references/recovery.md index c55c894c..1009446e 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/recovery.md +++ b/.claude/skills/reconstruct-with-nextflow/references/recovery.md @@ -105,7 +105,12 @@ session. What `-resume` does after a change: `script:` block in `nextflow/modules/*.nf` does. To force one step to recompute, delete that step's output for the affected -positions and their work dirs — not the whole `.nextflow/` cache. +positions and their work dirs — not the whole `.nextflow/` cache. Track's +output is inside the assembled plate, not under `5-track/`: per position, +`labels/`, `tracks.geff` and `tracks___.csv`, +plus its Ultrack database at `5-track/__/`. A rerun overwrites +all of these in place, so for track deleting the work dirs is what forces it; +never delete the assembled plate to clear tracking output. ## When to escalate to the user diff --git a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py index 41515edd..bc6b1c4e 100644 --- a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py +++ b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py @@ -4,19 +4,22 @@ Exits 1 and names every miss, 0 if every filter resolves. -WHY PREDICT RATHER THAN READ. The stores QC runs over — the assembled plate and -the tracking plate — do not exist when configs are scaffolded; the pipeline's -init phase creates them. So this derives their channel lists from the configs -that will produce them, which is exactly the coupling that goes wrong: +WHY PREDICT RATHER THAN READ. The store QC runs over — the assembled plate — +does not exist when configs are scaffolded; the pipeline's init phase creates +it. So this derives its channel list from the configs that will produce it, +which is exactly the coupling that goes wrong: assembled = raw channels (flat-field and deskew preserve names) + reconstruct's output_channel_names (waveorder derives these, so they are read from the settings model, not the YAML text) + virtual-stain's data.init_args.target_channel - tracking = ["_labels"] + labels = (/labels/) -`qc_track.yaml` naming `nuclei_prediction_labels` only matches if `track.yml` -says `target_channel: nuclei_prediction`. Two files, no shared source of truth. +Tracking writes its labels INTO the assembled plate, so `qc.yaml`'s +`instance_count.label` — and its one anchor channel, the channel the labels +were segmented from — only resolve if both equal `track.yml`'s +`target_channel`. Two files, no shared source of truth. Checked when `qc.yaml` +has an `instance_count` block and `track.yml` exists. RELATIONSHIP TO imaging-qc#226. Since that landed, `plan-stage` refuses an unknown channel name itself, so the pipeline fails at init rather than skipping @@ -71,26 +74,41 @@ def assembled_channels(raw_zarr, reconstruct_cfg, virtual_stain_cfg): return raw + list(rc.output_channel_names) + tgt -def tracking_channels(track_cfg): - t = yaml.safe_load(Path(track_cfg).read_text()) - return [f"{t['target_channel']}_labels"] +def tracking_entries(qc_cfg, track_cfg): + """`instance_count`'s label and anchor channels, and track's target_channel. + + None when there is nothing to compare: no `instance_count` block, or no track.yml. + """ + ic = (yaml.safe_load(Path(qc_cfg).read_text()) or {}).get("instance_count") + if not isinstance(ic, dict) or not Path(track_cfg).exists(): + return None + target = yaml.safe_load(Path(track_cfg).read_text())["target_channel"] + entries = [("instance_count.label", ic.get("label"))] + entries += [("instance_count.channels", c) for c in ic.get("channels") or []] + return target, entries if __name__ == "__main__": + if sys.argv[1:2] in (["-h"], ["--help"]): + sys.exit(print(__doc__)) + if len(sys.argv) != 3: + sys.exit(__doc__.split("\n\n")[1]) raw, cfgs = sys.argv[1], Path(sys.argv[2]) - checks = [ - ( - "qc.yaml", - assembled_channels(raw, cfgs / "reconstruct.yml", cfgs / "virtual_stain.yml"), - ) - ] - if (cfgs / "qc_track.yaml").exists(): - checks.append(("qc_track.yaml", tracking_channels(cfgs / "track.yml"))) bad = 0 - for name, predicted in checks: - print(f"{name}: store will have {predicted}") - for owner, ch in qc_filters(cfgs / name): - ok = ch in predicted + + predicted = assembled_channels(raw, cfgs / "reconstruct.yml", cfgs / "virtual_stain.yml") + print(f"qc.yaml: store will have {predicted}") + for owner, ch in qc_filters(cfgs / "qc.yaml"): + ok = ch in predicted + bad += not ok + print(f" {'ok ' if ok else 'MISS'} {owner}.channels: {ch!r}") + + tracking = tracking_entries(cfgs / "qc.yaml", cfgs / "track.yml") + if tracking: + target, entries = tracking + print(f"qc.yaml vs track.yml: tracking writes labels/{target}") + for key, value in entries: + ok = value == target bad += not ok - print(f" {'ok ' if ok else 'MISS'} {owner}.channels: {ch!r}") + print(f" {'ok ' if ok else 'MISS'} {key}: {value!r}") sys.exit(1 if bad else 0) diff --git a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh index 450b384a..36484c66 100755 --- a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh +++ b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh @@ -11,6 +11,9 @@ # # --output is the project root, so each step writes a sibling directory: # 0-flatfield/ 1-deskew/ 2-reconstruct/ 3-virtual-stain/ 4-assemble/ 5-track/ +# 5-track/ is tracking's work directory (per-FOV Ultrack databases), not a +# store: tracking writes labels/, tracks.geff and a tracks CSV +# into each position of 4-assemble/.zarr. # The number is the step's POSITION among the steps this run performs, so a run # without tracking ends at 4-assemble and one that also skipped a step earlier # shifts everything after it down. @@ -190,12 +193,14 @@ fi # Harmless when unset already, and nothing in the pipeline reads it. unset CLAUDECODE -# The four config flags from --concatenate_config down are the OPTIONAL steps: +# The three config flags from --concatenate_config down are the OPTIONAL steps: # a step runs only if its config is passed (biahub#306). Defaults by family — # -# A549 / cell-line assemble + track + QC (all four, as written) -# zebrafish / neuromast assemble + QC (delete --track_config and -# --qc_track_config) +# A549 / cell-line assemble + track + QC (all three, as written) +# zebrafish / neuromast assemble + QC (delete --track_config) +# +# --qc_track_config was removed: qc.yaml counts the labels tracking writes into +# the assembled store, so passing it errors at launch. # # TO SKIP A STEP, DELETE ITS LINE. Do not comment it out: `#` inside a # backslash-continued command does not start a comment line — the continuation @@ -215,7 +220,6 @@ nextflow run "${PIPELINE}" \ --concatenate_config "${CONFIGS}/concatenate.yml" \ --track_config "${CONFIGS}/track.yml" \ --qc_config "${CONFIGS}/qc.yaml" \ - --qc_track_config "${CONFIGS}/qc_track.yaml" \ --cleanup_intermediates "${CLEANUP_INTERMEDIATES}" \ -resume \ "$@" diff --git a/biahub/settings.py b/biahub/settings.py index e1ba9196..dec9dd4d 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -126,8 +126,9 @@ class TrackingSettings(MyBaseModel): target_channel: str = "nuclei_prediction" fov: str = "*/*/*" blank_frames_path: Path | None = None - # 2D writes an output plate with Z=1 (input must be projected); 3D keeps the - # selected z-window. Does not itself project the data. + # 2D tracks a projection (input must be projected) and writes its labels into the + # middle plane of the selected z-window; 3D writes them across the window. Does not + # itself project the data. output_mode: Literal["2D", "3D"] = "2D" # Which Z-planes to select for tracking. See ZSlicing. z_slicing: ZSlicing = ZSlicing() @@ -135,8 +136,6 @@ class TrackingSettings(MyBaseModel): tracking_config: dict[str, Any] = {} segmentation_method: Literal["foreground_contour", "cellpose"] = "foreground_contour" cellpose_config: CellposeConfig | None = None - # When None, preserve the OME-Zarr version of the input store. - output_ome_zarr_version: Literal["0.4", "0.5"] | None = None @field_validator("blank_frames_path") @classmethod diff --git a/biahub/track.py b/biahub/track.py index e80b8dfe..c3ba83d7 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -9,6 +9,7 @@ import numpy as np import pandas as pd import submitit +import zarr from iohub import open_ome_zarr from iohub.ngff.utils import create_empty_plate @@ -42,7 +43,6 @@ ) from biahub.utils.cluster import echo_resources, estimate_resources, get_submitit_cluster from biahub.utils.config import update_model, yaml_to_model -from biahub.utils.ngff import PROVENANCE_METADATA_KEYS, resolve_ome_zarr_version logger = logging.getLogger(__name__) @@ -51,6 +51,9 @@ NA_DET = 1.35 LAMBDA_ILL = 0.500 +# Provenance key on each label image: the track settings, plus where the labels sit in Z. +PROVENANCE_KEY = "biahub-track" + # Lazy imports for ultrack - imported only when needed in specific functions _ultrack_patched = False @@ -373,7 +376,7 @@ def resolve_z_slice(z: ZSlicing, z_shape: int) -> tuple[slice, int]: def apply_focus_slicing( data_dict: dict[str, ArrayLike], z_slicing: ZSlicing, pixel_size: float -) -> dict[str, ArrayLike]: +) -> tuple[dict[str, ArrayLike], slice]: """Slice every loaded channel to a focus window (``method='focus'``). The in-focus plane is found once per FOV on ``z_slicing.focus_channel`` (or the @@ -382,6 +385,9 @@ def apply_focus_slicing( Channel arrays are (T, Z, Y, X). Focus finding reads the focus channel one timepoint at a time; every channel stays lazy, so only the resolved window is materialised downstream. + + Returns the sliced ``data_dict`` and the resolved window, which locates the tracked + planes in the full stack (see :func:`tracked_z_window`). """ focus_channel = z_slicing.focus_channel or next(iter(data_dict)) if focus_channel not in data_dict: @@ -401,12 +407,30 @@ def apply_focus_slicing( ) click.echo(f"Focus-resolved z-slice: {z_slices}") - return {name: arr[:, z_slices] for name, arr in data_dict.items()} + return {name: arr[:, z_slices] for name, arr in data_dict.items()}, z_slices + + +def tracked_z_window( + read_slice: slice, focus_slice: slice | None, z_shape: int +) -> tuple[int, int]: + """Locate the tracked planes in the full stack as a ``[start, stop)`` window. + + ``read_slice`` is the read-time slice from :func:`resolve_z_slice`; + ``focus_slice`` is the per-FOV window from :func:`apply_focus_slicing`, relative to + the planes that were read, or None when ``method`` is not ``focus``. + """ + start, stop, _ = read_slice.indices(z_shape) + if focus_slice is not None: + focus_start, focus_stop, _ = focus_slice.indices(stop - start) + start, stop = start + focus_start, start + focus_stop + return start, stop def run_ultrack( tracking_config, database_path, + geff_path: Path | None = None, + geff_kwargs: dict | None = None, **track_kwargs, ): """ @@ -426,6 +450,11 @@ def run_ultrack( Ultrack configuration object defining segmentation, linking, and optimization parameters. database_path : Path Directory where tracking results, configuration files, and output data will be saved. + geff_path : Path, optional + If set, the solved tracks are also exported there as GEFF via + ``Tracker.to_geff`` (overwriting any previous export). + geff_kwargs : dict, optional + Extra keyword arguments for ``Tracker.to_geff``. **track_kwargs Keyword arguments forwarded to ``Tracker.track()``. E.g. ``detection``, ``edges`` (foreground+contour mode), or ``labels``, ``sigma`` (cellpose mode), plus ``scale`` @@ -467,6 +496,8 @@ def run_ultrack( labels = tracker.to_zarr( tracks_df=tracks_df, ) + if geff_path is not None: + tracker.to_geff(str(geff_path), overwrite=True, **(geff_kwargs or {})) with open(database_path / "config.toml", mode="w") as f: toml.dump(cfg.dict(by_alias=True), f) @@ -665,10 +696,11 @@ def _load_and_preprocess( z_slicing: ZSlicing | None, pixel_size: float | None, visualize: bool = False, -) -> dict[str, ArrayLike]: +) -> tuple[dict[str, ArrayLike], slice | None]: """Load, focus-slice, run the processing pipeline, and fill blank frames for one FOV. - Shared by the foreground+contour and cellpose segmentation paths. + Shared by the foreground+contour and cellpose segmentation paths. Also returns the + per-FOV focus window (None unless ``z_slicing.method`` is ``focus``). """ fov = "/".join(position_key) data_dict = load_data( @@ -677,10 +709,11 @@ def _load_and_preprocess( z_slices=z_slices, visualize=visualize, ) + focus_slice = None if z_slicing is not None and z_slicing.method == "focus": - data_dict = apply_focus_slicing(data_dict, z_slicing, pixel_size) + data_dict, focus_slice = apply_focus_slicing(data_dict, z_slicing, pixel_size) data_dict = run_preprocessing_pipeline(data_dict, input_images, visualize=visualize) - return fill_empty_frames_from_csv(fov, data_dict, blank_frames_path) + return fill_empty_frames_from_csv(fov, data_dict, blank_frames_path), focus_slice def detect_foreground_segmentation( @@ -787,11 +820,70 @@ def cellpose_segmentation( return cellpose_labels +# Label chunks, TZYX. 16 x 1 x 128 x 128 uint32 is 1 MB, the DCA v0.2 SHOULD (512 KB +# is the MUST) and its "2D+time" example; growing T keeps per-plane reads cheap. +LABEL_CHUNKS = (16, 1, 128, 128) +# Shards span one time chunk (16 timepoints, the DCA SHOULD) and the whole YX plane, +# capped at 16 chunks (2048 px, the DCA spatial-shard SHOULD). +MAX_LABEL_SHARD_CHUNKS = 16 + + +def label_chunks_and_shards( + shape: tuple[int, int, int, int], +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Chunks and shards ratio for a (T, Z, Y, X) label array of ``shape``.""" + chunks = tuple(min(c, s) for c, s in zip(LABEL_CHUNKS, shape, strict=True)) + shards_ratio = (1, 1) + tuple( + min(-(-s // c), MAX_LABEL_SHARD_CHUNKS) + for c, s in zip(chunks[2:], shape[2:], strict=True) + ) + return chunks, shards_ratio + + +def write_tracking_labels( + label_array, + labels: ArrayLike, + z_window: tuple[int, int], + output_mode: str, +) -> int | slice: + """Write tracked labels into the full-Z (T, Z, Y, X) label array. + + In 2D mode the (T, Y, X) labels go in the middle plane of the tracked window: the + middle of the volume for a full-stack projection, the plane itself for a one-plane + window, the middle of the section for a projected sub-volume. In 3D mode the + (T, Z, Y, X) labels fill the window. Returns the z index or slice written. + """ + start, stop = z_window + labels = np.asarray(labels, dtype=np.uint32) + if output_mode == "2D": + # A one-plane window that wasn't projected still has its Z axis. + if labels.ndim == 4 and labels.shape[1] == 1: + labels = labels[:, 0] + if labels.ndim != 3: + raise ValueError( + f"output_mode='2D' expects (T, Y, X) labels but tracking produced " + f"shape {labels.shape}. Ensure input_images projects Z (e.g. np.mean)." + ) + z_index = start + (stop - start) // 2 + label_array[:, z_index] = labels + return z_index + + if labels.ndim != 4 or labels.shape[1] != stop - start: + raise ValueError( + f"output_mode='3D' expects (T, {stop - start}, Y, X) labels for z-window " + f"[{start}, {stop}) but tracking produced shape {labels.shape}." + ) + label_array[:, start:stop] = labels + return slice(start, stop) + + def track_one_position( - position_key: str, + position_key: tuple[str, str, str], + position_dirpath: Path, input_images: list[ProcessingInputChannel], output_dirpath: Path, tracking_config, + label_name: str, blank_frames_path: Path | None = None, z_slices: slice | None = None, scale: tuple[float, ...] = (1, 1, 1, 1, 1), @@ -810,15 +902,20 @@ def track_one_position( Parameters ---------- - position_key : str - A string identifier for the field of view (e.g., "A_1_3"), typically composed of - Plate, Well, and Position joined by underscores. + position_key : tuple of str + (row, column, fov) of the field of view. + position_dirpath : Path + The ``-i`` position the outputs are written into. ``track --init`` has already + created its empty ``labels/`` image. input_images : list of ProcessingInputChannel Configuration describing which channels to load and how to preprocess them. output_dirpath : Path - Output directory where labeled Zarr volumes and track CSVs will be stored. + Working directory; the Ultrack database and config for this FOV go to + ``{output_dirpath}/{row}_{col}_{fov}/``. tracking_config : MainConfig Ultrack configuration containing segmentation, linking, and optimization settings. + label_name : str + Name of the label image under ``{position_dirpath}/labels/``. blank_frames_path : Path, optional Path to CSV file indicating empty frames for the current FOV. If None, blank frame filling is skipped. @@ -834,139 +931,173 @@ def track_one_position( Z-plane selection config. For ``method='focus'`` the in-focus window is resolved per-FOV from the loaded data (see :func:`apply_focus_slicing`). output_mode : str, optional - "2D" writes (T, Y, X) labels into the Z=0 plane; "3D" writes (T, Z, Y, X) - labels across the z-window. Default is "2D". + "2D" writes (T, Y, X) labels into the middle plane of the tracked z-window; "3D" + writes (T, Z, Y, X) labels across the window. Default is "2D". Returns ------- None - Outputs are saved directly to disk: - - Tracked object labels in Zarr format - - CSV file containing the track graph (IDs, positions, parents) + Outputs are saved into ``position_dirpath``: + - ``labels/``: the tracked labels, painted with ``track_id`` + - ``tracks.geff``: the track graph, linked to the labels by ``seg_id`` + - ``tracks___.csv``: the same tracks as a CSV, kept for now Notes ----- - - Output is saved to `{output_dirpath}/{position_key}/tracks_{position_key}.csv`. - - The Ultrack config is also saved as TOML in a `_config_tracking/{FOV}/` subdirectory. - In foreground+contour mode, missing "foreground"/"contour" channels raise a ValueError. """ if z_slicing is None: z_slicing = ZSlicing() + position_dirpath = Path(position_dirpath) fov = "_".join(position_key) click.echo(f"Processing FOV: {fov.replace('_', '/')}") - # Define path to save the tracking database and graph - filename = output_dirpath.stem - database_path = output_dirpath.parent / f"{filename}_config_tracking" / f"{fov}" + database_path = Path(output_dirpath) / fov os.makedirs(database_path, exist_ok=True) + with open_ome_zarr(position_dirpath, mode="r") as position: + z_shape = position.data.shape[2] + time_scale = position.scale[0] + units = {axis.name.lower(): getattr(axis, "unit", None) for axis in position.axes} + # Load, focus-slice (method='focus'), run the processing pipeline, and fill blank # frames once; both segmentation methods consume the resulting data_dict. # pixel_size is the transverse pixel size = last element of the (trimmed) scale. pixel_size = scale[-1] - data_dict = _load_and_preprocess( + data_dict, focus_slice = _load_and_preprocess( position_key, input_images, z_slices, blank_frames_path, z_slicing, pixel_size ) + z_window = tracked_z_window(z_slices or slice(None), focus_slice, z_shape) + + geff_kwargs = { + "zarr_format": 3, + "include_masks": False, + "include_overlaps": False, + "segmentation_path": f"../labels/{label_name}", + "scale": tuple(scale), + "spatial_unit": units.get("x"), + "time_scale": time_scale if units.get("t") else None, + "time_unit": units.get("t"), + } + ultrack_kwargs = { + "tracking_config": tracking_config, + "database_path": database_path, + "geff_path": position_dirpath / "tracks.geff", + "geff_kwargs": geff_kwargs, + "scale": scale, + "overwrite": True, + } if cellpose_config is not None: cellpose_labels = cellpose_segmentation(data_dict, cellpose_config) click.echo("Tracking with cellpose labels...") tracking_labels, tracks_df, _ = run_ultrack( - tracking_config=tracking_config, - database_path=database_path, labels=cellpose_labels, sigma=cellpose_config.labels_sigma, - scale=scale, - overwrite=True, + **ultrack_kwargs, ) else: foreground_mask, contour_gradient_map = detect_foreground_segmentation(data_dict) click.echo("Tracking with foreground + contour...") tracking_labels, tracks_df, _ = run_ultrack( - tracking_config=tracking_config, - database_path=database_path, detection=foreground_mask, edges=contour_gradient_map, - scale=scale, - overwrite=True, + **ultrack_kwargs, + ) + + # Kept alongside tracks.geff for downstream readers that still expect a CSV. + tracks_df.to_csv(position_dirpath / f"tracks_{fov}.csv", index=False) + + with open_ome_zarr(position_dirpath, mode="r+") as position: + label_image = position.get_label(label_name) + z_written = write_tracking_labels( + label_image["0"], tracking_labels, z_window, output_mode ) - # Save the tracks graph to a CSV file - csv_path = output_dirpath / Path(*position_key) / f"tracks_{fov}.csv" - os.makedirs(csv_path.parent, exist_ok=True) - - tracks_df.to_csv(csv_path, index=False) - - click.echo(f"Saved tracks to: {output_dirpath / Path(*position_key)}") - - # Save the tracking labels. The output plate stores (T, 1, Z_out, Y, X); in 2D - # mode Z_out is 1 and labels are (T, Y, X); in 3D mode labels are (T, Z, Y, X). - labels = np.asarray(tracking_labels, dtype=np.uint32) - with open_ome_zarr(output_dirpath / Path(*position_key), mode="r+") as output_dataset: - if output_mode == "2D": - if labels.ndim != 3: - raise ValueError( - f"output_mode='2D' expects (T, Y, X) labels but tracking produced " - f"shape {labels.shape}. Ensure input_images projects Z (e.g. np.mean)." - ) - output_dataset[0][:, 0, 0] = labels - else: - if labels.ndim != 4: - raise ValueError( - f"output_mode='3D' expects (T, Z, Y, X) labels but tracking produced " - f"shape {labels.shape}." - ) - output_dataset[0][:, 0] = labels + # Record where the labels sit in Z: the 2D GEFF has no z axis. + label_group = zarr.open_group(position_dirpath / "labels" / label_name, mode="r+") + provenance = dict(label_group.attrs.get(PROVENANCE_KEY, {})) + provenance["z_window"] = list(z_window) + if isinstance(z_written, int): + provenance["z_index"] = z_written + label_group.attrs[PROVENANCE_KEY] = provenance + + click.echo(f"Saved labels and tracks to: {position_dirpath}") return tracking_labels, tracks_df -def _init_output_plate( +def _tracked_shape( + input_position_dirpath: Path, settings: TrackingSettings +) -> tuple[int, int, int, int, int]: + """Return the (T, C, Z_tracked, Y, X) shape one position's tracking works on.""" + with open_ome_zarr(str(input_position_dirpath), mode="r") as dataset: + T, C, Z, Y, X = dataset.data.shape + # Data-free resolution: for focus this returns the fixed window size (min(total, Z)). + _, Z_win = resolve_z_slice(settings.z_slicing, Z) + return (T, C, 1 if settings.output_mode == "2D" else Z_win, Y, X) + + +def _labels_exist(input_position_dirpaths: list[Path], label_name: str) -> bool: + return all( + (Path(p) / "labels" / label_name / "zarr.json").exists() + or (Path(p) / "labels" / label_name / ".zattrs").exists() + for p in input_position_dirpaths + ) + + +def _init_labels( input_position_dirpaths: list[Path], - output_dirpath: Path, settings: TrackingSettings, -) -> tuple[int, int, int, int, int]: - """Create the empty tracking output plate. +) -> None: + """Create the empty ``labels/`` image in every input position. - Returns the (T, C, Z_out, Y, X) shape used for resource estimation. + The labels are (T, Z, Y, X) with the image's full Z, so both 2D and 3D tracking + write into the planes they tracked. Positions that already have the label are left + untouched, and the track settings are recorded on each label image. """ with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as dataset: T, C, Z, Y, X = dataset.data.shape scale = dataset.scale + channel_names = dataset.channel_names + dtype = dataset.data.dtype + version = dataset.version - # Data-free resolution: for focus this returns the fixed window size (min(total, Z)), - # so the plate Z matches the per-FOV tracked Z. - _, Z_win = resolve_z_slice(settings.z_slicing, Z) - - if settings.output_mode == "2D": - output_shape = (T, 1, 1, Y, X) - else: - output_shape = (T, 1, Z_win, Y, X) - + input_plate = Path(input_position_dirpaths[0]).parents[2] position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] + label_chunks, label_shards_ratio = label_chunks_and_shards((T, Z, Y, X)) - input_plate = Path(input_position_dirpaths[0]).parents[2] create_empty_plate( - store_path=output_dirpath, + store_path=input_plate, position_keys=position_keys, - channel_names=[f"{settings.target_channel}_labels"], - shape=output_shape, - chunks=None, + channel_names=channel_names, + shape=(T, C, Z, Y, X), scale=scale, - version=resolve_ome_zarr_version( - input_position_dirpaths[0], settings.output_ome_zarr_version - ), - dtype=np.uint32, - metadata_sources=input_plate, - metadata_keys=PROVENANCE_METADATA_KEYS, - extra_metadata={"biahub-track": settings.model_dump(mode="json")}, + version=version, + dtype=dtype, + label_names=[settings.target_channel], + label_dtype=np.uint32, + label_chunks=label_chunks, + label_shards_ratio=label_shards_ratio if version == "0.5" else None, ) - click.echo(f"Created {output_dirpath} ({len(position_keys)} positions)") + provenance = settings.model_dump(mode="json") + for position_dirpath in input_position_dirpaths: + label_group = zarr.open_group( + Path(position_dirpath) / "labels" / settings.target_channel, mode="r+" + ) + # Merge so a re-run init keeps the z placement the workers recorded. + label_group.attrs[PROVENANCE_KEY] = { + **label_group.attrs.get(PROVENANCE_KEY, {}), + **provenance, + } - return (T, C, output_shape[2], Y, X) + click.echo( + f"Created labels/{settings.target_channel} in {input_plate} " + f"({len(position_keys)} positions)" + ) def track( @@ -988,7 +1119,8 @@ def track( config_filepath : Path Path to the tracking configuration YAML. output_dirpath : Path - Path to the output Zarr store. + Working directory for the per-FOV Ultrack databases and SLURM logs. The + labels and tracks themselves are written into the ``-i`` positions. sbatch_filepath : str, optional Path to a SLURM batch file to override defaults. cluster : str, optional @@ -996,13 +1128,16 @@ def track( monitor : bool, optional If True, monitor submitted jobs. init_only : bool, optional - Only initialize the output store and exit. + Only create the empty label images in the ``-i`` positions and exit. input_images_path : str, optional Explicit pixel-data source filling the first null ``input_images`` path (Nextflow). If unset, that null path falls back to the ``-i`` plate (see Notes). Notes ----- + Outputs go into each ``-i`` position: ``labels/`` (OME-NGFF label + image), ``tracks.geff`` and ``tracks___.csv``. + Two input sources: ``-i`` (``input_position_dirpaths``) gives the plate structure and shape/scale, while ``input_images[].path`` gives the pixel data. The first null ``path`` is the primary source (``--input-images-path`` if given, else the ``-i`` @@ -1011,7 +1146,7 @@ def track( from ultrack import MainConfig output_dirpath = Path(output_dirpath) - slurm_out_path = output_dirpath.parent / "slurm_output" + slurm_out_path = output_dirpath / "slurm_output" settings = yaml_to_model(config_filepath, TrackingSettings) @@ -1024,8 +1159,10 @@ def track( image.path = primary_path break - output_shape = _init_output_plate(input_position_dirpaths, output_dirpath, settings) - T, C, Z_out, Y, X = output_shape + # Workers skip the scaffold --init already made: they all share the -i plate. + if init_only or not _labels_exist(input_position_dirpaths, settings.target_channel): + _init_labels(input_position_dirpaths, settings) + T, C, Z_out, Y, X = _tracked_shape(input_position_dirpaths[0], settings) _, num_cpus, gb_ram_per_cpu = estimate_resources( shape=[T, C, Z_out, Y, X], ram_multiplier=16, max_num_cpus=16 @@ -1046,7 +1183,7 @@ def track( # Fail here, once, rather than in every worker of the fan-out. check_cellpose_model_name(settings.cellpose_config.pretrained_model) warm_cellpose_weights(settings.cellpose_config.pretrained_model) - click.echo(f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)") + click.echo(f"Initialized tracking ({len(input_position_dirpaths)} positions)") return # Read shape/scale from the first input position for tracking parameters @@ -1065,6 +1202,7 @@ def track( tracking_cfg = update_model(default_config, tracking_cfg) position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] + output_dirpath.mkdir(parents=True, exist_ok=True) cellpose_cfg = ( settings.cellpose_config if settings.segmentation_method == "cellpose" else None @@ -1095,12 +1233,16 @@ def track( click.echo("Submitting jobs...") jobs = [] with submitit.helpers.clean_env(), executor.batch(): - for position_key in position_keys: + for position_key, position_dirpath in zip( + position_keys, input_position_dirpaths, strict=True + ): job = executor.submit( track_one_position, position_key=position_key, + position_dirpath=Path(position_dirpath), output_dirpath=output_dirpath, tracking_config=tracking_cfg, + label_name=settings.target_channel, input_images=settings.input_images, blank_frames_path=settings.blank_frames_path, z_slices=z_slices, @@ -1158,18 +1300,22 @@ def track_cli( ): """Track objects in 2D or 3D time-lapse microscopy data using configurable preprocessing. + Writes into each -i position: the labels as an OME-NGFF label image + (labels/), the tracks as tracks.geff, and a tracks CSV. -o is a + working directory for the Ultrack databases and SLURM logs. + \b Full SLURM fan-out: - >>> biahub track -i ./reconstruct.zarr/*/*/* -o ./track.zarr -c config.yml + >>> biahub track -i ./assemble.zarr/*/*/* -o ./track -c config.yml \b - Initialize the output plate only (Nextflow init step): - >>> biahub track --init -i ./reconstruct.zarr/*/*/* -o ./track.zarr -c config.yml + Create the empty label images only (Nextflow init step): + >>> biahub track --init -i ./assemble.zarr/*/*/* -o ./track -c config.yml \b In-process run of a single position (Nextflow per-position worker): - >>> biahub track --cluster debug -i ./reconstruct.zarr/B/3/000000 \\ - -o ./track.zarr -c config.yml --input-images-path ./virtual-stain.zarr + >>> biahub track --cluster debug -i ./assemble.zarr/B/3/000000 \\ + -o ./track -c config.yml """ # noqa: D301 track( input_position_dirpaths=input_position_dirpaths, diff --git a/nextflow/configs/a549/qc.yaml b/nextflow/configs/a549/qc.yaml index 6861b555..a998861f 100644 --- a/nextflow/configs/a549/qc.yaml +++ b/nextflow/configs/a549/qc.yaml @@ -1,5 +1,8 @@ # Pixel-value QC for the assembled intensity store — the plate assemble writes, -# holding the acquired channels plus the derived ones. +# holding the acquired channels plus the derived ones — and instance counts of the +# segmentation tracking writes into it at `/labels/nuclei_prediction`. +# The counts need that label image, so this config is for runs that track: the +# pipeline plans it after `track --init` and computes it after tracking. # # SELF-CONTAINED, with no Hydra `defaults:`. It used to inherit a base.yaml, which # existed to be shared and was shared by this one file; and the directory layout @@ -23,7 +26,6 @@ backend: iohub positions: null metric_group: default log_level: INFO -upstream_qc_dir: null metric_groups: default: @@ -32,6 +34,7 @@ metric_groups: - signal_to_noise - blank - tenengrad + - instance_count intensity_stats: grid: 16 @@ -101,6 +104,25 @@ tenengrad: n_std: 3.0 tail: lower +# Cells per frame, from the label image tracking writes (`label:`, named after the +# track config's `target_channel`). `channels:` is the ONE anchor channel the labels +# were segmented from: the count rows carry its C, so `blank` on that channel masks +# the count's verdict on a blank frame instead of failing it twice. `volume` +# granularity because 2D tracking writes its labels into a single Z plane of the +# full-Z label image; a per-slice count would read zero on every other plane. +instance_count: + label: nuclei_prediction + channels: + - nuclei_prediction + granularity: volume + gate: + # A count of zero on a frame that is not blank: the segmentation found nothing. + # The degenerate structural boundary, not a constant calibrated from a dataset. + hard_floor: 1 + # Put the verdict on the anchor channel (`low_instance_count__nuclei_prediction`) + # rather than repeating a channel-averaged flag on every channel. + granularity: channel + report: max_timepoints: 5 compact: false @@ -132,3 +154,4 @@ report: blank: "Blank Detection" tenengrad: "Tenengrad" intensity_stats: "Intensity Statistics" + instance_count: "Cell Count" diff --git a/nextflow/configs/a549/qc_track.yaml b/nextflow/configs/a549/qc_track.yaml deleted file mode 100644 index 2c6d9606..00000000 --- a/nextflow/configs/a549/qc_track.yaml +++ /dev/null @@ -1,51 +0,0 @@ -# Cell-count QC over a segmentation/tracking store (`_tracking.zarr`) -# rather than an intensity store. In the mantis/DynaCell layout the label classes -# ARE the channels, so `channels:` selects the label class by name and the C -# column in the output means the label class. Outputs land in the store's own -# `tables/qc/` group, since the pipeline passes no --output-dir. -# -# Self-contained on purpose: an external --config resolves Hydra defaults only -# against its own directory, so this cannot reference imaging-qc's internal -# /common or /segmentation groups — their keys are inlined below. - -backend: iohub -positions: null -metric_group: masks -log_level: INFO -upstream_qc_dir: null - -stage: 6 -display_name: "Tracking — Cell Count" -metrics: null # derived from metric_groups; override via CLI to run a subset - -metric_groups: - masks: - metrics: - - instance_count - -instance_count: - # `slice` suits a 2D segmentation (Z=1), which is what the tracking subsets hold - # (one label plane per (T, C)). For a 3D segmentation set `volume` so each - # instance is counted once instead of once per plane. - granularity: slice - channels: - - nuclei_prediction_labels - gate: - # A count of zero is an empty frame — a dropped/blank timepoint the - # segmentation had nothing to find in. This is the degenerate structural - # boundary, not a constant calibrated from any dataset. - hard_floor: 1 - # Per-(T,C,Z) verdicts, so one empty timepoint lands on the slice drop list - # instead of failing the whole position. - gate_granularity: slice - # No relative outlier fence, deliberately: at short T a zscore fence cannot - # fire and lower_quantile flags the wrong frames. Add one only with enough - # rows per population (n >= 11 for n_std 3.0) or pool with pooling_scope: store. - -report: - max_timepoints: 5 - compact: false - metric_archetypes: - instance_count: timeseries - metric_labels: - instance_count: "Instance Count" diff --git a/nextflow/configs/zebrafish/qc.yaml b/nextflow/configs/zebrafish/qc.yaml index 87eb3ec3..705f3ad8 100644 --- a/nextflow/configs/zebrafish/qc.yaml +++ b/nextflow/configs/zebrafish/qc.yaml @@ -23,7 +23,6 @@ backend: iohub positions: null metric_group: default log_level: INFO -upstream_qc_dir: null metric_groups: default: diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf index 3292a3c8..d2e41847 100644 --- a/nextflow/mantis-v2.nf +++ b/nextflow/mantis-v2.nf @@ -43,13 +43,13 @@ params.virtual_stain_config = null params.track_config = null params.concatenate_config = null params.max_positions = 0 -// QC, off by default. Each param points at a stage config for one store; set -// either, both, or neither. Both QC'd stores become tabs of ONE report. -// --qc_config nextflow/configs//qc.yaml (assembled store) -// --qc_track_config nextflow/configs//qc_track.yaml (tracking store) -// One self-contained file per store, beside that family's other step configs. +// QC, off by default. Points at the stage config for the assembled store: +// --qc_config nextflow/configs//qc.yaml +// Tracking writes its labels into that store, so a family that tracks counts +// instances in the same config (a549/qc.yaml) — there is no separate tracking +// store to QC any more. params.qc_config = null -params.qc_track_config = null +params.qc_track_config = null // removed; refused below include { collect_positions; dataset_name; check_environment } from './modules/common' include { flat_field_init_wf; flat_field_run_wf } from './modules/flat_field' @@ -128,9 +128,7 @@ workflow { // `4-assemble` where an A549 run also has `5-track` after it. def assemble_on = params.concatenate_config as boolean def track_on = params.track_config as boolean - def qc_image_on = params.qc_config as boolean - def qc_track_on = params.qc_track_config as boolean - def qc_on = qc_image_on || qc_track_on + def qc_on = params.qc_config as boolean // A step cannot outlive the step whose output it reads. Refuse the // combination at launch, naming the config to add or the one to drop, rather @@ -138,11 +136,12 @@ workflow { if (track_on && !assemble_on) { error "--track_config needs --concatenate_config: tracking reads the assembled plate." } - if (qc_image_on && !assemble_on) { + if (qc_on && !assemble_on) { error "--qc_config needs --concatenate_config: it QCs the assembled store." } - if (qc_track_on && !track_on) { - error "--qc_track_config needs --track_config: it QCs the tracking store." + if (params.qc_track_config) { + error "--qc_track_config was removed: tracking writes its labels into the assembled " + + "store, and --qc_config counts them (see nextflow/configs/a549/qc.yaml)." } // Whether this run deletes its intermediates once it finishes. Resolved now @@ -188,7 +187,9 @@ workflow { reconstruct_output = "${out}/${layout.reconstruct}/${ds}.zarr" virtual_stain_output = "${out}/${layout.virtual_stain}/${ds}.zarr" assemble_output = assemble_on ? "${out}/${layout.assemble}/${ds}.zarr" : null - track_output = track_on ? "${out}/${layout.track}/${ds}.zarr" : null + // Tracking writes its labels and tracks INTO the assembled plate; its own + // directory holds only the per-FOV Ultrack databases and SLURM placeholders. + track_dir = track_on ? "${out}/${layout.track}" : null // ======================================================================== // INIT PHASE — every config parsed, every output store scaffolded, before @@ -255,28 +256,26 @@ workflow { // already carries the phase and virtual-stain channels (concatenate // preserves channel names, so the track config's channel names resolve // unchanged), so the plate structure and the image data come from the - // same store. Its init also warms the shared cellpose weights cache, - // which is better done here than with N GPU workers racing for it. - tk_init = track_init_wf(assemble_output, track_output, params.track_config, + // same store. Its init creates the empty label image tracking fills in + // every position of that plate, and warms the shared cellpose weights + // cache, which is better done here than with N GPU workers racing for it. + tk_init = track_init_wf(assemble_output, track_dir, params.track_config, as_init.done) init_signals << tk_init.done } // QC's planning verbs belong in this phase for the same reason the step // inits do — `plan-stage` and `estimate-resources` read the store's - // structure and validate the config, and read no pixels. Each store's plan - // is gated on the init that scaffolded it. The tab label is the step - // directory, which is unique per store by construction. + // structure and validate the config, and read no pixels. The plan is gated + // on the init that scaffolded what it reads: with tracking on that is + // track's init, because a config counting instances in `labels/` is refused + // at plan time until the label image exists. The tab label is the step + // directory. def qc_stores = [] if (qc_on) { - if (qc_image_on) { - qc_stores << [label: layout.assemble, zarr: assemble_output, - config: params.qc_config, scaffolded: as_init.done] - } - if (qc_track_on) { - qc_stores << [label: layout.track, zarr: track_output, - config: params.qc_track_config, scaffolded: tk_init.done] - } + qc_stores << [label: layout.assemble, zarr: assemble_output, + config: params.qc_config, + scaffolded: track_on ? tk_init.done : as_init.done] // The trigger is mapped, NOT combined. `combine` concatenates the two // items, so what the producer emits leaks into the tuple's arity: an @@ -361,16 +360,16 @@ workflow { // virtual_stain_output, and gate on virtual_stain_done. if (track_on) { track_done = track_run_wf(all_positions, assemble_output, assemble_output, - track_output, params.track_config, + track_dir, params.track_config, tk_init.resources, assemble_done.done) } // ----- QC compute ------------------------------------------------------- - // Planned in the init phase; the compute it planned is gated on the step - // that wrote each store — the ASSEMBLED plate on assemble_done (the same - // signal track waits on, so image QC runs CONCURRENTLY with tracking), and - // the tracking store on track_done. Neither extends the critical path ahead - // of itself. + // Planned in the init phase; the compute it planned is gated on the last + // step that writes the assembled plate — track_done when tracking is on + // (its labels are counted in the same pass as the pixel metrics), else + // assemble_done. Waiting for tracking puts QC after it rather than beside + // it, the price of one QC pass and one table per store. // // A QC verdict cannot fail the pipeline: `imaging-qc gate` exits 0 whether // positions pass or fail, recording the verdict in each store's own @@ -387,15 +386,8 @@ workflow { // Keyed [zarr, config] to match what qc_plan_wf emitted, so each store's // planned work is released by its own producer and no other. Mapped, not // combined, for the reason spelled out at qc_plan_inputs above. - qc_compute_ready = channel.empty() - if (qc_image_on) { - qc_compute_ready = qc_compute_ready.mix( - assemble_done.done.map { tuple(assemble_output, params.qc_config) } ) - } - if (qc_track_on) { - qc_compute_ready = qc_compute_ready.mix( - track_done.done.map { tuple(track_output, params.qc_track_config) } ) - } + def plate_written = track_on ? track_done.done : assemble_done.done + qc_compute_ready = plate_written.map { tuple(assemble_output, params.qc_config) } qc = qc_compute_wf(qc_plan.items, qc_plan.stores, qc_compute_ready) qc_report = qc_report_wf(qc.done, spec, qc_report_dir) @@ -426,7 +418,7 @@ workflow { .collect { zarr -> new File(zarr).parent } def assemble_dir = new File(assemble_output).parent cleanup_paths << "${assemble_dir}/slurm_output" << "${assemble_dir}/.iohub-progress" - if (track_on) cleanup_paths << "${new File(track_output).parent}/slurm_output" + if (track_on) cleanup_paths << "${track_dir}/slurm_output" def cleanup_list = cleanup_targets(cleanup_paths, out) cleanup_record = "${out}/nextflow/intermediates_cleaned.txt" @@ -476,7 +468,7 @@ workflow { [label: 'virtual staining', done: virtual_stain_done.done, output: virtual_stain_output], ] if (assemble_on) step_events << [label: 'assemble', done: assemble_done.done, output: assemble_output] - if (track_on) step_events << [label: 'track', done: track_done.done, output: track_output] + if (track_on) step_events << [label: 'track', done: track_done.done, output: assemble_output] if (qc_on) step_events << [label: 'QC', done: qc_report.done, output: qc_report_dir] if (cleanup_on) step_events << [label: 'cleanup intermediates', done: cleanup_run.done, output: cleanup_record] diff --git a/nextflow/modules/tracking.nf b/nextflow/modules/tracking.nf index bf7d32d5..8fb24d96 100644 --- a/nextflow/modules/tracking.nf +++ b/nextflow/modules/tracking.nf @@ -1,20 +1,23 @@ -// Tracking subworkflows: init (scaffold + resources + weights cache) and run -// (fan-out × N positions). +// Tracking subworkflows: init (label scaffold + resources + weights cache) and +// run (fan-out × N positions). // // This module is PATH-AGNOSTIC. Callers pass the input zarr (plate structure), -// input images zarr (image data for tracking), output zarr, and config +// input images zarr (image data for tracking), work directory, and config // explicitly. // // Tracking is a 2-input step: `input_zarr` supplies the plate structure and // `input_images_zarr` the pixels. mantis-v2.nf passes the assembled plate for -// both. +// both. The outputs go INTO `input_zarr`'s positions — `labels/`, +// `tracks.geff` and a tracks CSV — so there is no tracking store; `work_dir` +// holds only the per-FOV Ultrack databases and SLURM placeholders. // // INIT AND RUN ARE SEPARATE SUBWORKFLOWS, and hoisting init is worth more here // than anywhere else: the tracking config is the LAST one a run would otherwise // parse, so a typo in it used to surface after every reconstruction step and // the assembly had finished. `biahub track --init` reads only the input plate's -// shape and scale (z-slice resolution is data-free), so it runs against the -// assembled plate as soon as `concatenate --init` has scaffolded it. It also +// shape and scale (z-slice resolution is data-free) and creates an empty label +// image in each position, so it runs against the assembled plate as soon as +// `concatenate --init` has scaffolded it. It also // warms the shared cellpose weights cache, which is strictly better done before // any GPU worker exists to race for it. @@ -26,7 +29,7 @@ process init_track { input: val input_zarr - val output_zarr + val work_dir val config path config_file // staged only for the task hash: see common.nf, #397 val trigger @@ -34,13 +37,15 @@ process init_track { output: stdout + // slurm_output_readme writes to /slurm_output, so + // "/." lands it in /slurm_output, beside track's placeholders. script: """ mkdir -p "${slurm_log_dir('track')}" - ${slurm_output_readme('track', output_zarr)} + ${slurm_output_readme('track', "${work_dir}/.")} biahub track --init \ -i "${input_zarr}"/*/*/* \ - -o "${output_zarr}" \ + -o "${work_dir}" \ -c "${config}" """ } @@ -59,7 +64,7 @@ process run_track { tuple val(position), val(meta) val input_zarr val input_images_zarr - val output_zarr + val work_dir val config path config_file // staged only for the task hash: see common.nf, #397 @@ -70,37 +75,37 @@ process run_track { """ biahub track --cluster debug \ -i "${input_zarr}/${position}" \ - -o "${output_zarr}" \ + -o "${work_dir}" \ -c "${config}" \ --input-images-path "${input_images_zarr}" # Ultrack's SQLite database is created 0644, and under the project dirs' # default ACL the group mask follows that mode — so the group can't write # the FOV's data.db. Grant it explicitly; -f because only the owner can. - chmod -Rf g+w "\$(dirname "${output_zarr}")/\$(basename "${output_zarr}" .zarr)_config_tracking/${position.replace('/', '_')}" || true + chmod -Rf g+w "${work_dir}/${position.replace('/', '_')}" || true """ } -// Validate the config, scaffold the output plate, warm the cellpose weights. +// Validate the config, create the empty label images, warm the cellpose weights. // Metadata-only and cheap, so it belongs in the pipeline's up-front init phase. // // take: -// input_zarr path to the input plate.zarr (plate structure) -// output_zarr path to the tracking output plate.zarr +// input_zarr path to the input plate.zarr (plate structure; receives the labels) +// work_dir tracking work directory (Ultrack databases, SLURM placeholders) // config path to the track settings YAML // trigger gating channel — init starts once this emits // emit: // resources the RESOURCES payload sizing one position's task -// done fires once the output plate exists +// done fires once every position has its empty label image workflow track_init_wf { take: input_zarr - output_zarr + work_dir config trigger main: - init_out = init_track(input_zarr, output_zarr, config, file(config), trigger.collect().map { 'done' }) + init_out = init_track(input_zarr, work_dir, config, file(config), trigger.collect().map { 'done' }) emit: resources = init_out.map { stdout_text -> parse_resources(stdout_text) } @@ -114,7 +119,7 @@ workflow track_init_wf { // positions collected channel of position keys // input_zarr plate structure store // input_images_zarr image data store -// output_zarr path to the tracking output plate.zarr +// work_dir tracking work directory // config path to the track settings YAML // resources RESOURCES payload from track_init_wf // prev_done gating channel — the input stores hold data @@ -123,7 +128,7 @@ workflow track_run_wf { positions input_zarr input_images_zarr - output_zarr + work_dir config resources prev_done @@ -142,7 +147,7 @@ workflow track_run_wf { .combine(prev_done.map { 'done' }) .map { pos, meta, _gate -> [pos, meta] } - tk_done = run_track(pos_meta, input_zarr, input_images_zarr, output_zarr, config, file(config)) | collect + tk_done = run_track(pos_meta, input_zarr, input_images_zarr, work_dir, config, file(config)) | collect emit: done = tk_done diff --git a/nextflow/nextflow.config b/nextflow/nextflow.config index ee5d2e12..9fee36ac 100644 --- a/nextflow/nextflow.config +++ b/nextflow/nextflow.config @@ -17,9 +17,9 @@ params { // of long-lived processes. Override per run: `--max_jobs 10`. max_jobs = 30 - // QC: drives the `imaging-qc` CLI. In mantis-v2.nf, --qc_config (assembled - // store) and --qc_track_config (tracking store) each turn on QC of one store; - // both become tabs of one report. qc-standalone.nf instead takes + // QC: drives the `imaging-qc` CLI. In mantis-v2.nf, --qc_config turns on QC + // of the assembled store, including the instance counts of the labels + // tracking writes into it. qc-standalone.nf instead takes // `stages_manifest`, a CSV (header: zarr_path,config_path) — one row runs one // config on one zarr, and every row becomes a tab. stages_manifest = null diff --git a/nextflow/qc-standalone.nf b/nextflow/qc-standalone.nf index b15c3f59..0d882dbc 100644 --- a/nextflow/qc-standalone.nf +++ b/nextflow/qc-standalone.nf @@ -49,7 +49,7 @@ workflow { .findAll { line -> line.trim() && !line.startsWith('zarr_path') } .collect { line -> line.split(',').collect { cell -> cell.trim() } } - // Label each tab by the store's parent directory (`4-assemble`, `5-track`), + // Label each tab by the store's parent directory (e.g. `4-assemble`), // which is what distinguishes stores of one dataset; fall back to the store // name when two stores would otherwise collide. def labels = rows.collect { row -> file(row[0]).parent.name } diff --git a/pyproject.toml b/pyproject.toml index cd17275c..ca2a0358 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,10 @@ segment = ["cellpose>=4.0.1"] # Deliberately CPU-only. ultrack has an optional GPU path when cupy and cucim # are installed, but it offers no speedup in tracking in our use cases. -track = ["ultrack>=0.8"] +# geff-spec>=1.3 records related_objects[].node_prop, which links tracks.geff to +# the label image; geff-spec 1.1 silently drops it, and geff 1.3 itself only caps +# geff-spec (<1.4), so the floor has to be stated here. +track = ["ultrack>=0.8", "geff>=1.3", "geff-spec>=1.3"] # PyQt6>=6.10 does not provide wheels gui = ["napari[pyqt6]", "PyQt6<6.10", "napari-animation"] @@ -109,22 +112,35 @@ index-strategy = "unsafe-best-match" # # git config --global url."ssh://git@github.com/".insteadOf "https://github.com/" # -# Pinned at #266: feat: plan-stage refuses an unknown channels:/per_channel: name -imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipeline", rev = "c5c51dc700d0aa52388d6f6a08eec5bfe16a5766" } - -# Pin iohub to main at the czbiohub-sf/iohub#476 merge, which syncs a write -# unit's shards before recording it as done. Without it a preemption or a -# write-back error between the write and the marker leaves a `.done` record -# beside a shard whose bytes never landed, and every later `--resume` skips it: -# the resume probe reads one element per shard, so corruption away from the -# origin reads back clean. That silently cost several pipeline runs on -# 2026_08_14_dynatrack. Also carries czbiohub-sf/iohub#472 -# (`create_empty_plate(extra_metadata=...)`) and czbiohub-sf/iohub#479, which +# Pinned at #243: feat: instance_count reads a label image inside the intensity +# store (label:), which counts the labels track writes into the assembled plate. +# Its #235 removed the `upstream_qc_dir` config key, so configs must not carry it. +imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipeline", rev = "f5ae656e0e37320ead74f7c1f95ab1ed304c1829" } + +# Pin ultrack to feat/geff-export-tracks on the ieivanov fork (royerlab/ultrack#288): +# Tracker.to_geff gains track_id/lineage_id/seg_id node props, related_objects +# pointing at the label image, axes with units and scale (no z in 2D), +# include_masks/include_overlaps and zarr_format=3, so track can write tracks.geff +# inside the assembled OME-NGFF 0.5 plate (#361). Repoint at a release once merged. +ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "3cd1b090307cfb13ac62d72df78ef4555ce0ffa3" } + +# Pin iohub to dev/dca-labels: main (d1f9e89, which carries czbiohub-sf/iohub#476, +# #472 and #479 — see below) merged with two unmerged branches that track needs to +# write OME-NGFF label images into the assembled plate (#360): +# - fix/labels-group-metadata (czbiohub-sf/iohub#485): the labels list goes on +# /labels/zarr.json, not the position, so ome-zarr-models (and so +# dca-validate) can open a field that has labels; +# - feat/create-empty-plate-labels (czbiohub-sf/iohub#486): +# create_empty_plate(label_names=...) scaffolds empty label images. +# Repoint at main once both merge, then at a PyPI release. +# +# Main's own reasons: #476 syncs a write unit's shards before recording it as +# done (without it a preemption between the write and the marker leaves a `.done` +# beside a shard whose bytes never landed, which silently cost several runs on +# 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's -# metadata per position (quadratic in FOVs per well: hours for a ~40k-FOV -# plate's init). Drop this back to a PyPI floor once iohub cuts a release -# containing all three. -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "d1f9e89294529d365719cf87a9b29622ca77a992" } +# metadata per position (hours for a ~40k-FOV plate's init). +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "2ab6667f93903f6e1036c1ddd1f35c74f58444de" } [tool.hatch.metadata] allow-direct-references = true diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 25911343..9120d22e 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -10,7 +10,17 @@ from biahub.cli.main import cli from biahub.settings import TrackingSettings, ZSlicing -from biahub.track import _init_output_plate, resolve_z_slice, track +from biahub.track import ( + _init_labels, + _tracked_shape, + label_chunks_and_shards, + resolve_z_slice, + track, + tracked_z_window, + write_tracking_labels, +) + +LABEL = "nuclei_prediction" @pytest.fixture(scope="function") @@ -81,6 +91,28 @@ def example_blank_frames_csv(tmp_path): yield csv_path +def _assert_tracking_outputs(position_path, z_index): + """Labels at z_index only, a valid GEFF linked to them, and the CSV.""" + import geff + + with open_ome_zarr(str(position_path), mode="r") as pos: + labels = pos.get_label(LABEL)["0"][:] + assert labels.dtype == np.uint32 + assert labels[:, z_index].any() + assert not np.delete(labels, z_index, axis=1).any() + + geff_path = position_path / "tracks.geff" + geff.validate_structure(str(geff_path)) + _, metadata = geff.read(str(geff_path), backend="networkx") + (related,) = metadata.related_objects + assert related.path == f"../labels/{LABEL}" + assert related.node_prop == "seg_id" + + (csv_path,) = position_path.glob("tracks_*.csv") + tracks = pd.read_csv(csv_path) + assert set(tracks["track_id"]) <= set(np.unique(labels)) - {0} + + def _make_tracking_config(plate_path, tmp_path): """Create a minimal tracking config pointing at the test plate.""" config_path = tmp_path / "track_config.yml" @@ -175,7 +207,7 @@ def test_track_cli_local( plate_path, _ = example_tracking_plate config_path = _make_tracking_config(plate_path, tmp_path) - output_path = tmp_path / "tracking_output.zarr" + output_path = tmp_path / "tracking_output" track( input_position_dirpaths=[ @@ -189,10 +221,10 @@ def test_track_cli_local( cluster="local", ) - assert output_path.exists() + # central z_slicing on Z=3 tracks planes [0, 3), so the 2D labels sit at z=1. for position in ["A/1/0", "B/1/0", "B/2/0"]: - position_path = output_path / position - assert position_path.exists() + _assert_tracking_outputs(plate_path / position, z_index=1) + assert (output_path / position.replace("/", "_")).is_dir() # Ultrack database def test_track_cli_with_blank_frames( @@ -207,7 +239,7 @@ def test_track_cli_with_blank_frames( plate_path, _ = example_tracking_plate config_path = _make_tracking_config(plate_path, tmp_path) - output_path = tmp_path / "tracking_output_blank_frames.zarr" + output_path = tmp_path / "tracking_output_blank_frames" # Add blank_frames_path to config with open(config_path) as f: @@ -228,7 +260,7 @@ def test_track_cli_with_blank_frames( cluster="local", ) - assert output_path.exists() + _assert_tracking_outputs(plate_path / "A" / "1" / "0", z_index=1) def test_track_cli_invalid_config(tmp_path, monkeypatch): @@ -266,9 +298,9 @@ def test_track_cli_missing_input_path(tmp_path, example_track_settings, monkeypa def test_track_cli_init_only(tmp_path, example_tracking_plate): - """Test that --init creates the output store and emits RESOURCES.""" + """--init creates empty full-Z label images in the -i positions and emits RESOURCES.""" plate_path, _ = example_tracking_plate - output_path = tmp_path / "track_output.zarr" + output_path = tmp_path / "track_output" config_path = _make_tracking_config(plate_path, tmp_path) runner = CliRunner() @@ -289,12 +321,26 @@ def test_track_cli_init_only(tmp_path, example_tracking_plate): ) assert result.exit_code == 0, result.output - assert output_path.exists() assert "RESOURCES:" in result.output - with open_ome_zarr(str(output_path / "A" / "1" / "0"), mode="r") as ds: - assert ds.data.dtype == np.uint32 - assert ds.channel_names == ["nuclei_prediction_labels"] + for position in ["A/1/0", "B/1/0", "B/2/0"]: + with open_ome_zarr(str(plate_path / position), mode="r") as pos: + assert pos.label_names() == [LABEL] + assert pos.channel_names == ["nuclei_prediction", "membrane_prediction"] + label = pos.get_label(LABEL) + assert label["0"].shape == (5, 3, 64, 64) + assert label["0"].dtype == np.uint32 + assert not label["0"][:].any() + assert [ax.name.lower() for ax in label.axes] == ["t", "z", "y", "x"] + label_attrs = yaml.safe_load( + (plate_path / position / "labels" / LABEL / "zarr.json").read_text() + )["attributes"] + assert label_attrs["biahub-track"]["target_channel"] == LABEL + # The labels listing lives on the labels group (OME-NGFF 0.5). + labels_attrs = yaml.safe_load( + (plate_path / position / "labels" / "zarr.json").read_text() + )["attributes"] + assert labels_attrs["ome"]["labels"] == [LABEL] def test_track_cli_debug_single_position(tmp_path, example_tracking_plate, monkeypatch): @@ -302,33 +348,34 @@ def test_track_cli_debug_single_position(tmp_path, example_tracking_plate, monke monkeypatch.setenv("ULTRACK_ARRAY_MODULE", "numpy") plate_path, _ = example_tracking_plate - output_path = tmp_path / "track_output.zarr" + output_path = tmp_path / "track_output" config_path = _make_tracking_config(plate_path, tmp_path) - runner = CliRunner() - result = runner.invoke( - cli, - [ - "track", - "-i", - str(plate_path / "A" / "1" / "0"), - "-o", - str(output_path), - "-c", - str(config_path), - "--cluster", - "debug", - ], - ) + def run(): + return CliRunner().invoke( + cli, + [ + "track", + "-i", + str(plate_path / "A" / "1" / "0"), + "-o", + str(output_path), + "-c", + str(config_path), + "--cluster", + "debug", + ], + ) + result = run() assert result.exit_code == 0, result.output - assert output_path.exists() assert "Tracking complete:" in result.output + _assert_tracking_outputs(plate_path / "A" / "1" / "0", z_index=1) - with open_ome_zarr(str(output_path / "A" / "1" / "0"), mode="r") as ds: - data = ds["0"][:] - assert data.dtype == np.uint32 - assert data.shape[0] > 0 + # A retry overwrites the same outputs. + result = run() + assert result.exit_code == 0, result.output + _assert_tracking_outputs(plate_path / "A" / "1" / "0", z_index=1) # --------------------------------------------------------------------------- @@ -394,9 +441,10 @@ def test_apply_focus_slicing_uniform_window(monkeypatch): "b": np.zeros((T, Z, Y, X)), } z = ZSlicing(method="focus", window_size=6, frac_below=1 / 3) - out = track_mod.apply_focus_slicing(data_dict, z, pixel_size=0.5) + out, window = track_mod.apply_focus_slicing(data_dict, z, pixel_size=0.5) # Same fixed window applied to every channel (center=15 -> slice(13, 19)). + assert window == slice(13, 19) assert out["a"].shape == (T, 6, Y, X) assert out["b"].shape == (T, 6, Y, X) assert np.array_equal(out["a"], data_dict["a"][:, 13:19]) @@ -449,22 +497,95 @@ def _minimal_settings(**overrides): [ ("2D", {"method": "central"}, 1), ("3D", {"method": "range", "range": (0, 2)}, 2), - # Regression: focus + 3D output Z equals the fixed focus window, not the - # full stack -- this is the shape-mismatch bug _init_output_plate had. + # Regression: focus + 3D tracked Z equals the fixed focus window, not the + # full stack. ("3D", {"method": "focus", "window_size": 2}, 2), ], ) -def test_init_output_plate_shape( +def test_init_labels_shape( tmp_path, example_tracking_plate, output_mode, z_slicing, expected_z ): plate_path, _ = example_tracking_plate - output_path = tmp_path / "init_shape.zarr" + position_path = plate_path / "A" / "1" / "0" settings = _minimal_settings(output_mode=output_mode, z_slicing=z_slicing) - shape = _init_output_plate([str(plate_path / "A" / "1" / "0")], output_path, settings) - assert shape[2] == expected_z - with open_ome_zarr(str(output_path / "A" / "1" / "0"), mode="r") as ds: - assert ds.data.shape[2] == expected_z + assert _tracked_shape(position_path, settings)[2] == expected_z + # The label image always spans the image's full Z. + _init_labels([position_path], settings) + with open_ome_zarr(str(position_path), mode="r") as pos: + assert pos.get_label(LABEL)["0"].shape == (5, 3, 64, 64) + + +def test_init_labels_leaves_existing_labels(tmp_path, example_tracking_plate): + plate_path, _ = example_tracking_plate + position_path = plate_path / "A" / "1" / "0" + settings = _minimal_settings() + + _init_labels([position_path], settings) + with open_ome_zarr(str(position_path), mode="r+") as pos: + pos.get_label(LABEL)["0"][0, 1, 0, 0] = 7 + image = pos["0"][:] + _init_labels([position_path], settings) + with open_ome_zarr(str(position_path), mode="r") as pos: + assert pos.get_label(LABEL)["0"][0, 1, 0, 0] == 7 + assert np.array_equal(pos["0"][:], image) + + +def test_label_chunks_and_shards(): + # 1 MB uint32 chunks (DCA SHOULD); shards cover YX up to 16 chunks. + chunks, shards = label_chunks_and_shards((67, 86, 1664, 1193)) + assert chunks == (16, 1, 128, 128) + assert shards == (1, 1, 13, 10) + assert label_chunks_and_shards((67, 86, 4096, 300))[1] == (1, 1, 16, 3) + # Clamped to small arrays. + assert label_chunks_and_shards((5, 3, 64, 64))[0] == (5, 1, 64, 64) + + +@pytest.mark.parametrize( + "read_slice, focus_slice, z_shape, expected", + [ + (slice(None), None, 30, (0, 30)), # all + (slice(9, 12), None, 21, (9, 12)), # central / range + (slice(10, 11), None, 30, (10, 11)), # one plane + (slice(None), slice(13, 19), 30, (13, 19)), # focus window + ], +) +def test_tracked_z_window(read_slice, focus_slice, z_shape, expected): + assert tracked_z_window(read_slice, focus_slice, z_shape) == expected + + +@pytest.mark.parametrize( + "z_window, expected_z", + [ + ((0, 30), 15), # projection of the whole volume -> middle of the volume + ((10, 11), 10), # one plane -> that plane + ((13, 19), 16), # projected section -> middle of the section + ], +) +def test_write_tracking_labels_2d_placement(z_window, expected_z): + label_array = np.zeros((2, 30, 4, 4), dtype=np.uint32) + labels = np.ones((2, 4, 4), dtype=np.int64) + + assert write_tracking_labels(label_array, labels, z_window, "2D") == expected_z + assert label_array[:, expected_z].all() + assert label_array.sum() == labels.sum() + + +def test_write_tracking_labels_2d_unprojected_single_plane(): + label_array = np.zeros((2, 30, 4, 4), dtype=np.uint32) + write_tracking_labels(label_array, np.ones((2, 1, 4, 4)), (10, 11), "2D") + assert label_array[:, 10].all() + + +def test_write_tracking_labels_3d_fills_window(): + label_array = np.zeros((2, 30, 4, 4), dtype=np.uint32) + labels = np.ones((2, 6, 4, 4), dtype=np.int64) + + assert write_tracking_labels(label_array, labels, (13, 19), "3D") == slice(13, 19) + assert label_array[:, 13:19].all() + assert label_array.sum() == labels.sum() + with pytest.raises(ValueError, match="3D"): + write_tracking_labels(label_array, np.ones((2, 5, 4, 4)), (13, 19), "3D") # --------------------------------------------------------------------------- @@ -477,7 +598,7 @@ def test_input_images_path_override(tmp_path, example_tracking_plate, monkeypatc monkeypatch.setenv("ULTRACK_ARRAY_MODULE", "numpy") plate_path, _ = example_tracking_plate - output_path = tmp_path / "override_output.zarr" + output_path = tmp_path / "override_output" config_path = _make_tracking_config(plate_path, tmp_path) runner = CliRunner() @@ -499,7 +620,7 @@ def test_input_images_path_override(tmp_path, example_tracking_plate, monkeypatc ) assert result.exit_code == 0, result.output - assert output_path.exists() + _assert_tracking_outputs(plate_path / "A" / "1" / "0", z_index=1) def test_track_init_rejects_an_unknown_cellpose_model( @@ -540,7 +661,7 @@ def __init__(self, gpu=False, pretrained_model="cpsam_v2", device=None): "-i", str(plate_path / "A" / "1" / "0"), "-o", - str(tmp_path / "out.zarr"), + str(tmp_path / "out"), "-c", str(config_path), "--init", diff --git a/uv.lock b/uv.lock index 8bda90ab..d5d9185d 100644 --- a/uv.lock +++ b/uv.lock @@ -315,6 +315,8 @@ dependencies = [ all = [ { name = "cellpose" }, { name = "cytoland" }, + { name = "geff" }, + { name = "geff-spec" }, { name = "napari", extra = ["pyqt6"] }, { name = "napari-animation" }, { name = "pyqt6" }, @@ -343,6 +345,8 @@ stain = [ { name = "cytoland" }, ] track = [ + { name = "geff" }, + { name = "geff-spec" }, { name = "ultrack" }, ] @@ -364,10 +368,14 @@ requires-dist = [ { name = "cytoland", marker = "extra == 'all'" }, { name = "cytoland", marker = "extra == 'stain'" }, { name = "dask", extras = ["array"] }, + { name = "geff", marker = "extra == 'all'", specifier = ">=1.3" }, + { name = "geff", marker = "extra == 'track'", specifier = ">=1.3" }, + { name = "geff-spec", marker = "extra == 'all'", specifier = ">=1.3" }, + { name = "geff-spec", marker = "extra == 'track'", specifier = ">=1.3" }, { name = "humanize" }, { name = "imageio-ffmpeg" }, - { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=c5c51dc700d0aa52388d6f6a08eec5bfe16a5766" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=d1f9e89294529d365719cf87a9b29622ca77a992" }, + { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=2ab6667f93903f6e1036c1ddd1f35c74f58444de" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -393,8 +401,8 @@ requires-dist = [ { name = "torch" }, { name = "tqdm" }, { name = "twine", marker = "extra == 'build'" }, - { name = "ultrack", marker = "extra == 'all'", specifier = ">=0.8" }, - { name = "ultrack", marker = "extra == 'track'", specifier = ">=0.8" }, + { name = "ultrack", marker = "extra == 'all'", git = "https://github.com/ieivanov/ultrack?rev=3cd1b090307cfb13ac62d72df78ef4555ce0ffa3" }, + { name = "ultrack", marker = "extra == 'track'", git = "https://github.com/ieivanov/ultrack?rev=3cd1b090307cfb13ac62d72df78ef4555ce0ffa3" }, { name = "waveorder", specifier = ">=3.0.4" }, { name = "zensical", marker = "extra == 'doc'", specifier = ">=0.0.32" }, ] @@ -1517,7 +1525,7 @@ s3 = [ [[package]] name = "geff" -version = "1.1.5.1.1" +version = "1.3.1.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "geff-spec" }, @@ -1528,14 +1536,14 @@ dependencies = [ { name = "typer" }, { name = "zarr" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e9/13/a1292643886905b34e9e492cb8c1dcdb5c5d553533ee3a5dec0f45567bbe/geff-1.1.5.1.1.tar.gz", hash = "sha256:7823ca1e82e06ee4931b0a9fca626021f75236c1831749fe477783662a550e3b", size = 126154, upload-time = "2026-04-06T14:07:09.271Z" } +sdist = { url = "https://files.pythonhosted.org/packages/86/45/da956fd26a35972f7ef549c8b2af9456d217315ef50282c4b3060de08de0/geff-1.3.1.1.3.tar.gz", hash = "sha256:318233ced75ffac29a6ecd813b3cf872ff57bb83e2f28bb64a81628a8d9ac2da", size = 133538, upload-time = "2026-08-27T20:53:03.702Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ce/12/6e8dcce315335269137157e9cbab82748e578e98f6ef2ede4332d1da0613/geff-1.1.5.1.1-py3-none-any.whl", hash = "sha256:584b048dc45c8df329e3675c125983c30ca7c50b89948e69a3c30de790e194cb", size = 64413, upload-time = "2026-04-06T14:07:07.872Z" }, + { url = "https://files.pythonhosted.org/packages/98/ae/dcc06a6dc25d52988ab8eb0bf2da34ed5adc7cc41c5cad0dd085797918b2/geff-1.3.1.1.3-py3-none-any.whl", hash = "sha256:d86d608f5a3eab5e7be3252d3a900f0d8ac0e2bf2a53e97ac546a91b675ff6b2", size = 69012, upload-time = "2026-08-27T20:53:02.416Z" }, ] [[package]] name = "geff-spec" -version = "1.1.1" +version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "annotated-types" }, @@ -1543,9 +1551,9 @@ dependencies = [ { name = "pydantic" }, { name = "zarr" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b9/59/1ac7b17c3594a64a992a77e94552203950bd771f9c769e92ed8c1264b591/geff_spec-1.1.1.tar.gz", hash = "sha256:e7baecfdcc1ebbca6eab716817315f22c7a65eecb14f7377735bea8ae30944d2", size = 19703, upload-time = "2025-11-20T18:55:43.553Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b3/51/3ecce25c6152b5bbdc1637b5490ad8e98304d7e8422bbf7ecbea601572ed/geff_spec-1.3.0.tar.gz", hash = "sha256:d4fc5de3759eff8b9ff888269b655f1cc6160eda859ff74f35354cee34180c66", size = 21114, upload-time = "2026-08-27T20:06:01.286Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/18/9d/4e52ab1a81556dcb6ed1b246d72ebde9773f23cf5848d0aa3ffba7c40ce7/geff_spec-1.1.1-py3-none-any.whl", hash = "sha256:2e1bd5ba5c6186cc5bac24f93ac74821f2c7e6ce2a7aa0af791c00b6203af80d", size = 14807, upload-time = "2025-11-20T18:55:42.065Z" }, + { url = "https://files.pythonhosted.org/packages/3c/79/5c5cfcbccbb83e32d8163cf26d998a6887d716b944f8599e6e14702e3cda/geff_spec-1.3.0-py3-none-any.whl", hash = "sha256:cc629ced0bf9ff2bbe6c1b2184fe9d89f134776f1d5d9f1daa2b0096e92e6fbd", size = 15602, upload-time = "2026-08-27T20:06:00.114Z" }, ] [[package]] @@ -1937,7 +1945,7 @@ wheels = [ [[package]] name = "imaging-qc-pipeline" version = "0.5.0" -source = { git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=c5c51dc700d0aa52388d6f6a08eec5bfe16a5766#c5c51dc700d0aa52388d6f6a08eec5bfe16a5766" } +source = { git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829#f5ae656e0e37320ead74f7c1f95ab1ed304c1829" } dependencies = [ { name = "dask", extra = ["array"] }, { name = "hydra-core" }, @@ -2010,8 +2018,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev16+d1f9e89" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=d1f9e89294529d365719cf87a9b29622ca77a992#d1f9e89294529d365719cf87a9b29622ca77a992" } +version = "0.3.12.dev20+2ab6667" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=2ab6667f93903f6e1036c1ddd1f35c74f58444de#2ab6667f93903f6e1036c1ddd1f35c74f58444de" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, @@ -3697,34 +3705,22 @@ wheels = [ [[package]] name = "numcodecs" -version = "0.17.0" +version = "0.15.1" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "deprecated" }, { name = "numpy" }, - { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/dd/ec/260cdb6304868de6db14eb31064bd2735c0200bcb3331d6b4c9e9be02a03/numcodecs-0.17.0.tar.gz", hash = "sha256:e8db2e337bdafd3bb5f891a2543b53b2b36a509ce9d587af2846db3715b6c8b9", size = 6288352, upload-time = "2026-09-17T18:12:42.262Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8f/e8/28cc96c77078ffcd08579211297cbf1f8ca6e76b4b53c8fbc029b879aaa8/numcodecs-0.17.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2e29732c5e3a83663e51b40007819d8fd0aae16a2322f7044ce13a2460a99e23", size = 1171637, upload-time = "2026-09-17T18:12:12.765Z" }, - { url = "https://files.pythonhosted.org/packages/96/59/1cde6df2f9baa26a1a21c36ac10312062acace29e5c95e029d8da9cf7c3d/numcodecs-0.17.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d30c69b4bdb1755af1022fa913e184eaadc4fc0cd38f736e483e8ad205e130d1", size = 978974, upload-time = "2026-09-17T18:12:14.496Z" }, - { url = "https://files.pythonhosted.org/packages/ef/86/15e1cc4e6644d7e33be613d17bb7cc939b1862ccd975fa2ce1055a1e3045/numcodecs-0.17.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1837d4d1d646cecd3ab2d1ba22956295d709edea0bddc952737c647bec1d03c4", size = 1384196, upload-time = "2026-09-17T18:12:16.327Z" }, - { url = "https://files.pythonhosted.org/packages/73/ca/b784745f189a12ccef60517c0c8526d579b40f35da4463c30c07a4677366/numcodecs-0.17.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1ebd63cdb8985c66257bc037fcdff5f38637aff72d7ef62612ec46f2299e8749", size = 1436978, upload-time = "2026-09-17T18:12:17.731Z" }, - { url = "https://files.pythonhosted.org/packages/7f/b0/f8b3852828c6712eae36e031d763cd52c2777290406066eae0b2a527c05f/numcodecs-0.17.0-cp312-cp312-win_amd64.whl", hash = "sha256:ecd0f6a10e3f8afbbb16ecc999d2b06aa2a31a2946f1c1a85d15d91a1ebcfef3", size = 1496300, upload-time = "2026-09-17T18:12:19.319Z" }, - { url = "https://files.pythonhosted.org/packages/11/f1/1d3d2bcb1240e5000f6647b5b0fd465b2b51ecef180bfa797a85df48cf2f/numcodecs-0.17.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:de2c66db238e74e66fe9be7e02b7e0129b75d3f812d38e4019eb0102cc2dcdf0", size = 1170875, upload-time = "2026-09-17T18:12:20.638Z" }, - { url = "https://files.pythonhosted.org/packages/64/81/64e2472a8b3a9fa26bccfc7d5fa876770a9027bd5cd77e5b4a7b807a0785/numcodecs-0.17.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:69b9b4685097c4d478a0c829debf4470555ec63e92cdd2c6b5f195460f1dc888", size = 976128, upload-time = "2026-09-17T18:12:21.856Z" }, - { url = "https://files.pythonhosted.org/packages/25/ea/2ab25f7e674cf1e78f123c5c2689d8a7dc85475554af0619bcce05cb32a9/numcodecs-0.17.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7065b3349b73d54785aa89e00d0b97d80f664e9056757929d28151f9208dc04c", size = 1379341, upload-time = "2026-09-17T18:12:23.159Z" }, - { url = "https://files.pythonhosted.org/packages/9d/96/b3bf9a31978d936654a73f2bb1036b92b6515164f170092d162419eb771c/numcodecs-0.17.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6c3342d91ed7cf59c1be84396edd364e936bb0ec9e366d24bb69689748d19625", size = 1430551, upload-time = "2026-09-17T18:12:24.97Z" }, - { url = "https://files.pythonhosted.org/packages/5c/ec/47515bea31725376aa6f061c335326cc7437f863ab704b8e167e80734bfd/numcodecs-0.17.0-cp313-cp313-win_amd64.whl", hash = "sha256:a854e9c89f58eeeb2453f3c1637d1916797edb6eaff26bc186a6cdb09d187092", size = 1492769, upload-time = "2026-09-17T18:12:26.702Z" }, - { url = "https://files.pythonhosted.org/packages/cc/e0/be0a4df898bd2cca26cf8071552925aa66601210c0e35be9c4070ede6467/numcodecs-0.17.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:0fc125d1c726c1937cde346e109e3662a2b4ff6be073289da7d124d172aceda5", size = 1172207, upload-time = "2026-09-17T18:12:28.061Z" }, - { url = "https://files.pythonhosted.org/packages/54/0f/9da01fd25953fc37273d7bab7a2174ff74520450a8d77abb837de5d5f040/numcodecs-0.17.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:6f1293581326e92293b142bd05b389f6682ed1ce333f36f116344bca340cfd10", size = 978019, upload-time = "2026-09-17T18:12:29.597Z" }, - { url = "https://files.pythonhosted.org/packages/31/ee/e306a14295a67f9852b9856077ed5ab6b67fae5797559ccdd22a3aa6d853/numcodecs-0.17.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a62e5a821ccfbe425bbdd9a079f8b6c41b7e796ff3c99324530561193a53047", size = 1387252, upload-time = "2026-09-17T18:12:31.065Z" }, - { url = "https://files.pythonhosted.org/packages/d9/13/107244147a8b2edd43ffc3f8bf229f07220aed200022d2e0b1dd761b68f8/numcodecs-0.17.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1cce4bf2278ed74841c2088acfd38e67c3e5aa77e3bc1962ef0fa2931becbb12", size = 1432514, upload-time = "2026-09-17T18:12:32.327Z" }, - { url = "https://files.pythonhosted.org/packages/dc/88/9460630aa3517f1745da87201a96ba84169f2ebee9558b223c0e5a35ab44/numcodecs-0.17.0-cp314-cp314-win_amd64.whl", hash = "sha256:4f43ba0d834ce012ed482996a7424df9077a47d5899ede2d1d54fe85e6eb12fa", size = 1534816, upload-time = "2026-09-17T18:12:33.566Z" }, - { url = "https://files.pythonhosted.org/packages/da/d7/c78d934e1baedd5c2e2d06ee087e23f9031ddf2dded75193195ab290139d/numcodecs-0.17.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:657b1f9aa4b1025aa0fa7d4bd8d7492900950a11f636dff622bd208c0b99e35e", size = 1196273, upload-time = "2026-09-17T18:12:34.952Z" }, - { url = "https://files.pythonhosted.org/packages/00/c1/bbc350003a32876a6a80ace17f7571ebd0110e39c523b8e07ecf64d431b2/numcodecs-0.17.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:4d83befe67a51ba6a988c562209bf13836438c1b6dce23049d84ff42854af32d", size = 1022050, upload-time = "2026-09-17T18:12:36.312Z" }, - { url = "https://files.pythonhosted.org/packages/ca/dc/fa7c6a1ce327093d04ddf0d0acf20fd1fe2a00a9b7eb43b6a7213eec1ae4/numcodecs-0.17.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3e4e351566b3ab2f6255a9d91c6c48e1d0f9ec6e2ae409a148e091a8fc0a80b0", size = 1380833, upload-time = "2026-09-17T18:12:37.739Z" }, - { url = "https://files.pythonhosted.org/packages/e6/38/33023f8771e8b7afb9c7f0dd50b0487dbef594617e245707cb965969edd3/numcodecs-0.17.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8697a4631fedded77a75d333e4926b1eb3a11bc7d3e30213e7e565d6910526d0", size = 1424937, upload-time = "2026-09-17T18:12:39.04Z" }, - { url = "https://files.pythonhosted.org/packages/86/43/a166898bd89ecc743762b0192b591a35aaecf442e47f85132a75113622f5/numcodecs-0.17.0-cp314-cp314t-win_amd64.whl", hash = "sha256:4c36f6fd14dc22939172145c24d3b3eab2410c34ed807906a5ece5f4541c7c43", size = 1523408, upload-time = "2026-09-17T18:12:40.72Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/63/fc/bb532969eb8236984ba65e4f0079a7da885b8ac0ce1f0835decbb3938a62/numcodecs-0.15.1.tar.gz", hash = "sha256:eeed77e4d6636641a2cc605fbc6078c7a8f2cc40f3dfa2b3f61e52e6091b04ff", size = 6267275, upload-time = "2025-02-10T10:23:33.254Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/7e/f12fc32d3beedc6a8f1ec69ea0ba72e93cb99c0350feed2cff5d04679bc3/numcodecs-0.15.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b0a9d9cd29a0088220682dda4a9898321f7813ff7802be2bbb545f6e3d2f10ff", size = 1691889, upload-time = "2025-02-10T10:23:12.934Z" }, + { url = "https://files.pythonhosted.org/packages/81/38/88e40d40288b73c3b3a390ed5614a34b0661d00255bdd4cfb91c32101364/numcodecs-0.15.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:a34f0fe5e5f3b837bbedbeb98794a6d4a12eeeef8d4697b523905837900b5e1c", size = 1189149, upload-time = "2025-02-10T10:23:15.803Z" }, + { url = "https://files.pythonhosted.org/packages/28/7d/7527d9180bc76011d6163c848c9cf02cd28a623c2c66cf543e1e86de7c5e/numcodecs-0.15.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c3a09e22140f2c691f7df26303ff8fa2dadcf26d7d0828398c0bc09b69e5efa3", size = 8879163, upload-time = "2025-02-10T10:23:18.582Z" }, + { url = "https://files.pythonhosted.org/packages/ab/bc/b6c3cde91c754860a3467a8c058dcf0b1a5ca14d82b1c5397c700cf8b1eb/numcodecs-0.15.1-cp312-cp312-win_amd64.whl", hash = "sha256:daed6066ffcf40082da847d318b5ab6123d69ceb433ba603cb87c323a541a8bc", size = 836785, upload-time = "2025-02-10T10:23:22.314Z" }, + { url = "https://files.pythonhosted.org/packages/78/57/acbc54b3419e5be65015e47177c76c0a73e037fd3ae2cde5808169194d4d/numcodecs-0.15.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e3d82b70500cf61e8d115faa0d0a76be6ecdc24a16477ee3279d711699ad85f3", size = 1688220, upload-time = "2025-02-10T10:23:23.79Z" }, + { url = "https://files.pythonhosted.org/packages/b6/56/9863fa6dc679f40a31bea5e9713ee5507a31dcd3ee82ea4b1a9268ce52e8/numcodecs-0.15.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1d471a1829ce52d3f365053a2bd1379e32e369517557c4027ddf5ac0d99c591e", size = 1180294, upload-time = "2025-02-10T10:23:25.533Z" }, + { url = "https://files.pythonhosted.org/packages/fa/91/d96999b41e3146b6c0ce6bddc5ad85803cb4d743c95394562c2a4bb8cded/numcodecs-0.15.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1dfdea4a67108205edfce99c1cb6cd621343bc7abb7e16a041c966776920e7de", size = 8834323, upload-time = "2025-02-10T10:23:27.46Z" }, + { url = "https://files.pythonhosted.org/packages/c3/32/233e5ede6568bdb044e6f99aaa9fa39827ff3109c6487fc137315f733586/numcodecs-0.15.1-cp313-cp313-win_amd64.whl", hash = "sha256:a4f7bdb26f1b34423cb56d48e75821223be38040907c9b5954eeb7463e7eb03c", size = 831955, upload-time = "2025-02-10T10:23:30.601Z" }, ] [[package]] @@ -6686,8 +6682,8 @@ wheels = [ [[package]] name = "ultrack" -version = "0.8.0" -source = { registry = "https://pypi.org/simple" } +version = "0.7.3" +source = { git = "https://github.com/ieivanov/ultrack?rev=3cd1b090307cfb13ac62d72df78ef4555ce0ffa3#3cd1b090307cfb13ac62d72df78ef4555ce0ffa3" } dependencies = [ { name = "blosc2" }, { name = "click" }, @@ -6696,6 +6692,7 @@ dependencies = [ { name = "fastapi" }, { name = "fasteners" }, { name = "geff" }, + { name = "geff-spec" }, { name = "gurobipy" }, { name = "higra" }, { name = "httpx" }, @@ -6728,10 +6725,6 @@ dependencies = [ { name = "websockets" }, { name = "zarr" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e2/1f/120af2a543438fb4f411b0cdea6b6fe055b41679c2826d0684e8c8ef820b/ultrack-0.8.0.tar.gz", hash = "sha256:9e11193dabd6398f8d27d5b024a51f92c0364e32493afaed1c470f5f7b70d511", size = 193136, upload-time = "2026-08-13T17:59:14.904Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/32/2b/811db2da0eb6c7f4b02630952216c7f84b2dd8e23a92b18fdf086dcd3b99/ultrack-0.8.0-py3-none-any.whl", hash = "sha256:7226f4e1c7cee3cdb0f61f358f1bb2a917d53d09eefe5b85221413d5ed40cab6", size = 274273, upload-time = "2026-08-13T17:59:13.611Z" }, -] [[package]] name = "urllib3" From d8e06a0c0c6ba73c5e22b4589bc4fc57c2148f65 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Thu, 8 Oct 2026 16:22:29 -0700 Subject: [PATCH 02/54] build(deps): repin iohub dev/dca-labels after #485 drops its legacy labels field Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ca2a0358..d39bea4c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -140,7 +140,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "3cd1b090307cfb13 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "2ab6667f93903f6e1036c1ddd1f35c74f58444de" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "6f29e64e6b9979328d38a3b08a3355732f59d7e9" } [tool.hatch.metadata] allow-direct-references = true diff --git a/uv.lock b/uv.lock index d5d9185d..e9bc9d6d 100644 --- a/uv.lock +++ b/uv.lock @@ -375,7 +375,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=2ab6667f93903f6e1036c1ddd1f35c74f58444de" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=6f29e64e6b9979328d38a3b08a3355732f59d7e9" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2018,8 +2018,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev20+2ab6667" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=2ab6667f93903f6e1036c1ddd1f35c74f58444de#2ab6667f93903f6e1036c1ddd1f35c74f58444de" } +version = "0.3.12.dev23+6f29e64" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=6f29e64e6b9979328d38a3b08a3355732f59d7e9#6f29e64e6b9979328d38a3b08a3355732f59d7e9" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 02e5f63cb80cd3d91a2c25a65049459f484a6604 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Thu, 8 Oct 2026 16:48:57 -0700 Subject: [PATCH 03/54] build(deps): repin iohub dev/dca-labels after iohub#486 label-default fix Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d39bea4c..700a851a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -140,7 +140,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "3cd1b090307cfb13 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "6f29e64e6b9979328d38a3b08a3355732f59d7e9" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "82ef216b059351a95f16d802f2667573dd59e1db" } [tool.hatch.metadata] allow-direct-references = true diff --git a/uv.lock b/uv.lock index e9bc9d6d..dc6713a4 100644 --- a/uv.lock +++ b/uv.lock @@ -375,7 +375,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=6f29e64e6b9979328d38a3b08a3355732f59d7e9" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=82ef216b059351a95f16d802f2667573dd59e1db" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2018,8 +2018,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev23+6f29e64" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=6f29e64e6b9979328d38a3b08a3355732f59d7e9#6f29e64e6b9979328d38a3b08a3355732f59d7e9" } +version = "0.3.12.dev25+82ef216" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=82ef216b059351a95f16d802f2667573dd59e1db#82ef216b059351a95f16d802f2667573dd59e1db" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From f863ad2407a4f4ad51b755cd9eaffeb1a5d40125 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Thu, 8 Oct 2026 17:12:05 -0700 Subject: [PATCH 04/54] feat(track): keep no Ultrack databases by default Each worker now builds its Ultrack database in a fresh scratch directory under its own $TMPDIR and deletes it once the FOV is exported. `-o` becomes optional and keeps the databases in {-o}/{row}_{col}_{fov} when given (e.g. to re-solve with other weights); the SLURM logs go to {-o}/slurm_output, else ./slurm_output, since submitit's folder must be on a shared filesystem. The scratch directory also holds ultrack's own temporary zarr arrays: create_zarr backs its default store with a TemporaryDirectory it lets go out of scope, so those directories were recreated by the first write and left in $TMPDIR, which on Bruno outlives the job. The full resolved Ultrack config (defaults included) is recorded under `biahub-track.ultrack_config` on each label image, so provenance no longer depends on config.toml beside the database. The worker returns nothing rather than a label array inside the deleted scratch directory. Nextflow passes no -o, so tracking has no step directory (no 5-track/, no README, no chmod of the SQLite file) and takes no number; its real logs stay in nextflow/slurm_output/track/. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../skills/reconstruct-with-nextflow/SKILL.md | 14 +- .../references/caveats.md | 9 +- .../references/datasets.md | 5 +- .../references/monitoring.md | 2 +- .../references/recovery.md | 10 +- .../templates/run_mantis_v2.sh | 12 +- biahub/cli/parsing.py | 11 +- biahub/track.py | 197 +++++++++++------- nextflow/mantis-v2.nf | 16 +- nextflow/modules/qc.nf | 2 +- nextflow/modules/tracking.nf | 31 +-- tests/test_cli/test_track_cli.py | 20 +- 12 files changed, 183 insertions(+), 146 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index a2e01e84..0429ab87 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -344,11 +344,10 @@ Do not run anything yet. Show the user: **State the directory numbers this run will produce.** The number is the step's position among the steps performed, not a fixed label, so a neuromast - run writes `4-assemble` as its last directory and an A549 run writes - `4-assemble` then `5-track` (tracking's work directory — its labels and - tracks go into the assembled plate). Older A549 runs on disk say `5-assemble` / - `4-track`, from when the numbers were fixed — say so if the user is comparing - against one. + run and an A549 run both end at `4-assemble`: tracking writes its labels and + tracks into the assembled plate and has no directory of its own. Older A549 + runs on disk also have `5-track` (or `5-assemble` / `4-track`, from when the + numbers were fixed) — say so if the user is comparing against one. 7. Known caveats that apply to this dataset. 8. Rough wall-time and that `-resume` is on. 9. Whether Slack notifications are on, and who will be @-mentioned at run end. @@ -363,9 +362,8 @@ Do not run anything yet. Show the user: pipeline deletes the flat-field, deskew, reconstruct and virtual-stain directories, the `slurm_output/` placeholders and `.iohub-progress/` resume markers beside the final stores, and the Nextflow work directory. Only the - assembled store (with tracking's labels and tracks inside it), tracking's - per-FOV Ultrack databases in `5-track/`, the QC report and the logs under - `nextflow/` remain, and the run is final: any rerun + assembled store (with tracking's labels and tracks inside it), the QC report + and the logs under `nextflow/` remain, and the run is final: any rerun recomputes from raw. The template's `CLEANUP_INTERMEDIATES="auto"` lets the pipeline decide at launch (`cleanup_decision` in `nextflow/modules/cleanup.nf`); check `concatenate.yml` now so the plan says diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 2be94790..492704f0 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -201,9 +201,12 @@ transfer. So simply do not run it: `track_wf` takes the assembled `.zarr` for *both* inputs, and writes its output back INTO it: per position, `labels/` (uint32, TZYX at the image's full Z; 2D-tracking masks sit in the middle plane of the tracked -z-window), `tracks.geff`, and `tracks___.csv`. `5-track/` is only -the work directory (`-o`): per-FOV Ultrack databases at `__/` and -`slurm_output/`. So any Z/Y/X crop or `time_indices` subset in +z-window), `tracks.geff`, and `tracks___.csv`. Nothing else is +kept: the pipeline passes no `-o`, so each worker builds its Ultrack database in +the node's `$TMPDIR` and deletes it after export, and track has no step +directory (pass `-o` to `biahub track` by hand to keep the databases, e.g. to +re-solve with other weights). The full Ultrack config is recorded under +`biahub-track.ultrack_config` on each label image. So any Z/Y/X crop or `time_indices` subset in `concatenate.yml` is what tracking sees, and tracking starts only after the whole plate assembles. `track --init` creates the empty label image in every position during the init phase, and with tracking on, QC is planned after that diff --git a/.claude/skills/reconstruct-with-nextflow/references/datasets.md b/.claude/skills/reconstruct-with-nextflow/references/datasets.md index 4c090817..f488bb7d 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/datasets.md +++ b/.claude/skills/reconstruct-with-nextflow/references/datasets.md @@ -48,7 +48,7 @@ Which root depends on the dataset family: | family | cues in the name | output root | deliverable | |---|---|---|---| | zebrafish / neuromast / dynatrack | `dynatrack`, `hpf`, `dpf`, `neuromast`, `cldnb`, `she`, `myo6b`, `zebrafish` | `/hpc/projects/tlg2_mantis/` | assemble (`4-assemble`) — **no tracking**, see `caveats.md` §4 | -| cell line / organelle / infection | `A549`, `HEK`, `iPSC`, organelle genes (`SEC61B`, `TOMM20`, `G3BP1`, `MAP4`, `CAAX`, `H2B`), viruses (`DENV`, `ZIKV`, `HSV1`) | `/hpc/projects/intracellular_dashboard/organelle_dynamics/` | assemble + track (`4-assemble`, `5-track`) | +| cell line / organelle / infection | `A549`, `HEK`, `iPSC`, organelle genes (`SEC61B`, `TOMM20`, `G3BP1`, `MAP4`, `CAAX`, `H2B`), viruses (`DENV`, `ZIKV`, `HSV1`) | `/hpc/projects/intracellular_dashboard/organelle_dynamics/` | assemble + track (`4-assemble`; tracking writes into it) | | instrument QC / calibration | `argolight`, `beads`, `psf`, `alignment`, `illumination`, `fluorescein`, `first_light`, `mantis_v2_*` | *no default* — ask. These are usually not reconstructed. | — | Neighbouring roots that exist and are **not** the default target — do not use @@ -73,9 +73,6 @@ layout is defined once in `nextflow/mantis-v2.nf` (`directory_layout()`): ├── 4-assemble/.zarr # the deliverable; with tracking, each position │ # also holds labels/, tracks.geff, │ # tracks___.csv -├── 5-track/ # A549 only — tracking's work dir, not a store -│ ├── __/ # per-FOV Ultrack database -│ └── slurm_output/ ├── qc/ │ ├── report/index.html # one report, one tab: the assembled store │ └── report_spec.yaml # written at launch, before any QC task diff --git a/.claude/skills/reconstruct-with-nextflow/references/monitoring.md b/.claude/skills/reconstruct-with-nextflow/references/monitoring.md index af5382d0..ca244e2d 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/monitoring.md +++ b/.claude/skills/reconstruct-with-nextflow/references/monitoring.md @@ -152,7 +152,7 @@ Steps run strictly in sequence; within a step, positions fan out with | `compute_transfer_function` + `run_apply_inv_tf` | one-shot + fan-out, CPU | TF is quick; apply is the bulk | | `run_virtual_stain_preprocess` + `run_virtual_stain` | one-shot + fan-out, **GPU** | `preempted` partition, any GPU node (A40 workstations included); a reclaimed task resumes at the timepoint it lost; longest per-position step | | `run_concatenate` | **single job**, one reserved node | longest single job; whole-plate I/O | -| `run_track` | fan-out, GPU (cellpose) | reads the assembled plate and writes `labels/`, `tracks.geff` and a tracks CSV into each position; `5-track/` holds only the Ultrack databases | +| `run_track` | fan-out, GPU (cellpose) | reads the assembled plate and writes `labels/`, `tracks.geff` and a tracks CSV into each position; nothing else is kept (the Ultrack database is temporary, in the node's `$TMPDIR`) | Per-step wall-time is derived from data volume (`estimate_time_minutes()` in `biahub/utils/cluster.py`) and passed through the init step's `RESOURCES:` payload, diff --git a/.claude/skills/reconstruct-with-nextflow/references/recovery.md b/.claude/skills/reconstruct-with-nextflow/references/recovery.md index 1009446e..dc109791 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/recovery.md +++ b/.claude/skills/reconstruct-with-nextflow/references/recovery.md @@ -106,11 +106,11 @@ session. What `-resume` does after a change: To force one step to recompute, delete that step's output for the affected positions and their work dirs — not the whole `.nextflow/` cache. Track's -output is inside the assembled plate, not under `5-track/`: per position, -`labels/`, `tracks.geff` and `tracks___.csv`, -plus its Ultrack database at `5-track/__/`. A rerun overwrites -all of these in place, so for track deleting the work dirs is what forces it; -never delete the assembled plate to clear tracking output. +output is inside the assembled plate: per position, `labels/`, +`tracks.geff` and `tracks___.csv` (its Ultrack database lives in +the node's `$TMPDIR` and is deleted after export). A rerun overwrites all of +these in place, so for track deleting the work dirs is what forces it; never +delete the assembled plate to clear tracking output. ## When to escalate to the user diff --git a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh index 36484c66..7302729d 100755 --- a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh +++ b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh @@ -10,13 +10,11 @@ # below. Keep it there: it is the run's provenance record of the exact command. # # --output is the project root, so each step writes a sibling directory: -# 0-flatfield/ 1-deskew/ 2-reconstruct/ 3-virtual-stain/ 4-assemble/ 5-track/ -# 5-track/ is tracking's work directory (per-FOV Ultrack databases), not a -# store: tracking writes labels/, tracks.geff and a tracks CSV -# into each position of 4-assemble/.zarr. -# The number is the step's POSITION among the steps this run performs, so a run -# without tracking ends at 4-assemble and one that also skipped a step earlier -# shifts everything after it down. +# 0-flatfield/ 1-deskew/ 2-reconstruct/ 3-virtual-stain/ 4-assemble/ +# Tracking has no directory: it writes labels/, tracks.geff and +# a tracks CSV into each position of 4-assemble/.zarr. +# The number is the step's POSITION among the steps this run performs that +# write a directory, so a run that skips a step shifts everything after it down. # The work dir defaults to /nextflow/work; override with `-work-dir`. # # Any extra arguments are forwarded to nextflow, e.g. diff --git a/biahub/cli/parsing.py b/biahub/cli/parsing.py index 85b0d594..b80a200c 100644 --- a/biahub/cli/parsing.py +++ b/biahub/cli/parsing.py @@ -24,8 +24,8 @@ def _validate_and_process_paths( return input_paths -def _str_to_path(ctx: click.Context, opt: click.Option, value: str) -> Path: - return Path(value) +def _str_to_path(ctx: click.Context, opt: click.Option, value: str | None) -> Path | None: + return None if value is None else Path(value) def _validate_and_process_config_paths(ctx, opt, value: tuple[str, ...]) -> list[Path]: @@ -125,14 +125,15 @@ def decorator(f: Callable) -> Callable: return decorator -def output_dirpath() -> Callable: +def output_dirpath(required: bool = True, help: str = "Path to output directory") -> Callable: def decorator(f: Callable) -> Callable: return click.option( "--output-dirpath", "-o", - required=True, + required=required, + default=None, type=click.Path(exists=False, file_okay=False, dir_okay=True), - help="Path to output directory", + help=help, callback=_str_to_path, )(f) diff --git a/biahub/track.py b/biahub/track.py index c3ba83d7..4116f960 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -1,7 +1,10 @@ import ast +import json import logging -import os +import shutil +import tempfile +from contextlib import contextmanager from pathlib import Path import click @@ -877,11 +880,48 @@ def write_tracking_labels( return slice(start, stop) +@contextmanager +def _fov_scratch(output_dirpath: Path | None, fov: str): + """Yield the Ultrack database directory for one FOV, inside a scratch directory. + + The scratch directory is fresh in this worker's ``$TMPDIR``, one per FOV because + ``$TMPDIR`` can be shared by all of a user's jobs on a node, and deleted on exit. + The database goes in it unless ``output_dirpath`` is given, in which case it is + kept in ``{output_dirpath}/{fov}``. + + While the context is open, ``tempfile`` also points into the scratch directory: + ultrack's ``create_zarr`` backs its default store with a ``TemporaryDirectory`` + that it lets go out of scope, so the directory is deleted, recreated by the first + chunk write, and never cleaned up. Kept in here, those arrays go with the rest. + """ + scratch = Path(tempfile.mkdtemp(prefix=f"biahub-track-{fov}-")) + previous_tempdir = tempfile.tempdir + tempfile.tempdir = str(scratch) + try: + if output_dirpath is not None: + database_path = Path(output_dirpath) / fov + else: + database_path = scratch / "database" + database_path.mkdir(parents=True, exist_ok=True) + yield database_path + finally: + tempfile.tempdir = previous_tempdir + shutil.rmtree(scratch, ignore_errors=True) + + +def _resolved_ultrack_config(tracking_config) -> dict: + """Return the full Ultrack config, defaults included, as JSON-safe provenance.""" + config = json.loads(json.dumps(tracking_config.dict(by_alias=True), default=str)) + # The working directory is the per-FOV database, often a deleted temp directory. + config.get("data", {}).pop("working_dir", None) + return config + + def track_one_position( position_key: tuple[str, str, str], position_dirpath: Path, input_images: list[ProcessingInputChannel], - output_dirpath: Path, + output_dirpath: Path | None, tracking_config, label_name: str, blank_frames_path: Path | None = None, @@ -909,9 +949,11 @@ def track_one_position( created its empty ``labels/`` image. input_images : list of ProcessingInputChannel Configuration describing which channels to load and how to preprocess them. - output_dirpath : Path - Working directory; the Ultrack database and config for this FOV go to - ``{output_dirpath}/{row}_{col}_{fov}/``. + output_dirpath : Path or None + Where to keep the Ultrack database and config for this FOV, in + ``{output_dirpath}/{row}_{col}_{fov}/``. If None they live in a scratch + directory in this worker's ``$TMPDIR`` and are deleted after export (see + :func:`_fov_scratch`). tracking_config : MainConfig Ultrack configuration containing segmentation, linking, and optimization settings. label_name : str @@ -953,80 +995,73 @@ def track_one_position( fov = "_".join(position_key) click.echo(f"Processing FOV: {fov.replace('_', '/')}") - database_path = Path(output_dirpath) / fov - os.makedirs(database_path, exist_ok=True) - with open_ome_zarr(position_dirpath, mode="r") as position: z_shape = position.data.shape[2] time_scale = position.scale[0] units = {axis.name.lower(): getattr(axis, "unit", None) for axis in position.axes} - # Load, focus-slice (method='focus'), run the processing pipeline, and fill blank - # frames once; both segmentation methods consume the resulting data_dict. - # pixel_size is the transverse pixel size = last element of the (trimmed) scale. - pixel_size = scale[-1] - data_dict, focus_slice = _load_and_preprocess( - position_key, input_images, z_slices, blank_frames_path, z_slicing, pixel_size - ) - z_window = tracked_z_window(z_slices or slice(None), focus_slice, z_shape) - - geff_kwargs = { - "zarr_format": 3, - "include_masks": False, - "include_overlaps": False, - "segmentation_path": f"../labels/{label_name}", - "scale": tuple(scale), - "spatial_unit": units.get("x"), - "time_scale": time_scale if units.get("t") else None, - "time_unit": units.get("t"), - } - ultrack_kwargs = { - "tracking_config": tracking_config, - "database_path": database_path, - "geff_path": position_dirpath / "tracks.geff", - "geff_kwargs": geff_kwargs, - "scale": scale, - "overwrite": True, - } - - if cellpose_config is not None: - cellpose_labels = cellpose_segmentation(data_dict, cellpose_config) - - click.echo("Tracking with cellpose labels...") - tracking_labels, tracks_df, _ = run_ultrack( - labels=cellpose_labels, - sigma=cellpose_config.labels_sigma, - **ultrack_kwargs, + with _fov_scratch(output_dirpath, fov) as database_path: + # Load, focus-slice (method='focus'), run the processing pipeline, and fill blank + # frames once; both segmentation methods consume the resulting data_dict. + # pixel_size is the transverse pixel size = last element of the (trimmed) scale. + pixel_size = scale[-1] + data_dict, focus_slice = _load_and_preprocess( + position_key, input_images, z_slices, blank_frames_path, z_slicing, pixel_size ) - else: - foreground_mask, contour_gradient_map = detect_foreground_segmentation(data_dict) + z_window = tracked_z_window(z_slices or slice(None), focus_slice, z_shape) + + geff_kwargs = { + "zarr_format": 3, + "include_masks": False, + "include_overlaps": False, + "segmentation_path": f"../labels/{label_name}", + "scale": tuple(scale), + "spatial_unit": units.get("x"), + "time_scale": time_scale if units.get("t") else None, + "time_unit": units.get("t"), + } + if cellpose_config is not None: + cellpose_labels = cellpose_segmentation(data_dict, cellpose_config) + detection_kwargs = { + "labels": cellpose_labels, + "sigma": cellpose_config.labels_sigma, + } + click.echo("Tracking with cellpose labels...") + else: + foreground_mask, contour_gradient_map = detect_foreground_segmentation(data_dict) + detection_kwargs = {"detection": foreground_mask, "edges": contour_gradient_map} + click.echo("Tracking with foreground + contour...") - click.echo("Tracking with foreground + contour...") tracking_labels, tracks_df, _ = run_ultrack( - detection=foreground_mask, - edges=contour_gradient_map, - **ultrack_kwargs, + tracking_config=tracking_config, + database_path=database_path, + geff_path=position_dirpath / "tracks.geff", + geff_kwargs=geff_kwargs, + scale=scale, + overwrite=True, + **detection_kwargs, ) - # Kept alongside tracks.geff for downstream readers that still expect a CSV. - tracks_df.to_csv(position_dirpath / f"tracks_{fov}.csv", index=False) + # Kept alongside tracks.geff for downstream readers that still expect a CSV. + tracks_df.to_csv(position_dirpath / f"tracks_{fov}.csv", index=False) - with open_ome_zarr(position_dirpath, mode="r+") as position: - label_image = position.get_label(label_name) - z_written = write_tracking_labels( - label_image["0"], tracking_labels, z_window, output_mode - ) + with open_ome_zarr(position_dirpath, mode="r+") as position: + label_image = position.get_label(label_name) + z_written = write_tracking_labels( + label_image["0"], tracking_labels, z_window, output_mode + ) - # Record where the labels sit in Z: the 2D GEFF has no z axis. + # Record where the labels sit in Z (the 2D GEFF has no z axis) and the full Ultrack + # config, so provenance does not depend on the (usually deleted) database directory. label_group = zarr.open_group(position_dirpath / "labels" / label_name, mode="r+") provenance = dict(label_group.attrs.get(PROVENANCE_KEY, {})) provenance["z_window"] = list(z_window) if isinstance(z_written, int): provenance["z_index"] = z_written + provenance["ultrack_config"] = _resolved_ultrack_config(tracking_config) label_group.attrs[PROVENANCE_KEY] = provenance click.echo(f"Saved labels and tracks to: {position_dirpath}") - return tracking_labels, tracks_df def _tracked_shape( @@ -1103,7 +1138,7 @@ def _init_labels( def track( input_position_dirpaths: list[Path], config_filepath: Path, - output_dirpath: Path, + output_dirpath: Path | None = None, sbatch_filepath: str | None = None, cluster: str = "slurm", monitor: bool = True, @@ -1118,9 +1153,12 @@ def track( Paths to input position directories (used for plate structure and metadata). config_filepath : Path Path to the tracking configuration YAML. - output_dirpath : Path - Working directory for the per-FOV Ultrack databases and SLURM logs. The - labels and tracks themselves are written into the ``-i`` positions. + output_dirpath : Path, optional + Where to keep the per-FOV Ultrack databases (``{output_dirpath}/{row}_{col}_{fov}/``) + and the SLURM logs (``{output_dirpath}/slurm_output``). If None, each worker + builds its database in its own ``$TMPDIR`` and deletes it after export, and the + SLURM logs go to ``./slurm_output``. The labels and tracks themselves are + always written into the ``-i`` positions. sbatch_filepath : str, optional Path to a SLURM batch file to override defaults. cluster : str, optional @@ -1145,8 +1183,13 @@ def track( """ from ultrack import MainConfig - output_dirpath = Path(output_dirpath) - slurm_out_path = output_dirpath / "slurm_output" + if output_dirpath is not None: + output_dirpath = Path(output_dirpath) + slurm_out_path = output_dirpath / "slurm_output" + else: + # Not $TMPDIR: submitit's folder must be visible from the compute nodes. + slurm_out_path = Path("slurm_output").absolute() + slurm_out_path.mkdir(parents=True, exist_ok=True) settings = yaml_to_model(config_filepath, TrackingSettings) @@ -1202,7 +1245,6 @@ def track( tracking_cfg = update_model(default_config, tracking_cfg) position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] - output_dirpath.mkdir(parents=True, exist_ok=True) cellpose_cfg = ( settings.cellpose_config if settings.segmentation_method == "cellpose" else None @@ -1254,7 +1296,6 @@ def track( jobs.append(job) job_ids = [job.job_id for job in jobs] - slurm_out_path.mkdir(exist_ok=True) log_path = Path(slurm_out_path / "submitit_jobs_ids.log") with log_path.open("w") as log_file: log_file.write("\n".join(job_ids)) @@ -1276,7 +1317,12 @@ def track( @click.command("track") @input_position_dirpaths() @config_filepath() -@output_dirpath() +@output_dirpath( + required=False, + help="Directory to keep the per-FOV Ultrack databases and the SLURM logs in. " + "If omitted, each worker builds its database in its own $TMPDIR and deletes it " + "after export, and the SLURM logs go to ./slurm_output.", +) @sbatch_filepath() @cluster() @monitor() @@ -1291,7 +1337,7 @@ def track( def track_cli( input_position_dirpaths: list[Path], config_filepath: Path, - output_dirpath: Path, + output_dirpath: Path | None = None, sbatch_filepath: str | None = None, cluster: str = "slurm", monitor: bool = False, @@ -1301,21 +1347,24 @@ def track_cli( """Track objects in 2D or 3D time-lapse microscopy data using configurable preprocessing. Writes into each -i position: the labels as an OME-NGFF label image - (labels/), the tracks as tracks.geff, and a tracks CSV. -o is a - working directory for the Ultrack databases and SLURM logs. + (labels/), the tracks as tracks.geff, and a tracks CSV. The + Ultrack databases are temporary unless -o names a directory to keep them in. \b Full SLURM fan-out: + >>> biahub track -i ./assemble.zarr/*/*/* -c config.yml + + \b + Keep the Ultrack databases (e.g. to re-solve with other weights): >>> biahub track -i ./assemble.zarr/*/*/* -o ./track -c config.yml \b Create the empty label images only (Nextflow init step): - >>> biahub track --init -i ./assemble.zarr/*/*/* -o ./track -c config.yml + >>> biahub track --init -i ./assemble.zarr/*/*/* -c config.yml \b In-process run of a single position (Nextflow per-position worker): - >>> biahub track --cluster debug -i ./assemble.zarr/B/3/000000 \\ - -o ./track -c config.yml + >>> biahub track --cluster debug -i ./assemble.zarr/B/3/000000 -c config.yml """ # noqa: D301 track( input_position_dirpaths=input_position_dirpaths, diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf index d2e41847..eb2e885d 100644 --- a/nextflow/mantis-v2.nf +++ b/nextflow/mantis-v2.nf @@ -96,7 +96,8 @@ def step_directories(performed) { reconstruct : 'reconstruct', virtual_stain : 'virtual-stain', assemble : 'assemble', - track : 'track', + // track writes into the assembled plate and keeps nothing else, so it + // has no directory and takes no number. ] def layout = [:] order.each { key, name -> @@ -124,8 +125,8 @@ workflow { // placeholder to author and no output to discard. // // Skipping a step DOES renumber the ones after it: the number is a position - // among the steps performed, so a neuromast run's assembled store is - // `4-assemble` where an A549 run also has `5-track` after it. + // among the steps that write a directory. Track writes into the assembled + // plate, so A549 and neuromast runs both end at `4-assemble`. def assemble_on = params.concatenate_config as boolean def track_on = params.track_config as boolean def qc_on = params.qc_config as boolean @@ -169,7 +170,6 @@ workflow { // numbering is derived from. Reconstruction proper is always in it. def performed = ['flat_field', 'deskew', 'reconstruct', 'virtual_stain'] if (assemble_on) performed << 'assemble' - if (track_on) performed << 'track' def layout = step_directories(performed) collect_positions(params.input) @@ -187,9 +187,6 @@ workflow { reconstruct_output = "${out}/${layout.reconstruct}/${ds}.zarr" virtual_stain_output = "${out}/${layout.virtual_stain}/${ds}.zarr" assemble_output = assemble_on ? "${out}/${layout.assemble}/${ds}.zarr" : null - // Tracking writes its labels and tracks INTO the assembled plate; its own - // directory holds only the per-FOV Ultrack databases and SLURM placeholders. - track_dir = track_on ? "${out}/${layout.track}" : null // ======================================================================== // INIT PHASE — every config parsed, every output store scaffolded, before @@ -259,7 +256,7 @@ workflow { // same store. Its init creates the empty label image tracking fills in // every position of that plate, and warms the shared cellpose weights // cache, which is better done here than with N GPU workers racing for it. - tk_init = track_init_wf(assemble_output, track_dir, params.track_config, + tk_init = track_init_wf(assemble_output, params.track_config, as_init.done) init_signals << tk_init.done } @@ -360,7 +357,7 @@ workflow { // virtual_stain_output, and gate on virtual_stain_done. if (track_on) { track_done = track_run_wf(all_positions, assemble_output, assemble_output, - track_dir, params.track_config, + params.track_config, tk_init.resources, assemble_done.done) } @@ -418,7 +415,6 @@ workflow { .collect { zarr -> new File(zarr).parent } def assemble_dir = new File(assemble_output).parent cleanup_paths << "${assemble_dir}/slurm_output" << "${assemble_dir}/.iohub-progress" - if (track_on) cleanup_paths << "${track_dir}/slurm_output" def cleanup_list = cleanup_targets(cleanup_paths, out) cleanup_record = "${out}/nextflow/intermediates_cleaned.txt" diff --git a/nextflow/modules/qc.nf b/nextflow/modules/qc.nf index 65fc74d4..f927d9e7 100644 --- a/nextflow/modules/qc.nf +++ b/nextflow/modules/qc.nf @@ -275,7 +275,7 @@ workflow qc_compute_wf { // // `tabs` is a list of maps: [label: , zarr: , config: ]. // Labels must be unique within a spec; a step directory name gives that for free, -// and it is also what identifies the tab to a reader — "4-assemble", "5-track" — +// and it is also what identifies the tab to a reader — e.g. "4-assemble" — // rather than the stage number, which is imaging-qc's internal table namespace. def qc_report_spec(tabs, spec_path, title) { def spec = file(spec_path) diff --git a/nextflow/modules/tracking.nf b/nextflow/modules/tracking.nf index 8fb24d96..77a2090d 100644 --- a/nextflow/modules/tracking.nf +++ b/nextflow/modules/tracking.nf @@ -2,14 +2,16 @@ // run (fan-out × N positions). // // This module is PATH-AGNOSTIC. Callers pass the input zarr (plate structure), -// input images zarr (image data for tracking), work directory, and config -// explicitly. +// input images zarr (image data for tracking), and config explicitly. // // Tracking is a 2-input step: `input_zarr` supplies the plate structure and // `input_images_zarr` the pixels. mantis-v2.nf passes the assembled plate for // both. The outputs go INTO `input_zarr`'s positions — `labels/`, -// `tracks.geff` and a tracks CSV — so there is no tracking store; `work_dir` -// holds only the per-FOV Ultrack databases and SLURM placeholders. +// `tracks.geff` and a tracks CSV — and nothing else is kept: each worker builds +// its Ultrack database in the node's $TMPDIR and deletes it after export (no +// -o), so tracking has no step directory. Its real logs are in +// nextflow/slurm_output/track/, like every step's; submitit's debug-mode +// placeholders land in the task's work directory. // // INIT AND RUN ARE SEPARATE SUBWORKFLOWS, and hoisting init is worth more here // than anywhere else: the tracking config is the LAST one a run would otherwise @@ -21,7 +23,7 @@ // warms the shared cellpose weights cache, which is strictly better done before // any GPU worker exists to race for it. -include { parse_resources; slurm_logs; slurm_log_dir; slurm_output_readme } from './common' +include { parse_resources; slurm_logs; slurm_log_dir } from './common' process init_track { @@ -29,7 +31,6 @@ process init_track { input: val input_zarr - val work_dir val config path config_file // staged only for the task hash: see common.nf, #397 val trigger @@ -37,15 +38,11 @@ process init_track { output: stdout - // slurm_output_readme writes to /slurm_output, so - // "/." lands it in /slurm_output, beside track's placeholders. script: """ mkdir -p "${slurm_log_dir('track')}" - ${slurm_output_readme('track', "${work_dir}/.")} biahub track --init \ -i "${input_zarr}"/*/*/* \ - -o "${work_dir}" \ -c "${config}" """ } @@ -64,7 +61,6 @@ process run_track { tuple val(position), val(meta) val input_zarr val input_images_zarr - val work_dir val config path config_file // staged only for the task hash: see common.nf, #397 @@ -75,13 +71,8 @@ process run_track { """ biahub track --cluster debug \ -i "${input_zarr}/${position}" \ - -o "${work_dir}" \ -c "${config}" \ --input-images-path "${input_images_zarr}" - # Ultrack's SQLite database is created 0644, and under the project dirs' - # default ACL the group mask follows that mode — so the group can't write - # the FOV's data.db. Grant it explicitly; -f because only the owner can. - chmod -Rf g+w "${work_dir}/${position.replace('/', '_')}" || true """ } @@ -91,7 +82,6 @@ process run_track { // // take: // input_zarr path to the input plate.zarr (plate structure; receives the labels) -// work_dir tracking work directory (Ultrack databases, SLURM placeholders) // config path to the track settings YAML // trigger gating channel — init starts once this emits // emit: @@ -100,12 +90,11 @@ process run_track { workflow track_init_wf { take: input_zarr - work_dir config trigger main: - init_out = init_track(input_zarr, work_dir, config, file(config), trigger.collect().map { 'done' }) + init_out = init_track(input_zarr, config, file(config), trigger.collect().map { 'done' }) emit: resources = init_out.map { stdout_text -> parse_resources(stdout_text) } @@ -119,7 +108,6 @@ workflow track_init_wf { // positions collected channel of position keys // input_zarr plate structure store // input_images_zarr image data store -// work_dir tracking work directory // config path to the track settings YAML // resources RESOURCES payload from track_init_wf // prev_done gating channel — the input stores hold data @@ -128,7 +116,6 @@ workflow track_run_wf { positions input_zarr input_images_zarr - work_dir config resources prev_done @@ -147,7 +134,7 @@ workflow track_run_wf { .combine(prev_done.map { 'done' }) .map { pos, meta, _gate -> [pos, meta] } - tk_done = run_track(pos_meta, input_zarr, input_images_zarr, work_dir, config, file(config)) | collect + tk_done = run_track(pos_meta, input_zarr, input_images_zarr, config, file(config)) | collect emit: done = tk_done diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 9120d22e..406712e5 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -1,4 +1,5 @@ import os +import tempfile import numpy as np import pandas as pd @@ -112,6 +113,12 @@ def _assert_tracking_outputs(position_path, z_index): tracks = pd.read_csv(csv_path) assert set(tracks["track_id"]) <= set(np.unique(labels)) - {0} + # The full Ultrack config is provenance on the label image, not a file beside it. + label_attrs = yaml.safe_load((position_path / "labels" / LABEL / "zarr.json").read_text()) + ultrack_config = label_attrs["attributes"]["biahub-track"]["ultrack_config"] + assert "solution_gap" in ultrack_config["tracking"] + assert "working_dir" not in ultrack_config["data"] + def _make_tracking_config(plate_path, tmp_path): """Create a minimal tracking config pointing at the test plate.""" @@ -300,7 +307,6 @@ def test_track_cli_missing_input_path(tmp_path, example_track_settings, monkeypa def test_track_cli_init_only(tmp_path, example_tracking_plate): """--init creates empty full-Z label images in the -i positions and emits RESOURCES.""" plate_path, _ = example_tracking_plate - output_path = tmp_path / "track_output" config_path = _make_tracking_config(plate_path, tmp_path) runner = CliRunner() @@ -312,8 +318,6 @@ def test_track_cli_init_only(tmp_path, example_tracking_plate): str(plate_path / "A" / "1" / "0"), str(plate_path / "B" / "1" / "0"), str(plate_path / "B" / "2" / "0"), - "-o", - str(output_path), "-c", str(config_path), "--init", @@ -348,8 +352,12 @@ def test_track_cli_debug_single_position(tmp_path, example_tracking_plate, monke monkeypatch.setenv("ULTRACK_ARRAY_MODULE", "numpy") plate_path, _ = example_tracking_plate - output_path = tmp_path / "track_output" config_path = _make_tracking_config(plate_path, tmp_path) + # No -o: the Ultrack database goes to $TMPDIR and the submitit folder to ./ + temp_dir = tmp_path / "tmpdir" + temp_dir.mkdir() + monkeypatch.setattr(tempfile, "tempdir", str(temp_dir)) + monkeypatch.chdir(tmp_path) def run(): return CliRunner().invoke( @@ -358,8 +366,6 @@ def run(): "track", "-i", str(plate_path / "A" / "1" / "0"), - "-o", - str(output_path), "-c", str(config_path), "--cluster", @@ -371,6 +377,8 @@ def run(): assert result.exit_code == 0, result.output assert "Tracking complete:" in result.output _assert_tracking_outputs(plate_path / "A" / "1" / "0", z_index=1) + assert not any(temp_dir.iterdir()) # the temporary database was deleted + assert (tmp_path / "slurm_output").is_dir() # A retry overwrites the same outputs. result = run() From 662e53f19a52f59d002fe8d0b1c886afb8ef206a Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Thu, 8 Oct 2026 17:33:30 -0700 Subject: [PATCH 05/54] build(deps): pin iohub to iohub#486 now that #485 is on main Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 17 +++++++---------- uv.lock | 6 +++--- 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 700a851a..c004e302 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -124,15 +124,12 @@ imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipelin # inside the assembled OME-NGFF 0.5 plate (#361). Repoint at a release once merged. ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "3cd1b090307cfb13ac62d72df78ef4555ce0ffa3" } -# Pin iohub to dev/dca-labels: main (d1f9e89, which carries czbiohub-sf/iohub#476, -# #472 and #479 — see below) merged with two unmerged branches that track needs to -# write OME-NGFF label images into the assembled plate (#360): -# - fix/labels-group-metadata (czbiohub-sf/iohub#485): the labels list goes on -# /labels/zarr.json, not the position, so ome-zarr-models (and so -# dca-validate) can open a field that has labels; -# - feat/create-empty-plate-labels (czbiohub-sf/iohub#486): -# create_empty_plate(label_names=...) scaffolds empty label images. -# Repoint at main once both merge, then at a PyPI release. +# Pin iohub to feat/create-empty-plate-labels (czbiohub-sf/iohub#486), which is +# main plus create_empty_plate(label_names=...): track --init scaffolds the empty +# label images with it (#360). Main already carries czbiohub-sf/iohub#485 (the +# labels list on /labels/zarr.json, so ome-zarr-models and dca-validate can +# open a field that has labels). Repoint at main once #486 merges, then at a PyPI +# release. # # Main's own reasons: #476 syncs a write unit's shards before recording it as # done (without it a preemption between the write and the marker leaves a `.done` @@ -140,7 +137,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "3cd1b090307cfb13 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "82ef216b059351a95f16d802f2667573dd59e1db" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "c18e5ff4be0d963078b15f359f1734d0f6ae3213" } [tool.hatch.metadata] allow-direct-references = true diff --git a/uv.lock b/uv.lock index dc6713a4..c52fae52 100644 --- a/uv.lock +++ b/uv.lock @@ -375,7 +375,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=82ef216b059351a95f16d802f2667573dd59e1db" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c18e5ff4be0d963078b15f359f1734d0f6ae3213" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2018,8 +2018,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev25+82ef216" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=82ef216b059351a95f16d802f2667573dd59e1db#82ef216b059351a95f16d802f2667573dd59e1db" } +version = "0.3.12.dev24+c18e5ff" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=c18e5ff4be0d963078b15f359f1734d0f6ae3213#c18e5ff4be0d963078b15f359f1734d0f6ae3213" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 8d473e8fbe52c58ab43e5c812e84935dcfe06430 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 9 Oct 2026 13:16:45 -0700 Subject: [PATCH 06/54] build(deps): repin iohub#486 after per-position label defaults fix Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c004e302..cd618869 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -137,7 +137,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "3cd1b090307cfb13 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "c18e5ff4be0d963078b15f359f1734d0f6ae3213" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "41964963348a2785f7339122c295893698c50eee" } [tool.hatch.metadata] allow-direct-references = true diff --git a/uv.lock b/uv.lock index c52fae52..025dc7d0 100644 --- a/uv.lock +++ b/uv.lock @@ -375,7 +375,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c18e5ff4be0d963078b15f359f1734d0f6ae3213" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=41964963348a2785f7339122c295893698c50eee" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2018,8 +2018,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev24+c18e5ff" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=c18e5ff4be0d963078b15f359f1734d0f6ae3213#c18e5ff4be0d963078b15f359f1734d0f6ae3213" } +version = "0.3.12.dev25+4196496" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=41964963348a2785f7339122c295893698c50eee#41964963348a2785f7339122c295893698c50eee" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 0c353758f74040c50d77eb516e7ecaa8cba60fad Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 9 Oct 2026 14:55:57 -0700 Subject: [PATCH 07/54] build(deps): repin iohub#486 after review fixes Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index cd618869..a774f694 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -137,7 +137,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "3cd1b090307cfb13 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "41964963348a2785f7339122c295893698c50eee" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "7bf91f1e7a4d4cf95f9c22cf32eb894cb6d28dbd" } [tool.hatch.metadata] allow-direct-references = true diff --git a/uv.lock b/uv.lock index 025dc7d0..eb3abbdc 100644 --- a/uv.lock +++ b/uv.lock @@ -375,7 +375,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=41964963348a2785f7339122c295893698c50eee" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=7bf91f1e7a4d4cf95f9c22cf32eb894cb6d28dbd" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2018,8 +2018,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev25+4196496" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=41964963348a2785f7339122c295893698c50eee#41964963348a2785f7339122c295893698c50eee" } +version = "0.3.12.dev27+7bf91f1" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=7bf91f1e7a4d4cf95f9c22cf32eb894cb6d28dbd#7bf91f1e7a4d4cf95f9c22cf32eb894cb6d28dbd" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From cee19fbccbbe6f4d492fd55b26443a6432c54f36 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 16 Oct 2026 13:40:45 -0700 Subject: [PATCH 08/54] build(deps): pin ultrack to royerlab main now that #288 and #289 merged Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 15 +++++++++------ uv.lock | 27 ++++++++------------------- 2 files changed, 17 insertions(+), 25 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a774f694..408a6dd5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -117,12 +117,15 @@ index-strategy = "unsafe-best-match" # Its #235 removed the `upstream_qc_dir` config key, so configs must not carry it. imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipeline", rev = "f5ae656e0e37320ead74f7c1f95ab1ed304c1829" } -# Pin ultrack to feat/geff-export-tracks on the ieivanov fork (royerlab/ultrack#288): -# Tracker.to_geff gains track_id/lineage_id/seg_id node props, related_objects -# pointing at the label image, axes with units and scale (no z in 2D), -# include_masks/include_overlaps and zarr_format=3, so track can write tracks.geff -# inside the assembled OME-NGFF 0.5 plate (#361). Repoint at a release once merged. -ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "3cd1b090307cfb13ac62d72df78ef4555ce0ffa3" } +# Pin ultrack to royerlab/ultrack main at the #289 merge. #288 gives +# Tracker.to_geff track_id/lineage_id/seg_id node props, related_objects pointing +# at the label image, axes with units and scale (no z in 2D), +# include_masks/include_overlaps and zarr_format=3, so track can write +# tracks.geff inside the assembled OME-NGFF 0.5 plate (#361); #289 stops +# create_zarr from leaking its temporary stores. Neither is in a PyPI release +# yet (0.8.0 predates both, and main's __version__ still reads 0.7.3, which this +# git pin overrides). Repoint at a release once one contains them. +ultrack = { git = "https://github.com/royerlab/ultrack", rev = "afbe2d6064aafa6dc25210353bb84239aace5102" } # Pin iohub to feat/create-empty-plate-labels (czbiohub-sf/iohub#486), which is # main plus create_empty_plate(label_names=...): track --init scaffolds the empty diff --git a/uv.lock b/uv.lock index eb3abbdc..7e0528f7 100644 --- a/uv.lock +++ b/uv.lock @@ -401,8 +401,8 @@ requires-dist = [ { name = "torch" }, { name = "tqdm" }, { name = "twine", marker = "extra == 'build'" }, - { name = "ultrack", marker = "extra == 'all'", git = "https://github.com/ieivanov/ultrack?rev=3cd1b090307cfb13ac62d72df78ef4555ce0ffa3" }, - { name = "ultrack", marker = "extra == 'track'", git = "https://github.com/ieivanov/ultrack?rev=3cd1b090307cfb13ac62d72df78ef4555ce0ffa3" }, + { name = "ultrack", marker = "extra == 'all'", git = "https://github.com/royerlab/ultrack?rev=afbe2d6064aafa6dc25210353bb84239aace5102" }, + { name = "ultrack", marker = "extra == 'track'", git = "https://github.com/royerlab/ultrack?rev=afbe2d6064aafa6dc25210353bb84239aace5102" }, { name = "waveorder", specifier = ">=3.0.4" }, { name = "zensical", marker = "extra == 'doc'", specifier = ">=0.0.32" }, ] @@ -1894,25 +1894,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/be/ea/416657fb1cf0cb93c2564922a94f8bd361ae3d73a591c4b84e257a2305d8/imagecodecs-2026.8.16-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:53f4b7d70a5554f8d5f6f09746e5e9569ce4932b12e6997d805b81c377621fed", size = 12928116, upload-time = "2026-08-15T16:06:08.041Z" }, { url = "https://files.pythonhosted.org/packages/21/cf/398d49c8b616bc997b58749812dab5719e68165caa0275d5a1e3e310a1ba/imagecodecs-2026.8.16-cp312-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:908fa48740d52acbaea200d422e2267623122a1987fc85e898c77b2bb02d3e00", size = 28094370, upload-time = "2026-08-15T16:06:11.642Z" }, { url = "https://files.pythonhosted.org/packages/67/1b/0771f9b93a8f43b1e98a20d33fab02a6eceb8bb75c32764b5d47ee4db36f/imagecodecs-2026.8.16-cp312-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:c1068e0db9bc09cbd32510e543caed50f231096aecbb4c9cf7bcfbd2ffc53a50", size = 28080673, upload-time = "2026-08-15T16:06:15.463Z" }, - { url = "https://files.pythonhosted.org/packages/66/57/876bf8245b672b1fc901b65b9d1d0a37c60f7ed36e6ae8815c4af49e2779/imagecodecs-2026.8.16-cp312-abi3-pyemscripten_2025_0_wasm32.whl", hash = "sha256:59e16b27641bd08b5162de01a8bc382cd32a5e448cb5c2f4a1b252ea96d37fed", size = 2280993, upload-time = "2026-08-15T16:06:17.576Z" }, - { url = "https://files.pythonhosted.org/packages/ac/a2/b0ea5fee8c2c4f3cf68bb2243f640ea087885598aedb2f98bde0298f73ec/imagecodecs-2026.8.16-cp312-abi3-pyemscripten_2026_0_wasm32.whl", hash = "sha256:fcb28f8e552f26c195710bddbf7e8685e789ee9309725df4a275df16d7a136bc", size = 2273594, upload-time = "2026-08-15T16:06:19.082Z" }, { url = "https://files.pythonhosted.org/packages/cf/b7/cae1b739d368bc17df669f43793b5e6a7d1f56e19438eaa41add21611919/imagecodecs-2026.8.16-cp312-abi3-win32.whl", hash = "sha256:82bc27758f25d181eb4fd4fd8997b3eb07118dc673dee93aa3e64df8485e89bc", size = 16562229, upload-time = "2026-08-15T16:06:21.677Z" }, { url = "https://files.pythonhosted.org/packages/3e/da/d168a0e706103db3b82bde1b088d493e84a20da26f58be42bb2f1fca87e7/imagecodecs-2026.8.16-cp312-abi3-win_amd64.whl", hash = "sha256:6444a421272665ebd0116b54ad96347b86b4a5e21b59ed184487145cd3937f4b", size = 21271669, upload-time = "2026-08-15T16:06:25.021Z" }, { url = "https://files.pythonhosted.org/packages/99/0b/31642c8c404d436496b76c7cf6df220617f2e1ae36ae25e6809aa0a36a10/imagecodecs-2026.8.16-cp312-abi3-win_arm64.whl", hash = "sha256:b05e3ae5fa190c9722a0865cc06a3efccf27b025cbeeb8e9fd6beae0de5ce1d2", size = 16870662, upload-time = "2026-08-15T16:06:28.055Z" }, - { url = "https://files.pythonhosted.org/packages/52/f4/cacc8eca9b6349e75b3fd01148d4b5e3cccae5b4873643fd821d53764b11/imagecodecs-2026.8.16-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:4af18875b6f9b9bed35b9905e2e00ef8e303981846e112e835583f52094581fe", size = 15419161, upload-time = "2026-08-15T16:06:31.213Z" }, - { url = "https://files.pythonhosted.org/packages/1f/15/93aa21ba529f61a988c2358e34ecb7f0a0a57f6735839bdb4902b285a86a/imagecodecs-2026.8.16-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1849ac820ca3073dfa9fe799d267b77cdf0f5118be4841274acf4c43e0608d3f", size = 13393687, upload-time = "2026-08-15T16:06:34.095Z" }, - { url = "https://files.pythonhosted.org/packages/77/bf/2963c729c8126821f90dfcfc8a22cf50d66ae82f5e9228e1261919cab020/imagecodecs-2026.8.16-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:8778982a9bb2d286793799d3e67715cc74479ff483572523e26316db10650e59", size = 33032087, upload-time = "2026-08-15T16:06:38.302Z" }, - { url = "https://files.pythonhosted.org/packages/6d/e3/92588bf8bddcf642fa4370adb6ad2aad92bed92c6f3b5a02f3c1c00f0199/imagecodecs-2026.8.16-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:8afd5d0802ee1ebcbdb809fe1758f8ca2b1900eabb7d737453297c7f218b4f6f", size = 32438276, upload-time = "2026-08-15T16:06:42.633Z" }, - { url = "https://files.pythonhosted.org/packages/1d/60/3d29f474719688346ec95dd4a7b3bcd799e31f4c7d15617f9c6dafa811cb/imagecodecs-2026.8.16-cp314-cp314t-win32.whl", hash = "sha256:6265622835428c21dc891261993f59c24f1e4255502c90aa36778996e91315f9", size = 17304429, upload-time = "2026-08-15T16:06:46.007Z" }, - { url = "https://files.pythonhosted.org/packages/74/95/90e45470ed4f0d8ca3e90174db4352bd40961315fc3ecffc15f0aed80683/imagecodecs-2026.8.16-cp314-cp314t-win_amd64.whl", hash = "sha256:a04b7f31fe4ec6416d2036b28b1df2f444ae1aa5163394059565e14677c9fc48", size = 22106210, upload-time = "2026-08-15T16:06:49.59Z" }, - { url = "https://files.pythonhosted.org/packages/78/1c/9ba4d8ede7829cabf7dd56ba806fe648e90dfd5a07d83f8afc28be034d79/imagecodecs-2026.8.16-cp314-cp314t-win_arm64.whl", hash = "sha256:fa7b9fd436b4e1a066fc71d42f506e028d2256f96950c03b1798e1fb257605b3", size = 17666610, upload-time = "2026-08-15T16:06:52.686Z" }, - { url = "https://files.pythonhosted.org/packages/e3/7d/183d18f5b3c28d9aefa2130f12d633a57636865453f1ce7433041e00d31a/imagecodecs-2026.8.16-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:f2b2f7f7a7c4918b3c1bc5a8dba7fd369c7456449fe7f61e522a72eec6905698", size = 15406581, upload-time = "2026-08-15T16:06:55.579Z" }, - { url = "https://files.pythonhosted.org/packages/62/41/f9cd37481294ba0eeecfb0063253406e79e1b2cb398d6198df26538c5e36/imagecodecs-2026.8.16-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:31537513f46db6822b5ccff57d02ea131e9f7786d52e2e0b69ce38792350a3fa", size = 13383849, upload-time = "2026-08-15T16:06:58.617Z" }, - { url = "https://files.pythonhosted.org/packages/37/dc/31ec12056945f18a6ddc1410325fa29b834eb52eb7bd3b63008db1bb1c66/imagecodecs-2026.8.16-cp315-cp315t-manylinux_2_28_aarch64.whl", hash = "sha256:207747fe4b1c65b46eedef0ede388cbc527d03e5de1dfc73873e305d8ecd6f06", size = 32775737, upload-time = "2026-08-15T16:07:02.631Z" }, - { url = "https://files.pythonhosted.org/packages/3e/ac/4ddd09edc70dccb68ba0ed0074b9389b0cb05862cfb9b351adc4f0adf1ae/imagecodecs-2026.8.16-cp315-cp315t-manylinux_2_28_x86_64.whl", hash = "sha256:70959750594ab4fb16cc966e344979109efac0d3d3897bc8bac20cf560a825bd", size = 32271785, upload-time = "2026-08-15T16:07:07.154Z" }, - { url = "https://files.pythonhosted.org/packages/15/28/aef2f6919cfface4354d21f5b2ba923ee6d5ccdbc673bfa68661902dea4a/imagecodecs-2026.8.16-cp315-cp315t-win32.whl", hash = "sha256:701645d424f8fdb9007da269888797b16c50b774f1720b68302466771c339a7d", size = 17304445, upload-time = "2026-08-15T16:07:10.479Z" }, - { url = "https://files.pythonhosted.org/packages/9b/a8/8716691a21dc3249ec5c8945a21010cdc0db68aaf21e21d23ad6700acbd6/imagecodecs-2026.8.16-cp315-cp315t-win_amd64.whl", hash = "sha256:8065dbb611562bb61b666a453d7c9b35533dd28b5afbeceb79dd69bb91162855", size = 22100436, upload-time = "2026-08-15T16:07:14.02Z" }, - { url = "https://files.pythonhosted.org/packages/56/a1/11a85d6243ef49da4cf141f0d52eb0a85028b0b0a1c82dda2ef8f1da3b57/imagecodecs-2026.8.16-cp315-cp315t-win_arm64.whl", hash = "sha256:64626fb7fc7a722136990dde9a5c3455dfcf579d5907db062eabb5c7d8c0e7cf", size = 17664176, upload-time = "2026-08-15T16:07:17.111Z" }, ] [[package]] @@ -5075,6 +5059,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1f/f9/023daa2a633cde62b2f7452e62d53273a8b23c3ce5dc31d7ecdc6b797897/pyqt6_sip-13.13.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:70199e2b37ac83a4799df2a2a2ec8358e1d20dfdf011f2baec17a8867c34fe12", size = 323559, upload-time = "2026-10-04T12:26:41.919Z" }, { url = "https://files.pythonhosted.org/packages/08/0b/1c8ebeb9184c9f017173ae952e89f5191571c680f3e18875217a1c00860b/pyqt6_sip-13.13.0-cp314-cp314-win_amd64.whl", hash = "sha256:26ee16a9b53aa7c8271a19154c35ab2e17514dbae85ac9a770ba40feac347be2", size = 55130, upload-time = "2026-10-04T12:26:44.552Z" }, { url = "https://files.pythonhosted.org/packages/35/b9/9bdfa5222e42aba61205847721d1b4c53145ffcb740952272a4fbe33289f/pyqt6_sip-13.13.0-cp314-cp314-win_arm64.whl", hash = "sha256:88e3bbc1a6bb1632ead7b15f77769f09b9a581eb77eb41a6862c32396f47b4a4", size = 49706, upload-time = "2026-10-04T12:26:45.471Z" }, + { url = "https://files.pythonhosted.org/packages/91/8c/a68a45e2a651795c9b5bcd95eec4ef9cb945c35af238f7ef0369a39a7e72/pyqt6_sip-13.13.0-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:e3b47eaab721871177148f4771a8b2259de5356e30cbcf04e88acfe55b57f257", size = 113961, upload-time = "2026-10-10T14:50:58.283Z" }, + { url = "https://files.pythonhosted.org/packages/f8/20/eebfef3e76175553f5049f2dd54d3e2e940b116ddb61d1ecdec12b93eb3a/pyqt6_sip-13.13.0-cp315-cp315-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f856e077035bd0316e26396fa70100b6ff13405d017464bbb501ae73070f3dce", size = 304973, upload-time = "2026-10-10T14:51:01.158Z" }, + { url = "https://files.pythonhosted.org/packages/09/44/a944b91a94c9745cfce05174d4a2dab4faa67b03b23c160948505993b1b0/pyqt6_sip-13.13.0-cp315-cp315-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:78fd781938d09797db6d0b23ac070a91d567c9170581dba3ca87630abb69055b", size = 327367, upload-time = "2026-10-10T14:50:59.831Z" }, + { url = "https://files.pythonhosted.org/packages/2e/b3/939df28667fd61f4d04abb7adab93f034fd4f62a7d3e0dac82ea58268532/pyqt6_sip-13.13.0-cp315-cp315-win_amd64.whl", hash = "sha256:67b37ae7d238462a89908b1c3e3a4bc9551fc9b8d3e8be05ff3314b7a24a3e53", size = 55311, upload-time = "2026-10-10T14:51:02.736Z" }, + { url = "https://files.pythonhosted.org/packages/89/48/a640f06a065655cc0286fc16a7fed8b09bfe3a8291e4223d483e834d0c05/pyqt6_sip-13.13.0-cp315-cp315-win_arm64.whl", hash = "sha256:a1cd42fdbb191d3ee0a0c8858348e18ddb5ea4bbc009af34da4a536db1f92e63", size = 49744, upload-time = "2026-10-10T14:51:04.017Z" }, ] [[package]] @@ -6683,7 +6672,7 @@ wheels = [ [[package]] name = "ultrack" version = "0.7.3" -source = { git = "https://github.com/ieivanov/ultrack?rev=3cd1b090307cfb13ac62d72df78ef4555ce0ffa3#3cd1b090307cfb13ac62d72df78ef4555ce0ffa3" } +source = { git = "https://github.com/royerlab/ultrack?rev=afbe2d6064aafa6dc25210353bb84239aace5102#afbe2d6064aafa6dc25210353bb84239aace5102" } dependencies = [ { name = "blosc2" }, { name = "click" }, From 324ff4da09535030fab3959be8815c21653b28b4 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 16 Oct 2026 13:52:16 -0700 Subject: [PATCH 09/54] refactor(track): drop the tempfile redirect now that ultrack#289 merged create_zarr now deletes its temporary store with the array, so the per-FOV temporary directory only needs to hold the Ultrack database again. Co-Authored-By: Claude Opus 5.5 (1M context) --- biahub/track.py | 41 +++++++++++++------------------- tests/test_cli/test_track_cli.py | 4 +++- 2 files changed, 19 insertions(+), 26 deletions(-) diff --git a/biahub/track.py b/biahub/track.py index 4116f960..97d67885 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -881,32 +881,23 @@ def write_tracking_labels( @contextmanager -def _fov_scratch(output_dirpath: Path | None, fov: str): - """Yield the Ultrack database directory for one FOV, inside a scratch directory. - - The scratch directory is fresh in this worker's ``$TMPDIR``, one per FOV because - ``$TMPDIR`` can be shared by all of a user's jobs on a node, and deleted on exit. - The database goes in it unless ``output_dirpath`` is given, in which case it is - kept in ``{output_dirpath}/{fov}``. - - While the context is open, ``tempfile`` also points into the scratch directory: - ultrack's ``create_zarr`` backs its default store with a ``TemporaryDirectory`` - that it lets go out of scope, so the directory is deleted, recreated by the first - chunk write, and never cleaned up. Kept in here, those arrays go with the rest. +def _database_dirpath(output_dirpath: Path | None, fov: str): + """Yield the Ultrack working directory (database and config) for one FOV. + + ``{output_dirpath}/{fov}`` when ``output_dirpath`` is given, and kept. Otherwise a + fresh directory in this worker's ``$TMPDIR``, deleted on exit: one per FOV, + because ``$TMPDIR`` can be shared by all of a user's jobs on a node. """ - scratch = Path(tempfile.mkdtemp(prefix=f"biahub-track-{fov}-")) - previous_tempdir = tempfile.tempdir - tempfile.tempdir = str(scratch) + if output_dirpath is not None: + dirpath = Path(output_dirpath) / fov + dirpath.mkdir(parents=True, exist_ok=True) + yield dirpath + return + dirpath = Path(tempfile.mkdtemp(prefix=f"biahub-track-{fov}-")) try: - if output_dirpath is not None: - database_path = Path(output_dirpath) / fov - else: - database_path = scratch / "database" - database_path.mkdir(parents=True, exist_ok=True) - yield database_path + yield dirpath finally: - tempfile.tempdir = previous_tempdir - shutil.rmtree(scratch, ignore_errors=True) + shutil.rmtree(dirpath, ignore_errors=True) def _resolved_ultrack_config(tracking_config) -> dict: @@ -953,7 +944,7 @@ def track_one_position( Where to keep the Ultrack database and config for this FOV, in ``{output_dirpath}/{row}_{col}_{fov}/``. If None they live in a scratch directory in this worker's ``$TMPDIR`` and are deleted after export (see - :func:`_fov_scratch`). + :func:`_database_dirpath`). tracking_config : MainConfig Ultrack configuration containing segmentation, linking, and optimization settings. label_name : str @@ -1000,7 +991,7 @@ def track_one_position( time_scale = position.scale[0] units = {axis.name.lower(): getattr(axis, "unit", None) for axis in position.axes} - with _fov_scratch(output_dirpath, fov) as database_path: + with _database_dirpath(output_dirpath, fov) as database_path: # Load, focus-slice (method='focus'), run the processing pipeline, and fill blank # frames once; both segmentation methods consume the resulting data_dict. # pixel_size is the transverse pixel size = last element of the (trimmed) scale. diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 406712e5..ac0be049 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -377,7 +377,9 @@ def run(): assert result.exit_code == 0, result.output assert "Tracking complete:" in result.output _assert_tracking_outputs(plate_path / "A" / "1" / "0", z_index=1) - assert not any(temp_dir.iterdir()) # the temporary database was deleted + # The temporary database was deleted. (ultrack's own temporary arrays go when + # they are garbage collected, at the latest when the worker process exits.) + assert not list(temp_dir.glob("biahub-track-*")) assert (tmp_path / "slurm_output").is_dir() # A retry overwrites the same outputs. From cbeedd8ca947e1dc1ec8241a6a75a77f7f5e711a Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 16 Oct 2026 14:46:17 -0700 Subject: [PATCH 10/54] build(deps): pin iohub to main now that #486 merged Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 13 ++++++------- uv.lock | 6 +++--- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 408a6dd5..4a889ca3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -127,12 +127,11 @@ imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipelin # git pin overrides). Repoint at a release once one contains them. ultrack = { git = "https://github.com/royerlab/ultrack", rev = "afbe2d6064aafa6dc25210353bb84239aace5102" } -# Pin iohub to feat/create-empty-plate-labels (czbiohub-sf/iohub#486), which is -# main plus create_empty_plate(label_names=...): track --init scaffolds the empty -# label images with it (#360). Main already carries czbiohub-sf/iohub#485 (the -# labels list on /labels/zarr.json, so ome-zarr-models and dca-validate can -# open a field that has labels). Repoint at main once #486 merges, then at a PyPI -# release. +# Pin iohub to main at the czbiohub-sf/iohub#486 merge, for +# create_empty_plate(label_names=...), with which track --init scaffolds the empty +# label images (#360), and czbiohub-sf/iohub#485, which puts the labels list on +# /labels/zarr.json so ome-zarr-models and dca-validate can open a field +# that has labels. Repoint at a PyPI release once one contains them. # # Main's own reasons: #476 syncs a write unit's shards before recording it as # done (without it a preemption between the write and the marker leaves a `.done` @@ -140,7 +139,7 @@ ultrack = { git = "https://github.com/royerlab/ultrack", rev = "afbe2d6064aafa6d # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "7bf91f1e7a4d4cf95f9c22cf32eb894cb6d28dbd" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "c9832a7675e0fb997a27db2796623d19b7dc9f1a" } [tool.hatch.metadata] allow-direct-references = true diff --git a/uv.lock b/uv.lock index 7e0528f7..c02d02c0 100644 --- a/uv.lock +++ b/uv.lock @@ -375,7 +375,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=7bf91f1e7a4d4cf95f9c22cf32eb894cb6d28dbd" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c9832a7675e0fb997a27db2796623d19b7dc9f1a" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2002,8 +2002,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev27+7bf91f1" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=7bf91f1e7a4d4cf95f9c22cf32eb894cb6d28dbd#7bf91f1e7a4d4cf95f9c22cf32eb894cb6d28dbd" } +version = "0.3.12.dev18+c9832a7" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=c9832a7675e0fb997a27db2796623d19b7dc9f1a#c9832a7675e0fb997a27db2796623d19b7dc9f1a" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From a98a267edbe1ec0e85f23933fddbb77a5d7a062b Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Sat, 17 Oct 2026 16:24:50 -0700 Subject: [PATCH 11/54] feat(project)!: biahub project CLI; segment and track no longer project Add `biahub project`, which turns a 3D plate into a 2D (Z=1) plate with one projection per output channel: max, mean, sum over a Z range or a window around a plane, or one plane (`slice`). The plane may be fixed or the in-focus plane found with the registration engine's focus criterion (FocusSettings subclass: centre crop, NA, wavelength), once per position (median over T) or per timepoint. Ranges are half-open like concatenate's Z_slice and relative to `center` when one is set, so the A549 60-plane window is `center: focus, range: [-20, 40]`. Unlisted input channels are dropped with a warning. The focus plane of every timepoint is recorded in the output position's `biahub-project` metadata. Segment and track stop selecting and projecting Z-planes. Segment loses `z_slice_2D`: a model without do_3D segments in 2D and needs a Z=1 input. Track loses `z_slicing` and the focus window; `output_mode` is checked against the input Z. Both raise, pointing at `biahub project`, when a 3D input meets a 2D request. Old keys are rejected with a migration message. Nextflow: a project step between assemble and track (`5-project`, `6-track` for A549); track reads the projected plate when project runs. The A549 project config reproduces the previous tracking input (mean over the 60-plane focus window on Phase3D), measured on the 800x800 centre crop. Co-Authored-By: Claude Fable 5.1 --- .../skills/reconstruct-with-nextflow/SKILL.md | 35 +- .../references/caveats.md | 39 +- .../references/datasets.md | 5 +- .../templates/run_mantis_v2.sh | 19 +- README.md | 1 + biahub/cli/main.py | 5 + biahub/project.py | 464 ++++++++++++++++++ biahub/segment.py | 56 ++- biahub/settings.py | 227 ++++++--- biahub/track.py | 258 ++-------- nextflow/configs/a549/project.yml | 22 + nextflow/configs/a549/segment.yml | 10 +- nextflow/configs/a549/track.yml | 43 +- nextflow/mantis-v2.nf | 100 ++-- nextflow/modules/project.nf | 137 ++++++ nextflow/modules/segment.nf | 10 +- settings/example_project_settings.yml | 60 +++ settings/example_segmentation_settings.yml | 8 +- settings/example_track_settings.yml | 44 +- tests/test_cellpose_inputs.py | 56 ++- tests/test_cli/test_cli.py | 1 + tests/test_cli/test_project_cli.py | 242 +++++++++ tests/test_cli/test_segment_cli.py | 88 ++-- tests/test_cli/test_track_cli.py | 216 +++----- tests/test_example_settings.py | 23 +- tests/test_project_settings.py | 117 +++++ tests/test_segment_settings.py | 51 +- 27 files changed, 1670 insertions(+), 667 deletions(-) create mode 100644 biahub/project.py create mode 100644 nextflow/configs/a549/project.yml create mode 100644 nextflow/modules/project.nf create mode 100644 settings/example_project_settings.yml create mode 100644 tests/test_cli/test_project_cli.py create mode 100644 tests/test_project_settings.py diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index 1c3ce329..47706440 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -3,7 +3,7 @@ name: reconstruct-with-nextflow description: >- Run a mantis-v2 microscope dataset through the biahub Nextflow reconstruction pipeline (nextflow/mantis-v2.nf: flat-field → deskew → reconstruct → - virtual-stain → assemble → track → QC) on the Bruno HPC cluster. Assemble, + virtual-stain → assemble → project → track → QC) on the Bruno HPC cluster. Assemble, track and QC are optional and selected per dataset family: A549 runs assemble + track + QC, neuromast runs assemble + QC and no tracking. Locates the raw acquisition under /hpc/instruments/cm.mantis, picks the output project @@ -323,13 +323,16 @@ Do not run anything yet. Show the user: what was mismatched and how it was resolved. 6. **The steps this run will perform, listed in order.** Reconstruction proper — flat-field → deskew → reconstruct → virtual-stain — always runs; assemble, - track and QC run only if their config is passed (biahub#306). Defaults by - family, which is what to present unless the user says otherwise: + project, track and QC run only if their config is passed (biahub#306). + Project turns the assembled 3D plate into the 2D (Z=1) plate that 2D + tracking reads (`biahub track` no longer selects or projects Z-planes), so + `--track_config` goes with `--project_config`. Defaults by family, which is + what to present unless the user says otherwise: | family | steps | |---|---| - | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → track → QC** | - | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → QC** (no tracking) | + | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → project → track → QC** | + | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → QC** (no projection or tracking) | Write the list out in the plan rather than naming the family, and say which optional steps are being skipped and why. If the user asks to skip anything @@ -342,9 +345,9 @@ Do not run anything yet. Show the user: **State the directory numbers this run will produce.** The number is the step's position among the steps performed, not a fixed label, so a neuromast run writes `4-assemble` as its last directory and an A549 run writes - `4-assemble` then `5-track`. Older A549 runs on disk say `5-assemble` / - `4-track`, from when the numbers were fixed — say so if the user is comparing - against one. + `4-assemble`, `5-project` then `6-track`. Older A549 runs on disk say + `5-track` (before the project step) or `5-assemble` / `4-track` (from when + the numbers were fixed) — say so if the user is comparing against one. 7. Known caveats that apply to this dataset. 8. Rough wall-time and that `-resume` is on. 9. Whether Slack notifications are on, and who will be @-mentioned at run end. @@ -396,13 +399,15 @@ copied file's group bits follow the *source* file's mode — `cp` from the repo (644) yields `rw-r--r--`. Hence the `chmod` after every copy. Plain `cp` only: `cp -p`/`-a` would also carry over the checkout's group. -The run script passes all four optional configs — `--concatenate_config`, -`--track_config`, `--qc_config` (`qc.yaml`), `--qc_track_config` -(`qc_track.yaml`) — so an A549 run needs no edit. **For a neuromast/zebrafish run, or any step the user asked to skip, -DELETE that flag's line** from the `nextflow run` call and note the skip in a -comment above it, so the script still records what this run did. A neuromast run -deletes `--track_config` and `--qc_track_config`; neither `track.yml` nor -`qc_track.yaml` need exist for that family, and neither does for zebrafish. +The run script passes all five optional configs — `--concatenate_config`, +`--project_config`, `--track_config`, `--qc_config` (`qc.yaml`), +`--qc_track_config` (`qc_track.yaml`) — so an A549 run needs no edit. **For a +neuromast/zebrafish run, or any step the user asked to skip, DELETE that flag's +line** from the `nextflow run` call and note the skip in a comment above it, so +the script still records what this run did. A neuromast run deletes +`--project_config`, `--track_config` and `--qc_track_config`; none of +`project.yml`, `track.yml` or `qc_track.yaml` need exist for that family, and +none does for zebrafish. Delete rather than comment: a `#` inside a backslash-continued command does not start a comment line — the continuation swallows it, every flag below is dropped diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 81dbd3f0..101572eb 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -135,27 +135,32 @@ per-position — leave it out. `target_channel` carries the `_prediction` suffix deliberately (§2). **`track.yml` — schema.** `TrackingSettings` forbids extra keys. Current -fields: `output_mode`, `z_slicing`, `target_channel`, `fov`, `input_images`, -`segmentation_method`, `cellpose_config`, `tracking_config`. Old configs under +fields: `output_mode`, `target_channel`, `fov`, `input_images`, +`segmentation_method`, `min_sibling_length`, `cellpose_config`, +`tracking_config`. Track no longer selects or projects Z-planes: it reads the +**projected** (Z=1) plate that `biahub project` writes from `project.yml`, and a +3D input with `output_mode: 2D` is an error. Old configs under `/hpc/projects/tlg2_mantis/` still carry rejected fields — migrate: ``` -mode -> output_mode -z_range + focus_config.z_window -> z_slicing.{method, window_size, focus_channel} +mode -> output_mode +z_slicing / z_range / focus_config -> project.yml: channels..z {center: focus, range: [below, above)}, focus.channel +input_images np.mean step -> project.yml: channels..method: mean ``` -(`focus_config`'s `NA_det`/`lambda_ill`/`pixel_size` have no equivalent.) -`target_channel` and `cellpose_config.input_channel` must match **pre-rename** -names in the **assembled** plate (§2, §5). `qc_track.yaml`'s `channels:` must -then match `_labels`, which is what the tracking store actually -holds — SKILL.md §5a checks both couplings before launch. +(`focus_config`'s `NA_det`/`lambda_ill` are `project.yml`'s `focus.na_det` / +`focus.lambda_ill`; the pixel size comes from the store.) `target_channel` and +`cellpose_config.input_channel` must match channels `project.yml` writes, which +carry the **pre-rename** names of the **assembled** plate (§2, §5). +`qc_track.yaml`'s `channels:` must then match `_labels`, which +is what the tracking store actually holds — SKILL.md §5a checks both couplings +before launch. -**`z_slicing.focus_channel` resolves against `input_images`, not the store.** -`apply_focus_slicing` raises if the focus channel is not among the loaded -channels, even when it exists in the plate — a template declaring -`focus_channel: Phase3D` while loading only `nuclei_prediction` failed all -positions after every other step had succeeded. Both shipped templates use -`method: all` (falls back to the first loaded channel), the safe default. +**`project.yml` — `focus.channel` must be in the assembled plate.** `biahub +project --init` resolves every channel name against the plate and fails there, +in the init phase, rather than after every other step. The A549 template finds +the focus on `Phase3D` and writes only `nuclei_prediction`; the focus channel +need not be an output. **`concatenate.yml` holds parameters only — no source paths.** The assemble subworkflow passes the three source stores on the command line, one `-i` per @@ -172,8 +177,8 @@ assembled store, and track's parameters (cellpose `diameter`, `min_area`/ `max_area`, linking `max_distance`) are tuned for A549 cells and do not transfer. So simply do not run it: -- Omit `--track_config` — delete that line, and `--qc_track_config` with it, - from the `nextflow run` call in the run script. A step runs only if its config +- Omit `--track_config` — delete that line, and `--project_config` and + `--qc_track_config` with it, from the `nextflow run` call in the run script. A step runs only if its config is passed ([biahub#306](https://github.com/czbiohub-sf/biahub/issues/306), fixed). Delete rather than comment: `#` inside a backslash-continued command drops every flag below it, `-resume` included. diff --git a/.claude/skills/reconstruct-with-nextflow/references/datasets.md b/.claude/skills/reconstruct-with-nextflow/references/datasets.md index 6ea22e0e..89c302df 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/datasets.md +++ b/.claude/skills/reconstruct-with-nextflow/references/datasets.md @@ -48,7 +48,7 @@ Which root depends on the dataset family: | family | cues in the name | output root | deliverable | |---|---|---|---| | zebrafish / neuromast / dynatrack | `dynatrack`, `hpf`, `dpf`, `neuromast`, `cldnb`, `she`, `myo6b`, `zebrafish` | `/hpc/projects/tlg2_mantis/` | assemble (`4-assemble`) — **no tracking**, see `caveats.md` §4 | -| cell line / organelle / infection | `A549`, `HEK`, `iPSC`, organelle genes (`SEC61B`, `TOMM20`, `G3BP1`, `MAP4`, `CAAX`, `H2B`), viruses (`DENV`, `ZIKV`, `HSV1`) | `/hpc/projects/intracellular_dashboard/organelle_dynamics/` | assemble + track (`4-assemble`, `5-track`) | +| cell line / organelle / infection | `A549`, `HEK`, `iPSC`, organelle genes (`SEC61B`, `TOMM20`, `G3BP1`, `MAP4`, `CAAX`, `H2B`), viruses (`DENV`, `ZIKV`, `HSV1`) | `/hpc/projects/intracellular_dashboard/organelle_dynamics/` | assemble + project + track (`4-assemble`, `5-project`, `6-track`) | | instrument QC / calibration | `argolight`, `beads`, `psf`, `alignment`, `illumination`, `fluorescein`, `first_light`, `mantis_v2_*` | *no default* — ask. These are usually not reconstructed. | — | Neighbouring roots that exist and are **not** the default target — do not use @@ -71,7 +71,8 @@ layout is defined once in `nextflow/mantis-v2.nf` (`directory_layout()`): ├── 2-reconstruct/.zarr ├── 3-virtual-stain/.zarr ├── 4-assemble/.zarr # the deliverable -├── 5-track/.zarr # A549 only — not run for neuromast +├── 5-project/.zarr # A549 only — the 2D (Z=1) plate track reads +├── 6-track/.zarr # A549 only — not run for neuromast ├── qc/ │ ├── report/index.html # one report, one tab per QC'd store │ └── report_spec.yaml # written at launch, before any QC task diff --git a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh index 450b384a..d351e050 100755 --- a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh +++ b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh @@ -1,16 +1,16 @@ #!/usr/bin/env bash # Run the mantis-v2 reconstruction pipeline (biahub nextflow/mantis-v2.nf): -# flat-field -> deskew -> reconstruct -> virtual-stain -> assemble -> track -> QC +# flat-field -> deskew -> reconstruct -> virtual-stain -> assemble -> project -> track -> QC # -# The last three are OPTIONAL: a step runs only if its config is passed. A549 -# runs assemble + track + QC; neuromast runs assemble + QC and no tracking. See -# the note above the `nextflow run` call at the bottom. +# The last four are OPTIONAL: a step runs only if its config is passed. A549 +# runs assemble + project + track + QC; neuromast runs assemble + QC and no +# projection or tracking. See the note above the `nextflow run` call at the bottom. # # Copy this into the project output directory and fill in the five variables # below. Keep it there: it is the run's provenance record of the exact command. # # --output is the project root, so each step writes a sibling directory: -# 0-flatfield/ 1-deskew/ 2-reconstruct/ 3-virtual-stain/ 4-assemble/ 5-track/ +# 0-flatfield/ 1-deskew/ 2-reconstruct/ 3-virtual-stain/ 4-assemble/ 5-project/ 6-track/ # The number is the step's POSITION among the steps this run performs, so a run # without tracking ends at 4-assemble and one that also skipped a step earlier # shifts everything after it down. @@ -190,12 +190,12 @@ fi # Harmless when unset already, and nothing in the pipeline reads it. unset CLAUDECODE -# The four config flags from --concatenate_config down are the OPTIONAL steps: +# The five config flags from --concatenate_config down are the OPTIONAL steps: # a step runs only if its config is passed (biahub#306). Defaults by family — # -# A549 / cell-line assemble + track + QC (all four, as written) -# zebrafish / neuromast assemble + QC (delete --track_config and -# --qc_track_config) +# A549 / cell-line assemble + project + track + QC (all five, as written) +# zebrafish / neuromast assemble + QC (delete --project_config, +# --track_config and --qc_track_config) # # TO SKIP A STEP, DELETE ITS LINE. Do not comment it out: `#` inside a # backslash-continued command does not start a comment line — the continuation @@ -213,6 +213,7 @@ nextflow run "${PIPELINE}" \ --reconstruct_config "${CONFIGS}/reconstruct.yml" \ --virtual_stain_config "${CONFIGS}/virtual_stain.yml" \ --concatenate_config "${CONFIGS}/concatenate.yml" \ + --project_config "${CONFIGS}/project.yml" \ --track_config "${CONFIGS}/track.yml" \ --qc_config "${CONFIGS}/qc.yaml" \ --qc_track_config "${CONFIGS}/qc_track.yaml" \ diff --git a/README.md b/README.md index a42bc6f9..b2df1d20 100644 --- a/README.md +++ b/README.md @@ -120,6 +120,7 @@ Some steps include an `estimate-*` command that writes a YAML of parameters you | `stitch` | Stitch positions within wells | | `concatenate` | Concatenate datasets channel-wise, with optional cropping | | | | +| `project` | Project a 3D plate to 2D: max, mean, sum or one plane per channel, optionally around the in-focus plane | | `segment` | Segment positions with a pretrained model or pipeline | | `virtual-stain` | Run VisCy/cytoland virtual staining | | `track` | Track objects in 2D/3D time-lapse data | diff --git a/biahub/cli/main.py b/biahub/cli/main.py index 48e0d9e6..d8d14104 100644 --- a/biahub/cli/main.py +++ b/biahub/cli/main.py @@ -155,6 +155,11 @@ def format_options(self, ctx, formatter): "import_path": "biahub.segment.segment_cli", "help": "Segment a position using pretrained model or pipeline", }, + { + "name": "project", + "import_path": "biahub.project.project_cli", + "help": "Project a 3D plate to 2D (max, mean, sum or one plane per channel)", + }, { "name": "virtual-stain", "import_path": "biahub.virtual_stain.virtual_stain_cli", diff --git a/biahub/project.py b/biahub/project.py new file mode 100644 index 00000000..ad9da6f0 --- /dev/null +++ b/biahub/project.py @@ -0,0 +1,464 @@ +"""Project a 3D plate to a 2D (Z=1) plate, one projection method per output channel. + +Each output channel is an input channel reduced over a selection of Z-planes (``max``, +``mean``, ``sum``) or one plane of it (``slice``). The selection is a fixed range, a +window around a plane, or a window around the in-focus plane, which is found per position +with the registration engine's focus criterion (see ``ProjectionSettings``). Segment and +track take the projected plate as their 2D input; neither projects by itself. +""" + +from pathlib import Path + +import click +import numpy as np +import submitit + +from iohub.ngff import open_ome_zarr +from iohub.ngff.utils import create_empty_plate, process_single_position + +from biahub.cli.monitor import monitor_jobs +from biahub.cli.parsing import ( + cluster, + config_filepath, + init_only, + input_position_dirpaths, + monitor, + output_dirpath, + resume, + sbatch_filepath, + sbatch_to_submitit, +) +from biahub.settings import ChannelProjection, ProjectionFocusSettings, ProjectionSettings +from biahub.utils.cluster import ( + apply_sbatch_overrides, + echo_resources, + estimate_resources, + get_submitit_cluster, + local_executor_parameters, +) +from biahub.utils.config import settings_fingerprint, yaml_to_model +from biahub.utils.ngff import ( + PROVENANCE_METADATA_KEYS, + get_output_paths, + resolve_ome_zarr_version, +) + +PROVENANCE_KEY = "biahub-project" + +# Methods whose result is not one of the input values, so the output is float32. +_FLOAT_METHODS = ("mean", "sum") + + +def resolve_z_indices( + projection: ChannelProjection, z_size: int, focus_z: int | None = None +) -> np.ndarray: + """Return the plane indices one channel's projection reads from a ``z_size``-plane stack. + + An unset ``z`` is the whole stack. A ``range`` without ``center`` is absolute; with a + ``center`` (``focus_z`` when it is ``"focus"``) the bounds are offsets from it, and + ``slice`` is the center alone. ``step`` thins the range. Indices outside the stack + are dropped; a selection that keeps none is an error. + """ + z = projection.z + if z is None: + return np.arange(z_size) + if z.center is None: + start, stop = z.range + else: + center = focus_z if z.center == "focus" else z.center + if center is None: + raise ValueError("z.center is 'focus' but no focus plane was given.") + if z.range is None: + start, stop = center, center + 1 + else: + start, stop = center + z.range[0], center + z.range[1] + indices = np.arange(start, stop, z.step) + indices = indices[(indices >= 0) & (indices < z_size)] + if indices.size == 0: + raise ValueError( + f"z planes {start}..{stop} (step {z.step}) select nothing from a {z_size}-plane " + "stack." + ) + return indices + + +def project_zyx(zyx: np.ndarray, method: str, z_indices: np.ndarray) -> np.ndarray: + """Reduce a (Z, Y, X) volume to (Y, X) over ``z_indices`` with ``method``.""" + if method == "slice": + return zyx[int(z_indices[0])] + planes = zyx[z_indices] + if method == "max": + return planes.max(axis=0) + if method == "mean": + return planes.mean(axis=0, dtype=np.float64).astype(np.float32) + if method == "sum": + return planes.sum(axis=0, dtype=np.float64).astype(np.float32) + raise ValueError(f"Unknown projection method {method!r}.") + + +def _project_czyx( + czyx_data: np.ndarray, + plans: list[tuple[int, ChannelProjection]], + focus_z: int | np.ndarray | None = None, + input_time_index: int | None = None, +) -> np.ndarray: + """Project one timepoint's (C, Z, Y, X) read to the (C_out, 1, Y, X) output. + + ``plans`` pairs each output channel with its index in ``czyx_data``. ``focus_z`` is + the position's in-focus plane, or one per timepoint (indexed with + ``input_time_index``, which iohub passes because it is in the signature). + """ + z_size = czyx_data.shape[1] + if isinstance(focus_z, np.ndarray): + focus_z = int(focus_z[input_time_index]) + projected = [ + project_zyx(czyx_data[c], plan.method, resolve_z_indices(plan, z_size, focus_z)) + for c, plan in plans + ] + return np.stack(projected)[:, np.newaxis] + + +def _resolve_channels( + settings: ProjectionSettings, channel_names: list[str], z_size: int +) -> tuple[list[tuple[int, ChannelProjection]], list[int], int | None]: + """Map the config's channel names onto the input plate. + + Returns the plans (index into the read, projection) in output order, the input + channel indices to read (the outputs, then the focus channel if it is not one), and + the focus channel's index in the read. Fixed selections are resolved against the + stack here, so a range outside it fails before any job is submitted. + """ + for name in settings.channels: + if name not in channel_names: + raise ValueError(f"Channel {name!r} is not in the input channels {channel_names}.") + read = [channel_names.index(name) for name in settings.channels] + plans = list(enumerate(settings.channels.values())) + + focus_index = None + if settings.uses_focus: + focus_channel = settings.focus.channel + if focus_channel not in channel_names: + raise ValueError( + f"focus.channel {focus_channel!r} is not in the input channels {channel_names}." + ) + if focus_channel not in settings.channels: + read.append(channel_names.index(focus_channel)) + focus_index = channel_names.index(focus_channel) + + for name, projection in settings.channels.items(): + if not projection.uses_focus: + try: + resolve_z_indices(projection, z_size) + except ValueError as e: + raise ValueError(f"Channel {name}: {e}") from e + + dropped = [name for name in channel_names if name not in settings.channels] + if dropped: + click.echo(f"Warning: input channels not in the config are dropped: {dropped}") + return plans, read, focus_index + + +def _init_output_plate( + input_position_dirpaths: list[Path], + output_dirpath: Path, + settings: ProjectionSettings, +) -> tuple[tuple[int, int, int, int, int], list[str], np.dtype]: + """Create the empty Z=1 output plate and record the settings as provenance. + + Returns the input (T, C, Z, Y, X) shape, its channel names and the output dtype. + """ + with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as input_dataset: + channel_names = input_dataset.channel_names + T, C, Z, Y, X = input_dataset.data.shape + scale = input_dataset.scale + input_dtype = input_dataset.data.dtype + + if Z == 1: + raise ValueError( + f"{input_position_dirpaths[0]} already has one Z-plane; nothing to project." + ) + # Resolve the config against the plate before scaffolding, so a bad channel name or + # range fails here. The dropped-channel warning is printed once, at init. + _resolve_channels(settings, channel_names, Z) + + methods = {projection.method for projection in settings.channels.values()} + dtype = np.dtype(np.float32) if methods & set(_FLOAT_METHODS) else np.dtype(input_dtype) + + input_plate = Path(input_position_dirpaths[0]).parents[2] + create_empty_plate( + store_path=output_dirpath, + position_keys=[Path(p).parts[-3:] for p in input_position_dirpaths], + channel_names=list(settings.channels), + shape=(T, len(settings.channels), 1, Y, X), + scale=scale, + version=resolve_ome_zarr_version( + input_position_dirpaths[0], settings.output_ome_zarr_version + ), + dtype=dtype, + metadata_sources=input_plate, + metadata_keys=PROVENANCE_METADATA_KEYS, + extra_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, + ) + return (T, C, Z, Y, X), channel_names, dtype + + +def _find_focus_per_t( + position_path: Path, channel_index: int, focus: ProjectionFocusSettings, pixel_size: float +) -> np.ndarray: + """Find the in-focus plane of every timepoint of one position, on ``focus.channel``. + + Measured on the centre crop with waveorder's transverse band like the engine's + FocusEstimator; an empty frame, or one without a focus, takes the middle plane. One + (Z, Y, X) volume is read at a time. + """ + from biahub.registration.estimators import EstimationError + from biahub.registration.methods.focus import center_crop_yx, find_focus + + crop_xy = (focus.center_crop_xy[0], focus.center_crop_xy[1]) + with open_ome_zarr(str(position_path), mode="r") as dataset: + T, _, Z, _, _ = dataset.data.shape + focus_z = np.empty(T, dtype=int) + for t in range(T): + zyx = np.asarray(dataset.data[t, channel_index]) + try: + focus_z[t] = find_focus( + center_crop_yx(zyx, crop_xy), pixel_size, focus.na_det, focus.lambda_ill + ) + except EstimationError: + focus_z[t] = Z // 2 + return np.clip(focus_z, 0, Z - 1) + + +def _project_position( + input_position_path: Path, + output_position_path: Path, + settings: ProjectionSettings, + plans: list[tuple[int, ChannelProjection]], + input_channel_indices: list[int], + focus_channel_index: int | None, + num_workers: int, + resume: bool, +) -> None: + """Project one position: the focus pre-pass, then every timepoint. + + The focus is found once per position for both ``focus.per`` modes and recorded under + the output position's ``biahub-project`` metadata, so the plane each (t) was + projected around can be read back. + """ + focus_z = None + if settings.uses_focus: + with open_ome_zarr(str(input_position_path), mode="r") as dataset: + pixel_size = dataset.scale[-1] + focus_per_t = _find_focus_per_t( + input_position_path, focus_channel_index, settings.focus, pixel_size + ) + if settings.focus.per == "fov": + focus_z = int(round(float(np.median(focus_per_t)))) + else: + focus_z = focus_per_t + click.echo( + f"Focus on {settings.focus.channel}: " + + ( + f"plane {focus_z} (median over T)" + if settings.focus.per == "fov" + else f"{focus_per_t.tolist()}" + ) + ) + with open_ome_zarr(str(output_position_path), mode="r+") as output_dataset: + record = dict(output_dataset.zattrs.get(PROVENANCE_KEY, {})) + record["focus"] = { + "channel": settings.focus.channel, + "per": settings.focus.per, + "focus_z": focus_z if settings.focus.per == "fov" else None, + "focus_z_per_t": focus_per_t.tolist(), + } + output_dataset.zattrs[PROVENANCE_KEY] = record + + process_single_position( + _project_czyx, + input_position_path, + output_position_path, + input_channel_indices=[input_channel_indices], + output_channel_indices=[list(range(len(plans)))], + num_workers=num_workers, + resume=resume, + resume_token=settings_fingerprint(settings), + plans=plans, + focus_z=focus_z, + ) + + +def project( + input_position_dirpaths: list[Path], + config_filepath: Path, + output_dirpath: Path, + sbatch_filepath: str | None = None, + cluster: str = "slurm", + monitor: bool = True, + init_only: bool = False, + resume: bool = False, +): + """Project a 3D plate to a 2D (Z=1) plate, per channel. + + Parameters + ---------- + input_position_dirpaths : list[Path] + Paths to the input position directories. + config_filepath : Path + Path to the projection configuration file (``ProjectionSettings``). + output_dirpath : Path + Path to the output plate. + sbatch_filepath : str, optional + Path to the SLURM batch file. + cluster : str, optional + Execution cluster: 'slurm' submits to a Slurm cluster, 'local' runs jobs as + subprocesses on this machine, 'debug' runs jobs in-process in the foreground. + monitor : bool, optional + If True, monitor the submitted jobs. + init_only : bool, optional + Only initialize the output store and exit; skip per-position processing. + resume : bool, optional + Skip the (time, channel) units a previous attempt already finished, rather + than recomputing the whole position; see ``iohub.ngff.utils.process_single_position``. + The focus pre-pass runs again either way. + """ + output_dirpath = Path(output_dirpath) + slurm_out_path = output_dirpath.parent / "slurm_output" + + settings = yaml_to_model(config_filepath, ProjectionSettings) + (T, C, Z, Y, X), channel_names, _ = _init_output_plate( + input_position_dirpaths, output_dirpath, settings + ) + plans, read_indices, focus_index = _resolve_channels(settings, channel_names, Z) + + # Not yet calibrated on a completed run. RAM: one timepoint's read (the listed + # channels plus the focus channel) per worker, with headroom for the float64 + # accumulation. Time per (t, c) volume: a reduction is cheap, so 0.1 min carries a + # wide margin; the focus pre-pass reads the focus channel once more per position + # (~40% on top of track's old window selection, #418), hence the higher rate. + time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( + shape=(T, len(read_indices), Z, Y, X), + ram_multiplier=4, + time_multiplier=0.15 if settings.uses_focus else 0.1, + max_num_cpus=8, + ) + mem_gb = num_cpus * gb_ram_per_cpu + echo_resources(num_cpus, mem_gb, time_minutes) + + if init_only: + click.echo(f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)") + return + + output_position_paths = get_output_paths(input_position_dirpaths, output_dirpath) + + slurm_args = { + "slurm_job_name": "project", + "slurm_mem": f"{mem_gb}G", + "slurm_cpus_per_task": num_cpus, + "slurm_array_parallelism": 100, # process up to 100 positions at a time + "slurm_time": time_minutes, + "slurm_partition": "preempted", + } + + if sbatch_filepath: + slurm_args = apply_sbatch_overrides(slurm_args, sbatch_to_submitit(sbatch_filepath)) + + resolved_cluster = get_submitit_cluster(cluster=cluster) + click.echo(f"Preparing jobs on cluster='{resolved_cluster}': {slurm_args}") + executor = submitit.AutoExecutor(folder=slurm_out_path, cluster=resolved_cluster) + if resolved_cluster == "local": + # Explicit --sbatch-filepath values (e.g. '#LOCAL --timeout-min') win over + # the computed local executor parameters. + slurm_args = {**local_executor_parameters(time_minutes), **slurm_args} + executor.update_parameters(**slurm_args) + + click.echo("Submitting jobs...") + jobs = [] + with submitit.helpers.clean_env(), executor.batch(): + for input_position_path, output_position_path in zip( + input_position_dirpaths, output_position_paths, strict=True + ): + jobs.append( + executor.submit( + _project_position, + input_position_path, + output_position_path, + settings=settings, + plans=plans, + input_channel_indices=read_indices, + focus_channel_index=focus_index, + num_workers=slurm_args["slurm_cpus_per_task"], + resume=resume, + ) + ) + + job_ids = [job.job_id for job in jobs] + slurm_out_path.mkdir(exist_ok=True) + log_path = slurm_out_path / "submitit_jobs_ids.log" + with log_path.open("w") as log_file: + log_file.write("\n".join(job_ids)) + + # submitit's DebugExecutor is lazy: .submit() wraps the callable in a DebugJob but + # execution only happens on .wait()/.done()/.result(). On the Nextflow path + # (--cluster debug) run each position in the foreground. + if resolved_cluster == "debug": + for job, path in zip(jobs, input_position_dirpaths, strict=True): + job.wait() + click.echo(f"Projection complete: {path}") + return + + if monitor: + monitor_jobs(jobs, input_position_dirpaths) + + +@click.command("project") +@input_position_dirpaths() +@config_filepath() +@output_dirpath() +@sbatch_filepath() +@cluster() +@monitor() +@init_only() +@resume() +def project_cli( + input_position_dirpaths: list[Path], + config_filepath: Path, + output_dirpath: Path, + sbatch_filepath: str | None = None, + cluster: str = "slurm", + monitor: bool = False, + init_only: bool = False, + resume: bool = False, +): + """Project a 3D plate to a 2D (Z=1) plate, with a projection method per channel. + + Each output channel is max, mean or sum over a Z range or window, or one plane of the + stack. The range or plane may be fixed or centred on the in-focus plane; see + settings/example_project_settings.yml. + + \b + SLURM fan-out of positions across a whole plate: + >>> biahub project -i ./input.zarr/*/*/* -c ./project.yml -o ./output.zarr + + \b + Initialize the output plate only (e.g. before running per-position Nextflow workers): + >>> biahub project --init -i ./input.zarr/*/*/* -c ./project.yml -o ./output.zarr + + \b + In-process run of a single position (e.g. from a Nextflow worker): + >>> biahub project --cluster debug -i ./input.zarr/A/1/0 -c ./project.yml -o ./output.zarr + """ # noqa: D301 + project( + input_position_dirpaths=input_position_dirpaths, + config_filepath=config_filepath, + output_dirpath=output_dirpath, + sbatch_filepath=sbatch_filepath, + cluster=cluster, + monitor=monitor, + init_only=init_only, + resume=resume, + ) + + +if __name__ == "__main__": + project_cli() diff --git a/biahub/segment.py b/biahub/segment.py index 87add8f8..8f035ce4 100644 --- a/biahub/segment.py +++ b/biahub/segment.py @@ -56,7 +56,7 @@ class ResolvedModel: pretrained_model: str channel_indices: list[int] eval_args: dict[str, Any] - z_slice_2D: int | None + is_3d: bool # (function, channel index, kwargs), applied in order preprocessing: list[tuple[Callable, int, dict[str, Any]]] = field(default_factory=list) @@ -69,9 +69,10 @@ def resolve_models( ) -> list[ResolvedModel]: """Resolve each configured model against the input plate, without touching ``settings``. - Channel names become dataset indices, 3D models get the dataset anisotropy unless the - config sets one, and ``z_slice_2D`` is checked against the stack. ``settings`` is left - exactly as written, so it can be recorded as provenance. + Channel names become dataset indices and 3D models get the dataset anisotropy unless + the config sets one. The models' dimensionality must match the input: 2D models need + a Z=1 plate (see `biahub project`), 3D models a stack. ``settings`` is left exactly as + written, so it can be recorded as provenance. """ def index(name: str, what: str) -> int: @@ -81,14 +82,22 @@ def index(name: str, what: str) -> int: ) return channel_names.index(name) + if settings.is_3d and z_size == 1: + raise ValueError( + "The input has one Z-plane but the models segment in 3D (eval_args.do_3D or " + "stitch_threshold); remove those arguments to segment in 2D." + ) + if not settings.is_3d and z_size > 1: + raise ValueError( + f"The input is 3D (Z={z_size}) but the models segment in 2D. Project the plate " + "with `biahub project` and segment the projected (Z=1) plate, or set " + "eval_args.do_3D: true to segment the stack." + ) + resolved = [] for name, model in settings.models.items(): - if model.z_slice_2D is not None and model.z_slice_2D >= z_size: - raise ValueError( - f"Model {name}: z_slice_2D={model.z_slice_2D} is outside the input stack (Z={z_size})." - ) eval_args = dict(model.eval_args) - if model.z_slice_2D is None and eval_args.get("anisotropy") is None: + if model.is_3d and eval_args.get("anisotropy") is None: eval_args["anisotropy"] = scale[-3] / scale[-1] resolved.append( ResolvedModel( @@ -96,7 +105,7 @@ def index(name: str, what: str) -> int: pretrained_model=model.pretrained_model, channel_indices=[index(c, f"Model {name}:") for c in model.channels], eval_args=eval_args, - z_slice_2D=model.z_slice_2D, + is_3d=model.is_3d, preprocessing=[ ( p.function, @@ -133,7 +142,7 @@ def _segment_czyx( ) -> np.ndarray: """Segment a CZYX image with each resolved cellpose model. - Returns an array of shape (n_models, Z or 1, Y, X). + Returns an array of shape (n_models, Z, Y, X); Z is 1 for 2D models. """ # Every job this step submits asks SLURM for a GPU, so an unusable one is a # broken allocation, not a reason to fall back to a ~130x slower CPU run. @@ -151,21 +160,21 @@ def _segment_czyx( czyx_data[c_idx] = func(czyx_data[c_idx], **kwargs) # Cellpose 4 refuses a z axis for 2D processing, so a 2D model gets the (C, Y, X) - # plane and a 3D model the (C, Z, Y, X) stack. It also ignores `channels` and keeps - # the first 3 channels it is given, so pass only the configured ones. + # plane of the Z=1 input and a 3D model the (C, Z, Y, X) stack. It also ignores + # `channels` and keeps the first 3 channels it is given, so pass only the + # configured ones. cellpose_model = _cached_model(model.pretrained_model, device) - if model.z_slice_2D is not None: - image, z_axis = czyx_data[model.channel_indices, model.z_slice_2D], None - else: + if model.is_3d: image, z_axis = czyx_data[model.channel_indices], 1 + else: + image, z_axis = czyx_data[model.channel_indices, 0], None segmentation, _, _ = cellpose_model.eval( image, channel_axis=0, z_axis=z_axis, **model.eval_args ) - # Cellpose squeezes singleton axes (e.g. a 3D model on a single-plane input); - # restore the (Z or 1, Y, X) shape the output plate expects. + # Cellpose squeezes singleton axes; restore the (Z, Y, X) shape the output + # plate expects (Z is 1 for a 2D model). segmentation = np.reshape( - segmentation, - (1, *image.shape[-2:]) if model.z_slice_2D is not None else image.shape[1:], + segmentation, image.shape[1:] if model.is_3d else (1, *image.shape[-2:]) ) czyx_segmentation.append(segmentation) return np.stack(czyx_segmentation, axis=0) @@ -190,13 +199,12 @@ def _init_output_plate( channel_names = input_dataset.channel_names models = resolve_models(settings, channel_names, scale, Z) - Z_out = 1 if models[0].z_slice_2D is not None else Z # Label channel i belongs to model i. On an existing output, create_empty_plate keeps # the stored channels and array shape, so a reordered config would write into swapped - # channels and a 3D -> 2D switch would keep the old Z planes. + # channels. The labels have the input's Z (1 for a 2D plate). label_channels = [m.name + "_labels" for m in models] - output_shape = (T, len(models), Z_out, Y, X) + output_shape = (T, len(models), Z, Y, X) position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] for key in position_keys: output_position = Path(output_dirpath, *key) @@ -303,7 +311,7 @@ def segment( # earlier conservative 2.5 min per frame. CPUs and RAM (the input timepoint, C volumes, # plus cellpose buffers, x up to 4 CPUs) deliberately keep a margin over the 2D run, # which also covers the unmeasured 3D path (~20-30 GB for an A549 plate). - is_2d = models[0].z_slice_2D is not None + is_2d = not settings.is_3d time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( shape=(T, len(models), Z, Y, X), ram_multiplier=C + 4, diff --git a/biahub/settings.py b/biahub/settings.py index ba66aa80..a523b4f1 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -92,47 +92,14 @@ def warn_model_type_ignored(cls, v): return v -class ZSlicing(MyBaseModel): - """How to SELECT the Z-planes used for tracking. - - The ``method`` decides which of the other fields are used; fields that belong to a - different method are simply ignored (all have defaults). The block selects a - z-window; the actual reduction to 2D is governed separately by - ``TrackingSettings.output_mode`` (plus any Z-projection step, e.g. ``np.mean``, in - ``input_images``). - - Methods - ------- - all - Use every plane (``slice(None)``). - central - Use an automatically centred window (see ``central_z_slice``). - range - Use the explicit ``range`` ``[start, stop]`` slice (falls back to all planes - if ``range`` is left unset). - focus - Detect the in-focus plane per-FOV (waveorder ``focus_from_transverse_band`` - on ``focus_channel``) and take a fixed window of ``window_size`` planes around it, - split ``frac_below``/``frac_above``. - """ - - method: Literal["all", "central", "range", "focus"] = "all" - range: tuple[int, int] | None = None # method: range - window_size: int = 48 # method: focus (fixed window size, in z-planes) - frac_below: float = 1 / 3 # method: focus - frac_above: float = 2 / 3 # method: focus - focus_channel: str | None = None # method: focus -- channel focus-finding runs on - - class TrackingSettings(MyBaseModel): target_channel: str = "nuclei_prediction" fov: str = "*/*/*" blank_frames_path: Path | None = None - # 2D writes an output plate with Z=1 (input must be projected); 3D keeps the - # selected z-window. Does not itself project the data. + # The dimensionality of the tracking. Track neither projects nor selects Z-planes: + # "2D" needs a Z=1 input (project the plate with `biahub project` first) and "3D" + # a Z>1 input; a mismatch is an error. output_mode: Literal["2D", "3D"] = "2D" - # Which Z-planes to select for tracking. See ZSlicing. - z_slicing: ZSlicing = ZSlicing() input_images: list[ProcessingInputChannel] tracking_config: dict[str, Any] = {} segmentation_method: Literal["foreground_contour", "cellpose"] = "foreground_contour" @@ -146,6 +113,17 @@ class TrackingSettings(MyBaseModel): # When None, preserve the OME-Zarr version of the input store. output_ome_zarr_version: Literal["0.4", "0.5"] | None = None + @model_validator(mode="before") + @classmethod + def reject_z_slicing(cls, data): + # Z-plane selection moved out of track into `biahub project`. + if isinstance(data, dict) and "z_slicing" in data: + raise ValueError( + "z_slicing was removed from track: select and project the Z-planes with " + "`biahub project` and track the projected (Z=1) plate." + ) + return data + @field_validator("blank_frames_path") @classmethod def validate_blank_frames_path(cls, v): @@ -331,6 +309,126 @@ def check_center_crop_xy(cls, v): return v +# -------------------------------------------------------------------------------------- +# Projection settings (`biahub project`) + + +class ZSelection(MyBaseModel): + """Which Z-planes of a stack feed one channel's projection. + + `range` is a half-open ``[start, stop)`` like concatenate's ``Z_slice``. With `center` + the bounds are offsets from the center (``[-20, 40]`` is 60 planes starting 20 below + it); `center` is a plane index or ``"focus"``, the in-focus plane found on + ``ProjectionSettings.focus.channel``. `step` keeps every step-th plane of the range. + A range that runs past the stack is clipped to it. + """ + + center: NonNegativeInt | Literal["focus"] | None = None + range: tuple[int, int] | None = None + step: PositiveInt = 1 + + @model_validator(mode="after") + def check_range(self): + if self.range is not None: + start, stop = self.range + if stop <= start: + raise ValueError(f"range {list(self.range)} is empty: stop must exceed start.") + if self.center is None and start < 0: + raise ValueError( + f"range {list(self.range)} is negative; offsets need a center." + ) + elif self.step != 1: + raise ValueError("step needs a range.") + return self + + +class ChannelProjection(MyBaseModel): + """How one output channel is projected from the input stack. + + ``max``, ``mean`` and ``sum`` reduce the selected planes (the whole stack when `z` is + unset); ``slice`` keeps the one plane ``z.center``. + """ + + method: Literal["max", "mean", "sum", "slice"] + z: ZSelection | None = None + + @model_validator(mode="after") + def check_z_for_method(self): + if self.method == "slice": + if self.z is None or self.z.center is None: + raise ValueError("method slice needs z.center (a plane index or 'focus').") + if self.z.range is not None: + raise ValueError( + "method slice selects one plane and takes no range; use max, mean or " + "sum to project a window." + ) + elif self.z is not None and self.z.center is not None and self.z.range is None: + raise ValueError( + f"method {self.method} needs z.range with z.center; use method slice to " + "keep the one plane." + ) + return self + + @property + def uses_focus(self) -> bool: + return self.z is not None and self.z.center == "focus" + + +class ProjectionFocusSettings(FocusSettings): + """Focus finding for `biahub project`: the engine's criterion, on one channel. + + `channel` is the input channel the focus is measured on (it need not be an output). + `per` is how the in-focus plane is resolved across time: ``fov`` takes one plane per + position, the median over timepoints; ``timepoint`` keeps one plane per timepoint, so + the window follows Z drift. + """ + + channel: str + per: Literal["fov", "timepoint"] = "fov" + # Projection has no drift axes to choose; the inherited field is fixed. + axes: Literal["z"] = "z" + + @model_validator(mode="before") + @classmethod + def reject_axes(cls, data): + if isinstance(data, dict) and "axes" in data: + raise ValueError("axes is a stabilization setting; projection has none.") + return data + + +class ProjectionSettings(MyBaseModel): + """`biahub project`: a 3D plate to a 2D (Z=1) plate, one projection per output channel. + + `channels` maps each OUTPUT channel (an input channel name, in output order) to its + projection; input channels not listed are dropped. `focus` is required when any + channel uses ``z.center: focus``. + """ + + channels: dict[str, ChannelProjection] = Field(min_length=1) + focus: ProjectionFocusSettings | None = None + # When None, preserve the OME-Zarr version of the input store. + output_ome_zarr_version: Literal["0.4", "0.5"] | None = None + + @model_validator(mode="after") + def check_focus(self): + focus_users = [name for name, c in self.channels.items() if c.uses_focus] + if focus_users and self.focus is None: + raise ValueError( + f"{focus_users} use z.center: focus; add a `focus` block naming the channel " + "the focus is found on." + ) + if self.focus is not None and not focus_users: + warnings.warn( + "`focus` is set but no channel uses z.center: focus; it is ignored.", + stacklevel=2, + ) + return self + + @property + def uses_focus(self) -> bool: + return any(c.uses_focus for c in self.channels.values()) + + class AffineTransformSettings(MyBaseModel): t_reference: Literal["first", "previous"] = "first" transform_type: Literal["euclidean", "similarity", "affine"] = "euclidean" @@ -755,15 +853,15 @@ class SegmentationModel(BaseModel): """One cellpose 4 model and how to apply it. ``channels`` are dataset channel names (1-3, any order: cellpose 4 takes up to 3 - channels). ``z_slice_2D`` selects the plane for 2D segmentation; without it the model - segments the Z stack, which needs ``eval_args.do_3D`` (or ``stitch_threshold``). - ``eval_args`` are the remaining ``CellposeModel.eval`` keyword arguments. + channels). The model segments in 2D unless ``eval_args.do_3D`` (or a + ``stitch_threshold``) is set; a 2D model needs a Z=1 input, which `biahub project` + makes from a stack. ``eval_args`` are the remaining ``CellposeModel.eval`` keyword + arguments. """ pretrained_model: str channels: list[str] = Field(min_length=1) eval_args: dict[str, Any] = {} - z_slice_2D: NonNegativeInt | None = None preprocessing: list[PreprocessingFunctions] = [] model_config = ConfigDict(extra="forbid", protected_namespaces=()) @@ -779,6 +877,12 @@ def reject_old_keys(cls, data): "eval_args.channels moved to the model's own `channels` field " "(dataset channel names)." ) + if "z_slice_2D" in data: + raise ValueError( + "z_slice_2D was removed: select or project the plane with " + "`biahub project` and segment the 2D (Z=1) plate. A model without " + "eval_args.do_3D segments in 2D." + ) return data @field_validator("channels") @@ -802,24 +906,17 @@ def validate_eval_args(cls, value): ) return value - @model_validator(mode="after") - def check_2d_3d_settings(self): - if self.z_slice_2D is not None and self.eval_args.get("do_3D"): - raise ValueError( - "If 'z_slice_2D' is provided, 'do_3D' in 'eval_args' must be False." - ) - # A model without z_slice_2D segments the Z stack, which cellpose 4 only accepts - # with do_3D or stitching (it raises "2D image processing selected, but z_axis..."). + @property + def is_3d(self) -> bool: + """Whether the model segments the Z stack. + + Cellpose 4 segments a stack only with do_3D or stitching; otherwise it refuses a + z axis ("2D image processing selected, but z_axis..."), so the model is 2D. + """ stitch_threshold = self.eval_args.get("stitch_threshold") - if self.z_slice_2D is None and not ( - self.eval_args.get("do_3D") - or (stitch_threshold is not None and stitch_threshold > 0) - ): - raise ValueError( - "A model without z_slice_2D segments in 3D: set eval_args.do_3D: true (or a " - "stitch_threshold), or set z_slice_2D for 2D segmentation." - ) - return self + return bool(self.eval_args.get("do_3D")) or ( + stitch_threshold is not None and stitch_threshold > 0 + ) class SegmentationSettings(BaseModel): @@ -829,16 +926,20 @@ class SegmentationSettings(BaseModel): model_config = {"extra": "forbid", "protected_namespaces": ()} @model_validator(mode="after") - def check_models_share_z(self): - # All models write into one output array, so its Z size must suit every model. - flat = [name for name, m in self.models.items() if m.z_slice_2D is not None] - if flat and len(flat) != len(self.models): + def check_models_share_dimensionality(self): + # One input plate is either 2D (Z=1) or 3D (Z>1), so it suits only one kind. + three_d = [name for name, m in self.models.items() if m.is_3d] + if three_d and len(three_d) != len(self.models): raise ValueError( - f"Cannot mix 2D and 3D models in one config: {flat} set z_slice_2D, the " - "others do not. Segment them in separate runs." + f"Cannot mix 2D and 3D models in one config: {three_d} set do_3D (or a " + "stitch_threshold), the others do not. Segment them in separate runs." ) return self + @property + def is_3d(self) -> bool: + return next(iter(self.models.values())).is_3d + # -------------------------------------------------------------------------------------- # Unified transform estimation / application settings (registration engine) diff --git a/biahub/track.py b/biahub/track.py index 4f1f9669..0bdcdb38 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -31,7 +31,6 @@ CellposeConfig, ProcessingInputChannel, TrackingSettings, - ZSlicing, ) from biahub.utils.cellpose import ( cellpose_device, @@ -52,11 +51,6 @@ logger = logging.getLogger(__name__) -# Optical parameters for waveorder focus finding (focus_from_transverse_band), -# the mantis label-free arm; the same values as biahub.settings.FocusSettings' defaults. -NA_DET = 1.35 -LAMBDA_ILL = 0.500 - # Lazy imports for ultrack - imported only when needed in specific functions _ultrack_patched = False @@ -259,156 +253,22 @@ def get_empty_frames_idx_from_csv(blank_frame_df: pd.DataFrame, fov: str) -> lis return None -def central_z_slice(z_shape: int) -> slice: - """ - Compute a centered Z-slice range from a 3D or 4D image volume. - - This function returns a slice object that selects a centered range of Z-planes - from a volumetric dataset, based on the total number of Z slices. It ensures - that the returned slice includes an odd number of planes (at least 3). - - Parameters - ---------- - z_shape : int - The total number of Z-planes (depth) in the dataset. - - Returns - ------- - slice - A slice object that extracts a centered Z-range from the dataset. - - Notes - ----- - - Ensures that at least 3 slices are returned. - - If the total number of slices is even, one extra slice is included to make it odd. - - This function is typically used when `z_slices=(0, 0)` is provided, - indicating that the user wants automatic central slicing. +def check_input_dimensionality(output_mode: str, z_size: int) -> None: + """Refuse an input whose Z does not match ``output_mode``. - Examples - -------- - >>> central_z_slice(21) - slice(10 - 1, 10 + 1 + 1) # slice(9, 12) - - >>> central_z_slice(8) - slice(3, 6) # (center=4, half_window=1) + Track neither projects nor selects Z-planes (that is `biahub project`): 2D tracking + needs a Z=1 plate and 3D tracking a stack. """ - n_slices = max(3, z_shape // 2) - if n_slices % 2 == 0: - n_slices += 1 - z_center = z_shape // 2 - half_window = n_slices // 2 - return slice(z_center - half_window, z_center + half_window + 1) - - -def _median_focus_plane(stack: ArrayLike, pixel_size: float) -> int: - """Median in-focus z-plane across timepoints for a (T, Z, Y, X) stack. - - Uses waveorder ``focus_from_transverse_band`` (with the module constants - ``NA_DET``/``LAMBDA_ILL``) per timepoint; empty frames fall back to the central plane. - """ - from waveorder.focus import focus_from_transverse_band - - z_shape = stack.shape[1] - z_focus = [] - for t in tqdm(range(stack.shape[0]), desc="Finding focus"): - zyx = np.asarray(stack[t]) - if zyx.sum() == 0: - z_focus.append(z_shape // 2) - continue - z_f = focus_from_transverse_band( - zyx, NA_det=NA_DET, lambda_ill=LAMBDA_ILL, pixel_size=pixel_size + if output_mode == "2D" and z_size > 1: + raise ValueError( + f"output_mode is 2D but the input is 3D (Z={z_size}). Project the plate with " + "`biahub project` and track the projected (Z=1) plate, or set output_mode: 3D." ) - z_focus.append(z_shape // 2 if z_f is None else int(np.clip(z_f, 0, z_shape - 1))) - return int(np.median(z_focus)) - - -def _focus_window( - center: int, window_size: int, z_shape: int, frac_below: float -) -> tuple[slice, int]: - """Build a fixed-size z-window of ``window_size`` planes around ``center``. - - The window is shifted (not clipped) to stay within ``[0, z_shape]``, so it always - spans exactly ``min(window_size, z_shape)`` planes. ``frac_below`` sets how far - below ``center`` it starts. - """ - size = min(window_size, z_shape) - start = center - int(round(frac_below * window_size)) - stop = start + size - if start < 0: - start, stop = 0, size - elif stop > z_shape: - start, stop = z_shape - size, z_shape - return slice(start, stop), size - - -def resolve_z_slice(z: ZSlicing, z_shape: int) -> tuple[slice, int]: - """Resolve the read-time Z-slice and its plane count from a :class:`ZSlicing`. - - ``z.method`` selects the strategy: ``all`` (every plane), ``central`` (see - :func:`central_z_slice`), or ``range`` (explicit ``z.range``). ``focus`` is resolved - later, per-FOV, from the loaded data (see :func:`apply_focus_slicing`); here it just - loads the full stack (``slice(None)``) but reports the fixed window count - ``min(z.window_size, z_shape)``, so the output plate Z still matches the tracked Z. - - Returns ``(slice, n_planes)``. - """ - if z.method == "all": - return slice(None), z_shape - - if z.method == "central": - z_slices = central_z_slice(z_shape) - return z_slices, z_slices.stop - z_slices.start - - if z.method == "range": - if z.range is None: - return slice(None), z_shape - start, stop = z.range - if stop <= start: - raise ValueError( - f"Invalid z_slicing.range {z.range}: must contain at least one slice " - "(stop > start)." - ) - return slice(start, stop), stop - start - - if z.method == "focus": - # Load the full stack now; the focus window is applied per-FOV after loading. - return slice(None), min(z.window_size, z_shape) - - raise ValueError(f"Unknown z_slicing.method: {z.method!r}") - - -def apply_focus_slicing( - data_dict: dict[str, ArrayLike], z_slicing: ZSlicing, pixel_size: float -) -> dict[str, ArrayLike]: - """Slice every loaded channel to a focus window (``method='focus'``). - - The in-focus plane is found once per FOV on ``z_slicing.focus_channel`` (or the - first loaded channel if unset), then the SAME fixed window of - ``z_slicing.window_size`` planes is applied to every channel so they stay Z-aligned. - Channel arrays are (T, Z, Y, X). Focus finding reads the focus channel one timepoint - at a time; every channel stays lazy, so only the resolved window is materialised - downstream. - """ - focus_channel = z_slicing.focus_channel or next(iter(data_dict)) - if focus_channel not in data_dict: + if output_mode == "3D" and z_size == 1: raise ValueError( - f"focus_channel '{focus_channel}' not in loaded channels {list(data_dict)}." + "output_mode is 3D but the input has one Z-plane; set output_mode: 2D." ) - # Keep the stack lazy: _median_focus_plane indexes one timepoint at a time, so a - # dask array reads a single (Z, Y, X) volume per step. Materialising the whole - # (T, Z, Y, X) focus channel here costs tens of GB on a full timelapse and OOM-kills - # the worker. - stack = data_dict[focus_channel] - - center = _median_focus_plane(stack, pixel_size) - z_slices, _ = _focus_window( - center, z_slicing.window_size, stack.shape[1], z_slicing.frac_below - ) - click.echo(f"Focus-resolved z-slice: {z_slices}") - - return {name: arr[:, z_slices] for name, arr in data_dict.items()} - def filter_short_siblings( tracks_df: pd.DataFrame, labels: ArrayLike, min_length: int @@ -590,15 +450,14 @@ def run_preprocessing_pipeline( def load_data( position_key: tuple[str, str, str], input_images: list[ProcessingInputChannel], - z_slices: slice, visualize: bool = False, ) -> dict[str, ArrayLike]: """ Load and extract specified channels from an OME-Zarr dataset for a given position. This function opens the OME-Zarr dataset corresponding to a specific position key, - extracts the required channels and Z-slices as defined in the `input_images` configuration, - and stores them in a dictionary for further processing. + extracts the required channels (every Z-plane) as defined in the `input_images` + configuration, and stores them in a dictionary for further processing. Parameters ---------- @@ -607,8 +466,6 @@ def load_data( in the dataset (e.g., (plate, well, position)). input_images : list of ProcessingInputChannel List of input image channel configurations. Each item defines the path and channels to load. - z_slices : slice - Slice object specifying which Z-planes to extract (e.g., central slices for 2D mode). visualize : bool, optional If True, opens a Napari viewer and displays each loaded channel. Default is False. @@ -616,7 +473,8 @@ def load_data( ------- dict of str to ArrayLike Dictionary mapping each channel name to its corresponding image data array. - Each array is a Dask array of shape (T, Z, Y, X), extracted from the OME-Zarr store. + Each array is a Dask array of shape (T, Z, Y, X), extracted from the OME-Zarr + store, or (T, Y, X) when the store has one Z-plane (a projected plate). Notes ----- @@ -637,9 +495,12 @@ def load_data( image_channel_names = dataset.channel_names for channel_name, _ in image.channels.items(): click.echo(f"Loading data for channel {channel_name} from {image.path}") - data_dict[channel_name] = dataset.data.dask_array()[ - :, image_channel_names.index(channel_name), z_slices, :, : + tzyx = dataset.data.dask_array()[ + :, image_channel_names.index(channel_name) ] + # A projected (Z=1) plate is a 2D series: drop the singleton axis so + # the processing functions and cellpose see (T, Y, X). + data_dict[channel_name] = tzyx[:, 0] if tzyx.shape[1] == 1 else tzyx if visualize: viewer = napari.Viewer() viewer.add_image(data_dict[channel_name], name=channel_name) @@ -689,25 +550,17 @@ def fill_empty_frames_from_csv( def _load_and_preprocess( position_key: str, input_images: list[ProcessingInputChannel], - z_slices: slice, blank_frames_path: Path | None, - z_slicing: ZSlicing | None, - pixel_size: float | None, visualize: bool = False, ) -> dict[str, ArrayLike]: - """Load, focus-slice, run the processing pipeline, and fill blank frames for one FOV. + """Load, run the processing pipeline, and fill blank frames for one FOV. Shared by the foreground+contour and cellpose segmentation paths. """ fov = "/".join(position_key) data_dict = load_data( - position_key=position_key, - input_images=input_images, - z_slices=z_slices, - visualize=visualize, + position_key=position_key, input_images=input_images, visualize=visualize ) - if z_slicing is not None and z_slicing.method == "focus": - data_dict = apply_focus_slicing(data_dict, z_slicing, pixel_size) data_dict = run_preprocessing_pipeline(data_dict, input_images, visualize=visualize) return fill_empty_frames_from_csv(fov, data_dict, blank_frames_path) @@ -771,8 +624,9 @@ def cellpose_segmentation( ) -> np.ndarray: """Run cellpose on a preprocessed ``data_dict`` channel. - Selects ``cellpose_config.input_channel``, mean-projects Z if needed, and returns an - integer label array of shape (T, Y, X) ready for ``run_ultrack(..., labels=...)``. + Selects ``cellpose_config.input_channel`` (a 2D series, (T, Y, X) or (T, 1, Y, X)) and + returns an integer label array of shape (T, Y, X) ready for + ``run_ultrack(..., labels=...)``. """ # Get the channel to segment channel_name = cellpose_config.input_channel @@ -784,11 +638,14 @@ def cellpose_segmentation( images = data_dict[channel_name] - # Project Z BEFORE materialising: mean-reducing a dask array streams chunk-wise, so - # peak memory is the (T, Y, X) result rather than the full (T, Z, Y, X) stack. + # Track does not project: a Z=1 input is squeezed, a stack is refused. if images.ndim == 4: - click.echo(f"Projecting Z-dimension via mean: {images.shape} -> (T, Y, X)") - images = images.mean(axis=1) + if images.shape[1] != 1: + raise ValueError( + f"Cellpose segmentation in track is 2D but channel '{channel_name}' has " + f"{images.shape[1]} Z-planes; project the plate with `biahub project` first." + ) + images = images[:, 0] if isinstance(images, da.Array): images = images.compute() @@ -822,10 +679,8 @@ def track_one_position( output_dirpath: Path, tracking_config, blank_frames_path: Path | None = None, - z_slices: slice | None = None, scale: tuple[float, ...] = (1, 1, 1, 1, 1), cellpose_config: CellposeConfig | None = None, - z_slicing: ZSlicing | None = None, output_mode: str = "2D", min_sibling_length: int | None = None, ) -> None: @@ -852,20 +707,15 @@ def track_one_position( blank_frames_path : Path, optional Path to CSV file indicating empty frames for the current FOV. If None, blank frame filling is skipped. - z_slices : slice, optional - Read-time Z-slice passed to ``load_data``. For ``method='focus'`` this is the - full stack and the window is applied per-FOV after loading. scale : tuple of float, optional Physical scale used for tracking, in (Z, Y, X) or (Y, X) order (already trimmed by the caller to match ``output_mode``). Default is (1, 1, 1, 1, 1). cellpose_config : CellposeConfig, optional If provided, uses cellpose segmentation instead of foreground+contour. - z_slicing : ZSlicing, optional - Z-plane selection config. For ``method='focus'`` the in-focus window is resolved - per-FOV from the loaded data (see :func:`apply_focus_slicing`). output_mode : str, optional - "2D" writes (T, Y, X) labels into the Z=0 plane; "3D" writes (T, Z, Y, X) - labels across the z-window. Default is "2D". + "2D" tracks a Z=1 input and writes (T, Y, X) labels; "3D" tracks the stack and + writes (T, Z, Y, X) labels. Track does not project (see `biahub project`). + Default is "2D". min_sibling_length : int, optional If set, remove daughter tracks of this many frames or fewer before writing (see :func:`filter_short_siblings`). Default None keeps every division. @@ -883,9 +733,6 @@ def track_one_position( - The Ultrack config is also saved as TOML in a `_config_tracking/{FOV}/` subdirectory. - In foreground+contour mode, missing "foreground"/"contour" channels raise a ValueError. """ - if z_slicing is None: - z_slicing = ZSlicing() - fov = "_".join(position_key) click.echo(f"Processing FOV: {fov.replace('_', '/')}") @@ -894,13 +741,9 @@ def track_one_position( database_path = output_dirpath.parent / f"{filename}_config_tracking" / f"{fov}" os.makedirs(database_path, exist_ok=True) - # Load, focus-slice (method='focus'), run the processing pipeline, and fill blank - # frames once; both segmentation methods consume the resulting data_dict. - # pixel_size is the transverse pixel size = last element of the (trimmed) scale. - pixel_size = scale[-1] - data_dict = _load_and_preprocess( - position_key, input_images, z_slices, blank_frames_path, z_slicing, pixel_size - ) + # Load, run the processing pipeline, and fill blank frames once; both segmentation + # methods consume the resulting data_dict. + data_dict = _load_and_preprocess(position_key, input_images, blank_frames_path) if cellpose_config is not None: cellpose_labels = cellpose_segmentation(data_dict, cellpose_config) @@ -940,15 +783,17 @@ def track_one_position( click.echo(f"Saved tracks to: {output_dirpath / Path(*position_key)}") - # Save the tracking labels. The output plate stores (T, 1, Z_out, Y, X); in 2D - # mode Z_out is 1 and labels are (T, Y, X); in 3D mode labels are (T, Z, Y, X). + # Save the tracking labels. The output plate stores (T, 1, Z, Y, X) with the input's + # Z: in 2D mode Z is 1 and labels are (T, Y, X); in 3D mode labels are (T, Z, Y, X). labels = np.asarray(tracking_labels, dtype=np.uint32) with open_ome_zarr(output_dirpath / Path(*position_key), mode="r+") as output_dataset: if output_mode == "2D": + if labels.ndim == 4 and labels.shape[1] == 1: + labels = labels[:, 0] if labels.ndim != 3: raise ValueError( f"output_mode='2D' expects (T, Y, X) labels but tracking produced " - f"shape {labels.shape}. Ensure input_images projects Z (e.g. np.mean)." + f"shape {labels.shape}." ) output_dataset[0][:, 0, 0] = labels else: @@ -968,20 +813,17 @@ def _init_output_plate( ) -> tuple[int, int, int, int, int]: """Create the empty tracking output plate. - Returns the (T, C, Z_out, Y, X) shape used for resource estimation. + The input's dimensionality must match ``settings.output_mode`` (see + :func:`check_input_dimensionality`); the labels take the input's Z. + + Returns the (T, C, Z, Y, X) shape used for resource estimation. """ with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as dataset: T, C, Z, Y, X = dataset.data.shape scale = dataset.scale - # Data-free resolution: for focus this returns the fixed window size (min(total, Z)), - # so the plate Z matches the per-FOV tracked Z. - _, Z_win = resolve_z_slice(settings.z_slicing, Z) - - if settings.output_mode == "2D": - output_shape = (T, 1, 1, Y, X) - else: - output_shape = (T, 1, Z_win, Y, X) + check_input_dimensionality(settings.output_mode, Z) + output_shape = (T, 1, Z, Y, X) position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] @@ -1092,8 +934,6 @@ def track( T, C, Z, Y, X = dataset.data.shape scale = dataset.scale - # Read-time slice (full stack for focus; the per-FOV window is applied after load). - z_slices, _ = resolve_z_slice(settings.z_slicing, Z) track_scale = scale[-2:] if settings.output_mode == "2D" else scale[-3:] tracking_cfg = settings.tracking_config @@ -1145,10 +985,8 @@ def track( tracking_config=tracking_cfg, input_images=settings.input_images, blank_frames_path=settings.blank_frames_path, - z_slices=z_slices, scale=track_scale, cellpose_config=cellpose_cfg, - z_slicing=settings.z_slicing, output_mode=settings.output_mode, min_sibling_length=settings.min_sibling_length, ) diff --git a/nextflow/configs/a549/project.yml b/nextflow/configs/a549/project.yml new file mode 100644 index 00000000..bd1fcf7d --- /dev/null +++ b/nextflow/configs/a549/project.yml @@ -0,0 +1,22 @@ +# Projection config for the mantis-v2 pipeline (`biahub project`). A549. +# +# Project reads the ASSEMBLED plate (-assemble/.zarr) and writes the 2D +# plate that track reads, so the channels below must exist THERE, and track.yml's +# target_channel / cellpose_config.input_channel must be among the channels written here. +# +# Mean of the virtual-stain nuclei over 60 planes around the in-focus plane of the +# phase channel: the window starts 20 planes below focus and ends 39 above it. This is +# the window track's own z_slicing selected before projection moved here (#418): against +# TrackEdit annotations (with min_sibling_length 3) division F1 rose 0.63 -> 0.68 on +# 2026_04_28, with fewer false divisions and fragmented tracks on 2026_04_21 at the same +# mitotic-cell detection as all planes (0.95); a narrower 40-plane window lost mitotic +# cells that drift in Z (0.89). The focus is found once per FOV (median over timepoints) +# on the 800x800 centre crop of the phase channel. +channels: + nuclei_prediction: + method: mean + z: {center: focus, range: [-20, 40]} + +focus: + channel: Phase3D + per: fov diff --git a/nextflow/configs/a549/segment.yml b/nextflow/configs/a549/segment.yml index 29fa2442..dba0b53f 100644 --- a/nextflow/configs/a549/segment.yml +++ b/nextflow/configs/a549/segment.yml @@ -1,12 +1,12 @@ # biahub segment: cellpose 4 nuclei labels from the virtual-stain nuclei channel of the -# assembled plate (mantis v2, 0.1133 um/px). Enable with --segment_config. +# PROJECTED plate (-project/.zarr, Z=1; mantis v2, 0.1133 um/px). Segment +# does not project: a model without eval_args.do_3D segments in 2D and needs a Z=1 +# input, which `biahub project` (project.yml) makes from the assembled plate. +# Not wired into mantis-v2.nf yet. models: nuclei: pretrained_model: cpsam_v2 # cellpose 4: cpsam_v2, cpsam, cpdino, cpdino-vitb, or a path - channels: [nuclei_prediction] # assembled-plate channel name(s), at most 3 - # 2D on one plane. Pick an in-focus plane of the virtual-stain stack for the dataset - # (selecting the plane per FOV from the data is planned, not implemented yet). - z_slice_2D: 30 + channels: [nuclei_prediction] # projected-plate channel name(s), at most 3 eval_args: diameter: null # native resolution (cellpose 4 resizes by 30 / diameter) cellprob_threshold: 0.0 diff --git a/nextflow/configs/a549/track.yml b/nextflow/configs/a549/track.yml index 3317cd4b..24a01f6d 100644 --- a/nextflow/configs/a549/track.yml +++ b/nextflow/configs/a549/track.yml @@ -1,35 +1,23 @@ # Tracking config for the mantis-v2 pipeline (`biahub track`). A549. # -# Track reads the ASSEMBLED plate (-assemble/.zarr), so target_channel -# and cellpose_config.input_channel must name channels that exist THERE. If you -# rename channels before tracking, update these to match. +# Track reads the PROJECTED plate (-project/.zarr, Z=1), so target_channel +# and cellpose_config.input_channel must name channels that `biahub project` wrote +# (project.yml's `channels`). If you rename channels before tracking, update these. # # CURRENT SCHEMA. TrackingSettings forbids extra keys; the old fields `mode`, -# `z_range` and `focus_config` are rejected. Mapping from the old schema: -# mode -> output_mode -# z_range + focus_config.z_window -> z_slicing.{method, window_size, focus_channel} -# focus_config's NA_det / lambda_ill / pixel_size have no equivalent — focus -# detection now runs waveorder's focus_from_transverse_band on -# z_slicing.focus_channel internally. +# `z_range`, `focus_config` and `z_slicing` are rejected. Z-plane selection and +# projection moved to `biahub project` (project.yml): the 60-plane focus window and +# the mean projection that used to be here live there now. Mapping from the old schema: +# mode -> output_mode +# z_slicing / z_range / focus_config -> project.yml: channels..z, focus +# input_images np.mean step -> project.yml: channels..method: mean -# 2D writes an output plate with Z=1; input_images below projects Z with np.mean. +# 2D tracks the Z=1 projected plate and writes a Z=1 label plate. Track does not +# project; a 3D input with output_mode 2D is an error. output_mode: "2D" fov: "*/*/*" target_channel: nuclei_prediction -# Segment a window of 60 planes around the in-focus plane rather than all planes: -# focus is found on the phase channel (median over timepoints), and the window is -# split 1/3 below / 2/3 above it (frac_below default). input_images below mean- -# projects the virtual-stain nuclei over that window. Against TrackEdit annotations -# (with min_sibling_length 3): division F1 0.63 -> 0.68 on 2026_04_28, with fewer -# false divisions and fragmented tracks on 2026_04_21 at the same mitotic-cell -# detection as all planes (0.95). A narrower 40-plane window lost mitotic cells -# that drift in Z (0.89). -z_slicing: - method: focus - window_size: 60 - focus_channel: Phase3D - segmentation_method: cellpose # Sibling filter (ultrack filter_short_sibling_tracks, same meaning): at a division @@ -51,16 +39,11 @@ cellpose_config: input_channel: nuclei_prediction labels_sigma: 5.0 +# The projected plate holds the nuclei already reduced to one plane; nothing to process. input_images: - path: null channels: - nuclei_prediction: - - function: np.mean - kwargs: - axis: 1 - per_timepoint: false - # Loaded only for focus finding (z_slicing.focus_channel); not processed. - Phase3D: [] + nuclei_prediction: [] tracking_config: segmentation_config: diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf index 3292a3c8..5e58bd2c 100644 --- a/nextflow/mantis-v2.nf +++ b/nextflow/mantis-v2.nf @@ -22,16 +22,17 @@ nextflow.enable.dsl = 2 // all the inits before any of the runs. See the INIT PHASE comment in the // workflow body for why that is possible and what it buys. // -// Flat-field → deskew → reconstruct → virtual-stain → assemble → track → QC is -// the full chain: assemble concatenates the deskew/reconstruct/virtual-stain -// channels into one plate, track reads that assembled plate as its single -// input, and QC reads the finished stores. Follow the chaining below for the -// pattern. +// Flat-field → deskew → reconstruct → virtual-stain → assemble → project → +// track → QC is the full chain: assemble concatenates the +// deskew/reconstruct/virtual-stain channels into one plate, project reduces +// that 3D plate to the 2D (Z=1) plate 2D tracking needs, track reads it as its +// single input, and QC reads the finished stores. Follow the chaining below for +// the pattern. // -// The last three are OPTIONAL and selected by whether their config is given, so -// a run performs the prefix it asks for: A549 wants assemble + track + QC, a -// neuromast run wants assemble + QC and no tracking (issue #306). A skipped step -// never renumbers the directories of the ones around it. +// The last four are OPTIONAL and selected by whether their config is given, so +// a run performs the prefix it asks for: A549 wants assemble + project + track +// + QC, a neuromast run wants assemble + QC and no tracking (issue #306). A +// skipped step never renumbers the directories of the ones around it. // --------------------------------------------------------------------------- params.input = null // raw source — may not be a zarr store @@ -40,6 +41,7 @@ params.deskew_config = null params.flat_field_config = null params.reconstruct_config = null params.virtual_stain_config = null +params.project_config = null params.track_config = null params.concatenate_config = null params.max_positions = 0 @@ -57,6 +59,7 @@ include { deskew_init_wf; deskew_run_wf } from './modules/deskew' include { reconstruct_init_wf; reconstruct_tf_wf; reconstruct_run_wf } from './modules/reconstruct' include { virtual_stain_init_wf; virtual_stain_run_wf } from './modules/virtual_stain' include { assemble_init_wf; assemble_run_wf } from './modules/assembly' +include { project_init_wf; project_run_wf } from './modules/project' include { track_init_wf; track_run_wf } from './modules/tracking' include { qc_plan_wf; qc_compute_wf; qc_report_wf; qc_report_spec } from './modules/qc' include { notify_step; notify_run_start; notify_run_end } from './modules/notify' @@ -96,6 +99,7 @@ def step_directories(performed) { reconstruct : 'reconstruct', virtual_stain : 'virtual-stain', assemble : 'assemble', + project : 'project', track : 'track', ] def layout = [:] @@ -125,8 +129,14 @@ workflow { // // Skipping a step DOES renumber the ones after it: the number is a position // among the steps performed, so a neuromast run's assembled store is - // `4-assemble` where an A549 run also has `5-track` after it. + // `4-assemble` where an A549 run also has `5-project` and `6-track` after it. + // + // Project turns the assembled 3D plate into the 2D (Z=1) plate that 2D + // tracking reads: `biahub track` no longer selects or projects Z-planes, so + // an A549 run names --project_config as well as --track_config. Tracking + // without it reads the assembled plate itself, which suits 3D tracking. def assemble_on = params.concatenate_config as boolean + def project_on = params.project_config as boolean def track_on = params.track_config as boolean def qc_image_on = params.qc_config as boolean def qc_track_on = params.qc_track_config as boolean @@ -135,8 +145,11 @@ workflow { // A step cannot outlive the step whose output it reads. Refuse the // combination at launch, naming the config to add or the one to drop, rather // than failing hours in with a missing store. + if (project_on && !assemble_on) { + error "--project_config needs --concatenate_config: project reads the assembled plate." + } if (track_on && !assemble_on) { - error "--track_config needs --concatenate_config: tracking reads the assembled plate." + error "--track_config needs --concatenate_config: tracking reads the assembled (or projected) plate." } if (qc_image_on && !assemble_on) { error "--qc_config needs --concatenate_config: it QCs the assembled store." @@ -170,6 +183,7 @@ workflow { // numbering is derived from. Reconstruction proper is always in it. def performed = ['flat_field', 'deskew', 'reconstruct', 'virtual_stain'] if (assemble_on) performed << 'assemble' + if (project_on) performed << 'project' if (track_on) performed << 'track' def layout = step_directories(performed) @@ -188,7 +202,11 @@ workflow { reconstruct_output = "${out}/${layout.reconstruct}/${ds}.zarr" virtual_stain_output = "${out}/${layout.virtual_stain}/${ds}.zarr" assemble_output = assemble_on ? "${out}/${layout.assemble}/${ds}.zarr" : null + project_output = project_on ? "${out}/${layout.project}/${ds}.zarr" : null track_output = track_on ? "${out}/${layout.track}/${ds}.zarr" : null + // What tracking reads: the projected (Z=1) plate when project runs, else the + // assembled plate. Both the plate structure and the pixels come from it. + track_input = project_on ? project_output : assemble_output // ======================================================================== // INIT PHASE — every config parsed, every output store scaffolded, before @@ -250,15 +268,25 @@ workflow { init_signals << as_init.done } + if (project_on) { + // Project reads the ASSEMBLED plate (the phase channel for focus finding, + // the virtual-stain channels to project) and scaffolds the Z=1 plate from + // its metadata, so a project config naming a channel the assembled plate + // lacks fails here. + pj_init = project_init_wf(assemble_output, project_output, params.project_config, + as_init.done) + init_signals << pj_init.done + } + if (track_on) { - // Track reads the ASSEMBLED plate for both of its inputs: assemble - // already carries the phase and virtual-stain channels (concatenate - // preserves channel names, so the track config's channel names resolve - // unchanged), so the plate structure and the image data come from the - // same store. Its init also warms the shared cellpose weights cache, - // which is better done here than with N GPU workers racing for it. - tk_init = track_init_wf(assemble_output, track_output, params.track_config, - as_init.done) + // Track reads ONE plate for both of its inputs (plate structure and image + // data): the projected plate when project runs, else the assembled plate + // (concatenate preserves channel names, so the track config's channel + // names resolve unchanged). Its init checks that the plate's Z matches + // output_mode and warms the shared cellpose weights cache, which is + // better done here than with N GPU workers racing for it. + tk_init = track_init_wf(track_input, track_output, params.track_config, + project_on ? pj_init.done : as_init.done) init_signals << tk_init.done } @@ -352,17 +380,25 @@ workflow { virtual_stain_done.done) } - // Tracking runs AFTER assemble rather than in parallel with it — the - // tradeoff for reading one store instead of two is that the whole plate must - // be assembled first. Two consequences: any Z/Y/X crop or time_indices subset - // in the concatenate config is what tracking sees, and the intermediate - // stores are no longer needed once assemble is verified. To go back to the - // parallel wiring, point the first input at reconstruct_output, the second at - // virtual_stain_output, and gate on virtual_stain_done. + // Projection reads the assembled plate, one task per position, and writes the + // 2D plate tracking reads. + if (project_on) { + project_done = project_run_wf(all_positions, assemble_output, project_output, + params.project_config, pj_init.resources, + assemble_done.done) + } + + // Tracking runs AFTER assemble (and project) rather than in parallel with it — + // the tradeoff for reading one store instead of two is that the whole plate + // must be assembled first. Two consequences: any Z/Y/X crop or time_indices + // subset in the concatenate config is what tracking sees, and the + // intermediate stores are no longer needed once assemble is verified. To go + // back to the parallel wiring, point the first input at reconstruct_output, + // the second at virtual_stain_output, and gate on virtual_stain_done. if (track_on) { - track_done = track_run_wf(all_positions, assemble_output, assemble_output, - track_output, params.track_config, - tk_init.resources, assemble_done.done) + track_done = track_run_wf(all_positions, track_input, track_input, + track_output, params.track_config, tk_init.resources, + project_on ? project_done.done : assemble_done.done) } // ----- QC compute ------------------------------------------------------- @@ -426,11 +462,16 @@ workflow { .collect { zarr -> new File(zarr).parent } def assemble_dir = new File(assemble_output).parent cleanup_paths << "${assemble_dir}/slurm_output" << "${assemble_dir}/.iohub-progress" + if (project_on) { + def project_dir = new File(project_output).parent + cleanup_paths << "${project_dir}/slurm_output" << "${project_dir}/.iohub-progress" + } if (track_on) cleanup_paths << "${new File(track_output).parent}/slurm_output" def cleanup_list = cleanup_targets(cleanup_paths, out) cleanup_record = "${out}/nextflow/intermediates_cleaned.txt" def final_signals = [assemble_done.done] + if (project_on) final_signals << project_done.done if (track_on) final_signals << track_done.done if (qc_on) final_signals << qc_report.done final_gate = channel.empty() @@ -476,6 +517,7 @@ workflow { [label: 'virtual staining', done: virtual_stain_done.done, output: virtual_stain_output], ] if (assemble_on) step_events << [label: 'assemble', done: assemble_done.done, output: assemble_output] + if (project_on) step_events << [label: 'project', done: project_done.done, output: project_output] if (track_on) step_events << [label: 'track', done: track_done.done, output: track_output] if (qc_on) step_events << [label: 'QC', done: qc_report.done, output: qc_report_dir] if (cleanup_on) step_events << [label: 'cleanup intermediates', done: cleanup_run.done, output: cleanup_record] diff --git a/nextflow/modules/project.nf b/nextflow/modules/project.nf new file mode 100644 index 00000000..c1994ff5 --- /dev/null +++ b/nextflow/modules/project.nf @@ -0,0 +1,137 @@ +// Projection subworkflows: init (scaffold + resources) and run (fan-out × N). +// +// This module is PATH-AGNOSTIC. Callers pass the input zarr, output zarr, and +// config explicitly; the module has no idea where it sits in the pipeline +// directory layout. The orchestrating pipeline (see mantis-v2.nf) owns the +// layout and the order of steps; this module just projects whatever 3D plate +// it is handed to a 2D (Z=1) plate. +// +// INIT AND RUN ARE SEPARATE SUBWORKFLOWS. The pipeline runs every step's init +// ahead of every step's compute, so that a bad config fails in the first +// minutes rather than hours in — which means it, not this module, decides when +// each phase happens. See the INIT PHASE comment in mantis-v2.nf. +// +// run_project uses `--cluster debug` so that submitit's DebugExecutor runs the +// work in-process. Nextflow already handles per-position fan-out and resource +// scheduling, so the CLI must NOT submit its own SLURM jobs. See +// modules/flat_field.nf, the template this follows. + +include { parse_resources; slurm_logs; slurm_log_dir; slurm_output_readme; retry_time } from './common' + + +process init_project { + label 'cpu_local' + + input: + val input_zarr + val output_zarr + val config + path config_file // staged only for the task hash: see common.nf, #397 + val trigger + + output: + stdout + + script: + """ + mkdir -p "${slurm_log_dir('project')}" + ${slurm_output_readme('project', output_zarr)} + biahub project --init \ + -i "${input_zarr}"/*/*/* \ + -o "${output_zarr}" \ + -c "${config}" + """ +} + +process run_project { + tag "${position}" + label 'cpu' + clusterOptions { slurm_logs('project') } + cpus { meta.cpus } + memory { "${meta.mem_gb} GB" } + time { retry_time(meta.time_minutes, task) } + + input: + tuple val(position), val(meta) + val input_zarr + val output_zarr + val config + path config_file // staged only for the task hash: see common.nf, #397 + + output: + val position + + script: + // --resume: a preempted task finishes the write it is in and stops early, so + // the retry (or a later `nextflow -resume`) recomputes only the (t, c) units + // this position had not finished. The focus pre-pass runs again either way. + """ + biahub project --cluster debug --resume \ + -i "${input_zarr}/${position}" \ + -o "${output_zarr}" \ + -c "${config}" + """ +} + + +// Validate the config and scaffold the output plate. Metadata-only and cheap, +// so it belongs in the pipeline's up-front init phase. +// +// take: +// input_zarr path to the input (3D) plate.zarr +// output_zarr path to the projected (Z=1) output plate.zarr +// config path to the project settings YAML +// trigger gating channel — init starts once this emits +// emit: +// resources the RESOURCES payload sizing one position's task +// done fires once the output plate exists +workflow project_init_wf { + take: + input_zarr + output_zarr + config + trigger + + main: + init_out = init_project(input_zarr, output_zarr, config, file(config), trigger.collect().map { 'done' }) + + emit: + // Value channels by the contract every step module emits on; see + // flat_field_init_wf for why the guarantee is made at the input. + resources = init_out.map { stdout_text -> parse_resources(stdout_text) } + done = init_out.map { 'done' } +} + + +// Fan out one projection task per position. +// +// take: +// positions collected channel of position keys (e.g. ['A/1/0', 'B/1/0']) +// input_zarr path to the input (3D) plate.zarr +// output_zarr path to the projected (Z=1) output plate.zarr +// config path to the project settings YAML +// resources RESOURCES payload from project_init_wf +// prev_done gating channel — compute starts once this emits +workflow project_run_wf { + take: + positions + input_zarr + output_zarr + config + resources + prev_done + + main: + // Gate channels are mapped to a token before combining, never combined raw; + // see flat_field_run_wf. + pos_meta = positions + .flatMap { items -> items } + .combine(resources) + .combine(prev_done.map { 'done' }) + .map { pos, meta, _gate -> [pos, meta] } + + pj_done = run_project(pos_meta, input_zarr, output_zarr, config, file(config)) | collect + + emit: + done = pj_done +} diff --git a/nextflow/modules/segment.nf b/nextflow/modules/segment.nf index 5867b7c1..f7ed7fdf 100644 --- a/nextflow/modules/segment.nf +++ b/nextflow/modules/segment.nf @@ -2,13 +2,13 @@ // (fan-out × N positions), the same split as tracking. // // This module is PATH-AGNOSTIC. Callers pass the input zarr, output zarr and -// config explicitly. mantis-v2.nf passes the assembled plate, which carries the -// virtual-stain channels the segment config names. +// config explicitly. For 2D models the input is the PROJECTED (Z=1) plate that +// `biahub project` writes from the assembled plate; segment does not project. // // `biahub segment --init` reads only the input plate's channels, shape and scale, -// so it runs as soon as the assembled plate is scaffolded: a config naming a -// channel the plate lacks, a z_slice_2D outside the stack, or an unknown cellpose -// model fails before any GPU job. It also warms the shared cellpose weights +// so it runs as soon as the input plate is scaffolded: a config naming a channel +// the plate lacks, 2D models on a 3D plate (or 3D models on a Z=1 plate), or an +// unknown cellpose model fails before any GPU job. It also warms the shared cellpose weights // cache, better done once here than by N GPU workers racing for it. include { parse_resources; slurm_logs; slurm_log_dir; slurm_output_readme; retry_time } from './common' diff --git a/settings/example_project_settings.yml b/settings/example_project_settings.yml new file mode 100644 index 00000000..b4af3299 --- /dev/null +++ b/settings/example_project_settings.yml @@ -0,0 +1,60 @@ +# biahub project: 3D plate -> 2D plate (Z=1). One entry per OUTPUT channel, in output +# order; the key is the input channel's name. Input channels not listed here are dropped +# (the command prints which). +# +# method: max | mean | sum | slice +# max/mean/sum reduce the selected planes; mean and sum write float32, max and slice +# keep the input dtype. slice keeps one plane. +# z: which planes feed the projection. Omit it for max/mean/sum over the whole stack. +# range: [start, stop) half-open, like Z_slice in concatenate. With `center` the bounds +# are offsets from the center. +# step: keep every step-th plane of the range (default 1) +# center: a plane index, or "focus" for the in-focus plane found on focus.channel. +# slice takes center alone; max/mean/sum take center with a range. +# A range that runs past the stack is clipped to it. + +channels: + # max over the whole stack + membrane_prediction: + method: max + + # max over an explicit range, every 2nd plane: z = 10, 12, ..., 48 + nuclei_prediction: + method: max + z: {range: [10, 50], step: 2} + + # sum over 21 planes around a manually chosen center: z = 20 .. 40 + GFP: + method: sum + z: {center: 30, range: [-10, 11]} + + # mean over 60 planes starting 20 below the in-focus plane: z = focus-20 .. focus+39 + RFP: + method: mean + z: {center: focus, range: [-20, 40]} + + # single plane, chosen manually + mCherry: + method: slice + z: {center: 42} + + # single plane, the in-focus plane + Phase3D: + method: slice + z: {center: focus} + +# Required when any channel uses `center: focus`. The focus criterion is the registration +# engine's (FocusSettings): waveorder's transverse band on a centre crop, with the +# detection NA and illumination wavelength in micrometres; the pixel size comes from the +# store. The plane of each timepoint is recorded in the output position's +# `biahub-project` metadata. +focus: + channel: Phase3D # input channel the focus is found on (need not be an output) + per: fov # fov: one plane per position, the median over timepoints + # timepoint: one plane per timepoint, so the window follows Z drift + center_crop_xy: [800, 800] # (X, Y) centre crop the focus is measured on + na_det: 1.35 + lambda_ill: 0.5 + +# When null, keep the OME-Zarr version of the input store. +output_ome_zarr_version: null diff --git a/settings/example_segmentation_settings.yml b/settings/example_segmentation_settings.yml index 15053759..5eaaf0b5 100644 --- a/settings/example_segmentation_settings.yml +++ b/settings/example_segmentation_settings.yml @@ -1,12 +1,13 @@ # Segmentation settings for cellpose 4 (`biahub segment`). # Model names become the output channel names (e.g. membrane -> membrane_labels). -# All models in one config must be either 2D (z_slice_2D set) or 3D (z_slice_2D: null with -# eval_args.do_3D: true), because they share one output array. Run 3D models separately. +# A model segments in 2D unless eval_args.do_3D (or a stitch_threshold) is set. Segment +# does not project: 2D models need a Z=1 input plate (make one from a stack with +# `biahub project`), 3D models a stack; a mismatch is an error. All models in one config +# must be either 2D or 3D, because they share one input and output. Run 3D models separately. models: membrane: pretrained_model: cpsam_v2 # cellpose 4 model (cpsam_v2, cpsam, cpdino, cpdino-vitb) or path to weights channels: ['mem', 'nuc'] # dataset channel names passed to cellpose (at most 3, any order) - z_slice_2D: 10 # plane for 2D segmentation; null = 3D (needs eval_args.do_3D: true) # Other CellposeModel.eval() arguments; see https://cellpose.readthedocs.io/en/latest/api.html eval_args: diameter: 65 # cellpose 4 resizes by 30 / diameter; null = native resolution @@ -23,7 +24,6 @@ models: nucleus: pretrained_model: cpsam_v2 channels: ['nuc'] - z_slice_2D: 10 eval_args: diameter: 60 cellprob_threshold: 0.0 diff --git a/settings/example_track_settings.yml b/settings/example_track_settings.yml index c34499b8..b7a8df46 100644 --- a/settings/example_track_settings.yml +++ b/settings/example_track_settings.yml @@ -1,26 +1,12 @@ -# Output layout. "2D" writes a Z=1 plate (the input must be projected, e.g. via the -# np.mean step below); "3D" keeps the selected z-window. output_mode does NOT project -# the data itself -- it only controls the output shape. +# The dimensionality of the tracking. "2D" tracks a Z=1 input plate and writes a Z=1 +# label plate; "3D" tracks the stack and writes labels of the same Z. Track does NOT +# project or select Z-planes: a 3D input with output_mode 2D is an error. To track a +# projection, a window or one plane of a stack, make the Z=1 plate first with +# `biahub project` (settings/example_project_settings.yml) and track that. output_mode: "2D" fov: "*/*/*" -# How to SELECT the Z-planes used for tracking. Set exactly one `method`; only the -# fields relevant to that method apply (others are rejected at load time): -# all - use every plane. -# central - auto-centered window. -# range - explicit [start, stop] (requires `range`). -# focus - fixed window of `window_size` planes around the per-FOV in-focus plane -# (`frac_below`/`frac_above` split it; `focus_channel` picks the channel -# focus-finding runs on). -z_slicing: - method: all - # range: [10, 25] # method: range - # window_size: 48 # method: focus (number of z-planes in the window; default 48) - # frac_below: 0.33 # method: focus - # frac_above: 0.67 # method: focus - # focus_channel: nuclei_prediction # method: focus - # Name of the target channel used to name the tracking output label channel. target_channel: nuclei_prediction blank_frames_path: blank_frames.csv @@ -59,17 +45,13 @@ cellpose_config: # virtual/computed channels (see the foreground_contour example below). input_images: - # Cellpose mode: load and Z-project the channel named by cellpose_config.input_channel. + # Cellpose mode: load the channel named by cellpose_config.input_channel (already a + # single plane in a projected plate; no processing needed). # path: null -> resolved from -i (or --input-images-path). Use an explicit - # /path/to/virtual_staining.zarr if the data lives elsewhere than the -i plate. + # /path/to/projected.zarr if the data lives elsewhere than the -i plate. - path: null channels: - nuclei_prediction: - # Z-projection using mean intensity - - function: np.mean - kwargs: - axis: 1 - per_timepoint: False + nuclei_prediction: [] # --------------------------------------------------------------------------- # Alternative: segmentation_method: foreground_contour @@ -79,18 +61,10 @@ input_images: # - path: null # channels: # nuclei_prediction: - # - function: np.mean - # kwargs: - # axis: 1 - # per_timepoint: False # - function: ultrack.imgproc.normalize # kwargs: # gamma: 0.7 # membrane_prediction: - # - function: np.mean - # kwargs: - # axis: 1 - # per_timepoint: False # - function: ultrack.imgproc.normalize # kwargs: # gamma: 0.7 diff --git a/tests/test_cellpose_inputs.py b/tests/test_cellpose_inputs.py index b9818881..4ccdabdc 100644 --- a/tests/test_cellpose_inputs.py +++ b/tests/test_cellpose_inputs.py @@ -30,7 +30,6 @@ def test_segment_uses_the_requested_model(fake_cellpose): channels=[0], pretrained_model="cpdino", eval_args={"diameter": None, "do_3D": False}, - z_slice_2D=0, ) assert fake_cellpose.built == ["cpdino"] @@ -46,7 +45,6 @@ def test_segment_passes_only_the_configured_channel(fake_cellpose): channels=[2], pretrained_model="cpsam_v2", eval_args={"diameter": None, "do_3D": False}, - z_slice_2D=0, ) (seen,) = fake_cellpose.evaluated @@ -77,16 +75,14 @@ def test_model_type_in_a_track_config_warns(): CellposeConfig() -def test_segment_slices_the_configured_plane(fake_cellpose): - # Channel 0 holds the plane index (x10) so the plane cellpose receives is visible. - czyx = np.stack([np.full((1, 8, 8), 10 * z, dtype=np.float32) for z in range(6)], axis=1) +def test_segment_2d_passes_the_single_plane_without_a_z_axis(fake_cellpose): + czyx = np.full((1, 1, 8, 8), 30, dtype=np.float32) _segment( czyx, channels=[0], pretrained_model="cpsam_v2", eval_args={"diameter": None, "do_3D": False}, - z_slice_2D=3, ) (seen,) = fake_cellpose.evaluated @@ -100,11 +96,7 @@ def test_mixing_2d_and_3d_models_is_refused(fake_cellpose): with pytest.raises(ValueError, match="2D and 3D"): SegmentationSettings( models={ - "flat": { - "pretrained_model": "cpsam_v2", - "channels": ["GFP"], - "z_slice_2D": 1, - }, + "flat": {"pretrained_model": "cpsam_v2", "channels": ["GFP"]}, "volume": { "pretrained_model": "cpsam_v2", "channels": ["RFP"], @@ -114,7 +106,8 @@ def test_mixing_2d_and_3d_models_is_refused(fake_cellpose): ) -def test_segment_cli_rejects_a_plane_outside_the_stack(fake_cellpose, example_plate, tmp_path): +def test_segment_cli_rejects_a_stack_for_a_2d_model(fake_cellpose, example_plate, tmp_path): + """Segment does not project: a 2D model needs the Z=1 plate `biahub project` writes.""" import yaml from click.testing import CliRunner @@ -131,7 +124,6 @@ def test_segment_cli_rejects_a_plane_outside_the_stack(fake_cellpose, example_pl "pretrained_model": "cpsam_v2", "channels": ["GFP"], "eval_args": {"do_3D": False}, - "z_slice_2D": 10, } } } @@ -140,7 +132,7 @@ def test_segment_cli_rejects_a_plane_outside_the_stack(fake_cellpose, example_pl result = CliRunner().invoke( cli, - # --init: the range check runs while creating the plate, before any job. + # --init: the dimensionality check runs while creating the plate, before any job. [ "segment", "-i", @@ -154,7 +146,7 @@ def test_segment_cli_rejects_a_plane_outside_the_stack(fake_cellpose, example_pl ) assert result.exit_code != 0 - assert "z_slice_2D" in str(result.exception) + result.output + assert "biahub project" in str(result.exception) + result.output assert not (tmp_path / "out.zarr").exists() @@ -167,12 +159,16 @@ def test_segment_3d_passes_the_z_stack(fake_cellpose): assert seen.shape == (1, 5, 8, 8) -def test_3d_model_without_do_3d_is_refused(fake_cellpose): - """Cellpose 4 would raise on a z stack with do_3D False and no stitching.""" - from biahub.settings import SegmentationModel - - with pytest.raises(ValueError, match="do_3D"): - SegmentationModel(pretrained_model="cpsam_v2", channels=["GFP"]) +def test_a_2d_model_is_refused_a_z_stack(fake_cellpose): + """Cellpose 4 would raise on a z stack with do_3D False and no stitching, so + resolve_models refuses it first (segment does not pick a plane; project does).""" + with pytest.raises(ValueError, match="input is 3D.*biahub project"): + _segment( + np.zeros((1, 5, 8, 8), dtype=np.float32), + channels=[0], + pretrained_model="cpsam_v2", + eval_args={"do_3D": False}, + ) def test_segment_loads_each_model_once_across_timepoints(fake_cellpose): @@ -182,7 +178,7 @@ def test_segment_loads_each_model_once_across_timepoints(fake_cellpose): from biahub.settings import SegmentationSettings settings = SegmentationSettings( - models={"nuc": {"pretrained_model": "cpdino", "channels": ["c0"], "z_slice_2D": 0}} + models={"nuc": {"pretrained_model": "cpdino", "channels": ["c0"]}} ) models = resolve_models(settings, ["c0"], scale=(1,) * 5, z_size=1) for _ in range(3): @@ -192,12 +188,22 @@ def test_segment_loads_each_model_once_across_timepoints(fake_cellpose): assert len(fake_cellpose.evaluated) == 3 -def test_segment_3d_keeps_z_on_a_single_plane_input(fake_cellpose): - """Cellpose squeezes singleton axes, so a 3D model on Z=1 returned (Y, X).""" +def test_segment_2d_restores_the_singleton_z(fake_cellpose): + """Cellpose returns (Y, X) for a 2D model; the label plate expects (1, Y, X).""" out = _segment( np.zeros((1, 1, 8, 8), dtype=np.float32), channels=[0], pretrained_model="cpsam_v2", - eval_args={"do_3D": True}, + eval_args={"do_3D": False}, ) assert out.shape == (1, 1, 8, 8) + + +def test_a_3d_model_is_refused_a_single_plane(fake_cellpose): + with pytest.raises(ValueError, match="one Z-plane but the models segment in 3D"): + _segment( + np.zeros((1, 1, 8, 8), dtype=np.float32), + channels=[0], + pretrained_model="cpsam_v2", + eval_args={"do_3D": True}, + ) diff --git a/tests/test_cli/test_cli.py b/tests/test_cli/test_cli.py index 04460918..848666ea 100644 --- a/tests/test_cli/test_cli.py +++ b/tests/test_cli/test_cli.py @@ -29,6 +29,7 @@ def test_main(): "deconvolve", "characterize-psf", "segment", + "project", "virtual-stain", "track", "process-with-config", diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py new file mode 100644 index 00000000..a337ff70 --- /dev/null +++ b/tests/test_cli/test_project_cli.py @@ -0,0 +1,242 @@ +"""biahub project: 3D plate -> 2D plate, one projection method per channel.""" + +import json + +import numpy as np +import pytest +import yaml + +from click.testing import CliRunner +from iohub.ngff import open_ome_zarr + +from biahub.cli.main import cli + +Z = 8 + + +@pytest.fixture +def plate(create_custom_plate, tmp_path): + """(T=3, C=3, Z=8, Y=5, X=6) uint16, channels GFP, RFP, Phase3D.""" + return create_custom_plate(tmp_path, z_size=Z) + + +def _config(tmp_path, channels, focus=None, **extra): + config = {"channels": channels, **extra} + if focus is not None: + config["focus"] = focus + path = tmp_path / "project.yml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) # channel order is output order + return path + + +def _run(*args): + return CliRunner().invoke(cli, ["project", *map(str, args)]) + + +def _read(position_path): + with open_ome_zarr(str(position_path), mode="r") as ds: + return np.asarray(ds.data), ds.channel_names, dict(ds.zattrs) + + +def _fake_focus(monkeypatch, planes): + """Make find_focus return the given plane for each timepoint, in order.""" + it = iter(planes) + monkeypatch.setattr( + "biahub.registration.methods.focus.find_focus", lambda *a, **k: next(it) + ) + + +def test_project_init_scaffolds_a_z1_plate_and_names_dropped_channels(plate, tmp_path): + plate_path, _ = plate + out = tmp_path / "proj.zarr" + config = _config( + tmp_path, {"RFP": {"method": "max"}, "GFP": {"method": "slice", "z": {"center": 2}}} + ) + + result = _run("--init", "-i", plate_path / "A" / "1" / "0", "-o", out, "-c", config) + + assert result.exit_code == 0, result.output + line = next(ln for ln in result.output.splitlines() if ln.startswith("RESOURCES:")) + assert set(json.loads(line.removeprefix("RESOURCES:"))) == { + "cpus", + "mem_gb", + "time_minutes", + } + assert "dropped: ['Phase3D']" in result.output + data, channel_names, zattrs = _read(out / "A" / "1" / "0") + assert channel_names == ["RFP", "GFP"] # config order, not input order + assert data.shape == (3, 2, 1, 5, 6) + assert data.dtype == np.uint16 # max and slice keep the input dtype + assert zattrs["biahub-project"]["channels"]["GFP"]["z"]["center"] == 2 + assert np.all(data == 0) # --init projects nothing + + +def test_project_methods_match_numpy(plate, tmp_path): + plate_path, _ = plate + out = tmp_path / "proj.zarr" + config = _config( + tmp_path, + { + "GFP": {"method": "max", "z": {"range": [2, 7], "step": 2}}, + "RFP": {"method": "mean"}, + "Phase3D": {"method": "slice", "z": {"center": 3}}, + }, + ) + position = plate_path / "B" / "1" / "0" + + result = _run("--cluster", "debug", "-i", position, "-o", out, "-c", config) + + assert result.exit_code == 0, result.output + assert "Projection complete" in result.output + source, _, _ = _read(position) + data, channel_names, _ = _read(out / "B" / "1" / "0") + assert channel_names == ["GFP", "RFP", "Phase3D"] + assert data.dtype == np.float32 # mean promotes the whole plate + assert np.array_equal(data[:, 0, 0], source[:, 0, 2:7:2].max(axis=1)) + np.testing.assert_allclose(data[:, 1, 0], source[:, 1].mean(axis=1, dtype=np.float64)) + assert np.array_equal(data[:, 2, 0], source[:, 2, 3]) + + +def test_project_sum_over_a_window_around_a_center(plate, tmp_path): + plate_path, _ = plate + out = tmp_path / "proj.zarr" + config = _config( + tmp_path, {"GFP": {"method": "sum", "z": {"center": 6, "range": [-2, 5]}}} + ) + position = plate_path / "A" / "1" / "0" + + result = _run("--cluster", "debug", "-i", position, "-o", out, "-c", config) + + assert result.exit_code == 0, result.output + source, _, _ = _read(position) + data, _, _ = _read(out / "A" / "1" / "0") + # 4..10 clipped to the 8-plane stack: 4..7 + np.testing.assert_allclose(data[:, 0, 0], source[:, 0, 4:8].sum(axis=1, dtype=np.float64)) + + +def test_project_focus_per_fov_takes_the_median_plane(plate, tmp_path, monkeypatch): + plate_path, _ = plate + out = tmp_path / "proj.zarr" + config = _config( + tmp_path, + { + "GFP": {"method": "mean", "z": {"center": "focus", "range": [-1, 2]}}, + "Phase3D": {"method": "slice", "z": {"center": "focus"}}, + }, + focus={"channel": "Phase3D", "per": "fov"}, + ) + position = plate_path / "A" / "1" / "0" + _fake_focus(monkeypatch, [1, 6, 2]) # one per timepoint; median 2 + + result = _run("--cluster", "debug", "-i", position, "-o", out, "-c", config) + + assert result.exit_code == 0, result.output + source, _, _ = _read(position) + data, _, zattrs = _read(out / "A" / "1" / "0") + np.testing.assert_allclose(data[:, 0, 0], source[:, 0, 1:4].mean(axis=1, dtype=np.float64)) + assert np.array_equal(data[:, 1, 0], source[:, 2, 2]) + assert zattrs["biahub-project"]["focus"] == { + "channel": "Phase3D", + "per": "fov", + "focus_z": 2, + "focus_z_per_t": [1, 6, 2], + } + + +def test_project_focus_per_timepoint_follows_the_plane(plate, tmp_path, monkeypatch): + plate_path, _ = plate + out = tmp_path / "proj.zarr" + config = _config( + tmp_path, + { + "GFP": {"method": "max", "z": {"center": "focus", "range": [-1, 2]}}, + "Phase3D": {"method": "slice", "z": {"center": "focus"}}, + }, + focus={"channel": "Phase3D", "per": "timepoint"}, + ) + position = plate_path / "A" / "1" / "0" + planes = [1, 6, 2] + _fake_focus(monkeypatch, planes) + + result = _run("--cluster", "debug", "-i", position, "-o", out, "-c", config) + + assert result.exit_code == 0, result.output + source, _, _ = _read(position) + data, _, zattrs = _read(out / "A" / "1" / "0") + for t, plane in enumerate(planes): + assert np.array_equal(data[t, 0, 0], source[t, 0, plane - 1 : plane + 2].max(axis=0)) + assert np.array_equal(data[t, 1, 0], source[t, 2, plane]) + assert zattrs["biahub-project"]["focus"]["focus_z"] is None + assert zattrs["biahub-project"]["focus"]["focus_z_per_t"] == planes + + +def test_project_focus_channel_need_not_be_an_output(plate, tmp_path, monkeypatch): + plate_path, _ = plate + out = tmp_path / "proj.zarr" + config = _config( + tmp_path, + {"GFP": {"method": "slice", "z": {"center": "focus"}}}, + focus={"channel": "Phase3D"}, + ) + _fake_focus(monkeypatch, [5, 5, 5]) + + result = _run("--cluster", "debug", "-i", plate_path / "A/1/0", "-o", out, "-c", config) + + assert result.exit_code == 0, result.output + source, _, _ = _read(plate_path / "A/1/0") + data, channel_names, _ = _read(out / "A" / "1" / "0") + assert channel_names == ["GFP"] + assert np.array_equal(data[:, 0, 0], source[:, 0, 5]) + + +@pytest.mark.parametrize( + "channels, message", + [ + ({"GFP": {"method": "max", "z": {"range": [10, 20]}}}, "select nothing"), + ({"GFP": {"method": "slice", "z": {"center": 42}}}, "select nothing"), + ({"DAPI": {"method": "max"}}, "not in the input channels"), + ], +) +def test_project_init_fails_on_a_config_the_plate_cannot_satisfy( + plate, tmp_path, channels, message +): + plate_path, _ = plate + out = tmp_path / "proj.zarr" + + result = _run( + "--init", "-i", plate_path / "A/1/0", "-o", out, "-c", _config(tmp_path, channels) + ) + + assert result.exit_code != 0 + assert message in str(result.exception) + assert not out.exists() + + +def test_project_refuses_a_plate_that_is_already_2d(create_custom_plate, tmp_path): + plate_path, _ = create_custom_plate(tmp_path, z_size=1) + config = _config(tmp_path, {"GFP": {"method": "max"}}) + + result = _run( + "--init", "-i", plate_path / "A/1/0", "-o", tmp_path / "o.zarr", "-c", config + ) + + assert result.exit_code != 0 + assert "already has one Z-plane" in str(result.exception) + + +def test_project_resume_skips_finished_timepoints(plate, tmp_path): + plate_path, _ = plate + out = tmp_path / "proj.zarr" + config = _config( + tmp_path, {"GFP": {"method": "max"}}, output_ome_zarr_version="0.5" + ) # resume needs Zarr v3 + position = plate_path / "A" / "1" / "0" + assert _run("--cluster", "debug", "-i", position, "-o", out, "-c", config).exit_code == 0 + before, _, _ = _read(out / "A" / "1" / "0") + + result = _run("--cluster", "debug", "--resume", "-i", position, "-o", out, "-c", config) + + assert result.exit_code == 0, result.output + assert "already written" in result.output + after, _, _ = _read(out / "A" / "1" / "0") + assert np.array_equal(before, after) diff --git a/tests/test_cli/test_segment_cli.py b/tests/test_cli/test_segment_cli.py index 20b44cf8..8a0c1449 100644 --- a/tests/test_cli/test_segment_cli.py +++ b/tests/test_cli/test_segment_cli.py @@ -21,7 +21,6 @@ def _write(pretrained_model="cpsam_v2", version=None): "nuc": { "pretrained_model": pretrained_model, "channels": ["GFP"], - "z_slice_2D": 1, "eval_args": {"diameter": None}, } } @@ -35,6 +34,12 @@ def _write(pretrained_model="cpsam_v2", version=None): return _write +@pytest.fixture +def plate_2d(create_custom_plate, tmp_path): + """A projected (Z=1) plate, what 2D models segment: (T=3, C=3, Z=1, Y=5, X=6).""" + return create_custom_plate(tmp_path, z_size=1) + + @pytest.fixture def cpu_cellpose(monkeypatch): """Workers ask for the GPU; run the (fake) model on the CPU instead.""" @@ -46,9 +51,9 @@ def _run(*args): def test_segment_init_prints_resources_and_stamps_provenance( - fake_cellpose, example_plate, segment_config, tmp_path + fake_cellpose, plate_2d, segment_config, tmp_path ): - plate_path, _ = example_plate # (T=3, C=6, Z=4, Y=5, X=6) + plate_path, _ = plate_2d # (T=3, C=3, Z=1, Y=5, X=6) out = tmp_path / "seg.zarr" result = _run( @@ -73,9 +78,9 @@ def test_segment_init_prints_resources_and_stamps_provenance( def test_segment_init_rejects_an_unknown_model( - fake_cellpose, example_plate, segment_config, tmp_path + fake_cellpose, plate_2d, segment_config, tmp_path ): - plate_path, _ = example_plate + plate_path, _ = plate_2d result = _run( "--init", "-i", @@ -90,9 +95,9 @@ def test_segment_init_rejects_an_unknown_model( def test_segment_debug_runs_one_position_in_process( - fake_cellpose, cpu_cellpose, example_plate, segment_config, tmp_path + fake_cellpose, cpu_cellpose, plate_2d, segment_config, tmp_path ): - plate_path, _ = example_plate + plate_path, _ = plate_2d out, config = tmp_path / "seg.zarr", segment_config() position = plate_path / "B" / "1" / "0" assert _run("--init", "-i", position, "-o", out, "-c", config).exit_code == 0 @@ -109,9 +114,9 @@ def test_segment_debug_runs_one_position_in_process( def test_segment_resume_skips_finished_work( - fake_cellpose, cpu_cellpose, example_plate, segment_config, tmp_path + fake_cellpose, cpu_cellpose, plate_2d, segment_config, tmp_path ): - plate_path, _ = example_plate + plate_path, _ = plate_2d out, config = tmp_path / "seg.zarr", segment_config(version="0.5") # resume needs Zarr v3 position = plate_path / "B" / "1" / "0" assert _run("--init", "-i", position, "-o", out, "-c", config).exit_code == 0 @@ -124,8 +129,8 @@ def test_segment_resume_skips_finished_work( assert len(fake_cellpose.evaluated) == n_first -def test_segment_legacy_local_flag_is_gone(example_plate, segment_config, tmp_path): - plate_path, _ = example_plate +def test_segment_legacy_local_flag_is_gone(plate_2d, segment_config, tmp_path): + plate_path, _ = plate_2d result = _run( "--local", "-i", @@ -167,13 +172,13 @@ def wait(self): @pytest.mark.parametrize("cluster, expected", [("local", 1), ("slurm", None), ("debug", None)]) def test_segment_gives_local_jobs_a_gpu( - fake_cellpose, example_plate, segment_config, tmp_path, monkeypatch, cluster, expected + fake_cellpose, plate_2d, segment_config, tmp_path, monkeypatch, cluster, expected ): """submitit's local executor only exposes GPUs through gpus_per_node; without it the local subprocess runs with CUDA_VISIBLE_DEVICES='' and cellpose_device refuses.""" monkeypatch.setattr("biahub.segment.submitit.AutoExecutor", _RecordingExecutor) monkeypatch.delenv("CI", raising=False) # CI forces cluster=debug - plate_path, _ = example_plate + plate_path, _ = plate_2d out = tmp_path / "seg.zarr" result = _run( @@ -193,7 +198,7 @@ def test_segment_gives_local_jobs_a_gpu( def test_segment_worker_stages_weights_before_cellpose_is_imported( - fake_cellpose, cpu_cellpose, example_plate, segment_config, tmp_path, monkeypatch + fake_cellpose, cpu_cellpose, plate_2d, segment_config, tmp_path, monkeypatch ): """Validating eval_args imports cellpose.models, which fixes the weights directory; staging after that is a no-op and every worker reads the checkpoint over NFS.""" @@ -210,7 +215,7 @@ def test_segment_worker_stages_weights_before_cellpose_is_imported( "get_valid_eval_args", lambda: calls.append("import cellpose") or real_valid_args(), ) - plate_path, _ = example_plate + plate_path, _ = plate_2d out, config = tmp_path / "seg.zarr", segment_config() position = plate_path / "B" / "1" / "0" assert _run("--init", "-i", position, "-o", out, "-c", config).exit_code == 0 @@ -226,15 +231,20 @@ def _resources(output): return json.loads(line.removeprefix("RESOURCES:")) -def test_segment_requests_less_time_for_2d_than_3d(fake_cellpose, example_plate, tmp_path): +def test_segment_requests_less_time_for_2d_than_3d( + fake_cellpose, plate_2d, example_plate, tmp_path +): """Calibrated on a real 2D run (~5 s/frame, 6 GB); 3D keeps the conservative rate.""" - plate_path, _ = example_plate - position = plate_path / "A" / "1" / "0" + positions = { + "2d": plate_2d[0] / "A" / "1" / "0", # Z=1 + "3d": example_plate[0] / "A" / "1" / "0", # Z=4 + } res = {} for kind, model in { - "2d": {"z_slice_2D": 1, "eval_args": {}}, + "2d": {"eval_args": {}}, "3d": {"eval_args": {"do_3D": True}}, }.items(): + position = positions[kind] cfg = tmp_path / f"{kind}.yml" cfg.write_text( yaml.safe_dump( @@ -253,7 +263,7 @@ def test_segment_requests_less_time_for_2d_than_3d(fake_cellpose, example_plate, def test_segment_slurm_head_does_not_stage_weights( - fake_cellpose, example_plate, segment_config, tmp_path, monkeypatch + fake_cellpose, plate_2d, segment_config, tmp_path, monkeypatch ): """Staging on the submitting node would export its CELLPOSE_LOCAL_MODELS_PATH to every SLURM worker, which then skips its own staging. Only in-process workers stage early.""" @@ -261,7 +271,7 @@ def test_segment_slurm_head_does_not_stage_weights( monkeypatch.setattr("biahub.segment.stage_cellpose_weights", lambda: staged.append(1)) monkeypatch.setattr("biahub.segment.submitit.AutoExecutor", _RecordingExecutor) monkeypatch.delenv("CI", raising=False) - plate_path, _ = example_plate + plate_path, _ = plate_2d result = _run( "--cluster", @@ -279,12 +289,12 @@ def test_segment_slurm_head_does_not_stage_weights( def test_segment_local_jobs_get_the_step_time_limit( - fake_cellpose, example_plate, segment_config, tmp_path, monkeypatch + fake_cellpose, plate_2d, segment_config, tmp_path, monkeypatch ): """submitit's local executor kills jobs after timeout_min (default 2 min).""" monkeypatch.setattr("biahub.segment.submitit.AutoExecutor", _RecordingExecutor) monkeypatch.delenv("CI", raising=False) - plate_path, _ = example_plate + plate_path, _ = plate_2d result = _run( "--cluster", @@ -305,9 +315,9 @@ def test_segment_local_jobs_get_the_step_time_limit( def test_segment_init_checks_models_before_creating_the_plate( - fake_cellpose, example_plate, segment_config, tmp_path + fake_cellpose, plate_2d, segment_config, tmp_path ): - plate_path, _ = example_plate + plate_path, _ = plate_2d out = tmp_path / "seg.zarr" result = _run( "--init", @@ -323,7 +333,7 @@ def test_segment_init_checks_models_before_creating_the_plate( def test_segment_local_sbatch_file_overrides_computed_parameters( - fake_cellpose, example_plate, segment_config, tmp_path, monkeypatch + fake_cellpose, plate_2d, segment_config, tmp_path, monkeypatch ): """A '#LOCAL --timeout-min' line and the computed local timeout used to collide (TypeError: multiple values for keyword argument 'timeout_min').""" @@ -331,7 +341,7 @@ def test_segment_local_sbatch_file_overrides_computed_parameters( monkeypatch.delenv("CI", raising=False) # conftest forces CI (debug) otherwise sbatch = tmp_path / "local.sbatch" sbatch.write_text("#LOCAL --timeout-min=5\n") - plate_path, _ = example_plate + plate_path, _ = plate_2d result = _run( "--cluster", @@ -353,8 +363,8 @@ def test_segment_local_sbatch_file_overrides_computed_parameters( def _two_model_config(path, order): models = { - "nuc": {"pretrained_model": "cpsam_v2", "channels": ["GFP"], "z_slice_2D": 1}, - "mem": {"pretrained_model": "cpsam_v2", "channels": ["RFP"], "z_slice_2D": 1}, + "nuc": {"pretrained_model": "cpsam_v2", "channels": ["GFP"]}, + "mem": {"pretrained_model": "cpsam_v2", "channels": ["RFP"]}, } config = { "models": {name: models[name] for name in order}, @@ -365,10 +375,10 @@ def _two_model_config(path, order): def test_segment_refuses_an_output_with_a_different_model_order( - fake_cellpose, cpu_cellpose, example_plate, tmp_path + fake_cellpose, cpu_cellpose, plate_2d, tmp_path ): """Reordering the models used to resume into swapped label channels.""" - plate_path, _ = example_plate + plate_path, _ = plate_2d out, position = tmp_path / "seg.zarr", plate_path / "B" / "1" / "0" first = _two_model_config(tmp_path / "first.yml", ["nuc", "mem"]) assert _run("--init", "-i", position, "-o", out, "-c", first).exit_code == 0 @@ -402,7 +412,7 @@ def test_segment_resume_token_depends_on_model_order(tmp_path): def test_segment_mem_per_cpu_override_drops_the_computed_total( - fake_cellpose, example_plate, segment_config, tmp_path, monkeypatch + fake_cellpose, plate_2d, segment_config, tmp_path, monkeypatch ): """The --sbatch-filepath help's own example, '#SBATCH --mem-per-cpu=16G', left --mem set too, and SLURM rejects jobs with both.""" @@ -410,7 +420,7 @@ def test_segment_mem_per_cpu_override_drops_the_computed_total( monkeypatch.delenv("CI", raising=False) # conftest forces CI (debug) otherwise sbatch = tmp_path / "mem.sbatch" sbatch.write_text("#SBATCH --mem-per-cpu=16G\n") - plate_path, _ = example_plate + plate_path, _ = plate_2d result = _run( "--cluster", @@ -431,21 +441,21 @@ def test_segment_mem_per_cpu_override_drops_the_computed_total( def test_segment_refuses_an_output_with_a_different_shape( - fake_cellpose, example_plate, tmp_path + fake_cellpose, plate_2d, example_plate, tmp_path ): """Switching a same-named model from 3D to 2D kept the old multi-plane array.""" - plate_path, _ = example_plate - out, position = tmp_path / "seg.zarr", plate_path / "B" / "1" / "0" + out = tmp_path / "seg.zarr" model = {"pretrained_model": "cpsam_v2", "channels": ["GFP"]} three_d = tmp_path / "3d.yml" three_d.write_text( yaml.safe_dump({"models": {"nuc": {**model, "eval_args": {"do_3D": True}}}}) ) two_d = tmp_path / "2d.yml" - two_d.write_text(yaml.safe_dump({"models": {"nuc": {**model, "z_slice_2D": 1}}})) - assert _run("--init", "-i", position, "-o", out, "-c", three_d).exit_code == 0 + two_d.write_text(yaml.safe_dump({"models": {"nuc": model}})) + stack = example_plate[0] / "B" / "1" / "0" # Z=4 + assert _run("--init", "-i", stack, "-o", out, "-c", three_d).exit_code == 0 - result = _run("--init", "-i", position, "-o", out, "-c", two_d) + result = _run("--init", "-i", plate_2d[0] / "B" / "1" / "0", "-o", out, "-c", two_d) assert result.exit_code != 0 assert "(3, 1, 4, 5, 6)" in str(result.exception) diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 297b7427..83812991 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -9,8 +9,8 @@ from iohub.ngff import open_ome_zarr from biahub.cli.main import cli -from biahub.settings import TrackingSettings, ZSlicing -from biahub.track import _init_output_plate, resolve_z_slice, track +from biahub.settings import TrackingSettings +from biahub.track import _init_output_plate, track @pytest.fixture(scope="function") @@ -33,21 +33,26 @@ def example_tracking_plate(tmp_path): channel_names=["nuclei_prediction", "membrane_prediction"], ) + # Deterministic spots: on random data ultrack occasionally finds no solution, and + # the failing in-process job drops submitit into pdb (which pytest reports as a + # read from stdin). + np.random.seed(0) for row, col, fov in position_list: position = plate_dataset.create_position(row, col, fov) - # Shape: (T, C, Z, Y, X) = (5, 2, 3, 64, 64) - data = np.random.uniform(0.1, 0.3, size=(5, 2, 3, 64, 64)).astype(np.float32) + # A projected (Z=1) plate, what 2D tracking reads: (T, C, Z, Y, X) = (5, 2, 1, 64, 64) + data = np.random.uniform(0.1, 0.3, size=(5, 2, 1, 64, 64)).astype(np.float32) - # Add some bright nuclei spots to channel 0 + # Add some bright nuclei spots to channel 0 (11 x 11 px, above the configs' + # min_area of 100; a single plane no longer merges spots across Z) for t in range(5): - for z in range(3): + for z in range(1): for _ in range(np.random.randint(3, 6)): y, x = np.random.randint(10, 54, 2) - data[t, 0, z, y - 3 : y + 4, x - 3 : x + 4] = np.random.uniform(0.7, 1.0) + data[t, 0, z, y - 5 : y + 6, x - 5 : x + 6] = np.random.uniform(0.7, 1.0) # Add some membrane boundaries to channel 1 for t in range(5): - for z in range(3): + for z in range(1): for _ in range(np.random.randint(2, 4)): y, x = np.random.randint(15, 49, 2) yy, xx = np.ogrid[:64, :64] @@ -87,28 +92,13 @@ def _make_tracking_config(plate_path, tmp_path): config = { "output_mode": "2D", "fov": "*/*/*", - "z_slicing": {"method": "central"}, "target_channel": "nuclei_prediction", "input_images": [ { # Primary data source left null -> resolved from the -i input plate. + # The plate is projected (Z=1), so the channels need no processing. "path": None, - "channels": { - "nuclei_prediction": [ - { - "function": "np.mean", - "kwargs": {"axis": 1}, - "per_timepoint": False, - }, - ], - "membrane_prediction": [ - { - "function": "np.mean", - "kwargs": {"axis": 1}, - "per_timepoint": False, - }, - ], - }, + "channels": {"nuclei_prediction": [], "membrane_prediction": []}, }, { "path": None, @@ -332,139 +322,81 @@ def test_track_cli_debug_single_position(tmp_path, example_tracking_plate, monke # --------------------------------------------------------------------------- -# Z-slicing resolution +# Dimensionality: track neither projects nor selects Z (that is `biahub project`) # --------------------------------------------------------------------------- -def test_resolve_z_slice_all(): - z_slices, n = resolve_z_slice(ZSlicing(method="all"), z_shape=30) - assert z_slices == slice(None) - assert n == 30 - - -def test_resolve_z_slice_central(): - z_slices, n = resolve_z_slice(ZSlicing(method="central"), z_shape=21) - assert n == z_slices.stop - z_slices.start - assert n >= 3 - - -def test_resolve_z_slice_range(): - z_slices, n = resolve_z_slice(ZSlicing(method="range", range=(5, 10)), z_shape=30) - assert z_slices == slice(5, 10) - assert n == 5 - - -def test_resolve_z_slice_range_invalid(): - with pytest.raises(ValueError): - resolve_z_slice(ZSlicing(method="range", range=(10, 5)), z_shape=30) - - -def test_resolve_z_slice_focus_loads_full_reports_window(): - # focus loads the full stack at read time; the count is the fixed window size. - z_slices, n = resolve_z_slice(ZSlicing(method="focus", window_size=15), z_shape=30) - assert z_slices == slice(None) - assert n == 15 - # window larger than the stack collapses the count to the full depth. - _, n = resolve_z_slice(ZSlicing(method="focus", window_size=50), z_shape=30) - assert n == 30 - - -def test_focus_window_fixed_size_and_shifts_at_edges(): - from biahub.track import _focus_window - - # centred window of the requested size. - assert _focus_window(15, 6, 30, 1 / 3) == (slice(13, 19), 6) - # window bigger than the stack collapses to the full range. - assert _focus_window(15, 50, 30, 1 / 3) == (slice(0, 30), 30) - # a window that would spill past an edge is shifted, not clipped. - z_slices, n = _focus_window(29, 6, 30, 0.0) - assert n == 6 - assert 0 <= z_slices.start and z_slices.stop <= 30 - - -def test_apply_focus_slicing_uniform_window(monkeypatch): - from biahub import track as track_mod - - # Deterministic focus centre so the test doesn't depend on waveorder. - monkeypatch.setattr(track_mod, "_median_focus_plane", lambda stack, pixel_size: 15) - - T, Z, Y, X = 4, 30, 8, 8 - data_dict = { - "a": np.arange(T * Z * Y * X).reshape(T, Z, Y, X).astype(float), - "b": np.zeros((T, Z, Y, X)), +def _minimal_settings(**overrides): + base = { + "target_channel": "nuclei_prediction", + "input_images": [{"path": None, "channels": {"nuclei_prediction": []}}], } - z = ZSlicing(method="focus", window_size=6, frac_below=1 / 3) - out = track_mod.apply_focus_slicing(data_dict, z, pixel_size=0.5) - - # Same fixed window applied to every channel (center=15 -> slice(13, 19)). - assert out["a"].shape == (T, 6, Y, X) - assert out["b"].shape == (T, 6, Y, X) - assert np.array_equal(out["a"], data_dict["a"][:, 13:19]) - - -def test_zslicing_focus_defaults_window_size(): - # focus works without an explicit window_size (defaults to 48). - z = ZSlicing(method="focus") - assert z.window_size == 48 + base.update(overrides) + return TrackingSettings(**base) -def test_zslicing_ignores_irrelevant_fields(): - # method decides which fields are used; the rest are ignored, not rejected. - z_slices, n = resolve_z_slice( - ZSlicing(method="all", range=(0, 5), window_size=10), z_shape=30 +def test_init_output_plate_takes_the_input_z( + tmp_path, example_tracking_plate, create_custom_plate +): + flat_path, _ = example_tracking_plate # Z=1 + shape = _init_output_plate( + [str(flat_path / "A" / "1" / "0")], tmp_path / "flat.zarr", _minimal_settings() ) - assert z_slices == slice(None) - assert n == 30 - + assert shape[2] == 1 + with open_ome_zarr(str(tmp_path / "flat.zarr" / "A" / "1" / "0"), mode="r") as ds: + assert ds.data.shape[2] == 1 + + stack_path, _ = create_custom_plate(tmp_path, z_size=3) + shape = _init_output_plate( + [str(stack_path / "A" / "1" / "0")], + tmp_path / "stack.zarr", + _minimal_settings(output_mode="3D"), + ) + assert shape[2] == 3 -def test_resolve_z_slice_range_unset_falls_back_to_all(): - z_slices, n = resolve_z_slice(ZSlicing(method="range"), z_shape=30) - assert z_slices == slice(None) - assert n == 30 +@pytest.mark.parametrize( + "output_mode, z_size, message", + [("2D", 3, "input is 3D.*biahub project"), ("3D", 1, "one Z-plane")], +) +def test_init_output_plate_refuses_a_dimensionality_mismatch( + tmp_path, create_custom_plate, output_mode, z_size, message +): + plate_path, _ = create_custom_plate(tmp_path, z_size=z_size) + with pytest.raises(ValueError, match=message): + _init_output_plate( + [str(plate_path / "A" / "1" / "0")], + tmp_path / "o.zarr", + _minimal_settings(output_mode=output_mode), + ) + assert not (tmp_path / "o.zarr").exists() -# --------------------------------------------------------------------------- -# Output plate shape matches tracked Z -# --------------------------------------------------------------------------- +def test_track_cli_refuses_a_3d_input_for_2d_tracking(tmp_path, create_custom_plate): + plate_path, _ = create_custom_plate(tmp_path, z_size=3) + config_path = _make_tracking_config(plate_path, tmp_path) -def _minimal_settings(**overrides): - base = { - "target_channel": "nuclei_prediction", - "input_images": [ - { - "path": None, - "channels": { - "nuclei_prediction": [{"function": "np.mean", "kwargs": {"axis": 1}}] - }, - } + result = CliRunner().invoke( + cli, + [ + "track", + "-i", + str(plate_path / "A/1/0"), + "-o", + str(tmp_path / "o.zarr"), + "-c", + str(config_path), + "--init", ], - } - base.update(overrides) - return TrackingSettings(**base) + ) + assert result.exit_code != 0 + assert "biahub project" in str(result.exception) -@pytest.mark.parametrize( - "output_mode, z_slicing, expected_z", - [ - ("2D", {"method": "central"}, 1), - ("3D", {"method": "range", "range": (0, 2)}, 2), - # Regression: focus + 3D output Z equals the fixed focus window, not the - # full stack -- this is the shape-mismatch bug _init_output_plate had. - ("3D", {"method": "focus", "window_size": 2}, 2), - ], -) -def test_init_output_plate_shape( - tmp_path, example_tracking_plate, output_mode, z_slicing, expected_z -): - plate_path, _ = example_tracking_plate - output_path = tmp_path / "init_shape.zarr" - settings = _minimal_settings(output_mode=output_mode, z_slicing=z_slicing) - shape = _init_output_plate([str(plate_path / "A" / "1" / "0")], output_path, settings) - assert shape[2] == expected_z - with open_ome_zarr(str(output_path / "A" / "1" / "0"), mode="r") as ds: - assert ds.data.shape[2] == expected_z +def test_z_slicing_is_rejected_with_a_pointer_to_project(): + with pytest.raises(ValueError, match="z_slicing was removed.*biahub project"): + _minimal_settings(z_slicing={"method": "focus"}) # --------------------------------------------------------------------------- diff --git a/tests/test_example_settings.py b/tests/test_example_settings.py index e84bdcec..2601b7cd 100644 --- a/tests/test_example_settings.py +++ b/tests/test_example_settings.py @@ -16,6 +16,7 @@ EstimateTransformSettings, FlatFieldCorrectionSettings, ProcessingImportFuncSettings, + ProjectionSettings, SegmentationSettings, StitchSettings, TrackingSettings, @@ -41,6 +42,7 @@ ("example_transform_settings.yml", TransformSettings), ("example_transform_settings_stabilization.yml", TransformSettings), ("example_process_with_config_settings.yml", ProcessingImportFuncSettings), + ("example_project_settings.yml", ProjectionSettings), ("example_segmentation_settings.yml", SegmentationSettings), ("example_stitch_settings.yml", StitchSettings), ("example_track_settings.yml", TrackingSettings), @@ -211,9 +213,24 @@ def test_nextflow_a549_track_config_validates(): with open("nextflow/configs/a549/track.yml") as f: settings = TrackingSettings(**yaml.safe_load(f)) - assert settings.z_slicing.method == "focus" - assert settings.z_slicing.window_size == 60 - assert settings.z_slicing.focus_channel in settings.input_images[0].channels + # Track reads the projected plate: 2D, no processing step on the nuclei channel. + assert settings.output_mode == "2D" + assert settings.input_images[0].channels == {"nuclei_prediction": []} + assert settings.cellpose_config.input_channel == "nuclei_prediction" + + +def test_nextflow_a549_project_config_validates(): + """The A549 projection keeps track's old 60-plane focus window (#418) and mean.""" + import yaml + + from biahub.settings import ProjectionSettings + + with open("nextflow/configs/a549/project.yml") as f: + settings = ProjectionSettings(**yaml.safe_load(f)) + nuclei = settings.channels["nuclei_prediction"] + assert nuclei.method == "mean" + assert nuclei.z.center == "focus" and nuclei.z.range == (-20, 40) + assert settings.focus.channel == "Phase3D" and settings.focus.per == "fov" def test_skill_config_templates_load(): diff --git a/tests/test_project_settings.py b/tests/test_project_settings.py new file mode 100644 index 00000000..69027516 --- /dev/null +++ b/tests/test_project_settings.py @@ -0,0 +1,117 @@ +"""Projection config schema and Z-plane resolution (`biahub project`).""" + +import numpy as np +import pytest + +from biahub.project import resolve_z_indices +from biahub.settings import ChannelProjection, ProjectionSettings + + +def _channel(method="max", **z): + return ChannelProjection(method=method, z=z or None) + + +# --------------------------------------------------------------------------- +# resolve_z_indices +# --------------------------------------------------------------------------- + + +def test_unset_z_is_the_whole_stack(): + assert np.array_equal(resolve_z_indices(_channel(), 5), np.arange(5)) + + +def test_absolute_range_with_step(): + indices = resolve_z_indices(_channel(range=[10, 50], step=2), 60) + assert indices.tolist() == list(range(10, 50, 2)) + + +def test_range_is_relative_to_a_center(): + indices = resolve_z_indices(_channel(center=30, range=[-20, 40]), 100) + assert indices.tolist() == list(range(10, 70)) + + +def test_focus_center_uses_the_given_plane(): + indices = resolve_z_indices(_channel(center="focus", range=[-2, 3]), 20, focus_z=5) + assert indices.tolist() == [3, 4, 5, 6, 7] + with pytest.raises(ValueError, match="no focus plane"): + resolve_z_indices(_channel(center="focus", range=[-2, 3]), 20) + + +def test_slice_is_one_plane(): + assert resolve_z_indices(_channel("slice", center=42), 60).tolist() == [42] + assert resolve_z_indices(_channel("slice", center="focus"), 60, focus_z=7).tolist() == [7] + + +def test_ranges_past_the_stack_are_clipped(): + assert resolve_z_indices(_channel(center=2, range=[-5, 5]), 10).tolist() == list(range(7)) + assert resolve_z_indices(_channel(range=[8, 50]), 10).tolist() == [8, 9] + + +def test_a_selection_outside_the_stack_is_an_error(): + with pytest.raises(ValueError, match="select nothing"): + resolve_z_indices(_channel("slice", center=42), 10) + with pytest.raises(ValueError, match="select nothing"): + resolve_z_indices(_channel(range=[10, 20]), 10) + + +# --------------------------------------------------------------------------- +# Schema +# --------------------------------------------------------------------------- + + +def test_slice_needs_a_center_and_takes_no_range(): + with pytest.raises(ValueError, match="slice needs z.center"): + _channel("slice") + with pytest.raises(ValueError, match="takes no range"): + _channel("slice", center=3, range=[0, 2]) + + +@pytest.mark.parametrize("method", ["max", "mean", "sum"]) +def test_a_center_needs_a_range_for_reductions(method): + with pytest.raises(ValueError, match="needs z.range"): + _channel(method, center=3) + + +def test_range_must_be_non_empty_and_absolute_without_a_center(): + with pytest.raises(ValueError, match="empty"): + _channel(range=[5, 5]) + with pytest.raises(ValueError, match="offsets need a center"): + _channel(range=[-5, 5]) + + +def test_step_needs_a_range(): + with pytest.raises(ValueError, match="step needs a range"): + _channel(step=2) + + +def test_focus_block_is_required_by_focus_users(): + with pytest.raises(ValueError, match="use z.center: focus"): + ProjectionSettings(channels={"a": {"method": "slice", "z": {"center": "focus"}}}) + s = ProjectionSettings( + channels={"a": {"method": "slice", "z": {"center": "focus"}}}, + focus={"channel": "Phase3D"}, + ) + assert s.uses_focus + assert s.focus.per == "fov" + assert s.focus.center_crop_xy == [800, 800] # FocusSettings' default crop + + +def test_an_unused_focus_block_warns(): + with pytest.warns(UserWarning, match="ignored"): + s = ProjectionSettings(channels={"a": {"method": "max"}}, focus={"channel": "x"}) + assert not s.uses_focus + + +def test_focus_has_no_axes(): + with pytest.raises(ValueError, match="axes"): + ProjectionSettings( + channels={"a": {"method": "slice", "z": {"center": "focus"}}}, + focus={"channel": "x", "axes": "xyz"}, + ) + + +def test_unknown_keys_are_refused(): + with pytest.raises(ValueError, match="Extra inputs"): + ProjectionSettings(channels={"a": {"method": "max", "mode": "max"}}) + with pytest.raises(ValueError, match="Extra inputs"): + ProjectionSettings(channels={"a": {"method": "max", "z": {"span": 3}}}) diff --git a/tests/test_segment_settings.py b/tests/test_segment_settings.py index 19363d98..217df4ec 100644 --- a/tests/test_segment_settings.py +++ b/tests/test_segment_settings.py @@ -41,7 +41,6 @@ def _settings(**model): base = { "pretrained_model": "cpsam_v2", "channels": ["nuc"], - "z_slice_2D": 2, "eval_args": {"diameter": None}, } return SegmentationSettings(models={"nucleus": {**base, **model}}) @@ -52,6 +51,12 @@ def test_new_schema_is_accepted(): m = s.models["nucleus"] assert m.pretrained_model == "cpsam_v2" assert m.channels == ["nuc"] + assert not m.is_3d and not s.is_3d # no do_3D / stitch_threshold: a 2D model + + +@pytest.mark.parametrize("eval_args", [{"do_3D": True}, {"stitch_threshold": 0.5}]) +def test_do_3d_or_stitching_makes_a_3d_model(eval_args): + assert _settings(eval_args=eval_args).is_3d def test_old_path_to_model_key_says_what_to_rename(): @@ -63,13 +68,27 @@ def test_old_path_to_model_key_says_what_to_rename(): "nucleus": { "path_to_model": "cpsam_v2", "channels": ["nuc"], - "z_slice_2D": 0, "eval_args": {}, } } ) +def test_removed_z_slice_2d_points_at_project(): + with pytest.raises(ValueError, match="z_slice_2D was removed.*biahub project"): + _settings(z_slice_2D=2) + + +def test_models_cannot_mix_2d_and_3d(): + from biahub.settings import SegmentationSettings + + model = {"pretrained_model": "cpsam_v2", "channels": ["nuc"]} + with pytest.raises(ValueError, match="Cannot mix 2D and 3D"): + SegmentationSettings( + models={"flat": model, "stack": {**model, "eval_args": {"do_3D": True}}} + ) + + def test_channels_inside_eval_args_says_where_they_moved(): with pytest.raises(ValueError, match="eval_args.channels"): _settings(eval_args={"channels": ["nuc"]}) @@ -100,39 +119,51 @@ def test_resolve_models_maps_names_without_mutating_the_settings(): before = s.model_dump(mode="json") (m,) = resolve_models( - s, channel_names=["BF", "nuc", "mem"], scale=(1, 1, 0.5, 0.1, 0.1), z_size=5 + s, channel_names=["BF", "nuc", "mem"], scale=(1, 1, 0.5, 0.1, 0.1), z_size=1 ) assert m.name == "nucleus" assert m.channel_indices == [2, 1] assert m.preprocessing[0][1] == 2 - assert m.z_slice_2D == 2 + assert not m.is_3d assert s.model_dump(mode="json") == before def test_resolve_models_fills_anisotropy_for_3d_only(): from biahub.segment import resolve_models - s3d = _settings(z_slice_2D=None, eval_args={"do_3D": True}) + s3d = _settings(eval_args={"do_3D": True}) (m,) = resolve_models(s3d, channel_names=["nuc"], scale=(1, 1, 0.5, 0.1, 0.1), z_size=5) + assert m.is_3d assert m.eval_args["anisotropy"] == pytest.approx(5.0) (m,) = resolve_models( - _settings(), channel_names=["nuc"], scale=(1, 1, 0.5, 0.1, 0.1), z_size=5 + _settings(), channel_names=["nuc"], scale=(1, 1, 0.5, 0.1, 0.1), z_size=1 ) assert "anisotropy" not in m.eval_args -def test_resolve_models_rejects_unknown_channels_and_planes(): +def test_resolve_models_rejects_unknown_channels(): from biahub.segment import resolve_models with pytest.raises(ValueError, match="not in the input"): resolve_models( - _settings(channels=["missing"]), channel_names=["nuc"], scale=(1,) * 5, z_size=5 + _settings(channels=["missing"]), channel_names=["nuc"], scale=(1,) * 5, z_size=1 ) - with pytest.raises(ValueError, match="z_slice_2D"): + + +def test_resolve_models_requires_the_input_dimensionality_to_match(): + """Segment does not project: 2D models need a Z=1 plate, 3D models a stack.""" + from biahub.segment import resolve_models + + with pytest.raises(ValueError, match="input is 3D.*biahub project"): + resolve_models(_settings(), channel_names=["nuc"], scale=(1,) * 5, z_size=5) + with pytest.raises(ValueError, match="one Z-plane but the models segment in 3D"): resolve_models( - _settings(z_slice_2D=9), channel_names=["nuc"], scale=(1,) * 5, z_size=5 + _settings(eval_args={"do_3D": True}), + channel_names=["nuc"], + scale=(1,) * 5, + z_size=1, ) From 1c82b7fdf1458a2884db418c047b0893594df1b0 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 14:40:38 -0700 Subject: [PATCH 12/54] feat(nextflow): A549 project.yml projects every assembled channel; skill checks it The A549 projection config lists every channel of a typical assembled plate under the pre-rename names project sees inside the pipeline: the in-focus plane of BF and Phase3D, the full-stack max of the raw fluorescence channels, and the mean of the virtual-stain predictions over the 60-plane focus window. The focus is found on Phase3D per timepoint, as biahub#400 asks for inference. The brightfield key is `BF`, matching reconstruct.yml; the fluorescence names and an older `BF - Oblique` are per-dataset edits. The reconstruct-with-nextflow check script also validates project.yml both ways (a listed channel the plate lacks is a MISS, an assembled channel the config omits is NOT PROJECTED) and track.yml against the projected plate. The skill carries unprojected channels into the pre-run plan for the user to decide, runs the rename on the projected plate too, and suggests the `BF - Oblique` edit in both reconstruct.yml and project.yml when the raw store names it so. Co-Authored-By: Claude Fable 5.1 --- .../skills/reconstruct-with-nextflow/SKILL.md | 31 ++++++++-- .../references/caveats.md | 50 ++++++++++----- .../templates/check_qc_channels.py | 62 ++++++++++++++++--- nextflow/configs/a549/project.yml | 57 +++++++++++++---- tests/test_example_settings.py | 14 ++++- 5 files changed, 176 insertions(+), 38 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index 47706440..eab3c0f5 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -191,7 +191,7 @@ the next dataset in the family would also need), fix it in a PR against `nextflow/configs//`, stating which dataset exposed it. Genuinely per-dataset values stay in the run directory. -### 5a. Check the QC channel filters against the stores they will QC +### 5a. Check the channel couplings: QC filters, project.yml, track.yml Once the configs are edited, run this from the biahub checkout: @@ -209,6 +209,20 @@ them: raw channel names for the deskewed data, waveorder's for the predictions, and `_labels` for the tracking store. +**`project.yml` against the assembled plate.** When `configs/project.yml` +exists the script also checks it both ways. A channel in `project.yml` that the +assembled plate will not have is a `MISS` (exit 1): `biahub project --init` +refuses it, so fix the name now — the raw fluorescence names are the camera +strings of this acquisition and are the usual mismatch. An assembled channel +that `project.yml` does not list is printed as `NOT PROJECTED`; that is a +warning, not a failure (`biahub project` drops it with a warning), **but the +intention of the pipeline is that every assembled channel is projected unless +the user decided otherwise**, so carry the list into the pre-run plan (step 6 +of §7) and ask the user whether to add each one or leave it out on purpose. +The script also checks that `track.yml`'s `target_channel` and +`cellpose_config.input_channel` are channels `project.yml` writes, since track +reads the projected plate. + **That last one is the coupling that actually breaks.** `qc_track.yaml` names `nuclei_prediction_labels`, which only resolves if `track.yml` says `target_channel: nuclei_prediction` — two files with no shared source of truth, @@ -342,6 +356,14 @@ Do not run anything yet. Show the user: that tracking is simply not run — there is no tracking by-product any more (`caveats.md` §4). + **When project runs, list what it projects.** State each assembled channel + with its `project.yml` method, and name any assembled channel the config + does not list (§5a's `NOT PROJECTED` lines). Those need a decision from the + user before launch: add the channel to `project.yml`, or confirm it is left + out on purpose. Do not launch with an unprojected channel the user has not + seen — the pipeline's intention is that every assembled channel is + projected unless the user explicitly decided otherwise. + **State the directory numbers this run will produce.** The number is the step's position among the steps performed, not a fixed label, so a neuromast run writes `4-assemble` as its last directory and an A549 run writes @@ -521,9 +543,10 @@ directory (`_rerun`) over relaunching in place. ## 11. Wrap up -1. Normalize channel names on the assembled plate with - `templates/rename_channels.py` (idempotent; `caveats.md` §2 — check first - whether a `rename-channels` CLI has landed on main, making this obsolete). +1. Normalize channel names on the assembled plate, and on the projected plate + when project ran, with `templates/rename_channels.py` (idempotent; + `caveats.md` §2 — check first whether a `rename-channels` CLI has landed on + main, making this obsolete). 2. Verify the assembled store (`-assemble/.zarr`, `4-assemble` unless earlier steps were skipped) opens with iohub; report shape, channels, size on disk. This is the deliverable for every family; tracking, when it diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 101572eb..6a32bad4 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -77,11 +77,14 @@ copy is group-writable like every other file the skill places in `` `biahub virtual-stain` names outputs verbatim from `target_channel`, but the rest of biahub keys off the `_prediction` suffix. -**⚠ Ordering.** tracking runs *inside* the pipeline and reads the assembled -plate, so **track sees the pre-rename names**. The shipped templates avoid the -trap by putting the suffix directly in the VS `target_channel` -(`nuclei_prediction`/`membrane_prediction`), making the rename a no-op for -those channels. +**⚠ Ordering.** project and tracking run *inside* the pipeline and read the +assembled plate, so **project.yml and track.yml see the pre-rename names** +(`project.yml`'s keys are `BF`, `GFP EX488 EM525-45`, ...). The +shipped templates avoid the trap for the predictions by putting the suffix +directly in the VS `target_channel` (`nuclei_prediction`/`membrane_prediction`), +making the rename a no-op for those channels. The projected plate inherits the +pre-rename names too: run `rename_channels.py` on `-project` as well as on +`-assemble` (same command, from the directory holding the store). **Stopgap:** PRs [#260](https://github.com/czbiohub-sf/biahub/pull/260) and [#250](https://github.com/czbiohub-sf/biahub/pull/250) carry a `rename-channels` @@ -97,10 +100,15 @@ lateral pixel size of the *input* store (no longer read from the zarr); `scan_step_um` comes from the acquisition — read `/hpc/instruments/cm.mantis//config.yaml`, do not assume 0.150. -**`reconstruct.yml` — `input_channel_names`.** Must name the brightfield -channel exactly as it appears in the deskewed store (`["BF - Oblique"]` -today). A mismatch fails at `compute_transfer_function`, after flat-field and -deskew have already burned hours. +**`reconstruct.yml` — `input_channel_names`, and `project.yml`'s brightfield +key.** Must name the brightfield channel exactly as it appears in the deskewed +store. The templates say `BF` (current acquisitions); older acquisitions name it +`BF - Oblique` (e.g. `2026_08_11_A549_SEC61B_DENV`, so a v9 rerun of it too). +Read the raw store's channel names and, when it says `BF - Oblique`, suggest +that edit in **both** `reconstruct.yml` and `project.yml` — §5a's check flags +the project.yml side; a reconstruct.yml mismatch fails at +`compute_transfer_function`, after flat-field and deskew have already burned +hours. **`reconstruct.yml` — `phase.transfer_function.{yx,z}_pixel_size`.** These are the voxel sizes of the *deskewed* store, so they follow from `deskew.yml`: @@ -156,11 +164,25 @@ carry the **pre-rename** names of the **assembled** plate (§2, §5). is what the tracking store actually holds — SKILL.md §5a checks both couplings before launch. -**`project.yml` — `focus.channel` must be in the assembled plate.** `biahub -project --init` resolves every channel name against the plate and fails there, -in the init phase, rather than after every other step. The A549 template finds -the focus on `Phase3D` and writes only `nuclei_prediction`; the focus channel -need not be an output. +**`project.yml` — every channel name must be in the assembled plate, and every +assembled channel should be in `project.yml`.** `biahub project --init` +resolves every listed channel (and `focus.channel`) against the plate and fails +there, in the init phase, rather than after every other step; a plate channel +the config does not list is dropped with a warning. The pipeline's intention is +that every assembled channel is projected unless the user decided otherwise, so +SKILL.md §5a's check prints the unlisted ones and the pre-run plan asks about +them. The A549 template lists a typical plate (brightfield, two raw +fluorescence channels, Phase3D, the two predictions) under the **PRE-rename +names**, because project runs inside the pipeline before the §2 rename: the +brightfield key is `BF`, the same name as `reconstruct.yml`'s +`input_channel_names` (edit both to `BF - Oblique` for the older acquisitions, +see above), and **the fluorescence names are the camera/filter strings** +(`GFP EX488 EM525-45`, `mCherry EX561 EM600-37` on the reference datasets, one +or two markers per dataset) — edit or remove them to match the raw store. +Label-free channels take the +one in-focus plane, fluorescence the full-stack max, the virtual-stain +predictions the mean over the 60-plane focus window; the focus is found per +timepoint on `Phase3D`. **`concatenate.yml` holds parameters only — no source paths.** The assemble subworkflow passes the three source stores on the command line, one `-i` per diff --git a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py index 41515edd..deae92a7 100644 --- a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py +++ b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py @@ -1,8 +1,20 @@ -"""Check a QC config's `channels:` against the store it will actually QC. +"""Check the channel couplings between the step configs and the stores they read. python check_qc_channels.py /configs -Exits 1 and names every miss, 0 if every filter resolves. +Exits 1 and names every miss, 0 if every filter resolves. Three couplings: + +1. Each QC config's `channels:` against the store it will QC. +2. `project.yml` (when present) against the assembled plate, both ways: a listed + channel the plate will not have is a MISS (`biahub project --init` refuses + it); an assembled channel the config does not list is printed as + NOT PROJECTED. That one is a warning, not a failure -- the CLI drops the + channel with a warning -- but the pipeline's intention is that every + assembled channel is projected unless the user decided otherwise, so the + skill carries these lines into the pre-run plan for the user to decide. +3. `track.yml`'s `target_channel` and `cellpose_config.input_channel` against + the plate track reads: the projected plate when `project.yml` exists, else + the assembled plate. WHY PREDICT RATHER THAN READ. The stores QC runs over — the assembled plate and the tracking plate — do not exist when configs are scaffolded; the pipeline's @@ -76,14 +88,25 @@ def tracking_channels(track_cfg): return [f"{t['target_channel']}_labels"] +def projected_channels(project_cfg): + """Return the channels `biahub project` writes, in output order (the config's keys).""" + return list(yaml.safe_load(Path(project_cfg).read_text())["channels"]) + + +def track_inputs(track_cfg): + """Return the channel names track reads from its input plate, keyed by the setting.""" + t = yaml.safe_load(Path(track_cfg).read_text()) + inputs = [("target_channel", t["target_channel"])] + cellpose = t.get("cellpose_config") or {} + if t.get("segmentation_method") == "cellpose" and "input_channel" in cellpose: + inputs.append(("cellpose_config.input_channel", cellpose["input_channel"])) + return inputs + + if __name__ == "__main__": raw, cfgs = sys.argv[1], Path(sys.argv[2]) - checks = [ - ( - "qc.yaml", - assembled_channels(raw, cfgs / "reconstruct.yml", cfgs / "virtual_stain.yml"), - ) - ] + assembled = assembled_channels(raw, cfgs / "reconstruct.yml", cfgs / "virtual_stain.yml") + checks = [("qc.yaml", assembled)] if (cfgs / "qc_track.yaml").exists(): checks.append(("qc_track.yaml", tracking_channels(cfgs / "track.yml"))) bad = 0 @@ -93,4 +116,27 @@ def tracking_channels(track_cfg): ok = ch in predicted bad += not ok print(f" {'ok ' if ok else 'MISS'} {owner}.channels: {ch!r}") + + # project.yml reads the assembled plate and writes the plate track reads. + track_reads = assembled + if (cfgs / "project.yml").exists(): + projected = projected_channels(cfgs / "project.yml") + track_reads = projected + print(f"project.yml: assembled plate will have {assembled}") + for ch in projected: + ok = ch in assembled + bad += not ok + print(f" {'ok ' if ok else 'MISS'} channels: {ch!r}") + for ch in assembled: + if ch not in projected: + # Not a failure: the CLI drops it with a warning. The skill asks the + # user about every such channel in the pre-run plan. + print(f" NOT PROJECTED {ch!r}: not in project.yml, will be dropped") + if (cfgs / "track.yml").exists(): + source = "projected" if (cfgs / "project.yml").exists() else "assembled" + print(f"track.yml: {source} plate will have {track_reads}") + for key, ch in track_inputs(cfgs / "track.yml"): + ok = ch in track_reads + bad += not ok + print(f" {'ok ' if ok else 'MISS'} {key}: {ch!r}") sys.exit(1 if bad else 0) diff --git a/nextflow/configs/a549/project.yml b/nextflow/configs/a549/project.yml index bd1fcf7d..0a53519b 100644 --- a/nextflow/configs/a549/project.yml +++ b/nextflow/configs/a549/project.yml @@ -1,22 +1,57 @@ # Projection config for the mantis-v2 pipeline (`biahub project`). A549. # # Project reads the ASSEMBLED plate (-assemble/.zarr) and writes the 2D -# plate that track reads, so the channels below must exist THERE, and track.yml's -# target_channel / cellpose_config.input_channel must be among the channels written here. +# (Z=1) plate that track reads and that downstream model training / embedding +# inference consume (biahub#400). The intention is that EVERY channel of the assembled +# plate is projected: a channel listed here that the plate lacks fails `--init`; a plate +# channel not listed here is dropped with a warning, so list it or decide to leave it +# out on purpose (the reconstruct-with-nextflow skill asks before launch). track.yml's +# target_channel / cellpose_config.input_channel must be among the channels here. # -# Mean of the virtual-stain nuclei over 60 planes around the in-focus plane of the -# phase channel: the window starts 20 planes below focus and ends 39 above it. This is -# the window track's own z_slicing selected before projection moved here (#418): against -# TrackEdit annotations (with min_sibling_length 3) division F1 rose 0.63 -> 0.68 on -# 2026_04_28, with fewer false divisions and fragmented tracks on 2026_04_21 at the same -# mitotic-cell detection as all planes (0.95); a narrower 40-plane window lost mitotic -# cells that drift in Z (0.89). The focus is found once per FOV (median over timepoints) -# on the 800x800 centre crop of the phase channel. +# CHANNEL NAMES ARE THE PRE-RENAME NAMES. Project runs inside the pipeline, before +# rename_channels.py normalises the assembled plate (caveats section 2), so the keys +# are the raw store's names. The brightfield channel is `BF` here, matching +# reconstruct.yml's input_channel_names; older acquisitions name it `BF - Oblique`, and +# then BOTH files need the edit. The fluorescence channels are the camera/filter +# strings (one or two markers per dataset). Edit them to match the raw store +# (`references/caveats.md` section 3). Phase3D, nuclei_prediction and +# membrane_prediction are fixed by reconstruct.yml and virtual_stain.yml. Run +# rename_channels.py on the projected plate as well. +# +# Focus: the in-focus plane of Phase3D is found for every position AND timepoint +# (`per: timepoint`, as biahub#400 asks for inference), on the 800x800 centre crop, and +# recorded per timepoint in each output position's `biahub-project` metadata. The +# light-sheet arm is registered to the label-free arm before assembly, so the phase +# focus applies to every channel. channels: + # Label-free: the one in-focus plane. + BF: + method: slice + z: {center: focus} + # Raw fluorescence (organelle markers): max over the whole stack. The marker signal + # is sparse and in focus wherever the organelle is, so the full-volume maximum keeps + # it through the cell's depth without the out-of-focus haze phase has. + GFP EX488 EM525-45: + method: max + mCherry EX561 EM600-37: + method: max + Phase3D: + method: slice + z: {center: focus} + # Virtual stain: mean over 60 planes starting 20 below the in-focus plane. This is + # the window track's own z_slicing selected before projection moved here (#418): + # against TrackEdit annotations (with min_sibling_length 3) division F1 rose 0.63 -> + # 0.68 on 2026_04_28, with fewer false divisions and fragmented tracks on 2026_04_21 + # at the same mitotic-cell detection as all planes (0.95); a narrower 40-plane window + # lost mitotic cells that drift in Z (0.89). That benchmark used one focus plane per + # FOV (the median over timepoints); per-timepoint focus follows Z drift instead. nuclei_prediction: method: mean z: {center: focus, range: [-20, 40]} + membrane_prediction: + method: mean + z: {center: focus, range: [-20, 40]} focus: channel: Phase3D - per: fov + per: timepoint diff --git a/tests/test_example_settings.py b/tests/test_example_settings.py index 2601b7cd..f6762b03 100644 --- a/tests/test_example_settings.py +++ b/tests/test_example_settings.py @@ -230,7 +230,19 @@ def test_nextflow_a549_project_config_validates(): nuclei = settings.channels["nuclei_prediction"] assert nuclei.method == "mean" assert nuclei.z.center == "focus" and nuclei.z.range == (-20, 40) - assert settings.focus.channel == "Phase3D" and settings.focus.per == "fov" + # Every channel of a typical A549 assembled plate, in plate order, under the + # PRE-rename names project sees inside the pipeline. + assert list(settings.channels) == [ + "BF", + "GFP EX488 EM525-45", + "mCherry EX561 EM600-37", + "Phase3D", + "nuclei_prediction", + "membrane_prediction", + ] + assert settings.channels["Phase3D"].method == "slice" + assert settings.channels["GFP EX488 EM525-45"].method == "max" + assert settings.focus.channel == "Phase3D" and settings.focus.per == "timepoint" def test_skill_config_templates_load(): From a8c6d1bf2ba7485b05ebb56e36abbce7b5329699 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 14:54:25 -0700 Subject: [PATCH 13/54] fix(project): float32 plate for any max/mean/sum channel; drop float64 accumulation NumPy already accumulates safely (float64 for integer input, pairwise float32 for float32 input), so mean and sum no longer force a float64 temporary. The output plate is float32 whenever a channel uses max, mean or sum -- float32 in and out is the common case -- and keeps the input dtype only when every channel is a slice. Co-Authored-By: Claude Fable 5.1 --- biahub/project.py | 11 ++++++----- settings/example_project_settings.yml | 4 ++-- tests/test_cli/test_project_cli.py | 3 ++- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/biahub/project.py b/biahub/project.py index ad9da6f0..222eb054 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -45,8 +45,9 @@ PROVENANCE_KEY = "biahub-project" -# Methods whose result is not one of the input values, so the output is float32. -_FLOAT_METHODS = ("mean", "sum") +# Methods that compute a new value; a plate with any of them is float32 (the common +# case is float32 in and out). Only an all-`slice` plate keeps the input dtype. +_FLOAT_METHODS = ("max", "mean", "sum") def resolve_z_indices( @@ -90,9 +91,9 @@ def project_zyx(zyx: np.ndarray, method: str, z_indices: np.ndarray) -> np.ndarr if method == "max": return planes.max(axis=0) if method == "mean": - return planes.mean(axis=0, dtype=np.float64).astype(np.float32) + return planes.mean(axis=0).astype(np.float32, copy=False) if method == "sum": - return planes.sum(axis=0, dtype=np.float64).astype(np.float32) + return planes.sum(axis=0).astype(np.float32, copy=False) raise ValueError(f"Unknown projection method {method!r}.") @@ -326,7 +327,7 @@ def project( slurm_out_path = output_dirpath.parent / "slurm_output" settings = yaml_to_model(config_filepath, ProjectionSettings) - (T, C, Z, Y, X), channel_names, _ = _init_output_plate( + (T, _, Z, Y, X), channel_names, _ = _init_output_plate( input_position_dirpaths, output_dirpath, settings ) plans, read_indices, focus_index = _resolve_channels(settings, channel_names, Z) diff --git a/settings/example_project_settings.yml b/settings/example_project_settings.yml index b4af3299..a9399919 100644 --- a/settings/example_project_settings.yml +++ b/settings/example_project_settings.yml @@ -3,8 +3,8 @@ # (the command prints which). # # method: max | mean | sum | slice -# max/mean/sum reduce the selected planes; mean and sum write float32, max and slice -# keep the input dtype. slice keeps one plane. +# max/mean/sum reduce the selected planes; slice keeps one plane. The output plate is +# float32 unless every channel is a slice, in which case it keeps the input dtype. # z: which planes feed the projection. Omit it for max/mean/sum over the whole stack. # range: [start, stop) half-open, like Z_slice in concatenate. With `center` the bounds # are offsets from the center. diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py index a337ff70..2a84c9d8 100644 --- a/tests/test_cli/test_project_cli.py +++ b/tests/test_cli/test_project_cli.py @@ -66,7 +66,7 @@ def test_project_init_scaffolds_a_z1_plate_and_names_dropped_channels(plate, tmp data, channel_names, zattrs = _read(out / "A" / "1" / "0") assert channel_names == ["RFP", "GFP"] # config order, not input order assert data.shape == (3, 2, 1, 5, 6) - assert data.dtype == np.uint16 # max and slice keep the input dtype + assert data.dtype == np.float32 # any max/mean/sum channel makes the plate float32 assert zattrs["biahub-project"]["channels"]["GFP"]["z"]["center"] == 2 assert np.all(data == 0) # --init projects nothing @@ -186,6 +186,7 @@ def test_project_focus_channel_need_not_be_an_output(plate, tmp_path, monkeypatc source, _, _ = _read(plate_path / "A/1/0") data, channel_names, _ = _read(out / "A" / "1" / "0") assert channel_names == ["GFP"] + assert data.dtype == np.uint16 # an all-slice plate keeps the input dtype assert np.array_equal(data[:, 0, 0], source[:, 0, 5]) From 1b240c27a72b77ef204dc3f258b858e5edf80ce6 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 15:15:07 -0700 Subject: [PATCH 14/54] refactor(project): focus record keys are z and z_per_t The focus record under the output position's biahub-project metadata is focus: {channel, per, z, z_per_t}; the focus_ prefix repeated the block name. Co-Authored-By: Claude Fable 5.1 --- biahub/project.py | 4 ++-- tests/test_cli/test_project_cli.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/biahub/project.py b/biahub/project.py index 222eb054..1518d2fe 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -270,8 +270,8 @@ def _project_position( record["focus"] = { "channel": settings.focus.channel, "per": settings.focus.per, - "focus_z": focus_z if settings.focus.per == "fov" else None, - "focus_z_per_t": focus_per_t.tolist(), + "z": focus_z if settings.focus.per == "fov" else None, + "z_per_t": focus_per_t.tolist(), } output_dataset.zattrs[PROVENANCE_KEY] = record diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py index 2a84c9d8..e47f40ea 100644 --- a/tests/test_cli/test_project_cli.py +++ b/tests/test_cli/test_project_cli.py @@ -138,8 +138,8 @@ def test_project_focus_per_fov_takes_the_median_plane(plate, tmp_path, monkeypat assert zattrs["biahub-project"]["focus"] == { "channel": "Phase3D", "per": "fov", - "focus_z": 2, - "focus_z_per_t": [1, 6, 2], + "z": 2, + "z_per_t": [1, 6, 2], } @@ -166,8 +166,8 @@ def test_project_focus_per_timepoint_follows_the_plane(plate, tmp_path, monkeypa for t, plane in enumerate(planes): assert np.array_equal(data[t, 0, 0], source[t, 0, plane - 1 : plane + 2].max(axis=0)) assert np.array_equal(data[t, 1, 0], source[t, 2, plane]) - assert zattrs["biahub-project"]["focus"]["focus_z"] is None - assert zattrs["biahub-project"]["focus"]["focus_z_per_t"] == planes + assert zattrs["biahub-project"]["focus"]["z"] is None + assert zattrs["biahub-project"]["focus"]["z_per_t"] == planes def test_project_focus_channel_need_not_be_an_output(plate, tmp_path, monkeypatch): From 67df18de0f4480b34431f9234fb9701ba3ab30e7 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 15:23:50 -0700 Subject: [PATCH 15/54] refactor(project)!: z selection is start/stop or center with below/above `range` meant absolute planes without a center and offsets with one. The two forms now have their own keys: `start`/`stop` are absolute, half-open like concatenate's time_indices; `below`/`above` count planes around `center`, selecting [center - below, center + above), so the A549 window is `center: focus, below: 20, above: 40`. The validator refuses mixing them. Co-Authored-By: Claude Fable 5.1 --- .../references/caveats.md | 2 +- biahub/project.py | 19 +++--- biahub/settings.py | 62 ++++++++++++------ nextflow/configs/a549/project.yml | 4 +- settings/example_project_settings.yml | 19 +++--- tests/test_cli/test_project_cli.py | 10 +-- tests/test_example_settings.py | 2 +- tests/test_project_settings.py | 63 ++++++++++++------- 8 files changed, 110 insertions(+), 71 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 6a32bad4..a3de9fcb 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -152,7 +152,7 @@ fields: `output_mode`, `target_channel`, `fov`, `input_images`, ``` mode -> output_mode -z_slicing / z_range / focus_config -> project.yml: channels..z {center: focus, range: [below, above)}, focus.channel +z_slicing / z_range / focus_config -> project.yml: channels..z {center: focus, below, above}, focus.channel input_images np.mean step -> project.yml: channels..method: mean ``` diff --git a/biahub/project.py b/biahub/project.py index 1518d2fe..bf050706 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -55,24 +55,25 @@ def resolve_z_indices( ) -> np.ndarray: """Return the plane indices one channel's projection reads from a ``z_size``-plane stack. - An unset ``z`` is the whole stack. A ``range`` without ``center`` is absolute; with a - ``center`` (``focus_z`` when it is ``"focus"``) the bounds are offsets from it, and - ``slice`` is the center alone. ``step`` thins the range. Indices outside the stack - are dropped; a selection that keeps none is an error. + An unset ``z`` is the whole stack. ``start``/``stop`` are absolute planes; with a + ``center`` (``focus_z`` when it is ``"focus"``) the planes are + ``[center - below, center + above)``, or the center alone for ``slice``. ``step`` + thins the selection. Indices outside the stack are dropped; a selection that keeps + none is an error. """ z = projection.z if z is None: return np.arange(z_size) - if z.center is None: - start, stop = z.range + if z.is_absolute: + start, stop = z.start, z.stop else: center = focus_z if z.center == "focus" else z.center if center is None: raise ValueError("z.center is 'focus' but no focus plane was given.") - if z.range is None: - start, stop = center, center + 1 + if z.is_window: + start, stop = center - z.below, center + z.above else: - start, stop = center + z.range[0], center + z.range[1] + start, stop = center, center + 1 indices = np.arange(start, stop, z.step) indices = indices[(indices >= 0) & (indices < z_size)] if indices.size == 0: diff --git a/biahub/settings.py b/biahub/settings.py index a523b4f1..b6acf745 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -316,29 +316,51 @@ def check_center_crop_xy(cls, v): class ZSelection(MyBaseModel): """Which Z-planes of a stack feed one channel's projection. - `range` is a half-open ``[start, stop)`` like concatenate's ``Z_slice``. With `center` - the bounds are offsets from the center (``[-20, 40]`` is 60 planes starting 20 below - it); `center` is a plane index or ``"focus"``, the in-focus plane found on - ``ProjectionSettings.focus.channel``. `step` keeps every step-th plane of the range. - A range that runs past the stack is clipped to it. + Two forms. Absolute: `start`/`stop`, a half-open ``[start, stop)`` like concatenate's + ``time_indices``. Around a plane: `center` with `below`/`above`, the planes + ``[center - below, center + above)`` (``below: 20, above: 40`` is 60 planes starting + 20 below the center; `above` counts the center plane itself). `center` is a plane + index or ``"focus"``, the in-focus plane found on ``ProjectionSettings.focus.channel``; + `center` alone is the one plane (``slice``). `step` keeps every step-th plane. A + selection that runs past the stack is clipped to it. """ center: NonNegativeInt | Literal["focus"] | None = None - range: tuple[int, int] | None = None + start: NonNegativeInt | None = None + stop: PositiveInt | None = None + below: NonNegativeInt | None = None + above: NonNegativeInt | None = None step: PositiveInt = 1 + @property + def is_absolute(self) -> bool: + return self.start is not None or self.stop is not None + + @property + def is_window(self) -> bool: + return self.below is not None or self.above is not None + @model_validator(mode="after") - def check_range(self): - if self.range is not None: - start, stop = self.range - if stop <= start: - raise ValueError(f"range {list(self.range)} is empty: stop must exceed start.") - if self.center is None and start < 0: + def check_form(self): + if self.is_absolute: + if self.start is None or self.stop is None: + raise ValueError("start and stop go together.") + if self.center is not None or self.is_window: raise ValueError( - f"range {list(self.range)} is negative; offsets need a center." + "start/stop are absolute planes; use center with below/above for a " + "window around a plane." ) + if self.stop <= self.start: + raise ValueError(f"start {self.start}, stop {self.stop} is empty.") + elif self.is_window: + if self.below is None or self.above is None: + raise ValueError("below and above go together.") + if self.center is None: + raise ValueError("below/above need a center.") + if self.below + self.above == 0: + raise ValueError("below and above select no plane.") elif self.step != 1: - raise ValueError("step needs a range.") + raise ValueError("step needs start/stop or below/above.") return self @@ -357,15 +379,15 @@ def check_z_for_method(self): if self.method == "slice": if self.z is None or self.z.center is None: raise ValueError("method slice needs z.center (a plane index or 'focus').") - if self.z.range is not None: + if self.z.is_window: raise ValueError( - "method slice selects one plane and takes no range; use max, mean or " - "sum to project a window." + "method slice selects one plane and takes no below/above; use max, " + "mean or sum to project a window." ) - elif self.z is not None and self.z.center is not None and self.z.range is None: + elif self.z is not None and self.z.center is not None and not self.z.is_window: raise ValueError( - f"method {self.method} needs z.range with z.center; use method slice to " - "keep the one plane." + f"method {self.method} needs z.below/above with z.center; use method slice " + "to keep the one plane." ) return self diff --git a/nextflow/configs/a549/project.yml b/nextflow/configs/a549/project.yml index 0a53519b..0fd39474 100644 --- a/nextflow/configs/a549/project.yml +++ b/nextflow/configs/a549/project.yml @@ -47,10 +47,10 @@ channels: # FOV (the median over timepoints); per-timepoint focus follows Z drift instead. nuclei_prediction: method: mean - z: {center: focus, range: [-20, 40]} + z: {center: focus, below: 20, above: 40} membrane_prediction: method: mean - z: {center: focus, range: [-20, 40]} + z: {center: focus, below: 20, above: 40} focus: channel: Phase3D diff --git a/settings/example_project_settings.yml b/settings/example_project_settings.yml index a9399919..d1dc0c6b 100644 --- a/settings/example_project_settings.yml +++ b/settings/example_project_settings.yml @@ -6,12 +6,13 @@ # max/mean/sum reduce the selected planes; slice keeps one plane. The output plate is # float32 unless every channel is a slice, in which case it keeps the input dtype. # z: which planes feed the projection. Omit it for max/mean/sum over the whole stack. -# range: [start, stop) half-open, like Z_slice in concatenate. With `center` the bounds -# are offsets from the center. -# step: keep every step-th plane of the range (default 1) -# center: a plane index, or "focus" for the in-focus plane found on focus.channel. -# slice takes center alone; max/mean/sum take center with a range. -# A range that runs past the stack is clipped to it. +# start, stop: absolute planes, half-open [start, stop) like time_indices in concatenate +# center: a plane index, or "focus" for the in-focus plane found on focus.channel. +# slice takes center alone; max/mean/sum take center with below/above. +# below, above: the window [center - below, center + above); `above` counts the +# center plane itself, so the window has below + above planes. +# step: keep every step-th plane of the selection (default 1) +# A selection that runs past the stack is clipped to it. channels: # max over the whole stack @@ -21,17 +22,17 @@ channels: # max over an explicit range, every 2nd plane: z = 10, 12, ..., 48 nuclei_prediction: method: max - z: {range: [10, 50], step: 2} + z: {start: 10, stop: 50, step: 2} # sum over 21 planes around a manually chosen center: z = 20 .. 40 GFP: method: sum - z: {center: 30, range: [-10, 11]} + z: {center: 30, below: 10, above: 11} # mean over 60 planes starting 20 below the in-focus plane: z = focus-20 .. focus+39 RFP: method: mean - z: {center: focus, range: [-20, 40]} + z: {center: focus, below: 20, above: 40} # single plane, chosen manually mCherry: diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py index e47f40ea..0a0e1c6d 100644 --- a/tests/test_cli/test_project_cli.py +++ b/tests/test_cli/test_project_cli.py @@ -77,7 +77,7 @@ def test_project_methods_match_numpy(plate, tmp_path): config = _config( tmp_path, { - "GFP": {"method": "max", "z": {"range": [2, 7], "step": 2}}, + "GFP": {"method": "max", "z": {"start": 2, "stop": 7, "step": 2}}, "RFP": {"method": "mean"}, "Phase3D": {"method": "slice", "z": {"center": 3}}, }, @@ -101,7 +101,7 @@ def test_project_sum_over_a_window_around_a_center(plate, tmp_path): plate_path, _ = plate out = tmp_path / "proj.zarr" config = _config( - tmp_path, {"GFP": {"method": "sum", "z": {"center": 6, "range": [-2, 5]}}} + tmp_path, {"GFP": {"method": "sum", "z": {"center": 6, "below": 2, "above": 5}}} ) position = plate_path / "A" / "1" / "0" @@ -120,7 +120,7 @@ def test_project_focus_per_fov_takes_the_median_plane(plate, tmp_path, monkeypat config = _config( tmp_path, { - "GFP": {"method": "mean", "z": {"center": "focus", "range": [-1, 2]}}, + "GFP": {"method": "mean", "z": {"center": "focus", "below": 1, "above": 2}}, "Phase3D": {"method": "slice", "z": {"center": "focus"}}, }, focus={"channel": "Phase3D", "per": "fov"}, @@ -149,7 +149,7 @@ def test_project_focus_per_timepoint_follows_the_plane(plate, tmp_path, monkeypa config = _config( tmp_path, { - "GFP": {"method": "max", "z": {"center": "focus", "range": [-1, 2]}}, + "GFP": {"method": "max", "z": {"center": "focus", "below": 1, "above": 2}}, "Phase3D": {"method": "slice", "z": {"center": "focus"}}, }, focus={"channel": "Phase3D", "per": "timepoint"}, @@ -193,7 +193,7 @@ def test_project_focus_channel_need_not_be_an_output(plate, tmp_path, monkeypatc @pytest.mark.parametrize( "channels, message", [ - ({"GFP": {"method": "max", "z": {"range": [10, 20]}}}, "select nothing"), + ({"GFP": {"method": "max", "z": {"start": 10, "stop": 20}}}, "select nothing"), ({"GFP": {"method": "slice", "z": {"center": 42}}}, "select nothing"), ({"DAPI": {"method": "max"}}, "not in the input channels"), ], diff --git a/tests/test_example_settings.py b/tests/test_example_settings.py index f6762b03..5d2f6b6d 100644 --- a/tests/test_example_settings.py +++ b/tests/test_example_settings.py @@ -229,7 +229,7 @@ def test_nextflow_a549_project_config_validates(): settings = ProjectionSettings(**yaml.safe_load(f)) nuclei = settings.channels["nuclei_prediction"] assert nuclei.method == "mean" - assert nuclei.z.center == "focus" and nuclei.z.range == (-20, 40) + assert nuclei.z.center == "focus" and (nuclei.z.below, nuclei.z.above) == (20, 40) # Every channel of a typical A549 assembled plate, in plate order, under the # PRE-rename names project sees inside the pipeline. assert list(settings.channels) == [ diff --git a/tests/test_project_settings.py b/tests/test_project_settings.py index 69027516..91c92a7a 100644 --- a/tests/test_project_settings.py +++ b/tests/test_project_settings.py @@ -20,21 +20,21 @@ def test_unset_z_is_the_whole_stack(): assert np.array_equal(resolve_z_indices(_channel(), 5), np.arange(5)) -def test_absolute_range_with_step(): - indices = resolve_z_indices(_channel(range=[10, 50], step=2), 60) +def test_absolute_start_stop_with_step(): + indices = resolve_z_indices(_channel(start=10, stop=50, step=2), 60) assert indices.tolist() == list(range(10, 50, 2)) -def test_range_is_relative_to_a_center(): - indices = resolve_z_indices(_channel(center=30, range=[-20, 40]), 100) - assert indices.tolist() == list(range(10, 70)) +def test_window_is_below_and_above_a_center(): + indices = resolve_z_indices(_channel(center=30, below=20, above=40), 100) + assert indices.tolist() == list(range(10, 70)) # 60 planes, 20 below the center def test_focus_center_uses_the_given_plane(): - indices = resolve_z_indices(_channel(center="focus", range=[-2, 3]), 20, focus_z=5) + indices = resolve_z_indices(_channel(center="focus", below=2, above=3), 20, focus_z=5) assert indices.tolist() == [3, 4, 5, 6, 7] with pytest.raises(ValueError, match="no focus plane"): - resolve_z_indices(_channel(center="focus", range=[-2, 3]), 20) + resolve_z_indices(_channel(center="focus", below=2, above=3), 20) def test_slice_is_one_plane(): @@ -42,16 +42,18 @@ def test_slice_is_one_plane(): assert resolve_z_indices(_channel("slice", center="focus"), 60, focus_z=7).tolist() == [7] -def test_ranges_past_the_stack_are_clipped(): - assert resolve_z_indices(_channel(center=2, range=[-5, 5]), 10).tolist() == list(range(7)) - assert resolve_z_indices(_channel(range=[8, 50]), 10).tolist() == [8, 9] +def test_selections_past_the_stack_are_clipped(): + assert resolve_z_indices(_channel(center=2, below=5, above=5), 10).tolist() == list( + range(7) + ) + assert resolve_z_indices(_channel(start=8, stop=50), 10).tolist() == [8, 9] def test_a_selection_outside_the_stack_is_an_error(): with pytest.raises(ValueError, match="select nothing"): resolve_z_indices(_channel("slice", center=42), 10) with pytest.raises(ValueError, match="select nothing"): - resolve_z_indices(_channel(range=[10, 20]), 10) + resolve_z_indices(_channel(start=10, stop=20), 10) # --------------------------------------------------------------------------- @@ -59,28 +61,41 @@ def test_a_selection_outside_the_stack_is_an_error(): # --------------------------------------------------------------------------- -def test_slice_needs_a_center_and_takes_no_range(): +def test_slice_needs_a_center_and_takes_no_window(): with pytest.raises(ValueError, match="slice needs z.center"): _channel("slice") - with pytest.raises(ValueError, match="takes no range"): - _channel("slice", center=3, range=[0, 2]) + with pytest.raises(ValueError, match="takes no below/above"): + _channel("slice", center=3, below=1, above=1) @pytest.mark.parametrize("method", ["max", "mean", "sum"]) -def test_a_center_needs_a_range_for_reductions(method): - with pytest.raises(ValueError, match="needs z.range"): +def test_a_center_needs_a_window_for_reductions(method): + with pytest.raises(ValueError, match="needs z.below/above"): _channel(method, center=3) -def test_range_must_be_non_empty_and_absolute_without_a_center(): - with pytest.raises(ValueError, match="empty"): - _channel(range=[5, 5]) - with pytest.raises(ValueError, match="offsets need a center"): - _channel(range=[-5, 5]) +def test_start_and_stop_go_together_and_are_not_mixed_with_a_center(): + with pytest.raises(ValueError, match="start and stop go together"): + _channel(start=5) + with pytest.raises(ValueError, match="is empty"): + _channel(start=5, stop=5) + with pytest.raises(ValueError, match="absolute planes"): + _channel(center=3, start=0, stop=5) + with pytest.raises(ValueError, match="absolute planes"): + _channel(start=0, stop=5, below=1, above=1) + + +def test_below_and_above_go_together_and_need_a_center(): + with pytest.raises(ValueError, match="below and above go together"): + _channel(center=3, below=2) + with pytest.raises(ValueError, match="need a center"): + _channel(below=2, above=2) + with pytest.raises(ValueError, match="select no plane"): + _channel(center=3, below=0, above=0) -def test_step_needs_a_range(): - with pytest.raises(ValueError, match="step needs a range"): +def test_step_needs_a_selection(): + with pytest.raises(ValueError, match="step needs"): _channel(step=2) @@ -114,4 +129,4 @@ def test_unknown_keys_are_refused(): with pytest.raises(ValueError, match="Extra inputs"): ProjectionSettings(channels={"a": {"method": "max", "mode": "max"}}) with pytest.raises(ValueError, match="Extra inputs"): - ProjectionSettings(channels={"a": {"method": "max", "z": {"span": 3}}}) + ProjectionSettings(channels={"a": {"method": "max", "z": {"range": [0, 3]}}}) From 1acc8929e40a48ec251f10eaa06ac8de894c748c Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 15:27:17 -0700 Subject: [PATCH 16/54] docs(nextflow): cleanup never targets the projected plate The projected store is a deliverable like the assembled plate: cleanup removes only its slurm_output/ and .iohub-progress/, as the wiring already did; say so in the comment. Co-Authored-By: Claude Fable 5.1 --- nextflow/mantis-v2.nf | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf index 5e58bd2c..a0350d7f 100644 --- a/nextflow/mantis-v2.nf +++ b/nextflow/mantis-v2.nf @@ -440,13 +440,15 @@ workflow { // ----- Cleanup of intermediates (--cleanup_intermediates) --------------- // Once the LAST step has finished, delete what only existed to feed the // assembled plate: the flat-field, deskew, reconstruct and virtual-stain step - // directories (store, slurm_output/ and .iohub-progress/ alike). Beside the - // final stores, delete what only served the run: slurm_output/ (submitit - // placeholders and a pointer README; the real logs are in - // nextflow/slurm_output/) and the assembled store's .iohub-progress/ resume - // markers, which left behind would make concatenate skip every write unit - // of a later run into this directory, keeping data built from the old - // intermediates. Tracking writes no resume markers. + // directories (store, slurm_output/ and .iohub-progress/ alike). The + // assembled and projected plates and the tracking store are DELIVERABLES and + // are never targets. Beside them, delete only what served the run: + // slurm_output/ (submitit placeholders and a pointer README; the real logs + // are in nextflow/slurm_output/) and the assembled and projected stores' + // .iohub-progress/ resume markers, which left behind would make concatenate + // or project skip every write unit of a later run into this directory, + // keeping data built from the old intermediates. Tracking writes no resume + // markers. // // What was removed is appended to nextflow/intermediates_cleaned.txt rather // than logged. Nextflow's `cleanup`, switched on above, empties the work From bcac86de0603a231b5f6e8d23979aef23a3f50b5 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 16:18:02 -0700 Subject: [PATCH 17/54] feat(project): chunks_czyx and shards_ratio; A549 writes 16-timepoint shards ProjectionSettings takes the output chunk shape and shards ratio with concatenate's names and semantics (shard = chunks * ratio, C not sharded) and passes them to create_empty_plate; None keeps iohub's defaults. The A549 config, following the Dynamic Cell Atlas array standard where a 2D float32 plate can: one-frame 256 x 256 chunks, the grid the assembled and tracking stores share, and shards of 16 timepoints by the full frame (the standard's T-shard floor; ~150 MB each, 30 files per position instead of 400). iohub writes one shard's timepoints per unit, so resume is 16 frames coarse. Co-Authored-By: Claude Fable 5.1 --- .../references/caveats.md | 6 ++++- biahub/project.py | 2 ++ biahub/settings.py | 23 +++++++++++++++++ nextflow/configs/a549/project.yml | 11 ++++++++ settings/example_project_settings.yml | 9 +++++++ tests/test_cli/test_project_cli.py | 25 +++++++++++++++++++ tests/test_example_settings.py | 3 +++ tests/test_project_settings.py | 15 +++++++++++ 8 files changed, 93 insertions(+), 1 deletion(-) diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index a3de9fcb..9c6356b9 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -232,7 +232,11 @@ recoverable per (t, c) unit. shape — changing T without recomputing ZYX silently changes the spatial extent. Larger T shards mean a larger RAM request and a worse blast radius when one tears; leave the T ratio alone without a measured reason. Channels cannot be -sharded — keep C at 1. Auto-sizing is +sharded — keep C at 1. The same two keys exist in `project.yml`: the A549 +template uses one-frame `256 x 256` chunks (the assembled and tracking grid) +and `shards_ratio: [16, 1, 1, 7, 5]`, 16-timepoint full-frame shards; its YX +entries are `ceil(Y/256), ceil(X/256)` and must be recomputed for a different +frame size, exactly as for `concatenate.yml`. Auto-sizing is [iohub#458](https://github.com/czbiohub-sf/iohub/issues/458); most of this section goes away when it lands. diff --git a/biahub/project.py b/biahub/project.py index bf050706..e4c37075 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -197,6 +197,8 @@ def _init_output_plate( input_position_dirpaths[0], settings.output_ome_zarr_version ), dtype=dtype, + chunks=(1, *settings.chunks_czyx) if settings.chunks_czyx else None, + shards_ratio=tuple(settings.shards_ratio) if settings.shards_ratio else None, metadata_sources=input_plate, metadata_keys=PROVENANCE_METADATA_KEYS, extra_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, diff --git a/biahub/settings.py b/biahub/settings.py index b6acf745..9852d48d 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -428,9 +428,32 @@ class ProjectionSettings(MyBaseModel): channels: dict[str, ChannelProjection] = Field(min_length=1) focus: ProjectionFocusSettings | None = None + # Output chunk shape (C, Z, Y, X) and shards ratio (T, C, Z, Y, X), as in concatenate: + # a shard is chunks * ratio, elementwise. None takes iohub's defaults (one-frame + # 256 x 256 chunks, one (t, c) frame per shard). Sharding needs OME-Zarr 0.5. + chunks_czyx: list[int] | None = None + shards_ratio: list[int] | None = None # When None, preserve the OME-Zarr version of the input store. output_ome_zarr_version: Literal["0.4", "0.5"] | None = None + @field_validator("chunks_czyx") + @classmethod + def check_chunk_size(cls, v): + if v is not None and (len(v) != 4 or any(i < 1 for i in v)): + raise ValueError("chunks_czyx must be 4 positive integers (C, Z, Y, X)") + return v + + @field_validator("shards_ratio") + @classmethod + def check_shards_ratio(cls, v): + if v is not None and (len(v) != 5 or any(i < 1 for i in v)): + raise ValueError("shards_ratio must be 5 positive integers (T, C, Z, Y, X)") + if v is not None and v[1] != 1: + raise ValueError( + "shards_ratio: iohub does not shard along C; the C entry must be 1" + ) + return v + @model_validator(mode="after") def check_focus(self): focus_users = [name for name, c in self.channels.items() if c.uses_focus] diff --git a/nextflow/configs/a549/project.yml b/nextflow/configs/a549/project.yml index 0fd39474..20c5af85 100644 --- a/nextflow/configs/a549/project.yml +++ b/nextflow/configs/a549/project.yml @@ -55,3 +55,14 @@ channels: focus: channel: Phase3D per: timepoint + +# Output layout, after the Dynamic Cell Atlas array standard (v0.2) where a 2D float32 +# plate can follow it. Chunks are one frame, 256 x 256: the same YX grid as the +# assembled and tracking stores, so a 256-aligned crop reads one chunk per channel in +# every store (the standard's 16-timepoint chunk suits uint8 frames; for float32 it +# would decode 16 frames per single-frame patch). Shards span 16 timepoints and the +# full frame (the standard's floor for the T shard; ~150 MB each here, 30 files per +# position instead of 400). iohub writes one shard's timepoints per unit, so --resume +# is 16 frames coarse and a torn shard loses 16 frames of one channel. +chunks_czyx: [1, 1, 256, 256] +shards_ratio: [16, 1, 1, 7, 5] # YX entries recomputed for THIS output shape: ceil(Y/256), ceil(X/256) diff --git a/settings/example_project_settings.yml b/settings/example_project_settings.yml index d1dc0c6b..c320ccf0 100644 --- a/settings/example_project_settings.yml +++ b/settings/example_project_settings.yml @@ -57,5 +57,14 @@ focus: na_det: 1.35 lambda_ill: 0.5 +# Output chunk shape [C, Z, Y, X] and shards ratio [T, C, Z, Y, X], as in concatenate: a +# shard is chunks * ratio, elementwise (the C entry must be 1). null takes iohub's +# defaults: one-frame 256 x 256 chunks, one (t, c) frame per shard. For a timelapse +# consider shards spanning 16 timepoints and the full frame, e.g. [16, 1, 1, 7, 5] for +# 1664 x 1193 frames (ceil(Y/256), ceil(X/256)); see nextflow/configs/a549/project.yml. +# Sharding needs OME-Zarr 0.5. +chunks_czyx: null +shards_ratio: null + # When null, keep the OME-Zarr version of the input store. output_ome_zarr_version: null diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py index 0a0e1c6d..284e289d 100644 --- a/tests/test_cli/test_project_cli.py +++ b/tests/test_cli/test_project_cli.py @@ -225,6 +225,31 @@ def test_project_refuses_a_plate_that_is_already_2d(create_custom_plate, tmp_pat assert "already has one Z-plane" in str(result.exception) +def test_project_writes_the_configured_chunks_and_shards(plate, tmp_path): + """Shards spanning timepoints: iohub writes one shard's timepoints per unit.""" + plate_path, _ = plate + out = tmp_path / "proj.zarr" + config = _config( + tmp_path, + {"GFP": {"method": "max"}, "Phase3D": {"method": "slice", "z": {"center": 1}}}, + output_ome_zarr_version="0.5", + chunks_czyx=[1, 1, 4, 4], + shards_ratio=[2, 1, 1, 2, 2], # (2, 1, 1, 8, 8): two timepoints per shard + ) + position = plate_path / "A" / "1" / "0" + + result = _run("--cluster", "debug", "-i", position, "-o", out, "-c", config) + + assert result.exit_code == 0, result.output + with open_ome_zarr(str(out / "A" / "1" / "0"), mode="r") as ds: + assert ds.data.chunks == (1, 1, 1, 4, 4) + assert ds.data.shards == (2, 1, 1, 8, 8) + source, _, _ = _read(position) + data, _, _ = _read(out / "A" / "1" / "0") + assert np.array_equal(data[:, 0, 0], source[:, 0].max(axis=1)) # all 3 timepoints + assert np.array_equal(data[:, 1, 0], source[:, 2, 1]) + + def test_project_resume_skips_finished_timepoints(plate, tmp_path): plate_path, _ = plate out = tmp_path / "proj.zarr" diff --git a/tests/test_example_settings.py b/tests/test_example_settings.py index 5d2f6b6d..94feb59d 100644 --- a/tests/test_example_settings.py +++ b/tests/test_example_settings.py @@ -241,6 +241,9 @@ def test_nextflow_a549_project_config_validates(): "membrane_prediction", ] assert settings.channels["Phase3D"].method == "slice" + # One-frame 256 x 256 chunks (the assembled/track grid), 16-timepoint full-frame shards. + assert settings.chunks_czyx == [1, 1, 256, 256] + assert settings.shards_ratio == [16, 1, 1, 7, 5] assert settings.channels["GFP EX488 EM525-45"].method == "max" assert settings.focus.channel == "Phase3D" and settings.focus.per == "timepoint" diff --git a/tests/test_project_settings.py b/tests/test_project_settings.py index 91c92a7a..0409c672 100644 --- a/tests/test_project_settings.py +++ b/tests/test_project_settings.py @@ -125,6 +125,21 @@ def test_focus_has_no_axes(): ) +def test_chunks_and_shards_are_validated(): + s = ProjectionSettings( + channels={"a": {"method": "max"}}, + chunks_czyx=[1, 1, 256, 256], + shards_ratio=[16, 1, 1, 7, 5], + ) + assert s.chunks_czyx == [1, 1, 256, 256] and s.shards_ratio == [16, 1, 1, 7, 5] + with pytest.raises(ValueError, match="4 positive integers"): + ProjectionSettings(channels={"a": {"method": "max"}}, chunks_czyx=[1, 256, 256]) + with pytest.raises(ValueError, match="5 positive integers"): + ProjectionSettings(channels={"a": {"method": "max"}}, shards_ratio=[16, 1, 7, 5]) + with pytest.raises(ValueError, match="C entry must be 1"): + ProjectionSettings(channels={"a": {"method": "max"}}, shards_ratio=[16, 2, 1, 7, 5]) + + def test_unknown_keys_are_refused(): with pytest.raises(ValueError, match="Extra inputs"): ProjectionSettings(channels={"a": {"method": "max", "mode": "max"}}) From dde9e933eb7b475dc6f04f54dda33415c6c081ab Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 16:22:29 -0700 Subject: [PATCH 18/54] perf(project): find the focus of the timepoints in a process pool The focus pre-pass ran the timepoints one after another while the job's other workers idled; with per-timepoint focus on a long timelapse it was the serial part of the step. The timepoints are independent, so with more than one worker they now run in a spawned ProcessPoolExecutor sized by the job's num_workers (the count the projection itself uses), each worker reading its own volume. One worker keeps the in-process loop. Co-Authored-By: Claude Fable 5.1 --- biahub/project.py | 76 +++++++++++++++++++++++------- tests/test_cli/test_project_cli.py | 17 +++++++ 2 files changed, 77 insertions(+), 16 deletions(-) diff --git a/biahub/project.py b/biahub/project.py index e4c37075..d89bdde5 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -7,6 +7,10 @@ track take the projected plate as their 2D input; neither projects by itself. """ +import multiprocessing as mp + +from concurrent.futures import ProcessPoolExecutor +from itertools import repeat from pathlib import Path import click @@ -206,30 +210,66 @@ def _init_output_plate( return (T, C, Z, Y, X), channel_names, dtype +def _focus_of_timepoint( + position_path: Path, + t: int, + channel_index: int, + crop_xy: tuple[int, int], + pixel_size: float, + na_det: float, + lambda_ill: float, +) -> int | None: + """Find the in-focus plane of one timepoint, or None when the frame has no focus. + + Top-level so a process pool can pickle it; each call reads its own (Z, Y, X) volume. + """ + from biahub.registration.estimators import EstimationError + from biahub.registration.methods.focus import center_crop_yx, find_focus + + with open_ome_zarr(str(position_path), mode="r") as dataset: + zyx = np.asarray(dataset.data[t, channel_index]) + try: + return find_focus(center_crop_yx(zyx, crop_xy), pixel_size, na_det, lambda_ill) + except EstimationError: + return None + + def _find_focus_per_t( - position_path: Path, channel_index: int, focus: ProjectionFocusSettings, pixel_size: float + position_path: Path, + channel_index: int, + focus: ProjectionFocusSettings, + pixel_size: float, + num_workers: int = 1, ) -> np.ndarray: """Find the in-focus plane of every timepoint of one position, on ``focus.channel``. Measured on the centre crop with waveorder's transverse band like the engine's - FocusEstimator; an empty frame, or one without a focus, takes the middle plane. One - (Z, Y, X) volume is read at a time. + FocusEstimator; an empty frame, or one without a focus, takes the middle plane. The + timepoints are independent, so with ``num_workers > 1`` they run in a process pool + (spawned, like iohub's, so a worker killed by the cgroup surfaces as an error); + each worker reads one (Z, Y, X) volume at a time. """ - from biahub.registration.estimators import EstimationError - from biahub.registration.methods.focus import center_crop_yx, find_focus - - crop_xy = (focus.center_crop_xy[0], focus.center_crop_xy[1]) with open_ome_zarr(str(position_path), mode="r") as dataset: T, _, Z, _, _ = dataset.data.shape - focus_z = np.empty(T, dtype=int) - for t in range(T): - zyx = np.asarray(dataset.data[t, channel_index]) - try: - focus_z[t] = find_focus( - center_crop_yx(zyx, crop_xy), pixel_size, focus.na_det, focus.lambda_ill + crop_xy = (focus.center_crop_xy[0], focus.center_crop_xy[1]) + constants = (channel_index, crop_xy, pixel_size, focus.na_det, focus.lambda_ill) + + if num_workers > 1 and T > 1: + with ProcessPoolExecutor( + max_workers=min(num_workers, T), mp_context=mp.get_context("spawn") + ) as pool: + planes = list( + pool.map( + _focus_of_timepoint, + repeat(position_path), + range(T), + *(repeat(value) for value in constants), ) - except EstimationError: - focus_z[t] = Z // 2 + ) + else: + planes = [_focus_of_timepoint(position_path, t, *constants) for t in range(T)] + + focus_z = np.array([Z // 2 if plane is None else plane for plane in planes], dtype=int) return np.clip(focus_z, 0, Z - 1) @@ -254,7 +294,11 @@ def _project_position( with open_ome_zarr(str(input_position_path), mode="r") as dataset: pixel_size = dataset.scale[-1] focus_per_t = _find_focus_per_t( - input_position_path, focus_channel_index, settings.focus, pixel_size + input_position_path, + focus_channel_index, + settings.focus, + pixel_size, + num_workers=num_workers, ) if settings.focus.per == "fov": focus_z = int(round(float(np.median(focus_per_t)))) diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py index 284e289d..10d4e8bd 100644 --- a/tests/test_cli/test_project_cli.py +++ b/tests/test_cli/test_project_cli.py @@ -225,6 +225,23 @@ def test_project_refuses_a_plate_that_is_already_2d(create_custom_plate, tmp_pat assert "already has one Z-plane" in str(result.exception) +def test_focus_pre_pass_in_parallel_matches_serial(plate): + """The timepoints are independent; the pool gives the same planes as the loop.""" + from biahub.project import _find_focus_per_t + from biahub.settings import ProjectionFocusSettings + + plate_path, _ = plate + position = plate_path / "A" / "1" / "0" + focus = ProjectionFocusSettings(channel="Phase3D") + + serial = _find_focus_per_t(position, 2, focus, pixel_size=0.1, num_workers=1) + parallel = _find_focus_per_t(position, 2, focus, pixel_size=0.1, num_workers=2) + + assert serial.shape == (3,) + assert np.array_equal(serial, parallel) + assert np.all((0 <= serial) & (serial < Z)) + + def test_project_writes_the_configured_chunks_and_shards(plate, tmp_path): """Shards spanning timepoints: iohub writes one shard's timepoints per unit.""" plate_path, _ = plate From 7cadafdde7847304d6ad7d4cefc349bf6a56f6c5 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 16:46:32 -0700 Subject: [PATCH 19/54] feat(nextflow): projected plate lives beside the assembled one as _proj.zarr The projected plate is the 2D companion of the assembled plate, not a step of its own: the pipeline now writes it to -assemble/_proj.zarr, so project has no step directory and track stays 5-track. Cleanup already covers the shared directory's slurm_output/ and .iohub-progress/; the slurm_output README notes that a directory can hold two steps' stores. The skill docs, run template and config comments follow. Co-Authored-By: Claude Fable 5.1 --- .../skills/reconstruct-with-nextflow/SKILL.md | 13 +++++++------ .../references/caveats.md | 5 +++-- .../references/datasets.md | 6 +++--- .../templates/run_mantis_v2.sh | 4 +++- nextflow/configs/a549/segment.yml | 2 +- nextflow/configs/a549/track.yml | 2 +- nextflow/mantis-v2.nf | 19 +++++++++++-------- nextflow/modules/common.nf | 4 +++- 8 files changed, 32 insertions(+), 23 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index eab3c0f5..b468a372 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -367,9 +367,10 @@ Do not run anything yet. Show the user: **State the directory numbers this run will produce.** The number is the step's position among the steps performed, not a fixed label, so a neuromast run writes `4-assemble` as its last directory and an A549 run writes - `4-assemble`, `5-project` then `6-track`. Older A549 runs on disk say - `5-track` (before the project step) or `5-assemble` / `4-track` (from when - the numbers were fixed) — say so if the user is comparing against one. + `4-assemble` then `5-track`. Project has no directory of its own: its plate + is `4-assemble/_proj.zarr`, beside the assembled plate. Older A549 + runs on disk say `5-assemble` / `4-track`, from when the numbers were fixed + — say so if the user is comparing against one. 7. Known caveats that apply to this dataset. 8. Rough wall-time and that `-resume` is on. 9. Whether Slack notifications are on, and who will be @-mentioned at run end. @@ -544,9 +545,9 @@ directory (`_rerun`) over relaunching in place. ## 11. Wrap up 1. Normalize channel names on the assembled plate, and on the projected plate - when project ran, with `templates/rename_channels.py` (idempotent; - `caveats.md` §2 — check first whether a `rename-channels` CLI has landed on - main, making this obsolete). + beside it when project ran (`DATASET=_proj`, same directory), with + `templates/rename_channels.py` (idempotent; `caveats.md` §2 — check first + whether a `rename-channels` CLI has landed on main, making this obsolete). 2. Verify the assembled store (`-assemble/.zarr`, `4-assemble` unless earlier steps were skipped) opens with iohub; report shape, channels, size on disk. This is the deliverable for every family; tracking, when it diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 9c6356b9..872b627d 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -83,8 +83,9 @@ assembled plate, so **project.yml and track.yml see the pre-rename names** shipped templates avoid the trap for the predictions by putting the suffix directly in the VS `target_channel` (`nuclei_prediction`/`membrane_prediction`), making the rename a no-op for those channels. The projected plate inherits the -pre-rename names too: run `rename_channels.py` on `-project` as well as on -`-assemble` (same command, from the directory holding the store). +pre-rename names too: it lives beside the assembled plate as +`-assemble/_proj.zarr`, so run `rename_channels.py` a second time +from the same directory with `DATASET=_proj`. **Stopgap:** PRs [#260](https://github.com/czbiohub-sf/biahub/pull/260) and [#250](https://github.com/czbiohub-sf/biahub/pull/250) carry a `rename-channels` diff --git a/.claude/skills/reconstruct-with-nextflow/references/datasets.md b/.claude/skills/reconstruct-with-nextflow/references/datasets.md index 89c302df..dbd866e7 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/datasets.md +++ b/.claude/skills/reconstruct-with-nextflow/references/datasets.md @@ -48,7 +48,7 @@ Which root depends on the dataset family: | family | cues in the name | output root | deliverable | |---|---|---|---| | zebrafish / neuromast / dynatrack | `dynatrack`, `hpf`, `dpf`, `neuromast`, `cldnb`, `she`, `myo6b`, `zebrafish` | `/hpc/projects/tlg2_mantis/` | assemble (`4-assemble`) — **no tracking**, see `caveats.md` §4 | -| cell line / organelle / infection | `A549`, `HEK`, `iPSC`, organelle genes (`SEC61B`, `TOMM20`, `G3BP1`, `MAP4`, `CAAX`, `H2B`), viruses (`DENV`, `ZIKV`, `HSV1`) | `/hpc/projects/intracellular_dashboard/organelle_dynamics/` | assemble + project + track (`4-assemble`, `5-project`, `6-track`) | +| cell line / organelle / infection | `A549`, `HEK`, `iPSC`, organelle genes (`SEC61B`, `TOMM20`, `G3BP1`, `MAP4`, `CAAX`, `H2B`), viruses (`DENV`, `ZIKV`, `HSV1`) | `/hpc/projects/intracellular_dashboard/organelle_dynamics/` | assemble + project + track (`4-assemble/.zarr`, `4-assemble/_proj.zarr`, `5-track`) | | instrument QC / calibration | `argolight`, `beads`, `psf`, `alignment`, `illumination`, `fluorescein`, `first_light`, `mantis_v2_*` | *no default* — ask. These are usually not reconstructed. | — | Neighbouring roots that exist and are **not** the default target — do not use @@ -71,8 +71,8 @@ layout is defined once in `nextflow/mantis-v2.nf` (`directory_layout()`): ├── 2-reconstruct/.zarr ├── 3-virtual-stain/.zarr ├── 4-assemble/.zarr # the deliverable -├── 5-project/.zarr # A549 only — the 2D (Z=1) plate track reads -├── 6-track/.zarr # A549 only — not run for neuromast +├── 4-assemble/_proj.zarr # A549 only — the 2D (Z=1) companion plate track reads +├── 5-track/.zarr # A549 only — not run for neuromast ├── qc/ │ ├── report/index.html # one report, one tab per QC'd store │ └── report_spec.yaml # written at launch, before any QC task diff --git a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh index d351e050..4b85ebbd 100755 --- a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh +++ b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh @@ -10,7 +10,9 @@ # below. Keep it there: it is the run's provenance record of the exact command. # # --output is the project root, so each step writes a sibling directory: -# 0-flatfield/ 1-deskew/ 2-reconstruct/ 3-virtual-stain/ 4-assemble/ 5-project/ 6-track/ +# 0-flatfield/ 1-deskew/ 2-reconstruct/ 3-virtual-stain/ 4-assemble/ 5-track/ +# The projected plate is written beside the assembled one as +# 4-assemble/_proj.zarr, not in a directory of its own. # The number is the step's POSITION among the steps this run performs, so a run # without tracking ends at 4-assemble and one that also skipped a step earlier # shifts everything after it down. diff --git a/nextflow/configs/a549/segment.yml b/nextflow/configs/a549/segment.yml index dba0b53f..ff4623ac 100644 --- a/nextflow/configs/a549/segment.yml +++ b/nextflow/configs/a549/segment.yml @@ -1,5 +1,5 @@ # biahub segment: cellpose 4 nuclei labels from the virtual-stain nuclei channel of the -# PROJECTED plate (-project/.zarr, Z=1; mantis v2, 0.1133 um/px). Segment +# PROJECTED plate (-assemble/_proj.zarr, Z=1; mantis v2, 0.1133 um/px). Segment # does not project: a model without eval_args.do_3D segments in 2D and needs a Z=1 # input, which `biahub project` (project.yml) makes from the assembled plate. # Not wired into mantis-v2.nf yet. diff --git a/nextflow/configs/a549/track.yml b/nextflow/configs/a549/track.yml index 24a01f6d..83490dbb 100644 --- a/nextflow/configs/a549/track.yml +++ b/nextflow/configs/a549/track.yml @@ -1,6 +1,6 @@ # Tracking config for the mantis-v2 pipeline (`biahub track`). A549. # -# Track reads the PROJECTED plate (-project/.zarr, Z=1), so target_channel +# Track reads the PROJECTED plate (-assemble/_proj.zarr, Z=1), so target_channel # and cellpose_config.input_channel must name channels that `biahub project` wrote # (project.yml's `channels`). If you rename channels before tracking, update these. # diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf index a0350d7f..5e0db652 100644 --- a/nextflow/mantis-v2.nf +++ b/nextflow/mantis-v2.nf @@ -99,7 +99,8 @@ def step_directories(performed) { reconstruct : 'reconstruct', virtual_stain : 'virtual-stain', assemble : 'assemble', - project : 'project', + // project writes beside the assembled plate (/_proj.zarr), + // not in a directory of its own, so it is not in this list. track : 'track', ] def layout = [:] @@ -129,12 +130,15 @@ workflow { // // Skipping a step DOES renumber the ones after it: the number is a position // among the steps performed, so a neuromast run's assembled store is - // `4-assemble` where an A549 run also has `5-project` and `6-track` after it. + // `4-assemble` where an A549 run also has `5-track` after it. // // Project turns the assembled 3D plate into the 2D (Z=1) plate that 2D // tracking reads: `biahub track` no longer selects or projects Z-planes, so // an A549 run names --project_config as well as --track_config. Tracking // without it reads the assembled plate itself, which suits 3D tracking. + // The projected plate is a companion of the assembled one and lives beside + // it as `-assemble/_proj.zarr`, so project has no step directory + // and does not renumber anything. def assemble_on = params.concatenate_config as boolean def project_on = params.project_config as boolean def track_on = params.track_config as boolean @@ -183,7 +187,6 @@ workflow { // numbering is derived from. Reconstruction proper is always in it. def performed = ['flat_field', 'deskew', 'reconstruct', 'virtual_stain'] if (assemble_on) performed << 'assemble' - if (project_on) performed << 'project' if (track_on) performed << 'track' def layout = step_directories(performed) @@ -202,7 +205,9 @@ workflow { reconstruct_output = "${out}/${layout.reconstruct}/${ds}.zarr" virtual_stain_output = "${out}/${layout.virtual_stain}/${ds}.zarr" assemble_output = assemble_on ? "${out}/${layout.assemble}/${ds}.zarr" : null - project_output = project_on ? "${out}/${layout.project}/${ds}.zarr" : null + // Beside the assembled plate, with a `_proj` suffix: the 2D companion of + // that store, not a step of its own. + project_output = project_on ? "${out}/${layout.assemble}/${ds}_proj.zarr" : null track_output = track_on ? "${out}/${layout.track}/${ds}.zarr" : null // What tracking reads: the projected (Z=1) plate when project runs, else the // assembled plate. Both the plate structure and the pixels come from it. @@ -463,11 +468,9 @@ workflow { def cleanup_paths = [ff_output, deskew_output, reconstruct_output, virtual_stain_output] .collect { zarr -> new File(zarr).parent } def assemble_dir = new File(assemble_output).parent + // The projected plate shares the assemble directory, so these two paths + // cover its slurm_output/ and .iohub-progress/ as well. cleanup_paths << "${assemble_dir}/slurm_output" << "${assemble_dir}/.iohub-progress" - if (project_on) { - def project_dir = new File(project_output).parent - cleanup_paths << "${project_dir}/slurm_output" << "${project_dir}/.iohub-progress" - } if (track_on) cleanup_paths << "${new File(track_output).parent}/slurm_output" def cleanup_list = cleanup_targets(cleanup_paths, out) cleanup_record = "${out}/nextflow/intermediates_cleaned.txt" diff --git a/nextflow/modules/common.nf b/nextflow/modules/common.nf index a3dae616..666d5434 100644 --- a/nextflow/modules/common.nf +++ b/nextflow/modules/common.nf @@ -88,7 +88,9 @@ The real per-task logs are in: ${log_dir}/ -one `_.out` / `.err` pair per task. +one `_.out` / `.err` pair per task. A directory that +holds more than one step's store (assemble and project) shares this +slurm_output/; each step's logs are in its own directory beside the one above. EOF""" } From c8f8e782f118341800640d50c6bc04da32a8ad23 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 17:12:33 -0700 Subject: [PATCH 20/54] Preserve focus metadata Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- biahub/project.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/biahub/project.py b/biahub/project.py index d89bdde5..fca20193 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -315,8 +315,7 @@ def _project_position( with open_ome_zarr(str(output_position_path), mode="r+") as output_dataset: record = dict(output_dataset.zattrs.get(PROVENANCE_KEY, {})) record["focus"] = { - "channel": settings.focus.channel, - "per": settings.focus.per, + **settings.focus.model_dump(mode="json", exclude={"axes"}), "z": focus_z if settings.focus.per == "fov" else None, "z_per_t": focus_per_t.tolist(), } From f06223b181866d982fcb2648d1cba42333061d8b Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 17:18:55 -0700 Subject: [PATCH 21/54] Validate existing output layout before reuse Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- biahub/project.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/biahub/project.py b/biahub/project.py index fca20193..cd3b3a4c 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -190,12 +190,29 @@ def _init_output_plate( methods = {projection.method for projection in settings.channels.values()} dtype = np.dtype(np.float32) if methods & set(_FLOAT_METHODS) else np.dtype(input_dtype) + output_shape = (T, len(settings.channels), 1, Y, X) + output_channels = list(settings.channels) + for input_path in input_position_dirpaths: + output_position = Path(output_dirpath, *Path(input_path).parts[-3:]) + if not output_position.exists(): + continue + with open_ome_zarr(str(output_position), mode="r") as existing: + if ( + existing.channel_names != output_channels + or tuple(existing.data.shape) != output_shape + or np.dtype(existing.data.dtype) != dtype + ): + raise ValueError( + f"{output_position} has incompatible channels, shape or dtype; " + "project into a new output store." + ) + input_plate = Path(input_position_dirpaths[0]).parents[2] create_empty_plate( store_path=output_dirpath, position_keys=[Path(p).parts[-3:] for p in input_position_dirpaths], - channel_names=list(settings.channels), - shape=(T, len(settings.channels), 1, Y, X), + channel_names=output_channels, + shape=output_shape, scale=scale, version=resolve_ome_zarr_version( input_position_dirpaths[0], settings.output_ome_zarr_version From 6e5b5e49432095feef19d8768433038b2692b14b Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 17:19:30 -0700 Subject: [PATCH 22/54] Include channel order in settings fingerprint Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- biahub/project.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/biahub/project.py b/biahub/project.py index cd3b3a4c..e0a9b107 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -346,7 +346,12 @@ def _project_position( output_channel_indices=[list(range(len(plans)))], num_workers=num_workers, resume=resume, - resume_token=settings_fingerprint(settings), + resume_token=settings_fingerprint( + { + "settings": settings.model_dump(mode="json"), + "channel_order": list(settings.channels), + } + ), plans=plans, focus_z=focus_z, ) From 0e504d518753e65d30a133666ad4e65e5aed25e5 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 17:21:25 -0700 Subject: [PATCH 23/54] fix(project): size RAM by the channels read iohub reads one timepoint's (C_read, Z, Y, X) block per worker, so the request must grow with the listed channels: the fixed ram_multiplier of 4 gave 4 GB/CPU whether a config read 2 channels or 6, and the A549 config (6 channels, 3.8 GiB per read, ~5-6 GiB peak) would have hit the cgroup. ram_multiplier is now len(read_indices) + 2 (the read, the window copy and the shard's output buffer): 6 GB/CPU for A549, unchanged for 2 channels. Co-Authored-By: Claude Fable 5.1 --- biahub/project.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/biahub/project.py b/biahub/project.py index e0a9b107..9218e896 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -400,14 +400,18 @@ def project( ) plans, read_indices, focus_index = _resolve_channels(settings, channel_names, Z) - # Not yet calibrated on a completed run. RAM: one timepoint's read (the listed - # channels plus the focus channel) per worker, with headroom for the float64 - # accumulation. Time per (t, c) volume: a reduction is cheap, so 0.1 min carries a - # wide margin; the focus pre-pass reads the focus channel once more per position - # (~40% on top of track's old window selection, #418), hence the higher rate. + # Not yet calibrated on a completed run. RAM per worker: iohub reads one timepoint's + # (C_read, Z, Y, X) block at once, so one ZYX volume per channel read (the listed + # channels plus the focus channel), plus the window copy a reduction makes and the + # shard's output buffer; +2 volumes covers those (6 channels of an A549 position: + # 3.8 GiB read, ~5-6 GiB peak, 6 GB/CPU requested). The focus pre-pass holds one + # volume and its crop per worker, less than that. Time per (t, c) volume: a + # reduction is cheap, so 0.1 min carries a wide margin; the focus pre-pass reads + # the focus channel once more per position (~40% on top of track's old window + # selection, #418), hence the higher rate. time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( shape=(T, len(read_indices), Z, Y, X), - ram_multiplier=4, + ram_multiplier=len(read_indices) + 2, time_multiplier=0.15 if settings.uses_focus else 0.1, max_num_cpus=8, ) From 1bf0071803aec538bea38eac75a3732048051607 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 17:23:03 -0700 Subject: [PATCH 24/54] test(project): focus record carries the focus settings Follows c8f8e782, which dumps the focus settings (minus axes) into the biahub-project focus record next to z and z_per_t. Co-Authored-By: Claude Fable 5.1 --- tests/test_cli/test_project_cli.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py index 10d4e8bd..2db37c06 100644 --- a/tests/test_cli/test_project_cli.py +++ b/tests/test_cli/test_project_cli.py @@ -135,9 +135,13 @@ def test_project_focus_per_fov_takes_the_median_plane(plate, tmp_path, monkeypat data, _, zattrs = _read(out / "A" / "1" / "0") np.testing.assert_allclose(data[:, 0, 0], source[:, 0, 1:4].mean(axis=1, dtype=np.float64)) assert np.array_equal(data[:, 1, 0], source[:, 2, 2]) + # The record carries the focus settings (minus the engine's `axes`) plus the result. assert zattrs["biahub-project"]["focus"] == { "channel": "Phase3D", "per": "fov", + "center_crop_xy": [800, 800], + "na_det": 1.35, + "lambda_ill": 0.5, "z": 2, "z_per_t": [1, 6, 2], } From 88f363e6e53f64b01fcf8c1eb05557609053bfc0 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 19 Oct 2026 17:41:19 -0700 Subject: [PATCH 25/54] fix(project): more RAM headroom per worker, calibrated on an A549 position One position of 2026_08_11 SEC61B v8 (6 channels, 86 x 1664 x 1193 float32, 5 shard units busy) peaked at 22 GB max RSS / 28 GB over the process tree, ~5-6 GiB per worker against the 6 GB/CPU requested. +4 volumes requests 7 GB/CPU there, a margin for a timelapse that keeps all 8 workers busy. The time request (70 min for a 3.6 min run) is left as is. Co-Authored-By: Claude Fable 5.1 --- biahub/project.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/biahub/project.py b/biahub/project.py index 9218e896..be1989af 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -400,18 +400,20 @@ def project( ) plans, read_indices, focus_index = _resolve_channels(settings, channel_names, Z) - # Not yet calibrated on a completed run. RAM per worker: iohub reads one timepoint's - # (C_read, Z, Y, X) block at once, so one ZYX volume per channel read (the listed - # channels plus the focus channel), plus the window copy a reduction makes and the - # shard's output buffer; +2 volumes covers those (6 channels of an A549 position: - # 3.8 GiB read, ~5-6 GiB peak, 6 GB/CPU requested). The focus pre-pass holds one - # volume and its crop per worker, less than that. Time per (t, c) volume: a - # reduction is cheap, so 0.1 min carries a wide margin; the focus pre-pass reads - # the focus channel once more per position (~40% on top of track's old window - # selection, #418), hence the higher rate. + # RAM per worker: iohub reads one timepoint's (C_read, Z, Y, X) block at once, so + # one ZYX volume per channel read (the listed channels plus the focus channel), plus + # the window copy a reduction makes and the shard's output buffer. Measured on an + # A549 position (2026_08_11 SEC61B v8, 6 channels of 86 x 1664 x 1193 float32, 5 + # shard units busy): 22 GB max RSS, 28 GB peak over the process tree, i.e. ~5-6 GiB + # per worker against a 3.8 GiB read. +4 volumes requests 7 GB/CPU there, a margin + # for a longer timelapse that keeps all 8 workers busy. The focus pre-pass holds + # one volume and its crop per worker, less than that. Time per (t, c) volume: the + # same position took 3.6 min against 70 requested, so 0.1 min carries a wide + # margin; the focus pre-pass reads the focus channel once more per position (~40% + # on top of track's old window selection, #418), hence the higher rate. time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( shape=(T, len(read_indices), Z, Y, X), - ram_multiplier=len(read_indices) + 2, + ram_multiplier=len(read_indices) + 4, time_multiplier=0.15 if settings.uses_focus else 0.1, max_num_cpus=8, ) From af5ec1f7bc88fc6db1c082aa76c424ddd2fd10eb Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Tue, 20 Oct 2026 10:35:21 -0700 Subject: [PATCH 26/54] build(deps): ultrack pin follows royerlab/ultrack#290's pyupgrade fix Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d571bfba..ddf44327 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -137,7 +137,7 @@ imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipelin # are painted from, so tracks.geff, the labels and the CSV agree after # `filter_short_sibling_tracks` (track's min_sibling_length). Repoint at # royerlab/ultrack main once #290 merges, and at a release once one contains them. -ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "81415ae53570b095c11dd62286063e45bf887ca5" } +ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5f84ccda3bdf072f5289f939c" } # Pin iohub to main at the czbiohub-sf/iohub#486 merge, for # create_empty_plate(label_names=...), with which track --init scaffolds the empty diff --git a/uv.lock b/uv.lock index 84dc29ae..8753730d 100644 --- a/uv.lock +++ b/uv.lock @@ -448,8 +448,8 @@ requires-dist = [ { name = "torch" }, { name = "tqdm" }, { name = "twine", marker = "extra == 'build'" }, - { name = "ultrack", marker = "extra == 'all'", git = "https://github.com/ieivanov/ultrack?rev=81415ae53570b095c11dd62286063e45bf887ca5" }, - { name = "ultrack", marker = "extra == 'track'", git = "https://github.com/ieivanov/ultrack?rev=81415ae53570b095c11dd62286063e45bf887ca5" }, + { name = "ultrack", marker = "extra == 'all'", git = "https://github.com/ieivanov/ultrack?rev=69e7589a050b0aa5f84ccda3bdf072f5289f939c" }, + { name = "ultrack", marker = "extra == 'track'", git = "https://github.com/ieivanov/ultrack?rev=69e7589a050b0aa5f84ccda3bdf072f5289f939c" }, { name = "waveorder", git = "https://github.com/mehta-lab/waveorder?rev=8d796509b88bafc46a31b81889f6b8034a9050db" }, { name = "zensical", marker = "extra == 'doc'", specifier = ">=0.0.32" }, ] @@ -6878,7 +6878,7 @@ wheels = [ [[package]] name = "ultrack" version = "0.7.3" -source = { git = "https://github.com/ieivanov/ultrack?rev=81415ae53570b095c11dd62286063e45bf887ca5#81415ae53570b095c11dd62286063e45bf887ca5" } +source = { git = "https://github.com/ieivanov/ultrack?rev=69e7589a050b0aa5f84ccda3bdf072f5289f939c#69e7589a050b0aa5f84ccda3bdf072f5289f939c" } dependencies = [ { name = "blosc2" }, { name = "click" }, From 108724a1195b3062220cabab3a2d9e4789bb0666 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Tue, 20 Oct 2026 17:34:55 -0700 Subject: [PATCH 27/54] feat(track)!: cellpose is the default segmentation method TrackingSettings.segmentation_method defaulted to foreground_contour, the legacy path no production config has used since April; a config that omitted the key fell back to it silently. The default is now cellpose, and cellpose without a cellpose_config is refused at load time instead of falling back. The unit tests, which need no weights or GPU, name foreground_contour explicitly. Co-Authored-By: Claude Fable 5.1 --- biahub/settings.py | 17 ++++++++++++++++- settings/example_track_settings.yml | 3 ++- tests/test_cli/test_track_cli.py | 14 ++++++++++++++ 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/biahub/settings.py b/biahub/settings.py index d7ff8dc8..db77ad4d 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -103,7 +103,10 @@ class TrackingSettings(MyBaseModel): output_mode: Literal["2D", "3D"] = "2D" input_images: list[ProcessingInputChannel] tracking_config: dict[str, Any] = {} - segmentation_method: Literal["foreground_contour", "cellpose"] = "foreground_contour" + # What feeds ultrack: cellpose labels (the default, with `cellpose_config`), or a + # foreground mask and contour map computed by the `input_images` processing + # functions (the legacy path; see settings/example_track_settings.yml). + segmentation_method: Literal["foreground_contour", "cellpose"] = "cellpose" # Post-processing (ultrack's filter_short_sibling_tracks, same meaning): at a division # where exactly one daughter lives this many frames or fewer and does not divide # again, drop that daughter and merge the other into the parent. A division whose two @@ -130,6 +133,18 @@ def validate_blank_frames_path(cls, v): return v return Path(v) + @model_validator(mode="after") + def check_segmentation_method(self): + # Without this a config that omits cellpose_config would silently fall back to + # the foreground/contour path at tracking time. + if self.segmentation_method == "cellpose" and self.cellpose_config is None: + raise ValueError( + "segmentation_method is cellpose (the default) but cellpose_config is not " + "set; add a cellpose_config, or set segmentation_method: foreground_contour " + "to track a foreground mask and contour map from input_images." + ) + return self + class EdgeGraphSettings(BaseModel): method: Literal["knn", "radius", "full"] = "knn" diff --git a/settings/example_track_settings.yml b/settings/example_track_settings.yml index b7a8df46..92155c7e 100644 --- a/settings/example_track_settings.yml +++ b/settings/example_track_settings.yml @@ -11,7 +11,8 @@ fov: "*/*/*" target_channel: nuclei_prediction blank_frames_path: blank_frames.csv -# Segmentation method used to feed ultrack: "cellpose" or "foreground_contour". +# Segmentation method used to feed ultrack: "cellpose" (the default; needs +# cellpose_config) or "foreground_contour" (the legacy path). # - cellpose: run a Cellpose model per frame to produce instance labels. # - foreground_contour: derive a foreground mask + contour map (ultrack segments). segmentation_method: cellpose diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index c0655b92..f975c096 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -124,6 +124,8 @@ def _make_tracking_config(plate_path, tmp_path): "output_mode": "2D", "fov": "*/*/*", "target_channel": "nuclei_prediction", + # The foreground/contour path needs no cellpose weights or GPU in the tests. + "segmentation_method": "foreground_contour", "input_images": [ { # Primary data source left null -> resolved from the -i input plate. @@ -378,12 +380,24 @@ def run(): def _minimal_settings(**overrides): base = { "target_channel": "nuclei_prediction", + "segmentation_method": "foreground_contour", "input_images": [{"path": None, "channels": {"nuclei_prediction": []}}], } base.update(overrides) return TrackingSettings(**base) +def test_cellpose_is_the_default_and_needs_its_config(): + base = { + "target_channel": "nuclei_prediction", + "input_images": [{"path": None, "channels": {"nuclei_prediction": []}}], + } + with pytest.raises(ValueError, match="cellpose_config is not set"): + TrackingSettings(**base) + settings = TrackingSettings(**base, cellpose_config={"input_channel": "nuclei_prediction"}) + assert settings.segmentation_method == "cellpose" + + def test_init_labels_take_the_input_z(tmp_path, example_tracking_plate, create_custom_plate): """The label image has the plate's Z: one plane for a projected plate, the stack in 3D.""" flat_path, _ = example_tracking_plate # Z=1 From 16063a14af633228eabe75b98593d777454e6577 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Tue, 20 Oct 2026 17:36:19 -0700 Subject: [PATCH 28/54] feat(pyramid): pyramids on images and label images; mantis-v2 adds them last (#359) `biahub pyramid` is rewritten to the processing-command template and computes pyramids on each position's image and on every label image under labels/, in place: --levels (default 3, the DCA v0.2 target), --method for the images and --label-method (max/stride/mode, default max) for the labels, so no level holds a label id level 0 does not. Spatial axes of size 1 are not downsampled, so the projected Z=1 plate gets factor 2 in Y and X only and keeps its Z scale. --cluster/--monitor/--init replace --local; --init refuses a plate carrying a pyramid of another depth and prints RESOURCES. PyramidSettings validates the options and is the biahub-pyramid provenance record on each position. mantis-v2 gains --pyramid_levels (default 3, 1 = off) and a pyramid module run once per plate, assembled and projected, as the LAST writer of each: after tracking and after QC's compute, since all three rewrite the position's zarr.json. Level 0 is never touched. iohub is pinned to czbiohub-sf/iohub#488 (PositionLabel.compute_pyramid via a shared MultiscaleNode; mode downsampler fix). Co-Authored-By: Claude Fable 5.1 --- .../skills/reconstruct-with-nextflow/SKILL.md | 11 +- .../references/caveats.md | 14 + .../templates/run_mantis_v2.sh | 4 + README.md | 2 +- biahub/cli/main.py | 2 +- biahub/pyramid.py | 288 +++++++++++++----- biahub/settings.py | 21 ++ nextflow/mantis-v2.nf | 90 +++++- nextflow/modules/pyramid.nf | 131 ++++++++ pyproject.toml | 23 +- tests/test_cli/test_pyramid_cli.py | 140 +++++++++ uv.lock | 20 +- 12 files changed, 628 insertions(+), 118 deletions(-) create mode 100644 nextflow/modules/pyramid.nf create mode 100644 tests/test_cli/test_pyramid_cli.py diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index 0d019a28..d5bf27fb 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -3,7 +3,7 @@ name: reconstruct-with-nextflow description: >- Run a mantis-v2 microscope dataset through the biahub Nextflow reconstruction pipeline (nextflow/mantis-v2.nf: flat-field → deskew → reconstruct → - virtual-stain → assemble → project → track → QC) on the Bruno HPC cluster. Assemble, + virtual-stain → assemble → project → track → QC → pyramid) on the Bruno HPC cluster. Assemble, track and QC are optional and selected per dataset family: A549 runs assemble + track + QC, neuromast runs assemble + QC and no tracking. Locates the raw acquisition under /hpc/instruments/cm.mantis, picks the output project @@ -349,8 +349,13 @@ Do not run anything yet. Show the user: | family | steps | |---|---| - | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → project → track → QC** | - | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → QC** (no projection or tracking) | + | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → project → track → QC → pyramid** | + | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → QC → pyramid** (no projection or tracking) | + + Pyramid is on by default whenever assemble runs: it adds two coarser levels + (3 in total, DCA v0.2) to every image and label image of the assembled and + projected plates, in place, once tracking and QC have finished writing into + them. `--pyramid_levels 1` skips it. Write the list out in the plan rather than naming the family, and say which optional steps are being skipped and why. If the user asks to skip anything diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 7cc81a96..5b3f9e5c 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -416,3 +416,17 @@ Consequences worth knowing: differ" means nothing here; `corr` and max abs diff are the measures that do. A real virtual-stain regression moves the distribution, not the last bit. Deskew, reconstruct and assemble are fully deterministic. + +## 15. Pyramids are written last, and are not a resume unit + +`biahub pyramid` rewrites each position's `zarr.json` (the multiscales list and +its `biahub-pyramid` record), as do tracking and imaging-qc, so the pipeline +runs it after both for each plate. A run that fails after the pyramid step +left the plates with 3 levels; a relaunch with `-resume` reuses the cached +pyramid tasks, and a relaunch that recomputes earlier steps rewrites level 0 +only — rerun `biahub pyramid -i /*/*/* --levels 3` by hand in that +case, since the coarser levels are not recomputed by anything else. A plate +that already carries a pyramid of another depth is refused at init +(`already has a pyramid of a different depth`); the fix is iohub's +`delete_pyramid` on each position or matching `--pyramid_levels`. + diff --git a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh index 5cd09c9f..6528641d 100755 --- a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh +++ b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh @@ -202,6 +202,10 @@ unset CLAUDECODE # --qc_track_config was removed: qc_project.yaml counts the labels tracking writes # into the projected plate, so passing the old flag errors at launch. # +# Pyramids are not a flag here: the assembled and projected plates get 3 +# resolution levels (images mean-pooled, labels max-pooled; DCA v0.2, biahub#359) +# after tracking and QC by default. Add `--pyramid_levels 1` to skip them. +# # TO SKIP A STEP, DELETE ITS LINE. Do not comment it out: `#` inside a # backslash-continued command does not start a comment line — the continuation # swallows it, every flag below it is dropped including `-resume`, and bash then diff --git a/README.md b/README.md index b2df1d20..69f873a9 100644 --- a/README.md +++ b/README.md @@ -99,7 +99,7 @@ Some steps include an `estimate-*` command that writes a YAML of parameters you | | | | `flat-field` | Apply flat-field correction to selected channels | | `flip` | Flip images in a dataset | -| `pyramid` | Create multiscale pyramid levels | +| `pyramid` | Compute multiscale pyramids of images and label images, in place | | `process-with-config` | Apply arbitrary YAML-defined functions to a dataset | | | | | `estimate-deskew` | Estimate deskewing parameters | diff --git a/biahub/cli/main.py b/biahub/cli/main.py index d8d14104..a9ed017b 100644 --- a/biahub/cli/main.py +++ b/biahub/cli/main.py @@ -93,7 +93,7 @@ def format_options(self, ctx, formatter): { "name": "pyramid", "import_path": "biahub.pyramid.pyramid_cli", - "help": "Create pyramid levels for a dataset", + "help": "Compute multiscale pyramids of a plate's images and label images, in place", }, { "name": "apply-transform", diff --git a/biahub/pyramid.py b/biahub/pyramid.py index e4aa84b0..7925986c 100644 --- a/biahub/pyramid.py +++ b/biahub/pyramid.py @@ -1,4 +1,12 @@ -import datetime +"""Compute multiscale pyramids in place, on a plate's images and its label images. + +Every position of the input plate gets ``levels`` resolution levels: level 0 is the +existing array and each further level is 2x coarser along the spatial axes longer than +one voxel, so a projected (Z=1) plate is downsampled in Y and X only. Images are pooled +with ``method`` and the label images under ``labels/`` with ``label_method``, one that +never invents a label id. The work is iohub's ``compute_pyramid``, run once per +position; this module fans it out and records what it did. +""" from pathlib import Path @@ -7,134 +15,246 @@ from iohub.ngff import open_ome_zarr +from biahub.cli.monitor import monitor_jobs from biahub.cli.parsing import ( + cluster, + init_only, input_position_dirpaths, - local, + monitor, sbatch_filepath, sbatch_to_submitit, ) +from biahub.settings import LABEL_PYRAMID_METHODS, PYRAMID_METHODS, PyramidSettings from biahub.utils.cluster import ( apply_sbatch_overrides, + echo_resources, estimate_resources, get_submitit_cluster, + local_executor_parameters, ) +PROVENANCE_KEY = "biahub-pyramid" + -def pyramid(fov_path: Path, levels: int, method: str) -> None: +def pyramid_dims(shape_zyx: tuple[int, int, int]) -> set[str]: + """Spatial axes a pyramid downsamples: those longer than one voxel. + + A singleton axis cannot be halved, and scaling its metadata would misreport the + voxel size, so the Z=1 plate ``biahub project`` writes is downsampled in Y and X + only. Time and channel are never downsampled. """ - Create pyramid levels for a single field of view. + return {name for name, size in zip("zyx", shape_zyx, strict=True) if size > 1} - Delegates to iohub's Position.compute_pyramid() which uses cascade downsampling, - where each level is downsampled from the previous level rather than from level 0. - This avoids aliasing artifacts and chunk boundary issues. - Parameters - ---------- - fov_path : Path - Path to the FOV position directory - levels : int - Total number of resolution levels (including level 0). - E.g., levels=4 creates arrays "0", "1", "2", "3". - method : str - Downsampling method (e.g., 'mean', 'max', 'min') +def _check_levels(position_dirpath: Path, levels: int) -> None: + """Refuse a position whose image or labels already have a different pyramid. + + ``compute_pyramid`` recomputes an existing pyramid of the same depth in place, which + is what a retry wants, but a different depth needs ``delete_pyramid`` first; catch + that before any job is submitted, on the first position. """ - click.echo(f"Computing pyramid for FOV: {fov_path}") + with open_ome_zarr(position_dirpath, mode="r") as position: + found = {"image": len(position.array_keys())} + for name, label in position.labels(): + found[f"labels/{name}"] = len(label.array_keys()) + mismatched = {key: n for key, n in found.items() if n not in (1, levels)} + if mismatched: + raise ValueError( + f"{position_dirpath} already has a pyramid of a different depth than the " + f"{levels} levels requested: {mismatched}. Delete it (iohub's delete_pyramid) " + "or request that depth." + ) - with open_ome_zarr(fov_path, mode="r+") as dataset: - dataset.compute_pyramid(levels=levels, method=method) +def _pyramid_position(position_dirpath: Path, settings: PyramidSettings) -> None: + """Compute the pyramid of one position's image and of each of its label images.""" + click.echo(f"Computing pyramid for {position_dirpath}") + with open_ome_zarr(position_dirpath, mode="r+") as position: + dims = pyramid_dims(position.data.shape[-3:]) + position.compute_pyramid(levels=settings.levels, method=settings.method, dims=dims) + for name, label in position.labels(): + click.echo(f"Computing pyramid for labels/{name} of {position_dirpath}") + label.compute_pyramid( + levels=settings.levels, method=settings.label_method, dims=dims + ) + position.zattrs[PROVENANCE_KEY] = {**settings.model_dump(), "dims": sorted(dims)} -@click.command("pyramid") -@input_position_dirpaths() -@sbatch_filepath() -@local() -@click.option( - "--levels", - "-lv", - type=int, - default=4, - show_default=True, - help="Total number of resolution levels including level 0. E.g., levels=4 creates 0, 1, 2, 3.", -) -@click.option( - "--method", - "-m", - type=click.Choice( - [ - "stride", - "median", - "mode", - "mean", - "min", - "max", - ] - ), - default="mean", - show_default=True, - help="Downsampling method to use.", -) -def pyramid_cli( + +def pyramid( input_position_dirpaths: list[Path], - levels: int = 4, + levels: int = 3, method: str = "mean", - sbatch_filepath: Path | None = None, - local: bool = False, + label_method: str = "max", + sbatch_filepath: str | None = None, + cluster: str = "slurm", + monitor: bool = True, + init_only: bool = False, ) -> None: - """Create multi-scale pyramids for OME-Zarr datasets. - - Uses cascade downsampling to generate progressively downscaled pyramid levels. - Each level is 2x downsampled from the previous (e.g., levels=4 creates the original - plus 2x, 4x, and 8x downsampled versions as arrays "0", "1", "2", "3"). + """Compute multiscale pyramids of the given positions, images and labels alike. - >>> biahub pyramid -i ./data.zarr/*/*/* --levels 4 --local + Parameters + ---------- + input_position_dirpaths : list[Path] + Paths to the position directories; the pyramids are written into them. + levels : int, optional + Total number of resolution levels, counting level 0. + method : str, optional + Downsampling method for the images, one of ``PYRAMID_METHODS``. + label_method : str, optional + Downsampling method for the label images, one of ``LABEL_PYRAMID_METHODS``. + sbatch_filepath : str, optional + Path to the SLURM batch file. + cluster : str, optional + Execution cluster: 'slurm' submits to a Slurm cluster, 'local' runs jobs as + subprocesses on this machine, 'debug' runs jobs in-process in the foreground. + monitor : bool, optional + If True, monitor the submitted jobs. + init_only : bool, optional + Only validate the request against the first position and report the resources + a position needs; skip per-position processing. """ - if levels <= 1: - click.echo("No pyramid levels to create (levels must be > 1).") - return + slurm_out_path = input_position_dirpaths[0].parents[3] / "slurm_output" + + settings = PyramidSettings(levels=levels, method=method, label_method=label_method) + _check_levels(input_position_dirpaths[0], settings.levels) - # Estimate resources based on first FOV data shape - with open_ome_zarr(input_position_dirpaths[0], mode="r") as dataset: - T, C, Z, Y, X = dataset.data.shape + with open_ome_zarr(input_position_dirpaths[0], mode="r") as position: + T, C, Z, Y, X = position.data.shape + dtype = position.data.dtype - _, num_cpus, gb_ram = estimate_resources(shape=(T, C, Z, Y, X), ram_multiplier=5) + # iohub downsamples one target shard at a time from the matching source region + # (16 timepoints by the full frame on a projected plate), single-threaded, so the + # working set is a few volumes, not the position. Time per (t, c) volume is well + # under the 0.05 min requested here. + time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( + shape=(T, C, Z, Y, X), + dtype=dtype, + ram_multiplier=4, + time_multiplier=0.05, + max_num_cpus=4, + ) + mem_gb = num_cpus * gb_ram_per_cpu + echo_resources(num_cpus, mem_gb, time_minutes) - cluster = get_submitit_cluster(local) + if init_only: + click.echo( + f"Validated {settings.levels}-level pyramids for {len(input_position_dirpaths)} " + "positions" + ) + return slurm_args = { "slurm_job_name": "pyramid", - "slurm_partition": "cpu", + "slurm_mem": f"{mem_gb}G", "slurm_cpus_per_task": num_cpus, - "slurm_mem_per_cpu": f"{gb_ram}G", - "slurm_time": 30, - "slurm_array_parallelism": 100, + "slurm_array_parallelism": 100, # process up to 100 positions at a time + "slurm_time": time_minutes, + "slurm_partition": "preempted", } - # Override with sbatch file parameters if provided if sbatch_filepath: slurm_args = apply_sbatch_overrides(slurm_args, sbatch_to_submitit(sbatch_filepath)) - slurm_out_path = Path("slurm_output") - slurm_out_path.mkdir(exist_ok=True) - - executor = submitit.AutoExecutor(folder=slurm_out_path, cluster=cluster) + resolved_cluster = get_submitit_cluster(cluster=cluster) + click.echo(f"Preparing jobs on cluster='{resolved_cluster}': {slurm_args}") + executor = submitit.AutoExecutor(folder=slurm_out_path, cluster=resolved_cluster) + if resolved_cluster == "local": + # Explicit --sbatch-filepath values (e.g. '#LOCAL --timeout-min') win over + # the computed local executor parameters. + slurm_args = {**local_executor_parameters(time_minutes), **slurm_args} executor.update_parameters(**slurm_args) - click.echo( - f"Submitting {len(input_position_dirpaths)} pyramid jobs with resources: {slurm_args}" - ) - + click.echo("Submitting jobs...") jobs = [] with submitit.helpers.clean_env(), executor.batch(): - for fov_path in input_position_dirpaths: - job = executor.submit(pyramid, fov_path=fov_path, levels=levels, method=method) - jobs.append(job) + for input_position_path in input_position_dirpaths: + jobs.append(executor.submit(_pyramid_position, input_position_path, settings)) job_ids = [job.job_id for job in jobs] - timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") - log_path = slurm_out_path / f"pyramid-jobs_{timestamp}.log" + slurm_out_path.mkdir(exist_ok=True) + log_path = slurm_out_path / "submitit_jobs_ids.log" with log_path.open("w") as log_file: log_file.write("\n".join(job_ids)) + # submitit's DebugExecutor is lazy: .submit() wraps the callable in a DebugJob but + # execution only happens on .wait()/.done()/.result(). On the Nextflow path + # (--cluster debug) run each position in the foreground. + if resolved_cluster == "debug": + for job, path in zip(jobs, input_position_dirpaths, strict=True): + job.wait() + click.echo(f"Pyramid complete: {path}") + return + + if monitor: + monitor_jobs(jobs, input_position_dirpaths) + + +@click.command("pyramid") +@input_position_dirpaths() +@sbatch_filepath() +@cluster() +@monitor() +@init_only() +@click.option( + "--levels", + type=int, + default=3, + show_default=True, + help="Total number of resolution levels, counting level 0: 3 adds 2x and 4x levels.", +) +@click.option( + "--method", + type=click.Choice(PYRAMID_METHODS), + default="mean", + show_default=True, + help="Downsampling method for the images.", +) +@click.option( + "--label-method", + type=click.Choice(LABEL_PYRAMID_METHODS), + default="max", + show_default=True, + help="Downsampling method for the label images under labels/; none invents a label id.", +) +def pyramid_cli( + input_position_dirpaths: list[Path], + sbatch_filepath: str | None = None, + cluster: str = "slurm", + monitor: bool = False, + init_only: bool = False, + levels: int = 3, + method: str = "mean", + label_method: str = "max", +) -> None: + """Compute multiscale pyramids of a plate's images and label images, in place. + + Each level is 2x coarser than the previous along the spatial axes longer than one + voxel, so a Z=1 plate is downsampled in Y and X only. Images use --method, the label + images under labels/ use --label-method. Rerunning recomputes the same levels. + + \b + SLURM fan-out of positions across a whole plate: + >>> biahub pyramid -i ./plate.zarr/*/*/* --levels 3 + + \b + Validate the request and report the per-position resources only: + >>> biahub pyramid --init -i ./plate.zarr/*/*/* --levels 3 + + \b + In-process run of a single position (e.g. from a Nextflow worker): + >>> biahub pyramid --cluster debug -i ./plate.zarr/A/1/0 --levels 3 + """ # noqa: D301 + pyramid( + input_position_dirpaths=input_position_dirpaths, + levels=levels, + method=method, + label_method=label_method, + sbatch_filepath=sbatch_filepath, + cluster=cluster, + monitor=monitor, + init_only=init_only, + ) + if __name__ == "__main__": pyramid_cli() diff --git a/biahub/settings.py b/biahub/settings.py index d7ff8dc8..7e6658fb 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -473,6 +473,27 @@ def uses_focus(self) -> bool: return any(c.uses_focus for c in self.channels.values()) +PYRAMID_METHODS = ("mean", "median", "mode", "min", "max", "stride") +# Methods that only ever pick a value present in the window, so no level of a +# label pyramid holds an id that level 0 does not (DCA v0.2: "maximum filter or +# decimation filter for label arrays"). +LABEL_PYRAMID_METHODS = ("max", "stride", "mode") + + +class PyramidSettings(MyBaseModel): + """`biahub pyramid`: the pyramid every image and label image of a position gets. + + ``levels`` counts level 0, so ``3`` adds two downsampled levels, each 2x coarser + than the previous along the spatial axes longer than one voxel (Z is left alone + on a Z=1 plate). ``method`` pools the images and ``label_method`` the label + images under ``labels/``. + """ + + levels: Annotated[int, Field(ge=2)] = 3 + method: Literal[PYRAMID_METHODS] = "mean" + label_method: Literal[LABEL_PYRAMID_METHODS] = "max" + + class AffineTransformSettings(MyBaseModel): t_reference: Literal["first", "previous"] = "first" transform_type: Literal["euclidean", "similarity", "affine"] = "euclidean" diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf index 8c637818..50aadc2e 100644 --- a/nextflow/mantis-v2.nf +++ b/nextflow/mantis-v2.nf @@ -23,16 +23,19 @@ nextflow.enable.dsl = 2 // workflow body for why that is possible and what it buys. // // Flat-field → deskew → reconstruct → virtual-stain → assemble → project → -// track → QC is the full chain: assemble concatenates the +// track → QC → pyramid is the full chain: assemble concatenates the // deskew/reconstruct/virtual-stain channels into one plate, project reduces // that 3D plate to the 2D (Z=1) plate 2D tracking needs, track reads it as its -// single input, and QC reads the finished stores. Follow the chaining below for -// the pattern. +// single input, QC reads the finished stores, and pyramid adds the coarser +// resolution levels to the assembled and projected plates once nothing else +// writes into them. Follow the chaining below for the pattern. // -// The last four are OPTIONAL and selected by whether their config is given, so -// a run performs the prefix it asks for: A549 wants assemble + project + track -// + QC, a neuromast run wants assemble + QC and no tracking (issue #306). A -// skipped step never renumbers the directories of the ones around it. +// Assemble, project, track and QC are OPTIONAL and selected by whether their +// config is given, so a run performs the prefix it asks for: A549 wants +// assemble + project + track + QC, a neuromast run wants assemble + QC and no +// tracking (issue #306). Pyramids are on by default for whichever of the two +// plates exist (--pyramid_levels 1 turns them off). A skipped step never +// renumbers the directories of the ones around it. // --------------------------------------------------------------------------- params.input = null // raw source — may not be a zarr store @@ -54,6 +57,11 @@ params.max_positions = 0 params.qc_config = null params.qc_project_config = null params.qc_track_config = null // removed; refused below +// Resolution levels of the assembled and projected plates, counting level 0: +// 3 is the DCA v0.2 target (issue #359) and adds a 2x and a 4x level, in Y and X +// (and Z when the plate has more than one plane), to every image and label +// image. 1 turns the step off. Images are mean-pooled, labels max-pooled. +params.pyramid_levels = 3 include { collect_positions; dataset_name; check_environment } from './modules/common' include { flat_field_init_wf; flat_field_run_wf } from './modules/flat_field' @@ -64,6 +72,12 @@ include { assemble_init_wf; assemble_run_wf } from './modules/assembly' include { project_init_wf; project_run_wf } from './modules/project' include { track_init_wf; track_run_wf } from './modules/tracking' include { qc_plan_wf; qc_compute_wf; qc_report_wf; qc_report_spec } from './modules/qc' +// One module, two plates: a process can only be invoked once per workflow +// context, so the assembled and projected plates get their own aliases. +include { pyramid_init_wf as pyramid_assembled_init_wf; + pyramid_run_wf as pyramid_assembled_run_wf } from './modules/pyramid' +include { pyramid_init_wf as pyramid_projected_init_wf; + pyramid_run_wf as pyramid_projected_run_wf } from './modules/pyramid' include { notify_step; notify_run_start; notify_run_end } from './modules/notify' include { cleanup_decision; cleanup_targets; cleanup_intermediates_wf } from './modules/cleanup' @@ -101,9 +115,9 @@ def step_directories(performed) { reconstruct : 'reconstruct', virtual_stain : 'virtual-stain', assemble : 'assemble', - // project writes beside the assembled plate (/_proj.zarr) - // and track writes into the plate it reads; neither has a directory of - // its own or takes a number. + // project writes beside the assembled plate (/_proj.zarr), + // and track and pyramid write into the plates they read; none has a + // directory of its own or takes a number. ] def layout = [:] order.each { key, name -> @@ -147,6 +161,13 @@ workflow { def qc_image_on = params.qc_config as boolean def qc_project_on = params.qc_project_config as boolean def qc_on = qc_image_on || qc_project_on + // Pyramids go on the deliverable plates, so without assemble there is + // nothing to pyramid and the default is silently moot rather than an error. + def pyramid_levels = params.pyramid_levels as int + if (pyramid_levels < 1) { + error "--pyramid_levels must be at least 1 (1 = no pyramid), got ${params.pyramid_levels}." + } + def pyramid_on = pyramid_levels > 1 && assemble_on // A step cannot outlive the step whose output it reads. Refuse the // combination at launch, naming the config to add or the one to drop, rather @@ -302,6 +323,19 @@ workflow { init_signals << tk_init.done } + if (pyramid_on) { + // Pyramid's init scaffolds nothing: it checks the plate's first position + // against the requested depth and sizes the task. Gated on the init that + // scaffolded each plate; the label images tracking adds later get their + // levels at compute time, from the same depth. + pya_init = pyramid_assembled_init_wf(assemble_output, pyramid_levels, as_init.done) + init_signals << pya_init.done + if (project_on) { + pyp_init = pyramid_projected_init_wf(project_output, pyramid_levels, pj_init.done) + init_signals << pyp_init.done + } + } + // QC's planning verbs belong in this phase for the same reason the step // inits do — `plan-stage` and `estimate-resources` read the store's // structure and validate the config, and read no pixels. Each store's plan @@ -454,6 +488,32 @@ workflow { qc_report = qc_report_wf(qc.done, spec, qc_report_dir) } + // ----- Pyramids --------------------------------------------------------- + // The LAST writer of each plate. `biahub pyramid` rewrites a position's + // multiscales metadata and adds its own provenance record to the position's + // zarr.json, and so does everything else that writes into the plate after + // assemble: tracking (labels/, tracks.geff) and imaging-qc, which stamps the + // tables it wrote into the same position attrs. Two writers of one zarr.json + // lose updates, so pyramids wait for all of them: the assembled plate's for + // tracking when it writes there (track without project), the projected + // plate's for tracking, and both for QC's compute when QC is on. The QC + // report only reads the tables, so it is not waited for. Level 0 is never + // touched, so no downstream reader sees a different image. + if (pyramid_on) { + def assembled_written = (track_on && !project_on) ? track_done.done : assemble_done.done + def assembled_ready = qc_on ? qc.done : assembled_written + pyramid_assembled_done = pyramid_assembled_run_wf(all_positions, assemble_output, + pyramid_levels, pya_init.resources, + assembled_ready) + if (project_on) { + def projected_written = track_on ? track_done.done : project_done.done + def projected_ready = qc_on ? qc.done : projected_written + pyramid_projected_done = pyramid_projected_run_wf(all_positions, project_output, + pyramid_levels, pyp_init.resources, + projected_ready) + } + } + // ----- Cleanup of intermediates (--cleanup_intermediates) --------------- // Once the LAST step has finished, delete what only existed to feed the // assembled plate: the flat-field, deskew, reconstruct and virtual-stain step @@ -490,6 +550,10 @@ workflow { if (project_on) final_signals << project_done.done if (track_on) final_signals << track_done.done if (qc_on) final_signals << qc_report.done + if (pyramid_on) { + final_signals << pyramid_assembled_done.done + if (project_on) final_signals << pyramid_projected_done.done + } final_gate = channel.empty() final_signals.each { signal -> final_gate = final_gate.mix(signal) } @@ -536,6 +600,12 @@ workflow { if (project_on) step_events << [label: 'project', done: project_done.done, output: project_output] if (track_on) step_events << [label: 'track', done: track_done.done, output: track_input] if (qc_on) step_events << [label: 'QC', done: qc_report.done, output: qc_report_dir] + if (pyramid_on) { + step_events << [label: 'pyramid (assembled)', done: pyramid_assembled_done.done, output: assemble_output] + if (project_on) { + step_events << [label: 'pyramid (projected)', done: pyramid_projected_done.done, output: project_output] + } + } if (cleanup_on) step_events << [label: 'cleanup intermediates', done: cleanup_run.done, output: cleanup_record] steps = step_events.collect { event -> event.label } diff --git a/nextflow/modules/pyramid.nf b/nextflow/modules/pyramid.nf new file mode 100644 index 00000000..345a6f9c --- /dev/null +++ b/nextflow/modules/pyramid.nf @@ -0,0 +1,131 @@ +// Pyramid subworkflows: init (validate + resources) and run (fan-out × N). +// +// This module is PATH-AGNOSTIC. Callers pass the plate and the number of levels +// explicitly; the module has no idea where the plate sits in the pipeline +// directory layout or which step wrote it. mantis-v2.nf includes it twice, +// under aliases, for the assembled and the projected plates. +// +// `biahub pyramid` writes IN PLACE: every position of the plate gets `levels` +// resolution levels (level 0 is the existing array), on its image and on every +// label image under `labels/`, so it must run after the LAST step that writes +// pixels or labels into that plate. Level 0 is never touched, so QC of the same +// plate can run concurrently with it. Rerunning recomputes the same levels, +// which is what a retried task wants. +// +// INIT AND RUN ARE SEPARATE SUBWORKFLOWS, like every step's (see the INIT PHASE +// comment in mantis-v2.nf). The init here scaffolds nothing: it checks the first +// position against the requested depth (a plate already carrying a pyramid of +// another depth is refused now rather than after tracking) and prints the +// RESOURCES line one position's task is sized from. +// +// run_pyramid uses `--cluster debug` so that submitit's DebugExecutor runs the +// work in-process; Nextflow owns the per-position fan-out. See modules/project.nf, +// the template this follows. + +include { parse_resources; slurm_logs; slurm_log_dir; retry_time } from './common' + + +process init_pyramid { + label 'cpu_local' + + input: + val input_zarr + val levels + val trigger + + output: + stdout + + script: + """ + mkdir -p "${slurm_log_dir('pyramid')}" + biahub pyramid --init \ + -i "${input_zarr}"/*/*/* \ + --levels ${levels} + """ +} + +process run_pyramid { + // Two plates share this process (under the aliased workflows), so the tag + // names the store as well as the position. + tag "${new File(input_zarr as String).name}/${position}" + label 'cpu' + clusterOptions { slurm_logs('pyramid') } + cpus { meta.cpus } + memory { "${meta.mem_gb} GB" } + time { retry_time(meta.time_minutes, task) } + + input: + tuple val(position), val(meta) + val input_zarr + val levels + + output: + val position + + script: + """ + biahub pyramid --cluster debug \ + -i "${input_zarr}/${position}" \ + --levels ${levels} + """ +} + + +// Validate the request against the plate and size one position's task. +// Metadata-only and cheap, so it belongs in the pipeline's up-front init phase. +// +// take: +// input_zarr path to the plate.zarr the pyramids are written into +// levels total number of resolution levels, counting level 0 +// trigger gating channel — init starts once this emits (the plate exists) +// emit: +// resources the RESOURCES payload sizing one position's task +// done fires once the request has been validated +workflow pyramid_init_wf { + take: + input_zarr + levels + trigger + + main: + init_out = init_pyramid(input_zarr, levels, trigger.collect().map { 'done' }) + + emit: + // Value channels by the contract every step module emits on; see + // flat_field_init_wf for why the guarantee is made at the input. + resources = init_out.map { stdout_text -> parse_resources(stdout_text) } + done = init_out.map { 'done' } +} + + +// Fan out one pyramid task per position. +// +// take: +// positions collected channel of position keys (e.g. ['A/1/0', 'B/1/0']) +// input_zarr path to the plate.zarr the pyramids are written into +// levels total number of resolution levels, counting level 0 +// resources RESOURCES payload from pyramid_init_wf +// prev_done gating channel — the last step writing into the plate has finished +workflow pyramid_run_wf { + take: + positions + input_zarr + levels + resources + prev_done + + main: + // Gate channels are mapped to a token before combining, never combined raw; + // see flat_field_run_wf. + pos_meta = positions + .flatMap { items -> items } + .combine(resources) + .combine(prev_done.map { 'done' }) + .map { pos, meta, _gate -> [pos, meta] } + + py_done = run_pyramid(pos_meta, input_zarr, levels) | collect + + emit: + done = py_done +} diff --git a/pyproject.toml b/pyproject.toml index ddf44327..4fe81902 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,10 +99,10 @@ biahub = "biahub.cli.main:cli" index-strategy = "unsafe-best-match" # waveorder's own [tool.uv.sources] pins iohub to an older main commit # (b7bc648, the czbiohub-sf/iohub#476 merge) than the one biahub needs below -# (d1f9e89, which adds #479). uv refuses two git URLs for one package, so -# force the whole graph onto biahub's rev. Drop together with the iohub pin. +# (the czbiohub-sf/iohub#488 branch). uv refuses two git URLs for one package, +# so force the whole graph onto biahub's rev. Drop together with the iohub pin. override-dependencies = [ - "iohub @ git+https://github.com/czbiohub-sf/iohub@d1f9e89294529d365719cf87a9b29622ca77a992", + "iohub @ git+https://github.com/czbiohub-sf/iohub@6862b10b1aa83426495c264f2a879730f0488103", ] [tool.uv.sources] @@ -139,11 +139,16 @@ imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipelin # royerlab/ultrack main once #290 merges, and at a release once one contains them. ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5f84ccda3bdf072f5289f939c" } -# Pin iohub to main at the czbiohub-sf/iohub#486 merge, for -# create_empty_plate(label_names=...), with which track --init scaffolds the empty -# label images (#360), and czbiohub-sf/iohub#485, which puts the labels list on -# /labels/zarr.json so ome-zarr-models and dca-validate can open a field -# that has labels. Repoint at a PyPI release once one contains them. +# Pin iohub to the czbiohub-sf/iohub#488 branch (feat/label-pyramid, on main at +# the #486 merge): PositionLabel.compute_pyramid, with which `biahub pyramid` +# gives the label images tracking writes the same levels as the image (#359), +# and a fix for the "mode" downsampler. Repoint at main once #488 merges, and at +# a PyPI release once one contains it. +# +# Main's reasons: #486 adds create_empty_plate(label_names=...), with which +# track --init scaffolds the empty label images (#360), and #485 puts the labels +# list on /labels/zarr.json so ome-zarr-models and dca-validate can open a +# field that has labels. # # Main's own reasons: #476 syncs a write unit's shards before recording it as # done (without it a preemption between the write and the marker leaves a `.done` @@ -151,7 +156,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "c9832a7675e0fb997a27db2796623d19b7dc9f1a" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "6862b10b1aa83426495c264f2a879730f0488103" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/tests/test_cli/test_pyramid_cli.py b/tests/test_cli/test_pyramid_cli.py new file mode 100644 index 00000000..1589ff39 --- /dev/null +++ b/tests/test_cli/test_pyramid_cli.py @@ -0,0 +1,140 @@ +"""biahub pyramid: multiscale levels on a plate's images and label images, in place.""" + +import json + +import numpy as np +import pytest + +from click.testing import CliRunner +from iohub.ngff import open_ome_zarr + +from biahub.cli.main import cli +from biahub.pyramid import PROVENANCE_KEY, pyramid_dims + +POSITIONS = (("A", "1", "0"), ("B", "1", "0")) + + +@pytest.fixture +def plate(create_custom_plate, tmp_path): + """(T=3, C=3, Z=4, Y=16, X=12) uint16 with a uint32 label image in each position.""" + plate_path, _ = create_custom_plate( + tmp_path, position_list=POSITIONS, z_size=4, y_size=16, x_size=12 + ) + for key in POSITIONS: + with open_ome_zarr(plate_path.joinpath(*key), mode="r+") as position: + labels = np.zeros((3, 4, 16, 12), dtype=np.uint32) + labels[:, :, 2:10, 1:7] = 5 + labels[:, :, 12:16, 8:12] = 9 + position.create_label("nuclei", labels) + return plate_path + + +def _run(*args): + return CliRunner().invoke(cli, ["pyramid", *map(str, args)]) + + +def _levels(position_path): + with open_ome_zarr(position_path, mode="r") as position: + label = position.get_label("nuclei") + return ( + {level: position[level].shape for level in position.array_keys()}, + {level: label[level].shape for level in label.array_keys()}, + dict(position.zattrs).get(PROVENANCE_KEY), + ) + + +def test_pyramid_dims_skips_singleton_axes(): + assert pyramid_dims((4, 16, 12)) == {"z", "y", "x"} + assert pyramid_dims((1, 16, 12)) == {"y", "x"} + + +def test_pyramid_writes_image_and_label_levels(plate): + positions = [plate.joinpath(*key) for key in POSITIONS] + + result = _run("--cluster", "debug", "-i", *positions, "--levels", "3") + + assert result.exit_code == 0, result.output + for position_path in positions: + image_levels, label_levels, record = _levels(position_path) + assert image_levels == { + "0": (3, 3, 4, 16, 12), + "1": (3, 3, 2, 8, 6), + "2": (3, 3, 1, 4, 3), + } + assert label_levels == {"0": (3, 4, 16, 12), "1": (3, 2, 8, 6), "2": (3, 1, 4, 3)} + assert record == { + "levels": 3, + "method": "mean", + "label_method": "max", + "dims": ["x", "y", "z"], + } + with open_ome_zarr(position_path, mode="r") as position: + label = position.get_label("nuclei") + for level in ("1", "2"): + # Max pooling keeps only ids that exist at level 0, and keeps both blobs. + assert set(np.unique(label[level][:])) == {0, 5, 9} + assert position.get_effective_scale(level)[2:] == [2.0 ** int(level)] * 3 + assert label.get_effective_scale(level)[1:] == [2.0 ** int(level)] * 3 + assert position[1][:].mean() > 0 + + +def test_pyramid_z1_plate_is_downsampled_in_yx_only(create_custom_plate, tmp_path): + plate_path, _ = create_custom_plate( + tmp_path, position_list=POSITIONS[:1], z_size=1, y_size=16, x_size=12 + ) + position_path = plate_path.joinpath(*POSITIONS[0]) + with open_ome_zarr(position_path, mode="r+") as position: + position.create_label("nuclei", np.ones((3, 1, 16, 12), dtype=np.uint32)) + + result = _run("--cluster", "debug", "-i", position_path, "--label-method", "stride") + + assert result.exit_code == 0, result.output + image_levels, label_levels, record = _levels(position_path) + assert image_levels["2"] == (3, 3, 1, 4, 3) + assert label_levels["2"] == (3, 1, 4, 3) + assert record["dims"] == ["x", "y"] + assert record["label_method"] == "stride" + with open_ome_zarr(position_path, mode="r") as position: + # Z keeps its size and its scale at every level. + assert position.get_effective_scale("2")[2] == 1.0 + assert position.get_label("nuclei").get_effective_scale("2")[1] == 1.0 + + +def test_pyramid_rerun_recomputes_and_different_depth_is_refused(plate): + position_path = plate.joinpath(*POSITIONS[0]) + assert _run("--cluster", "debug", "-i", position_path).exit_code == 0 + + rerun = _run("--cluster", "debug", "-i", position_path, "--method", "max") + assert rerun.exit_code == 0, rerun.output + assert _levels(position_path)[2]["method"] == "max" + + deeper = _run("--cluster", "debug", "-i", position_path, "--levels", "4") + assert deeper.exit_code != 0 + assert "different depth" in str(deeper.exception) + assert "'image': 3" in str(deeper.exception) + + +def test_pyramid_init_validates_and_reports_resources(plate): + positions = [plate.joinpath(*key) for key in POSITIONS] + + result = _run("--init", "-i", *positions) + + assert result.exit_code == 0, result.output + line = next(ln for ln in result.output.splitlines() if ln.startswith("RESOURCES:")) + assert set(json.loads(line.removeprefix("RESOURCES:"))) == { + "cpus", + "mem_gb", + "time_minutes", + } + assert "Validated 3-level pyramids for 2 positions" in result.output + # Nothing was written. + image_levels, label_levels, record = _levels(positions[0]) + assert list(image_levels) == ["0"] and list(label_levels) == ["0"] and record is None + + +def test_pyramid_refuses_mean_for_labels_and_one_level(plate): + position_path = plate.joinpath(*POSITIONS[0]) + assert _run("--init", "-i", position_path, "--label-method", "mean").exit_code == 2 + one_level = _run("--init", "-i", position_path, "--levels", "1") + assert one_level.exit_code != 0 + assert "levels" in str(one_level.exception) diff --git a/uv.lock b/uv.lock index 4051eb08..052c3d2a 100644 --- a/uv.lock +++ b/uv.lock @@ -3,25 +3,25 @@ revision = 5 requires-python = ">=3.12" resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", - "python_full_version == '3.14.*' and sys_platform == 'win32'", "python_full_version >= '3.15' and sys_platform == 'emscripten'", - "python_full_version == '3.14.*' and sys_platform == 'emscripten'", "python_full_version >= '3.15' and sys_platform == 'darwin'", - "python_full_version == '3.14.*' and sys_platform == 'darwin'", "python_full_version >= '3.15' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.14.*' and sys_platform == 'win32'", + "python_full_version == '3.14.*' and sys_platform == 'emscripten'", + "python_full_version == '3.14.*' and sys_platform == 'darwin'", "python_full_version == '3.14.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version < '3.13' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version < '3.13' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform == 'darwin'", - "python_full_version < '3.13' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.13' and sys_platform == 'win32'", + "python_full_version < '3.13' and sys_platform == 'emscripten'", + "python_full_version < '3.13' and sys_platform == 'darwin'", "python_full_version < '3.13' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c9832a7675e0fb997a27db2796623d19b7dc9f1a" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=6862b10b1aa83426495c264f2a879730f0488103" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c9832a7675e0fb997a27db2796623d19b7dc9f1a" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=6862b10b1aa83426495c264f2a879730f0488103" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev18+c9832a7" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=c9832a7675e0fb997a27db2796623d19b7dc9f1a#c9832a7675e0fb997a27db2796623d19b7dc9f1a" } +version = "0.3.12.dev19+6862b10" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=6862b10b1aa83426495c264f2a879730f0488103#6862b10b1aa83426495c264f2a879730f0488103" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From c8d2b60de552c58c00a71b5a0107ee9510421c2d Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Tue, 20 Oct 2026 17:51:04 -0700 Subject: [PATCH 29/54] feat(nextflow): QC is the last step; pyramids run after tracking and ahead of it QC's compute is gated on each plate's pyramid (or on tracking / the step that created the plate when pyramids are off) rather than running concurrently with projection and tracking, so imaging-qc's position-attr stamps never race another writer of the same zarr.json and the report describes the finished plates. Co-Authored-By: Claude Fable 5.1 --- .../skills/reconstruct-with-nextflow/SKILL.md | 11 +-- .../references/caveats.md | 4 +- .../templates/run_mantis_v2.sh | 3 +- nextflow/mantis-v2.nf | 81 ++++++++++--------- nextflow/modules/pyramid.nf | 6 +- 5 files changed, 54 insertions(+), 51 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index d5bf27fb..517c0ec0 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -3,7 +3,7 @@ name: reconstruct-with-nextflow description: >- Run a mantis-v2 microscope dataset through the biahub Nextflow reconstruction pipeline (nextflow/mantis-v2.nf: flat-field → deskew → reconstruct → - virtual-stain → assemble → project → track → QC → pyramid) on the Bruno HPC cluster. Assemble, + virtual-stain → assemble → project → track → pyramid → QC) on the Bruno HPC cluster. Assemble, track and QC are optional and selected per dataset family: A549 runs assemble + track + QC, neuromast runs assemble + QC and no tracking. Locates the raw acquisition under /hpc/instruments/cm.mantis, picks the output project @@ -349,13 +349,14 @@ Do not run anything yet. Show the user: | family | steps | |---|---| - | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → project → track → QC → pyramid** | - | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → QC → pyramid** (no projection or tracking) | + | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → project → track → pyramid → QC** | + | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → pyramid → QC** (no projection or tracking) | Pyramid is on by default whenever assemble runs: it adds two coarser levels (3 in total, DCA v0.2) to every image and label image of the assembled and - projected plates, in place, once tracking and QC have finished writing into - them. `--pyramid_levels 1` skips it. + projected plates, in place, once tracking has finished writing into them. + QC is the very last step and waits for the pyramids. `--pyramid_levels 1` + skips them. Write the list out in the plan rather than naming the family, and say which optional steps are being skipped and why. If the user asks to skip anything diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 5b3f9e5c..d8cb5636 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -417,11 +417,11 @@ Consequences worth knowing: real virtual-stain regression moves the distribution, not the last bit. Deskew, reconstruct and assemble are fully deterministic. -## 15. Pyramids are written last, and are not a resume unit +## 15. Pyramids follow tracking and precede QC, and are not a resume unit `biahub pyramid` rewrites each position's `zarr.json` (the multiscales list and its `biahub-pyramid` record), as do tracking and imaging-qc, so the pipeline -runs it after both for each plate. A run that fails after the pyramid step +runs it after tracking and holds QC, the very last step, until it has finished. A run that fails after the pyramid step left the plates with 3 levels; a relaunch with `-resume` reuses the cached pyramid tasks, and a relaunch that recomputes earlier steps rewrites level 0 only — rerun `biahub pyramid -i /*/*/* --levels 3` by hand in that diff --git a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh index 6528641d..d5b583b4 100755 --- a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh +++ b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh @@ -204,7 +204,8 @@ unset CLAUDECODE # # Pyramids are not a flag here: the assembled and projected plates get 3 # resolution levels (images mean-pooled, labels max-pooled; DCA v0.2, biahub#359) -# after tracking and QC by default. Add `--pyramid_levels 1` to skip them. +# after tracking by default, and QC runs last, after them. Add +# `--pyramid_levels 1` to skip them. # # TO SKIP A STEP, DELETE ITS LINE. Do not comment it out: `#` inside a # backslash-continued command does not start a comment line — the continuation diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf index 50aadc2e..3ce9bb96 100644 --- a/nextflow/mantis-v2.nf +++ b/nextflow/mantis-v2.nf @@ -23,12 +23,13 @@ nextflow.enable.dsl = 2 // workflow body for why that is possible and what it buys. // // Flat-field → deskew → reconstruct → virtual-stain → assemble → project → -// track → QC → pyramid is the full chain: assemble concatenates the +// track → pyramid → QC is the full chain: assemble concatenates the // deskew/reconstruct/virtual-stain channels into one plate, project reduces // that 3D plate to the 2D (Z=1) plate 2D tracking needs, track reads it as its -// single input, QC reads the finished stores, and pyramid adds the coarser -// resolution levels to the assembled and projected plates once nothing else -// writes into them. Follow the chaining below for the pattern. +// single input, pyramid adds the coarser resolution levels to the assembled +// and projected plates once nothing else writes pixels or labels into them, +// and QC, the very last step, reads the finished stores. Follow the chaining +// below for the pattern. // // Assemble, project, track and QC are OPTIONAL and selected by whether their // config is given, so a run performs the prefix it asks for: A549 wants @@ -451,12 +452,35 @@ workflow { project_on ? project_done.done : assemble_done.done) } + // ----- Pyramids --------------------------------------------------------- + // After the LAST step that writes pixels or labels into each plate: the + // assembled plate's tracking when it writes there (track without project), + // else assemble; the projected plate's tracking, else project. `biahub + // pyramid` rewrites a position's multiscales metadata and adds its own + // record to the position's zarr.json, so it must not overlap another writer + // of that file, and QC (which stamps the tables it wrote into the same + // attrs) is held until it has finished. Level 0 is never touched. + if (pyramid_on) { + def assembled_written = (track_on && !project_on) ? track_done.done : assemble_done.done + pyramid_assembled_done = pyramid_assembled_run_wf(all_positions, assemble_output, + pyramid_levels, pya_init.resources, + assembled_written) + if (project_on) { + def projected_written = track_on ? track_done.done : project_done.done + pyramid_projected_done = pyramid_projected_run_wf(all_positions, project_output, + pyramid_levels, pyp_init.resources, + projected_written) + } + } + // ----- QC compute ------------------------------------------------------- - // Planned in the init phase; the compute it planned is gated on the last - // step that writes each store: the ASSEMBLED plate on assemble_done (so its - // pixel QC runs concurrently with projection and tracking), and the - // PROJECTED plate on track_done when tracking is on (its labels are counted - // there), else project_done. + // THE VERY LAST STEP. Planned in the init phase; the compute it planned is + // gated on the last writer of each store: its pyramid when pyramids are on, + // else tracking where it writes into that store (the projected plate, or the + // assembled one for 3D tracking), else the step that created it. QC used to + // run concurrently with projection and tracking; it now waits, so that + // imaging-qc's stamps on the position attrs never race another writer of the + // same zarr.json and the report describes the finished plates. // // A QC verdict cannot fail the pipeline: `imaging-qc gate` exits 0 whether // positions pass or fail, recording the verdict in each store's own @@ -475,45 +499,22 @@ workflow { // combined, for the reason spelled out at qc_plan_inputs above. qc_compute_ready = channel.empty() if (qc_image_on) { + def assembled_final = pyramid_on ? pyramid_assembled_done.done + : (track_on && !project_on) ? track_done.done : assemble_done.done qc_compute_ready = qc_compute_ready.mix( - assemble_done.done.map { tuple(assemble_output, params.qc_config) } ) + assembled_final.map { tuple(assemble_output, params.qc_config) } ) } if (qc_project_on) { - def proj_written = track_on ? track_done.done : project_done.done + def projected_final = pyramid_on ? pyramid_projected_done.done + : track_on ? track_done.done : project_done.done qc_compute_ready = qc_compute_ready.mix( - proj_written.map { tuple(project_output, params.qc_project_config) } ) + projected_final.map { tuple(project_output, params.qc_project_config) } ) } qc = qc_compute_wf(qc_plan.items, qc_plan.stores, qc_compute_ready) qc_report = qc_report_wf(qc.done, spec, qc_report_dir) } - // ----- Pyramids --------------------------------------------------------- - // The LAST writer of each plate. `biahub pyramid` rewrites a position's - // multiscales metadata and adds its own provenance record to the position's - // zarr.json, and so does everything else that writes into the plate after - // assemble: tracking (labels/, tracks.geff) and imaging-qc, which stamps the - // tables it wrote into the same position attrs. Two writers of one zarr.json - // lose updates, so pyramids wait for all of them: the assembled plate's for - // tracking when it writes there (track without project), the projected - // plate's for tracking, and both for QC's compute when QC is on. The QC - // report only reads the tables, so it is not waited for. Level 0 is never - // touched, so no downstream reader sees a different image. - if (pyramid_on) { - def assembled_written = (track_on && !project_on) ? track_done.done : assemble_done.done - def assembled_ready = qc_on ? qc.done : assembled_written - pyramid_assembled_done = pyramid_assembled_run_wf(all_positions, assemble_output, - pyramid_levels, pya_init.resources, - assembled_ready) - if (project_on) { - def projected_written = track_on ? track_done.done : project_done.done - def projected_ready = qc_on ? qc.done : projected_written - pyramid_projected_done = pyramid_projected_run_wf(all_positions, project_output, - pyramid_levels, pyp_init.resources, - projected_ready) - } - } - // ----- Cleanup of intermediates (--cleanup_intermediates) --------------- // Once the LAST step has finished, delete what only existed to feed the // assembled plate: the flat-field, deskew, reconstruct and virtual-stain step @@ -549,11 +550,11 @@ workflow { def final_signals = [assemble_done.done] if (project_on) final_signals << project_done.done if (track_on) final_signals << track_done.done - if (qc_on) final_signals << qc_report.done if (pyramid_on) { final_signals << pyramid_assembled_done.done if (project_on) final_signals << pyramid_projected_done.done } + if (qc_on) final_signals << qc_report.done final_gate = channel.empty() final_signals.each { signal -> final_gate = final_gate.mix(signal) } @@ -599,13 +600,13 @@ workflow { if (assemble_on) step_events << [label: 'assemble', done: assemble_done.done, output: assemble_output] if (project_on) step_events << [label: 'project', done: project_done.done, output: project_output] if (track_on) step_events << [label: 'track', done: track_done.done, output: track_input] - if (qc_on) step_events << [label: 'QC', done: qc_report.done, output: qc_report_dir] if (pyramid_on) { step_events << [label: 'pyramid (assembled)', done: pyramid_assembled_done.done, output: assemble_output] if (project_on) { step_events << [label: 'pyramid (projected)', done: pyramid_projected_done.done, output: project_output] } } + if (qc_on) step_events << [label: 'QC', done: qc_report.done, output: qc_report_dir] if (cleanup_on) step_events << [label: 'cleanup intermediates', done: cleanup_run.done, output: cleanup_record] steps = step_events.collect { event -> event.label } diff --git a/nextflow/modules/pyramid.nf b/nextflow/modules/pyramid.nf index 345a6f9c..ee10bdbf 100644 --- a/nextflow/modules/pyramid.nf +++ b/nextflow/modules/pyramid.nf @@ -8,9 +8,9 @@ // `biahub pyramid` writes IN PLACE: every position of the plate gets `levels` // resolution levels (level 0 is the existing array), on its image and on every // label image under `labels/`, so it must run after the LAST step that writes -// pixels or labels into that plate. Level 0 is never touched, so QC of the same -// plate can run concurrently with it. Rerunning recomputes the same levels, -// which is what a retried task wants. +// pixels or labels into that plate, and before QC, which stamps the same +// position attrs it rewrites. Level 0 is never touched. Rerunning recomputes +// the same levels, which is what a retried task wants. // // INIT AND RUN ARE SEPARATE SUBWORKFLOWS, like every step's (see the INIT PHASE // comment in mantis-v2.nf). The init here scaffolds nothing: it checks the first From 936c1474ad125d40bdce9732563042b3915cfaf7 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Tue, 20 Oct 2026 18:14:15 -0700 Subject: [PATCH 30/54] build(deps): iohub pin follows the #488 review fixes (label-safe defaults, pyramid_dims) Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 4 ++-- uv.lock | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4fe81902..d48c7413 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,7 +102,7 @@ index-strategy = "unsafe-best-match" # (the czbiohub-sf/iohub#488 branch). uv refuses two git URLs for one package, # so force the whole graph onto biahub's rev. Drop together with the iohub pin. override-dependencies = [ - "iohub @ git+https://github.com/czbiohub-sf/iohub@6862b10b1aa83426495c264f2a879730f0488103", + "iohub @ git+https://github.com/czbiohub-sf/iohub@876b50eeef29b41a19e88832b3aea8019993d978", ] [tool.uv.sources] @@ -156,7 +156,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "6862b10b1aa83426495c264f2a879730f0488103" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "876b50eeef29b41a19e88832b3aea8019993d978" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index 052c3d2a..fb76f388 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=6862b10b1aa83426495c264f2a879730f0488103" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=876b50eeef29b41a19e88832b3aea8019993d978" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=6862b10b1aa83426495c264f2a879730f0488103" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=876b50eeef29b41a19e88832b3aea8019993d978" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev19+6862b10" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=6862b10b1aa83426495c264f2a879730f0488103#6862b10b1aa83426495c264f2a879730f0488103" } +version = "0.3.12.dev20+876b50e" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=876b50eeef29b41a19e88832b3aea8019993d978#876b50eeef29b41a19e88832b3aea8019993d978" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 098b16f4e36bcd47dc7ec9333bb07ac69605da62 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Tue, 20 Oct 2026 18:46:55 -0700 Subject: [PATCH 31/54] build(deps): iohub pin follows the second #488 review round Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 4 ++-- uv.lock | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d48c7413..91470495 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,7 +102,7 @@ index-strategy = "unsafe-best-match" # (the czbiohub-sf/iohub#488 branch). uv refuses two git URLs for one package, # so force the whole graph onto biahub's rev. Drop together with the iohub pin. override-dependencies = [ - "iohub @ git+https://github.com/czbiohub-sf/iohub@876b50eeef29b41a19e88832b3aea8019993d978", + "iohub @ git+https://github.com/czbiohub-sf/iohub@97bdc5c73f182f9204e0987d15c308379150278e", ] [tool.uv.sources] @@ -156,7 +156,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "876b50eeef29b41a19e88832b3aea8019993d978" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "97bdc5c73f182f9204e0987d15c308379150278e" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index fb76f388..7fe69d55 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=876b50eeef29b41a19e88832b3aea8019993d978" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=97bdc5c73f182f9204e0987d15c308379150278e" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=876b50eeef29b41a19e88832b3aea8019993d978" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=97bdc5c73f182f9204e0987d15c308379150278e" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev20+876b50e" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=876b50eeef29b41a19e88832b3aea8019993d978#876b50eeef29b41a19e88832b3aea8019993d978" } +version = "0.3.12.dev21+97bdc5c" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=97bdc5c73f182f9204e0987d15c308379150278e#97bdc5c73f182f9204e0987d15c308379150278e" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From dccb1f8cd96d2c18e314920e89b5b3b32740adb5 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Tue, 20 Oct 2026 19:02:09 -0700 Subject: [PATCH 32/54] build(deps): iohub pin follows the third #488 review round Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 4 ++-- uv.lock | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 91470495..bc068623 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,7 +102,7 @@ index-strategy = "unsafe-best-match" # (the czbiohub-sf/iohub#488 branch). uv refuses two git URLs for one package, # so force the whole graph onto biahub's rev. Drop together with the iohub pin. override-dependencies = [ - "iohub @ git+https://github.com/czbiohub-sf/iohub@97bdc5c73f182f9204e0987d15c308379150278e", + "iohub @ git+https://github.com/czbiohub-sf/iohub@d8e649e2fe9de3b126854430916149dde745c05c", ] [tool.uv.sources] @@ -156,7 +156,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "97bdc5c73f182f9204e0987d15c308379150278e" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "d8e649e2fe9de3b126854430916149dde745c05c" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index 7fe69d55..7279c2c7 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=97bdc5c73f182f9204e0987d15c308379150278e" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=d8e649e2fe9de3b126854430916149dde745c05c" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=97bdc5c73f182f9204e0987d15c308379150278e" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=d8e649e2fe9de3b126854430916149dde745c05c" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev21+97bdc5c" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=97bdc5c73f182f9204e0987d15c308379150278e#97bdc5c73f182f9204e0987d15c308379150278e" } +version = "0.3.12.dev22+d8e649e" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=d8e649e2fe9de3b126854430916149dde745c05c#d8e649e2fe9de3b126854430916149dde745c05c" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 1b5b1752e3ceaa7f9d156c1da652f8a5f6dd0aa1 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Tue, 20 Oct 2026 19:10:13 -0700 Subject: [PATCH 33/54] feat(segment,track)!: segmentation moves out of track into the plate The pipeline becomes assemble -> project -> segment -> track: `biahub segment` writes one OME-NGFF label image per model into the plate it reads (`labels/`, no -o; -o keeps the separate label plate) through iohub's `process_single_position(output_label=...)` (czbiohub-sf/iohub#489), so a preempted position resumes. `biahub track` takes that label image as its input (`labels: `, `labels_sigma`) and writes the tracked labels beside it as `labels/_tracked`, with tracks.geff and the CSV. The cellpose code, `cellpose_config` and `target_channel` leave track (both keys are rejected with a pointer to the segment config); `foreground_contour` stays as the legacy path and keeps `target_channel` and `input_images`. Blank frames listed in the CSV are filled with the nearest frame's labels, as the images were before. Track no longer needs a GPU: the `cpu` partition (a new `cpu_reserved` Nextflow label; ultrack cannot resume). mantis-v2.nf gains --segment_config between project and track; segment's init creates the label images so track's init can check that its input exists, and qc_project.yaml counts `labels/nuclei_tracked`. The A549 segment.yml takes the cellpose values tracking was validated with (cpsam_v2, diameter 80, min_size 2500). Label chunking helpers move to biahub.utils.ngff. Co-Authored-By: Claude Fable 5.1 --- .../skills/reconstruct-with-nextflow/SKILL.md | 92 ++--- .../references/caveats.md | 41 ++- .../references/datasets.md | 12 +- .../templates/check_qc_channels.py | 101 ++++-- .../templates/run_mantis_v2.sh | 25 +- README.md | 12 +- biahub/segment.py | 189 ++++++++-- biahub/settings.py | 111 +++--- biahub/track.py | 329 ++++++------------ biahub/utils/ngff.py | 89 ++++- nextflow/configs/a549/qc_project.yaml | 12 +- nextflow/configs/a549/segment.yml | 25 +- nextflow/configs/a549/track.yml | 41 +-- nextflow/mantis-v2.nf | 141 +++++--- nextflow/modules/segment.nf | 49 ++- nextflow/modules/tracking.nf | 47 +-- nextflow/nextflow.config | 13 +- pyproject.toml | 15 +- settings/example_segmentation_settings.yml | 4 +- settings/example_track_settings.yml | 124 +++---- tests/test_cellpose_inputs.py | 25 -- tests/test_cli/test_segment_cli.py | 144 ++++++++ tests/test_cli/test_track_cli.py | 240 ++++++++----- tests/test_example_settings.py | 10 +- tests/test_segment_settings.py | 22 ++ uv.lock | 20 +- 26 files changed, 1183 insertions(+), 750 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index 0d019a28..c78497a6 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -3,9 +3,10 @@ name: reconstruct-with-nextflow description: >- Run a mantis-v2 microscope dataset through the biahub Nextflow reconstruction pipeline (nextflow/mantis-v2.nf: flat-field → deskew → reconstruct → - virtual-stain → assemble → project → track → QC) on the Bruno HPC cluster. Assemble, - track and QC are optional and selected per dataset family: A549 runs assemble - + track + QC, neuromast runs assemble + QC and no tracking. Locates the raw + virtual-stain → assemble → project → segment → track → QC) on the Bruno HPC + cluster. Assemble, project, segment, track and QC are optional and selected per + dataset family: A549 runs all of them, neuromast runs assemble + QC and no + segmentation or tracking. Locates the raw acquisition under /hpc/instruments/cm.mantis, picks the output project directory, scaffolds configs, launches the run in a tmux session in the foreground so it can be watched, recovers from Lustre/torn-shard I/O errors, @@ -174,8 +175,9 @@ say which you think it is and let the user decide. A fresh reprocess goes to a The templates in `/nextflow/configs/{a549,zebrafish}/` are the source of truth — they version with the pipeline and are reviewed. The QC config sits with -them, one self-contained `qc.yaml` for the assembled store; A549's also counts -cells in the label image tracking writes into that store. It carries `.yaml` +them, one self-contained `qc.yaml` for the assembled store; A549's +`qc_project.yaml` also counts cells in the tracked label image tracking writes +into the projected plate. It carries `.yaml` rather than `.yml` because imaging-qc loads a stage config through Hydra, which strips the extension and looks for `.yaml`. A previous run's `configs/` directory is a fallback only: an unreviewed snapshot that drifts @@ -191,7 +193,7 @@ the next dataset in the family would also need), fix it in a PR against `nextflow/configs//`, stating which dataset exposed it. Genuinely per-dataset values stay in the run directory. -### 5a. Check the channel couplings: QC filters, project.yml, track.yml +### 5a. Check the channel couplings: QC filters, project.yml, segment.yml, track.yml Once the configs are edited, run this from the biahub checkout: @@ -218,17 +220,19 @@ warning, not a failure (`biahub project` drops it with a warning), **but the intention of the pipeline is that every assembled channel is projected unless the user decided otherwise**, so carry the list into the pre-run plan (step 6 of §7) and ask the user whether to add each one or leave it out on purpose. -The script also checks that `track.yml`'s `target_channel` and -`cellpose_config.input_channel` are channels `project.yml` writes, since track -reads the projected plate. - -**`qc_project.yaml` against the projected plate and `track.yml` — the coupling -that actually breaks.** Tracking writes its labels into each position of the -plate it reads, at `labels/`, and `qc_project.yaml`'s -`instance_count` names them twice: `label: nuclei_prediction` and its anchor -`channels: [nuclei_prediction]`, the channel they were segmented from. Both only -resolve if `track.yml` says `target_channel: nuclei_prediction`, and the anchor -must be a channel `project.yml` writes — three files with no shared source of +The script also checks that `segment.yml`'s model `channels` are channels +`project.yml` writes, since segment reads the projected plate, and that +`track.yml`'s `labels` is a `segment.yml` model name, since track reads the +label image segment wrote (`labels/`). + +**`qc_project.yaml` against the projected plate, `segment.yml` and `track.yml` +— the coupling that actually breaks.** Segment writes `labels/` into each +position of the plate it reads and tracking writes `labels/_tracked` +beside it. `qc_project.yaml`'s `instance_count` names the tracked one, `label: +nuclei_tracked`, and its anchor `channels: [nuclei_prediction]`, the channel the +labels were segmented from. The label only resolves if `track.yml` says `labels: +nuclei` and `segment.yml` has a model `nuclei`; the anchor must be that model's +channel and a channel `project.yml` writes — four files with no shared source of truth, and changing one is the ordinary way to desynchronise them. The script checks all of it when `qc_project.yaml` exists. @@ -341,16 +345,18 @@ Do not run anything yet. Show the user: what was mismatched and how it was resolved. 6. **The steps this run will perform, listed in order.** Reconstruction proper — flat-field → deskew → reconstruct → virtual-stain — always runs; assemble, - project, track and QC run only if their config is passed (biahub#306). - Project turns the assembled 3D plate into the 2D (Z=1) plate that 2D - tracking reads (`biahub track` no longer selects or projects Z-planes), so - `--track_config` goes with `--project_config`. Defaults by family, which is + project, segment, track and QC run only if their config is passed + (biahub#306). Project turns the assembled 3D plate into the 2D (Z=1) plate + that 2D segmentation and tracking read (neither selects or projects + Z-planes); segment writes the cellpose labels into that plate and track + links them (`biahub track` no longer segments), so `--segment_config` and + `--track_config` go with `--project_config`. Defaults by family, which is what to present unless the user says otherwise: | family | steps | |---|---| - | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → project → track → QC** | - | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → QC** (no projection or tracking) | + | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → project → segment → track → QC** | + | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → QC** (no projection, segmentation or tracking) | Write the list out in the plan rather than naming the family, and say which optional steps are being skipped and why. If the user asks to skip anything @@ -371,9 +377,10 @@ Do not run anything yet. Show the user: **State the directory numbers this run will produce.** The number is the step's position among the steps performed, not a fixed label, so a neuromast run and an A549 run both end at `4-assemble`: project writes its plate - beside the assembled one as `4-assemble/_proj.zarr`, and tracking - writes its labels and tracks into that plate; neither has a directory of - its own. Older A549 runs on disk also have `5-track` (or `5-assemble` / + beside the assembled one as `4-assemble/_proj.zarr`, segment + writes its label images into that plate and tracking its tracked labels and + tracks; none of the three has a directory of its own. Older A549 runs on + disk also have `5-track` (or `5-assemble` / `4-track`, from when the numbers were fixed) — say so if the user is comparing against one. 7. Known caveats that apply to this dataset. @@ -390,7 +397,8 @@ Do not run anything yet. Show the user: pipeline deletes the flat-field, deskew, reconstruct and virtual-stain directories, the `slurm_output/` placeholders and `.iohub-progress/` resume markers beside the final stores, and the Nextflow work directory. Only the - assembled store (with tracking's labels and tracks inside it), the QC report + assembled and projected stores (with the segmentation and tracking labels + and the tracks inside the latter), the QC report and the logs under `nextflow/` remain, and the run is final: any rerun recomputes from raw. The template's `CLEANUP_INTERMEDIATES="auto"` lets the pipeline decide at launch (`cleanup_decision` in @@ -427,15 +435,16 @@ copied file's group bits follow the *source* file's mode — `cp` from the repo (644) yields `rw-r--r--`. Hence the `chmod` after every copy. Plain `cp` only: `cp -p`/`-a` would also carry over the checkout's group. -The run script passes all five optional configs — `--concatenate_config`, -`--project_config`, `--track_config`, `--qc_config` (`qc.yaml`), -`--qc_project_config` (`qc_project.yaml`) — so an A549 run needs no edit. **For a -neuromast/zebrafish run, or any step the user asked to skip, DELETE that flag's -line** from the `nextflow run` call and note the skip in a comment above it, so -the script still records what this run did. A neuromast run deletes -`--project_config`, `--track_config` and `--qc_project_config`; none of -`project.yml`, `track.yml` or `qc_project.yaml` need exist for that family, and -none does for zebrafish. `--qc_track_config` is gone: passing it errors at launch. +The run script passes all six optional configs — `--concatenate_config`, +`--project_config`, `--segment_config`, `--track_config`, `--qc_config` +(`qc.yaml`), `--qc_project_config` (`qc_project.yaml`) — so an A549 run needs no +edit. **For a neuromast/zebrafish run, or any step the user asked to skip, +DELETE that flag's line** from the `nextflow run` call and note the skip in a +comment above it, so the script still records what this run did. A neuromast +run deletes `--project_config`, `--segment_config`, `--track_config` and +`--qc_project_config`; none of `project.yml`, `segment.yml`, `track.yml` or +`qc_project.yaml` need exist for that family, and none does for zebrafish. +`--qc_track_config` is gone: passing it errors at launch. Delete rather than comment: a `#` inside a backslash-continued command does not start a comment line — the continuation swallows it, every flag below is dropped @@ -555,11 +564,12 @@ directory (`_rerun`) over relaunching in place. whether a `rename-channels` CLI has landed on main, making this obsolete). 2. Verify the assembled store (`-assemble/.zarr`, `4-assemble` unless earlier steps were skipped) opens with iohub; report shape, channels, - size on disk. This is the deliverable for every family. When tracking ran, - its output is inside it: check a position has `labels/nuclei_prediction` - (the track config's `target_channel`; uint32 TZYX at the image's full Z), - `tracks.geff` and `tracks___.csv`, and report them beside - the image data. + size on disk. This is the deliverable for every family. When segmentation + and tracking ran, their output is inside the projected plate: check a + position has `labels/nuclei` (the segment config's model name) and + `labels/nuclei_tracked` (the track config's `labels` plus `_tracked`; both + uint32 TZYX at the image's full Z), `tracks.geff` and + `tracks___.csv`, and report them beside the image data. 3. If QC ran, report its verdict: the `QC_SUMMARY` line from the pipeline log (`pass=`/`fail=`/`gates_fail=`), and point at the report at `/qc/report/index.html` — one page, one tab: the assembled store diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 7cc81a96..02a6b9ec 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -143,30 +143,45 @@ model silently stalls. Top level may contain **only** `model`, `data`, per-position — leave it out. `target_channel` carries the `_prediction` suffix deliberately (§2). +**`segment.yml` — schema.** `SegmentationSettings`: `models.` with +`pretrained_model`, `channels` (projected-plate channel names, at most 3) and +`eval_args` (cellpose `CellposeModel.eval` arguments: `diameter`, +`cellprob_threshold`, `flow_threshold`, `min_size`, ...). Segment reads the +**projected** (Z=1) plate and writes one label image per model INTO it, +`labels/`; it does not project (a model without `eval_args.do_3D` needs a +Z=1 input). The A549 template carries the cellpose values tracking was validated +with (cpsam_v2, diameter 80, min_size 2500) — they used to live in `track.yml`'s +`cellpose_config`. + **`track.yml` — schema.** `TrackingSettings` forbids extra keys. Current -fields: `output_mode`, `target_channel`, `fov`, `input_images`, -`blank_frames_path`, `segmentation_method`, `min_sibling_length`, -`cellpose_config`, `tracking_config`. Track no longer selects or projects -Z-planes: it reads the **projected** (Z=1) plate that `biahub project` writes -from `project.yml`, writes its labels and tracks into that plate, and a 3D input -with `output_mode: 2D` is an error. Old configs under +fields: `output_mode`, `fov`, `blank_frames_path`, `segmentation_method` +(`labels`, the default, or the legacy `foreground_contour`), `labels`, +`labels_sigma`, `min_sibling_length`, `tracking_config` (plus `target_channel` +and `input_images` on the legacy path). Track neither segments nor selects or +projects Z-planes: it reads the label image `biahub segment` wrote into the +**projected** plate (`labels: `), writes the tracked +labels beside it as `labels/_tracked` with `tracks.geff` and a tracks +CSV, and a 3D input with `output_mode: 2D` is an error. Old configs under `/hpc/projects/tlg2_mantis/` still carry rejected fields — migrate: ``` mode -> output_mode z_slicing / z_range / focus_config -> project.yml: channels..z {center: focus, below, above}, focus.channel input_images np.mean step -> project.yml: channels..method: mean +cellpose_config (model, diameter, thresholds, min_size, input_channel) + -> segment.yml: models. +cellpose_config.labels_sigma -> labels_sigma +target_channel -> labels: output_ome_zarr_version -> (delete: labels take the projected plate's version) ``` (`focus_config`'s `NA_det`/`lambda_ill` are `project.yml`'s `focus.na_det` / -`focus.lambda_ill`; the pixel size comes from the store.) `target_channel` and -`cellpose_config.input_channel` must match channels `project.yml` writes, which -carry the **pre-rename** names of the **assembled** plate (§2, §5). -`target_channel` also names the label image tracking writes into each position -of the projected plate, `labels/` (no `_labels` suffix), so -`qc_project.yaml`'s `instance_count.label` and its anchor `channels:` entry must -equal it — SKILL.md §5a checks these couplings before launch. +`focus.lambda_ill`; the pixel size comes from the store.) `segment.yml`'s model +`channels` must match channels `project.yml` writes, which carry the +**pre-rename** names of the **assembled** plate (§2, §5). `qc_project.yaml`'s +`instance_count.label` must be `_tracked` and its anchor +`channels:` entry the segment model's channel — SKILL.md §5a checks these +couplings before launch. **QC configs reject `upstream_qc_dir`.** The imaging-qc pin (f5ae656) dropped the key; a `qc.yaml` copied from an older run that still carries diff --git a/.claude/skills/reconstruct-with-nextflow/references/datasets.md b/.claude/skills/reconstruct-with-nextflow/references/datasets.md index 5c2948b3..1f0a65b8 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/datasets.md +++ b/.claude/skills/reconstruct-with-nextflow/references/datasets.md @@ -48,7 +48,7 @@ Which root depends on the dataset family: | family | cues in the name | output root | deliverable | |---|---|---|---| | zebrafish / neuromast / dynatrack | `dynatrack`, `hpf`, `dpf`, `neuromast`, `cldnb`, `she`, `myo6b`, `zebrafish` | `/hpc/projects/tlg2_mantis/` | assemble (`4-assemble`) — **no tracking**, see `caveats.md` §4 | -| cell line / organelle / infection | `A549`, `HEK`, `iPSC`, organelle genes (`SEC61B`, `TOMM20`, `G3BP1`, `MAP4`, `CAAX`, `H2B`), viruses (`DENV`, `ZIKV`, `HSV1`) | `/hpc/projects/intracellular_dashboard/organelle_dynamics/` | assemble + project + track (`4-assemble/.zarr` and `4-assemble/_proj.zarr`; tracking writes into the latter) | +| cell line / organelle / infection | `A549`, `HEK`, `iPSC`, organelle genes (`SEC61B`, `TOMM20`, `G3BP1`, `MAP4`, `CAAX`, `H2B`), viruses (`DENV`, `ZIKV`, `HSV1`) | `/hpc/projects/intracellular_dashboard/organelle_dynamics/` | assemble + project + segment + track (`4-assemble/.zarr` and `4-assemble/_proj.zarr`; segment and track write into the latter) | | instrument QC / calibration | `argolight`, `beads`, `psf`, `alignment`, `illumination`, `fluorescein`, `first_light`, `mantis_v2_*` | *no default* — ask. These are usually not reconstructed. | — | Neighbouring roots that exist and are **not** the default target — do not use @@ -72,9 +72,10 @@ layout is defined once in `nextflow/mantis-v2.nf` (`directory_layout()`): ├── 3-virtual-stain/.zarr ├── 4-assemble/.zarr # the deliverable ├── 4-assemble/_proj.zarr # A549 only — the 2D (Z=1) companion plate; with -│ # tracking, each position also holds -│ # labels/, tracks.geff and -│ # tracks___.csv +│ # segmentation and tracking, each position also +│ # holds labels/ (segment), +│ # labels/_tracked, tracks.geff and +│ # tracks___.csv (track) ├── qc/ │ ├── report/index.html # one report, one tab per QC'd store │ └── report_spec.yaml # written at launch, before any QC task @@ -84,7 +85,10 @@ layout is defined once in `nextflow/mantis-v2.nf` (`directory_layout()`): │ ├── reconstruct.yml │ ├── virtual_stain.yml # older runs name this predict.yml │ ├── concatenate.yml +│ ├── project.yml # A549 only +│ ├── segment.yml # A549 only │ ├── track.yml # A549 only +│ ├── qc_project.yaml # A549 only │ └── qc.yaml # QC config for the assembled store ├── nextflow/ │ ├── work/ # Nextflow work dir (default: /nextflow/work) diff --git a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py index c492424e..73ece064 100644 --- a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py +++ b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py @@ -12,9 +12,10 @@ channel with a warning -- but the pipeline's intention is that every assembled channel is projected unless the user decided otherwise, so the skill carries these lines into the pre-run plan for the user to decide. -3. `track.yml`'s `target_channel` and `cellpose_config.input_channel` against - the plate track reads: the projected plate when `project.yml` exists, else - the assembled plate. +3. `segment.yml`'s model channels against the plate segment reads (the + projected plate when `project.yml` exists, else the assembled plate), and + `track.yml`'s `labels` against `segment.yml`'s model names, since track reads + the label image segment wrote (`labels/`). WHY PREDICT RATHER THAN READ. The store QC runs over — the assembled plate — does not exist when configs are scaffolded; the pipeline's init phase creates @@ -25,12 +26,14 @@ + reconstruct's output_channel_names (waveorder derives these, so they are read from the settings model, not the YAML text) + virtual-stain's data.init_args.target_channel - labels = (/labels/) - -Tracking writes its labels INTO the assembled plate, so `qc.yaml`'s -`instance_count.label` — and its one anchor channel, the channel the labels -were segmented from — only resolve if both equal `track.yml`'s -`target_channel`. Two files, no shared source of truth. Checked when `qc.yaml` + labels = (/labels/) + + _tracked (/labels/_tracked) + +Segment and track write their labels INTO the plate they read, so +`qc_project.yaml`'s `instance_count.label` only resolves if it is +`_tracked` (or a `segment.yml` model name), and its one anchor +channel, the channel the labels were segmented from, must be that model's +channel. Four files, no shared source of truth. Checked when `qc_project.yaml` has an `instance_count` block and `track.yml` exists. RELATIONSHIP TO imaging-qc#226. Since that landed, `plan-stage` refuses an @@ -86,18 +89,14 @@ def assembled_channels(raw_zarr, reconstruct_cfg, virtual_stain_cfg): return raw + list(rc.output_channel_names) + tgt -def tracking_entries(qc_cfg, track_cfg): - """`instance_count`'s label and anchor channels, and track's target_channel. - - None when there is nothing to compare: no `instance_count` block, or no track.yml. - """ +def instance_count_entries(qc_cfg): + """`instance_count`'s label and anchor channels, or None without such a block.""" ic = (yaml.safe_load(Path(qc_cfg).read_text()) or {}).get("instance_count") - if not isinstance(ic, dict) or not Path(track_cfg).exists(): + if not isinstance(ic, dict): return None - target = yaml.safe_load(Path(track_cfg).read_text())["target_channel"] entries = [("instance_count.label", ic.get("label"))] entries += [("instance_count.channels", c) for c in ic.get("channels") or []] - return target, entries + return entries def projected_channels(project_cfg): @@ -105,14 +104,18 @@ def projected_channels(project_cfg): return list(yaml.safe_load(Path(project_cfg).read_text())["channels"]) -def track_inputs(track_cfg): - """Return the channel names track reads from its input plate, keyed by the setting.""" +def segment_models(segment_cfg): + """Return {model name: [channels]} of a segment config: the label images it writes.""" + models = yaml.safe_load(Path(segment_cfg).read_text())["models"] + return {name: list(model["channels"]) for name, model in models.items()} + + +def tracked_label(track_cfg): + """Return (input label, output label) of a track config, or None on the legacy path.""" t = yaml.safe_load(Path(track_cfg).read_text()) - inputs = [("target_channel", t["target_channel"])] - cellpose = t.get("cellpose_config") or {} - if t.get("segmentation_method") == "cellpose" and "input_channel" in cellpose: - inputs.append(("cellpose_config.input_channel", cellpose["input_channel"])) - return inputs + if t.get("segmentation_method", "labels") != "labels": + return None + return t["labels"], f"{t['labels']}_tracked" if __name__ == "__main__": @@ -145,29 +148,51 @@ def track_inputs(track_cfg): # Not a failure: the CLI drops it with a warning. The skill asks the # user about every such channel in the pre-run plan. print(f" NOT PROJECTED {ch!r}: not in project.yml, will be dropped") + source = "projected" if (cfgs / "project.yml").exists() else "assembled" + + # segment.yml reads the plate project (or assemble) writes and adds labels/. + models = {} + if (cfgs / "segment.yml").exists(): + models = segment_models(cfgs / "segment.yml") + print(f"segment.yml: {source} plate will have {track_reads}") + for name, channels in models.items(): + for ch in channels: + ok = ch in track_reads + bad += not ok + print(f" {'ok ' if ok else 'MISS'} models.{name}.channels: {ch!r}") + + # track.yml reads labels/ of that plate and writes labels/_tracked. + tracked = None if (cfgs / "track.yml").exists(): - source = "projected" if (cfgs / "project.yml").exists() else "assembled" - print(f"track.yml: {source} plate will have {track_reads}") - for key, ch in track_inputs(cfgs / "track.yml"): - ok = ch in track_reads + tracked = tracked_label(cfgs / "track.yml") + if tracked: + print(f"track.yml: segment.yml writes labels/{{{', '.join(models) or '-'}}}") + ok = tracked[0] in models bad += not ok - print(f" {'ok ' if ok else 'MISS'} {key}: {ch!r}") + print(f" {'ok ' if ok else 'MISS'} labels: {tracked[0]!r}") - # qc_project.yaml QCs the plate track reads and writes its labels into: its - # `channels:` against that plate, and its `instance_count` against the label - # image `labels/`. + # qc_project.yaml QCs the plate segment and track write into: its `channels:` + # against that plate, and its `instance_count` against the label images there + # (the tracked one, labels/_tracked, or a segment model's) with + # the anchor channel being the one the model segmented. if (cfgs / "qc_project.yaml").exists(): - print(f"qc_project.yaml: projected plate will have {track_reads}") + print(f"qc_project.yaml: {source} plate will have {track_reads}") for owner, ch in qc_filters(cfgs / "qc_project.yaml"): ok = ch in track_reads bad += not ok print(f" {'ok ' if ok else 'MISS'} {owner}.channels: {ch!r}") - tracking = tracking_entries(cfgs / "qc_project.yaml", cfgs / "track.yml") - if tracking: - target, entries = tracking - print(f"qc_project.yaml vs track.yml: tracking writes labels/{target}") + entries = instance_count_entries(cfgs / "qc_project.yaml") + if entries and (models or tracked): + label_images = dict(models) + if tracked and tracked[0] in models: + label_images[tracked[1]] = models[tracked[0]] + print(f"qc_project.yaml vs segment.yml/track.yml: label images {sorted(label_images)}") + label = dict(entries).get("instance_count.label") for key, value in entries: - ok = value == target + if key == "instance_count.label": + ok = value in label_images + else: + ok = value in label_images.get(label, []) bad += not ok print(f" {'ok ' if ok else 'MISS'} {key}: {value!r}") sys.exit(1 if bad else 0) diff --git a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh index 5cd09c9f..86b39dfe 100755 --- a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh +++ b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh @@ -1,19 +1,21 @@ #!/usr/bin/env bash # Run the mantis-v2 reconstruction pipeline (biahub nextflow/mantis-v2.nf): -# flat-field -> deskew -> reconstruct -> virtual-stain -> assemble -> project -> track -> QC +# flat-field -> deskew -> reconstruct -> virtual-stain -> assemble -> project +# -> segment -> track -> QC # -# The last four are OPTIONAL: a step runs only if its config is passed. A549 -# runs assemble + project + track + QC; neuromast runs assemble + QC and no -# projection or tracking. See the note above the `nextflow run` call at the bottom. +# The last five are OPTIONAL: a step runs only if its config is passed. A549 +# runs assemble + project + segment + track + QC; neuromast runs assemble + QC and +# no projection, segmentation or tracking. See the note above the `nextflow run` +# call at the bottom. # # Copy this into the project output directory and fill in the five variables # below. Keep it there: it is the run's provenance record of the exact command. # # --output is the project root, so each step writes a sibling directory: # 0-flatfield/ 1-deskew/ 2-reconstruct/ 3-virtual-stain/ 4-assemble/ -# Project writes its plate beside the assembled one as 4-assemble/_proj.zarr, -# and tracking writes labels/, tracks.geff and a tracks CSV into -# each position of that plate; neither has a directory of its own. +# Project writes its plate beside the assembled one as 4-assemble/_proj.zarr; +# segment writes labels/ into each position of that plate and tracking +# labels/_tracked, tracks.geff and a tracks CSV; none has a directory of its own. # The number is the step's POSITION among the steps this run performs that # write a directory, so a run that skips a step shifts everything after it down. # The work dir defaults to /nextflow/work; override with `-work-dir`. @@ -192,12 +194,12 @@ fi # Harmless when unset already, and nothing in the pipeline reads it. unset CLAUDECODE -# The five config flags from --concatenate_config down are the OPTIONAL steps: +# The six config flags from --concatenate_config down are the OPTIONAL steps: # a step runs only if its config is passed (biahub#306). Defaults by family — # -# A549 / cell-line assemble + project + track + QC (all five, as written) -# zebrafish / neuromast assemble + QC (delete --project_config, --track_config -# and --qc_project_config) +# A549 / cell-line assemble + project + segment + track + QC (all six, as written) +# zebrafish / neuromast assemble + QC (delete --project_config, --segment_config, +# --track_config and --qc_project_config) # # --qc_track_config was removed: qc_project.yaml counts the labels tracking writes # into the projected plate, so passing the old flag errors at launch. @@ -219,6 +221,7 @@ nextflow run "${PIPELINE}" \ --virtual_stain_config "${CONFIGS}/virtual_stain.yml" \ --concatenate_config "${CONFIGS}/concatenate.yml" \ --project_config "${CONFIGS}/project.yml" \ + --segment_config "${CONFIGS}/segment.yml" \ --track_config "${CONFIGS}/track.yml" \ --qc_config "${CONFIGS}/qc.yaml" \ --qc_project_config "${CONFIGS}/qc_project.yaml" \ diff --git a/README.md b/README.md index b2df1d20..3ce56ba6 100644 --- a/README.md +++ b/README.md @@ -121,9 +121,9 @@ Some steps include an `estimate-*` command that writes a YAML of parameters you | `concatenate` | Concatenate datasets channel-wise, with optional cropping | | | | | `project` | Project a 3D plate to 2D: max, mean, sum or one plane per channel, optionally around the in-focus plane | -| `segment` | Segment positions with a pretrained model or pipeline | +| `segment` | Segment positions with cellpose 4; writes label images into the plate (or a separate label plate with `-o`) | | `virtual-stain` | Run VisCy/cytoland virtual staining | -| `track` | Track objects in 2D/3D time-lapse data | +| `track` | Track the instance labels of a label image (ultrack); writes tracked labels and tracks.geff into the plate | | | | | `nf list-positions` | List position keys of a plate zarr (used for Nextflow fan-out) | @@ -224,7 +224,7 @@ With `--biahub_project`, each task runs as `uv run --project biahub ...`; Profiles select *where* work runs: - `-profile local` — everything on the current machine. Good for debugging and single-node runs. -- `-profile slurm` — per-position work goes to the `preempted` partition, virtual-stain prediction to the `preempted` partition on GPU nodes (it resumes where a reclaimed task stopped), tracking to the `gpu` partition, and lightweight init steps stay local. Preempted, timed-out, and OOM-killed tasks are retried automatically (up to 5 times, with escalating time/memory); genuine errors fail fast instead of burning retries. +- `-profile slurm` — per-position work goes to the `preempted` partition, virtual-stain prediction and segmentation to the `preempted` partition on GPU nodes (both resume where a reclaimed task stopped), tracking to the `cpu` partition (ultrack cannot resume), and lightweight init steps stay local. Preempted, timed-out, and OOM-killed tasks are retried automatically (up to 5 times, with escalating time/memory); genuine errors fail fast instead of burning retries. Nextflow's own flags still apply — most usefully `-resume` to reuse completed tasks after a failure or a config tweak, and `-w /fast/scratch` to move the work directory off the output filesystem. @@ -235,7 +235,7 @@ Each step writes `.zarr` into its own numbered subdirectory of `--outpu ``` output/ ├── 0-flatfield/ 1-deskew/ 2-reconstruct/ -├── 3-virtual-stain/ 4-track/ 5-assemble/ +├── 3-virtual-stain/ 4-assemble/ # .zarr and, with --project_config, _proj.zarr │ └── slurm_output/README.md # (in each step) points to the real logs └── nextflow/ ├── report.html timeline.html dag.html trace.txt @@ -245,9 +245,9 @@ output/ Per-task logs (progress and tracebacks) are in `nextflow/slurm_output//`. The `slurm_output/` inside each step directory is created by the CLI, but because tasks run it with `--cluster debug`, it only holds submitit placeholder files; each step's init writes a `README.md` there that points to the real logs. -With cleanup on (`--cleanup_intermediates true`, or `auto` when `concatenate.yml` crops nothing), once the last step (QC, when it runs) has finished, the pipeline deletes the flat-field, deskew, reconstruct and virtual-stain directories, the `slurm_output/` placeholders beside the assembled and tracking stores, the assembled store's `.iohub-progress/` resume markers, and the task directories in `nextflow/work`. What it removed is appended to `nextflow/intermediates_cleaned.txt`. Assemble copies those stores' channels, so when `concatenate.yml` keeps all the data this only removes the duplicate copy; if it crops, the cropped-out data is lost. A cleaned run is final: with no cached tasks and no resume markers left, rerunning it, even with `-resume`, recomputes everything from the raw input. +With cleanup on (`--cleanup_intermediates true`, or `auto` when `concatenate.yml` crops nothing), once the last step (QC, when it runs) has finished, the pipeline deletes the flat-field, deskew, reconstruct and virtual-stain directories, the `slurm_output/` placeholders beside the assembled and projected plates, their `.iohub-progress/` resume markers (assemble, project and segment write them; track does not), and the task directories in `nextflow/work`. What it removed is appended to `nextflow/intermediates_cleaned.txt`. Assemble copies those stores' channels, so when `concatenate.yml` keeps all the data this only removes the duplicate copy; if it crops, the cropped-out data is lost. A cleaned run is final: with no cached tasks and no resume markers left, rerunning it, even with `-resume`, recomputes everything from the raw input. -The pipeline runs flat-field → deskew → reconstruct → virtual-stain → assemble → track; `5-assemble` concatenates the deskew, reconstruct, and virtual-stain channels into one plate, which tracking then consumes. +The pipeline runs flat-field → deskew → reconstruct → virtual-stain → assemble → project → segment → track → QC; `4-assemble` concatenates the deskew, reconstruct, and virtual-stain channels into one plate, project reduces it to the 2D companion plate beside it, segment writes cellpose label images into that plate, and track links them and writes the tracked labels and `tracks.geff` beside them. Assemble, project, segment, track and QC each run only when their config is passed. ### Adapting a pipeline diff --git a/biahub/segment.py b/biahub/segment.py index 8f035ce4..1b2085c4 100644 --- a/biahub/segment.py +++ b/biahub/segment.py @@ -40,6 +40,7 @@ from biahub.utils.config import settings_fingerprint, yaml_to_model from biahub.utils.ngff import ( PROVENANCE_METADATA_KEYS, + create_label_images, get_output_paths, resolve_ome_zarr_version, ) @@ -47,6 +48,9 @@ if TYPE_CHECKING: import torch +# Provenance key: on each label image (no -o) or on each output position (-o). +PROVENANCE_KEY = "biahub-segment" + @dataclass(frozen=True) class ResolvedModel: @@ -180,12 +184,56 @@ def _segment_czyx( return np.stack(czyx_segmentation, axis=0) +def _init_label_images( + input_position_dirpaths: list[Path], + settings: SegmentationSettings, +) -> tuple[tuple[int, int, int, int, int], list[ResolvedModel]]: + """Create the empty ``labels/`` images in the input positions (no ``-o``). + + One label image per model, named after it, with the image's shape (its Z: one + plane for a projected plate). Existing label images are left untouched, so both + --init and per-position workers can call this; one whose shape differs from the + image's (a model renamed from 3D to 2D, say) is refused. The settings are recorded + on each label image exactly as written (names, not indices). + + Returns the input (T, C, Z, Y, X) shape and the resolved models. + """ + with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as input_dataset: + T, C, Z, Y, X = input_dataset.data.shape + scale = input_dataset.scale + channel_names = input_dataset.channel_names + + models = resolve_models(settings, channel_names, scale, Z) + label_names = [m.name for m in models] + + for position_dirpath in input_position_dirpaths: + with open_ome_zarr(str(position_dirpath), mode="r") as position: + for name in label_names: + if name not in position.label_names(): + continue + existing_shape = tuple(position.get_label(name)["0"].shape) + if existing_shape != (T, Z, Y, X): + raise ValueError( + f"{position_dirpath} already has a label image {name!r} of shape " + f"{existing_shape}, but the image is {(T, Z, Y, X)} (T, Z, Y, X). " + "Rename the model or remove the old label image." + ) + + create_label_images( + input_position_dirpaths, + label_names, + PROVENANCE_KEY, + settings.model_dump(mode="json"), + ) + return (T, C, Z, Y, X), models + + def _init_output_plate( input_position_dirpaths: list[Path], output_dirpath: Path, settings: SegmentationSettings, ) -> tuple[tuple[int, int, int, int, int], list[ResolvedModel]]: - """Create (or extend) the empty label plate and resolve the models against the input. + """Create (or extend) a separate label plate (``-o``) and resolve the models. create_empty_plate is idempotent: re-running with the same positions is a no-op and new positions are appended, so both --init and per-position workers can call it. @@ -234,26 +282,76 @@ def _init_output_plate( ), metadata_sources=input_plate, metadata_keys=PROVENANCE_METADATA_KEYS, - extra_metadata={"biahub-segment": settings.model_dump(mode="json")}, + extra_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, ) return (T, C, Z, Y, X), models -def _resume_token(settings: SegmentationSettings) -> str: +def _resume_token(settings: SegmentationSettings, model: str | None = None) -> str: """Fingerprint the settings, including the order of the models. - settings_fingerprint sorts keys, but model order decides which label channel each - model writes, so finished work from another order must not be skipped. + settings_fingerprint sorts keys, but in a separate label plate (``-o``) the model + order decides which label channel each model writes, so finished work from another + order must not be skipped. In the input plate each model has its own label image, + so the token names the model instead and the order does not matter. """ + if model is not None: + return settings_fingerprint( + {"settings": settings.model_dump(mode="json"), "model": model} + ) return settings_fingerprint( {"settings": settings.model_dump(mode="json"), "model_order": list(settings.models)} ) +def _segment_position( + input_position_path: Path, + output_position_path: Path | None, + models: list[ResolvedModel], + n_channels: int, + settings: SegmentationSettings, + resume: bool, +) -> None: + """Segment one position: into its own label images, or into the ``-o`` plate. + + Without ``-o`` each model gets one ``process_single_position`` call writing the + label image named after it (iohub's ``output_label``); the cellpose model cache + makes the repeated calls free. With ``-o`` one call writes all the models as + channels of the output position. + """ + if output_position_path is None: + for model in models: + process_single_position( + _segment_czyx, + input_position_path, + input_position_path, + input_channel_indices=[list(range(n_channels))], + output_label=model.name, + # One process per position: timepoints share the loaded model and + # the GPU, which is the bottleneck. + num_workers=1, + resume=resume, + resume_token=_resume_token(settings, model.name), + models=[model], + ) + return + process_single_position( + _segment_czyx, + input_position_path, + output_position_path, + input_channel_indices=[list(range(n_channels))], + output_channel_indices=[list(range(len(models)))], + num_workers=1, + resume=resume, + resume_token=_resume_token(settings), + models=models, + ) + + def segment( input_position_dirpaths: list[Path], config_filepath: Path, - output_dirpath: Path, + output_dirpath: Path | None = None, sbatch_filepath: str | None = None, cluster: str = "slurm", monitor: bool = True, @@ -269,8 +367,11 @@ def segment( or "input.zarr/*/*/*". config_filepath : Path Path to the segmentation YAML configuration file. - output_dirpath : Path - Path to the "output.zarr" label plate. + output_dirpath : Path, optional + Path to a separate "output.zarr" label plate, with one channel per model + (``_labels``). If None (the default), the labels are written into the + ``-i`` positions as OME-NGFF label images, ``labels/``, which is what + `biahub track` reads, and the SLURM logs go to ``./slurm_output``. sbatch_filepath : str, optional SBATCH filepath that contains slurm parameters to overwrite defaults. cluster : str, optional @@ -284,8 +385,12 @@ def segment( Skip the (time, channel) units a previous attempt already finished; see ``iohub.ngff.utils.process_single_position``. """ - output_dirpath = Path(output_dirpath) - slurm_out_path = output_dirpath.parent / "slurm_output" + if output_dirpath is not None: + output_dirpath = Path(output_dirpath) + slurm_out_path = output_dirpath.parent / "slurm_output" + else: + # Not $TMPDIR: submitit's folder must be visible from the compute nodes. + slurm_out_path = Path("slurm_output").absolute() resolved_cluster = get_submitit_cluster(cluster=cluster) if not init_only and resolved_cluster == "debug": @@ -300,9 +405,12 @@ def segment( # Fail on a bad model name before scaffolding anything. for name in dict.fromkeys(m.pretrained_model for m in settings.models.values()): check_cellpose_model_name(name) - (T, C, Z, Y, X), models = _init_output_plate( - input_position_dirpaths, output_dirpath, settings - ) + if output_dirpath is None: + (T, C, Z, Y, X), models = _init_label_images(input_position_dirpaths, settings) + else: + (T, C, Z, Y, X), models = _init_output_plate( + input_position_dirpaths, output_dirpath, settings + ) # Calibrated on a real 2D run (2026_04_28 SEC61B, 67 x 1664 x 1193, 6 input channels, # cpsam_v2 on one GPU): ~5 s per frame (6-10 min per position, incl. weight staging) and @@ -327,10 +435,22 @@ def segment( # than N workers racing for it. Workers must not (staging needs cellpose unimported). for name in dict.fromkeys(m.pretrained_model for m in models): warm_cellpose_weights(name) - click.echo(f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)") + if output_dirpath is None: + click.echo( + f"Created labels/{{{', '.join(m.name for m in models)}}} in " + f"{Path(input_position_dirpaths[0]).parents[2]} " + f"({len(input_position_dirpaths)} positions)" + ) + else: + click.echo( + f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)" + ) return - output_position_paths = get_output_paths(input_position_dirpaths, output_dirpath) + if output_dirpath is None: + output_position_paths = [None] * len(input_position_dirpaths) + else: + output_position_paths = get_output_paths(input_position_dirpaths, output_dirpath) slurm_args = { "slurm_job_name": "segment", @@ -363,18 +483,13 @@ def segment( ): jobs.append( executor.submit( - process_single_position, - _segment_czyx, + _segment_position, input_position_path, output_position_path, - input_channel_indices=[list(range(C))], - output_channel_indices=[list(range(len(models)))], - # One process per position: timepoints share the loaded model and - # the GPU, which is the bottleneck. - num_workers=1, - resume=resume, - resume_token=_resume_token(settings), models=models, + n_channels=C, + settings=settings, + resume=resume, ) ) @@ -399,7 +514,13 @@ def segment( @click.command("segment") @input_position_dirpaths() @config_filepath() -@output_dirpath() +@output_dirpath( + required=False, + help="Separate label plate to write, one channel per model (_labels). " + "If omitted, the labels are written into the -i positions as label images, " + "labels/, which `biahub track` reads, and the SLURM logs go to " + "./slurm_output.", +) @sbatch_filepath() @cluster() @monitor() @@ -408,7 +529,7 @@ def segment( def segment_cli( input_position_dirpaths: list[Path], config_filepath: Path, - output_dirpath: Path, + output_dirpath: Path | None = None, sbatch_filepath: str | None = None, cluster: str = "slurm", monitor: bool = False, @@ -417,18 +538,24 @@ def segment_cli( ): """Segment positions with cellpose 4 models configured in a YAML file. + Writes one label image per model into each -i position (labels/), where + `biahub track` reads them; -o writes a separate label plate instead. + \b SLURM fan-out of positions across a whole plate: - >>> biahub segment -i ./input.zarr/*/*/* -c ./segment.yml -o ./segment.zarr + >>> biahub segment -i ./input.zarr/*/*/* -c ./segment.yml \b - Initialize the output plate only (Nextflow init step): - >>> biahub segment --init -i ./input.zarr/*/*/* -c ./segment.yml -o ./segment.zarr + Create the empty label images only (Nextflow init step): + >>> biahub segment --init -i ./input.zarr/*/*/* -c ./segment.yml \b In-process run of a single position (Nextflow per-position worker): - >>> biahub segment --cluster debug -i ./input.zarr/B/3/000000 -c ./segment.yml \\ - -o ./segment.zarr --resume + >>> biahub segment --cluster debug -i ./input.zarr/B/3/000000 -c ./segment.yml --resume + + \b + Into a separate label plate: + >>> biahub segment -i ./input.zarr/*/*/* -c ./segment.yml -o ./segment.zarr """ # noqa: D301 segment( input_position_dirpaths=input_position_dirpaths, diff --git a/biahub/settings.py b/biahub/settings.py index db77ad4d..b1a17dd6 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -62,38 +62,12 @@ def validate_path_not_plate(cls, v): return v -class CellposeConfig(MyBaseModel): - """Configuration for Cellpose segmentation used as input to tracking.""" - - # cellpose 4 model: cpsam_v2, cpsam, cpdino, cpdino-vitb, or a path to weights. - pretrained_model: str = "cpsam_v2" - # Deprecated: cellpose 4 ignores it (``model_type: nuclei`` always ran cpsam_v2). - model_type: str | None = None - # Cellpose 4 resizes each frame by 30 / diameter before inference (80 -> 37.5%), - # which shrinks small objects such as rounded mitotic nuclei. None = native size. - diameter: float | None = 80 - cellprob_threshold: float = 0.0 - flow_threshold: float = 0.4 - gpu: bool = True - min_size: int = 500 - input_channel: str = "nuclei_prediction" - labels_sigma: float = 5.0 - - @field_validator("model_type") - @classmethod - def warn_model_type_ignored(cls, v): - if v is not None: - warnings.warn( - f"cellpose_config.model_type={v!r} is ignored by cellpose 4, which runs " - "pretrained_model (default cpsam_v2). Remove model_type from the config.", - UserWarning, - stacklevel=2, - ) - return v +#: Suffix of the label image tracking writes next to the one it reads: +#: ``labels/`` (from `biahub segment`) -> ``labels/_tracked``. +TRACKED_LABEL_SUFFIX = "_tracked" class TrackingSettings(MyBaseModel): - target_channel: str = "nuclei_prediction" fov: str = "*/*/*" blank_frames_path: Path | None = None # The dimensionality of the tracking. Track neither projects nor selects Z-planes: @@ -101,29 +75,52 @@ class TrackingSettings(MyBaseModel): # a Z>1 input; a mismatch is an error. The labels are written into the plate track # reads, with its Z, so they take that store's OME-Zarr version. output_mode: Literal["2D", "3D"] = "2D" - input_images: list[ProcessingInputChannel] + # What feeds ultrack: the instance labels of a label image of the plate track reads + # (`labels`, written by `biahub segment`; the default), or a foreground mask and + # contour map computed by the `input_images` processing functions (the legacy + # path; see settings/example_track_settings.yml). Track does not segment. + segmentation_method: Literal["labels", "foreground_contour"] = "labels" + # segmentation_method: labels -- the label image, `labels/` in each -i + # position, named after the segment config's model. The tracked labels go to + # `labels/_tracked`. + labels: str | None = None + # Smoothing (pixels) ultrack applies to the labels' contours before it builds its + # own segmentation hypotheses from them (ultrack `labels_to_contours`). + labels_sigma: float = 5.0 + # segmentation_method: foreground_contour only -- the name of the label image the + # tracks are written to. With `labels` the name is derived (_tracked). + target_channel: str | None = None + # Image channels to load and process per FOV; needed by foreground_contour (its + # processing functions compute the mask and contour) and optional otherwise. + input_images: list[ProcessingInputChannel] = [] tracking_config: dict[str, Any] = {} - # What feeds ultrack: cellpose labels (the default, with `cellpose_config`), or a - # foreground mask and contour map computed by the `input_images` processing - # functions (the legacy path; see settings/example_track_settings.yml). - segmentation_method: Literal["foreground_contour", "cellpose"] = "cellpose" # Post-processing (ultrack's filter_short_sibling_tracks, same meaning): at a division # where exactly one daughter lives this many frames or fewer and does not divide # again, drop that daughter and merge the other into the parent. A division whose two # daughters are both that short is kept. Removes short false divisions, e.g. a # virtual-staining hallucination linked as a daughter. None (default) keeps all. min_sibling_length: PositiveInt | None = None - cellpose_config: CellposeConfig | None = None @model_validator(mode="before") @classmethod - def reject_z_slicing(cls, data): + def reject_removed_keys(cls, data): + if not isinstance(data, dict): + return data # Z-plane selection moved out of track into `biahub project`. - if isinstance(data, dict) and "z_slicing" in data: + if "z_slicing" in data: raise ValueError( "z_slicing was removed from track: select and project the Z-planes with " "`biahub project` and track the projected (Z=1) plate." ) + # Segmentation moved out of track into `biahub segment`. + if "cellpose_config" in data or data.get("segmentation_method") == "cellpose": + raise ValueError( + "Track no longer runs cellpose (cellpose_config, segmentation_method: " + "cellpose): segment the plate with `biahub segment` (its config takes the " + "cellpose settings; see settings/example_segmentation_settings.yml), which " + "writes labels/ into the plate, then set `labels: ` here " + "(and `labels_sigma` for the contour smoothing)." + ) return data @field_validator("blank_frames_path") @@ -135,16 +132,40 @@ def validate_blank_frames_path(cls, v): @model_validator(mode="after") def check_segmentation_method(self): - # Without this a config that omits cellpose_config would silently fall back to - # the foreground/contour path at tracking time. - if self.segmentation_method == "cellpose" and self.cellpose_config is None: - raise ValueError( - "segmentation_method is cellpose (the default) but cellpose_config is not " - "set; add a cellpose_config, or set segmentation_method: foreground_contour " - "to track a foreground mask and contour map from input_images." - ) + if self.segmentation_method == "labels": + if self.labels is None: + raise ValueError( + "segmentation_method is labels (the default) but `labels` is not set: " + "name the label image `biahub segment` wrote (its model name), or set " + "segmentation_method: foreground_contour to track a foreground mask " + "and contour map from input_images." + ) + if self.target_channel is not None: + raise ValueError( + "target_channel is for segmentation_method: foreground_contour; with " + f"labels the tracks are written to labels/{self.labels}" + f"{TRACKED_LABEL_SUFFIX}." + ) + else: + if self.target_channel is None: + raise ValueError( + "segmentation_method: foreground_contour needs target_channel, the name " + "of the label image the tracks are written to." + ) + if not self.input_images: + raise ValueError( + "segmentation_method: foreground_contour needs input_images whose " + "processing functions produce the foreground mask and contour map." + ) return self + @property + def output_label(self) -> str: + """Name of the label image tracking writes, ``labels/``.""" + if self.segmentation_method == "labels": + return f"{self.labels}{TRACKED_LABEL_SUFFIX}" + return self.target_channel + class EdgeGraphSettings(BaseModel): method: Literal["knn", "radius", "full"] = "knn" diff --git a/biahub/track.py b/biahub/track.py index e96afb83..85f78d4e 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -15,9 +15,7 @@ import zarr from iohub import open_ome_zarr -from iohub.ngff.utils import create_empty_plate from numpy.typing import ArrayLike -from tqdm import tqdm from biahub.cli.monitor import monitor_jobs from biahub.cli.parsing import ( @@ -31,18 +29,7 @@ sbatch_to_submitit, ) from biahub.cli.resolve_function import resolve_function -from biahub.settings import ( - CellposeConfig, - ProcessingInputChannel, - TrackingSettings, -) -from biahub.utils.cellpose import ( - cellpose_device, - check_cellpose_model_name, - load_cellpose_model, - stage_cellpose_weights, - warm_cellpose_weights, -) +from biahub.settings import ProcessingInputChannel, TrackingSettings from biahub.utils.cluster import ( apply_sbatch_overrides, echo_resources, @@ -51,6 +38,7 @@ local_executor_parameters, ) from biahub.utils.config import update_model, yaml_to_model +from biahub.utils.ngff import create_label_images, labels_exist logger = logging.getLogger(__name__) @@ -335,7 +323,7 @@ def run_ultrack( GEFF export. **track_kwargs Keyword arguments forwarded to ``Tracker.track()``. E.g. ``detection``, ``edges`` - (foreground+contour mode), or ``labels``, ``sigma`` (cellpose mode), plus ``scale`` + (foreground+contour mode), or ``labels``, ``sigma`` (label-image mode), plus ``scale`` and ``overwrite``. Returns @@ -526,7 +514,7 @@ def load_data( :, image_channel_names.index(channel_name) ] # A projected (Z=1) plate is a 2D series: drop the singleton axis so - # the processing functions and cellpose see (T, Y, X). + # the processing functions see (T, Y, X). data_dict[channel_name] = tzyx[:, 0] if tzyx.shape[1] == 1 else tzyx if visualize: viewer = napari.Viewer() @@ -582,7 +570,7 @@ def _load_and_preprocess( ) -> dict[str, ArrayLike]: """Load, run the processing pipeline, and fill blank frames for one FOV. - Shared by the foreground+contour and cellpose segmentation paths. + Used by the foreground+contour path (the labels path reads a label image). """ fov = "/".join(position_key) data_dict = load_data( @@ -608,116 +596,40 @@ def detect_foreground_segmentation( raise ValueError("Foreground and contour channels are required for tracking.") -def run_cellpose_per_frame( - images: np.ndarray, - pretrained_model: str = "cpsam_v2", - diameter: float | None = 80, - cellprob_threshold: float = 0.0, - flow_threshold: float = 0.4, - gpu: bool = True, - min_size: int = 500, -) -> np.ndarray: - """Run cellpose on each 2D frame of a (T, Y, X) array. - - Returns an integer label array of shape (T, Y, X). - - Note: cellpose is imported lazily so that importing ``biahub.track`` does not - require the ``segment`` extra unless cellpose segmentation is actually used. - """ - # Resolve the device before staging, so a host with no usable GPU fails in - # seconds rather than after copying 1.2 GB of weights, and stage before the - # import, which is when cellpose freezes the directory it loads them from. - device = cellpose_device(gpu) - stage_cellpose_weights() - - model = load_cellpose_model(pretrained_model, device) - - T = images.shape[0] - labels = np.zeros_like(images, dtype=np.int32) - for t in tqdm(range(T), desc="Cellpose segmentation"): - mask, _, _ = model.eval( - images[t], - diameter=diameter, - cellprob_threshold=cellprob_threshold, - flow_threshold=flow_threshold, - min_size=min_size, - ) - labels[t] = np.asarray(mask) - return labels - - -def cellpose_segmentation( - data_dict: dict[str, ArrayLike], cellpose_config: CellposeConfig +def load_labels( + position_dirpath: Path, + label_name: str, + output_mode: str, + blank_frames_path: Path | None = None, ) -> np.ndarray: - """Run cellpose on a preprocessed ``data_dict`` channel. + """Read the instance labels ``labels/`` of one position. - Selects ``cellpose_config.input_channel`` (a 2D series, (T, Y, X) or (T, 1, Y, X)) and - returns an integer label array of shape (T, Y, X) ready for - ``run_ultrack(..., labels=...)``. + Returns (T, Y, X) labels for 2D tracking (the single plane of a projected plate) + and (T, Z, Y, X) for 3D. Frames listed as blank for this FOV in + ``blank_frames_path`` are filled with the nearest non-blank frame's labels, as + the images used to be before segmentation moved out of track, so a dropped + timepoint does not break every track. """ - # Get the channel to segment - channel_name = cellpose_config.input_channel - if channel_name not in data_dict: - raise ValueError( - f"Cellpose input channel '{channel_name}' not found in data. " - f"Available: {list(data_dict.keys())}" - ) - - images = data_dict[channel_name] - - # Track does not project: a Z=1 input is squeezed, a stack is refused. - if images.ndim == 4: - if images.shape[1] != 1: + with open_ome_zarr(str(position_dirpath), mode="r") as position: + if label_name not in position.label_names(): raise ValueError( - f"Cellpose segmentation in track is 2D but channel '{channel_name}' has " - f"{images.shape[1]} Z-planes; project the plate with `biahub project` first." + f"{position_dirpath} has no label image {label_name!r} (labels: " + f"{position.label_names()}); segment the plate with `biahub segment` first." ) - images = images[:, 0] - - if isinstance(images, da.Array): - images = images.compute() - images = np.asarray(images) - - click.echo( - f"Running cellpose ({cellpose_config.pretrained_model}, " - f"diameter={cellpose_config.diameter}) on channel '{channel_name}'..." - ) - cellpose_labels = run_cellpose_per_frame( - images, - pretrained_model=cellpose_config.pretrained_model, - diameter=cellpose_config.diameter, - cellprob_threshold=cellpose_config.cellprob_threshold, - flow_threshold=cellpose_config.flow_threshold, - gpu=cellpose_config.gpu, - min_size=cellpose_config.min_size, - ) - - n_cells = [len(np.unique(cellpose_labels[t])) - 1 for t in range(cellpose_labels.shape[0])] + labels = np.asarray(position.get_label(label_name)["0"][:]) + if output_mode == "2D": + labels = labels[:, 0] + if blank_frames_path: + fov = "/".join(Path(position_dirpath).parts[-3:]) + empty_frames_idx = get_empty_frames_idx_from_csv(pd.read_csv(blank_frames_path), fov) + if empty_frames_idx: + labels = fill_empty_frames(labels, empty_frames_idx) + n_cells = [len(np.unique(frame)) - 1 for frame in labels] click.echo( - f"Cellpose cells per frame: mean={np.mean(n_cells):.1f}, " + f"Labels {label_name!r}: cells per frame mean={np.mean(n_cells):.1f}, " f"min={np.min(n_cells)}, max={np.max(n_cells)}" ) - return cellpose_labels - - -# Label chunks, TZYX. 16 x 1 x 128 x 128 uint32 is 1 MB, the DCA v0.2 SHOULD (512 KB -# is the MUST) and its "2D+time" example; growing T keeps per-plane reads cheap. -LABEL_CHUNKS = (16, 1, 128, 128) -# Shards span one time chunk (16 timepoints, the DCA SHOULD) and the whole YX plane, -# capped at 16 chunks (2048 px, the DCA spatial-shard SHOULD). -MAX_LABEL_SHARD_CHUNKS = 16 - - -def label_chunks_and_shards( - shape: tuple[int, int, int, int], -) -> tuple[tuple[int, ...], tuple[int, ...]]: - """Chunks and shards ratio for a (T, Z, Y, X) label array of ``shape``.""" - chunks = tuple(min(c, s) for c, s in zip(LABEL_CHUNKS, shape, strict=True)) - shards_ratio = (1, 1) + tuple( - min(-(-s // c), MAX_LABEL_SHARD_CHUNKS) - for c, s in zip(chunks[2:], shape[2:], strict=True) - ) - return chunks, shards_ratio + return labels def write_tracking_labels(label_array, labels: ArrayLike, output_mode: str) -> None: @@ -784,18 +696,20 @@ def track_one_position( label_name: str, blank_frames_path: Path | None = None, scale: tuple[float, ...] = (1, 1, 1, 1, 1), - cellpose_config: CellposeConfig | None = None, + input_label: str | None = None, + labels_sigma: float = 5.0, output_mode: str = "2D", min_sibling_length: int | None = None, ) -> None: """ Run tracking on a single field of view. - Supports two segmentation modes: - - **foreground+contour** (default): the preprocessing pipeline must produce a binary + Supports two segmentation inputs: + - **labels** (``input_label`` set): the instance labels of ``labels/`` + in the position (written by `biahub segment`) are handed to ultrack, which builds + its own segmentation hypotheses from their contours (smoothed by ``labels_sigma``). + - **foreground+contour** (legacy): the preprocessing pipeline must produce a binary foreground mask and a corresponding contour map, which are handed to ultrack. - - **cellpose**: runs cellpose per-frame on a specified channel and passes the resulting - instance labels directly to ultrack. Parameters ---------- @@ -805,7 +719,8 @@ def track_one_position( The ``-i`` position the outputs are written into. ``track --init`` has already created its empty ``labels/`` image. input_images : list of ProcessingInputChannel - Configuration describing which channels to load and how to preprocess them. + Configuration describing which channels to load and how to preprocess them + (foreground+contour); may be empty with ``input_label``. output_dirpath : Path or None Where to keep the Ultrack database and config for this FOV, in ``{output_dirpath}/{row}_{col}_{fov}/``. If None they live in a scratch @@ -821,8 +736,12 @@ def track_one_position( scale : tuple of float, optional Physical scale used for tracking, in (Z, Y, X) or (Y, X) order (already trimmed by the caller to match ``output_mode``). Default is (1, 1, 1, 1, 1). - cellpose_config : CellposeConfig, optional - If provided, uses cellpose segmentation instead of foreground+contour. + input_label : str, optional + Name of the label image to track, ``labels/`` in the position. If + None, the foreground+contour path is used. + labels_sigma : float, optional + Smoothing of the labels' contours before ultrack re-segments them (ultrack's + ``labels_to_contours`` sigma). Default is 5.0. output_mode : str, optional "2D" tracks a Z=1 input and writes (T, Y, X) labels into its one plane; "3D" tracks the stack and writes (T, Z, Y, X) labels. Track does not project (see @@ -853,10 +772,6 @@ def track_one_position( units = {axis.name.lower(): getattr(axis, "unit", None) for axis in position.axes} with _database_dirpath(output_dirpath, fov) as database_path: - # Load, run the processing pipeline, and fill blank frames once; both - # segmentation methods consume the resulting data_dict. - data_dict = _load_and_preprocess(position_key, input_images, blank_frames_path) - geff_kwargs = { "zarr_format": 3, "include_masks": False, @@ -867,14 +782,13 @@ def track_one_position( "time_scale": time_scale if units.get("t") else None, "time_unit": units.get("t"), } - if cellpose_config is not None: - cellpose_labels = cellpose_segmentation(data_dict, cellpose_config) - detection_kwargs = { - "labels": cellpose_labels, - "sigma": cellpose_config.labels_sigma, - } - click.echo("Tracking with cellpose labels...") + if input_label is not None: + labels = load_labels(position_dirpath, input_label, output_mode, blank_frames_path) + detection_kwargs = {"labels": labels, "sigma": labels_sigma} + click.echo(f"Tracking the labels of labels/{input_label}...") else: + # Load, run the processing pipeline, and fill blank frames. + data_dict = _load_and_preprocess(position_key, input_images, blank_frames_path) foreground_mask, contour_gradient_map = detect_foreground_segmentation(data_dict) detection_kwargs = {"detection": foreground_mask, "edges": contour_gradient_map} click.echo("Tracking with foreground + contour...") @@ -907,71 +821,49 @@ def track_one_position( click.echo(f"Saved labels and tracks to: {position_dirpath}") -def _labels_exist(input_position_dirpaths: list[Path], label_name: str) -> bool: - return all( - (Path(p) / "labels" / label_name / "zarr.json").exists() - or (Path(p) / "labels" / label_name / ".zattrs").exists() - for p in input_position_dirpaths - ) - - def _init_labels( input_position_dirpaths: list[Path], settings: TrackingSettings, ) -> tuple[int, int, int, int, int]: - """Create the empty ``labels/`` image in every input position. + """Create the empty ``labels/`` image in every input position. The input's dimensionality must match ``settings.output_mode`` (see - :func:`check_input_dimensionality`). The labels are (T, Z, Y, X) with the plate's Z: - one plane for a projected plate, the stack for 3D tracking. Positions that already - have the label are left untouched, and the track settings are recorded on each - label image. + :func:`check_input_dimensionality`), and with ``settings.labels`` every position + must already have that label image (``biahub segment --init`` creates it, so a + pipeline missing its segment step fails here). The tracked labels are (T, Z, Y, X) + with the plate's Z: one plane for a projected plate, the stack for 3D tracking. + Positions that already have the output label are left untouched, and the track + settings are recorded on each label image. Returns the input (T, C, Z, Y, X) shape used for resource estimation. """ with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as dataset: - T, C, Z, Y, X = dataset.data.shape - scale = dataset.scale - channel_names = dataset.channel_names - dtype = dataset.data.dtype - version = dataset.version - + Z = dataset.data.shape[2] check_input_dimensionality(settings.output_mode, Z) - input_plate = Path(input_position_dirpaths[0]).parents[2] - position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] - label_chunks, label_shards_ratio = label_chunks_and_shards((T, Z, Y, X)) - - create_empty_plate( - store_path=input_plate, - position_keys=position_keys, - channel_names=channel_names, - shape=(T, C, Z, Y, X), - scale=scale, - version=version, - dtype=dtype, - label_names=[settings.target_channel], - label_dtype=np.uint32, - label_chunks=label_chunks, - label_shards_ratio=label_shards_ratio if version == "0.5" else None, - ) - - provenance = settings.model_dump(mode="json") - for position_dirpath in input_position_dirpaths: - label_group = zarr.open_group( - Path(position_dirpath) / "labels" / settings.target_channel, mode="r+" - ) - # Merge so a re-run init keeps what the workers recorded (the Ultrack config). - label_group.attrs[PROVENANCE_KEY] = { - **label_group.attrs.get(PROVENANCE_KEY, {}), - **provenance, - } + if settings.labels is not None: + missing = [ + str(p) for p in input_position_dirpaths if not labels_exist([p], settings.labels) + ] + if missing: + raise ValueError( + f"labels: {settings.labels!r} names a label image that " + f"{len(missing)} position(s) do not have (first: {missing[0]}); run " + "`biahub segment` on the plate first (its model name is the label name)." + ) + shape = create_label_images( + input_position_dirpaths, + [settings.output_label], + PROVENANCE_KEY, + settings.model_dump(mode="json"), + ) click.echo( - f"Created labels/{settings.target_channel} in {input_plate} " - f"({len(position_keys)} positions)" + f"Created labels/{settings.output_label} in " + f"{Path(input_position_dirpaths[0]).parents[2]} " + f"({len(input_position_dirpaths)} positions)" ) - return (T, C, Z, Y, X) + return shape def track( @@ -1012,14 +904,17 @@ def track( Notes ----- - Outputs go into each ``-i`` position: ``labels/`` (OME-NGFF label - image), ``tracks.geff`` and ``tracks___.csv``. The ``-i`` plate is - the plate track reads, so 2D tracking writes into the projected (Z=1) plate. - - Two input sources: ``-i`` (``input_position_dirpaths``) gives the plate structure - and shape/scale, while ``input_images[].path`` gives the pixel data. The first null - ``path`` is the primary source (``--input-images-path`` if given, else the ``-i`` - plate); later null entries are virtual/computed channels. + Outputs go into each ``-i`` position: ``labels/_tracked`` (OME-NGFF label + image; ``labels/`` on the foreground+contour path), + ``tracks.geff`` and ``tracks___.csv``. The ``-i`` plate is the plate + track reads, so 2D tracking writes into the projected (Z=1) plate, next to the + ``labels/`` image `biahub segment` wrote there. + + Two input sources: ``-i`` (``input_position_dirpaths``) gives the plate structure, + shape/scale and the label image, while ``input_images[].path`` gives pixel data + (foreground+contour). The first null ``path`` is the primary source + (``--input-images-path`` if given, else the ``-i`` plate); later null entries are + virtual/computed channels. """ from ultrack import MainConfig @@ -1043,7 +938,7 @@ def track( break # Workers skip the scaffold --init already made: they all share the -i plate. - if init_only or not _labels_exist(input_position_dirpaths, settings.target_channel): + if init_only or not labels_exist(input_position_dirpaths, settings.output_label): _init_labels(input_position_dirpaths, settings) with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as dataset: T, C, Z, Y, X = dataset.data.shape @@ -1060,15 +955,6 @@ def track( echo_resources(num_cpus, mem_gb, time_minutes) if init_only: - # --init runs once, on the head node, before the per-position fan-out, so it - # is the one place that can populate the shared weights cache without every - # worker racing to do it. Deliberately NOT done outside this branch: workers - # reach this code too, and importing cellpose here would fix the weights - # directory before stage_cellpose_weights could redirect it. - if settings.segmentation_method == "cellpose": - # Fail here, once, rather than in every worker of the fan-out. - check_cellpose_model_name(settings.cellpose_config.pretrained_model) - warm_cellpose_weights(settings.cellpose_config.pretrained_model) click.echo(f"Initialized tracking ({len(input_position_dirpaths)} positions)") return @@ -1082,21 +968,16 @@ def track( position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] - cellpose_cfg = ( - settings.cellpose_config if settings.segmentation_method == "cellpose" else None - ) - slurm_args = { "slurm_job_name": "track", "slurm_mem": f"{mem_gb}G", "slurm_cpus_per_task": num_cpus, "slurm_array_parallelism": 100, # process up to 100 positions at a time "slurm_time": time_minutes, - # Tracking runs cellpose on the GPU; use the non-preemptible `gpu` partition - # so long per-FOV jobs aren't evicted mid-run (unlike `preempted`). - # Override via --sbatch-filepath if a different partition is needed. - "slurm_partition": "gpu", - "slurm_gpus_per_node": 1, + # Ultrack is CPU-bound (segmentation is `biahub segment`'s, on the GPU) and + # cannot resume, so the non-preemptible `cpu` partition rather than + # `preempted`. Override via --sbatch-filepath if a different partition is needed. + "slurm_partition": "cpu", "slurm_use_srun": False, } @@ -1109,7 +990,7 @@ def track( if resolved_cluster == "local": # Explicit --sbatch-filepath values (e.g. '#LOCAL --timeout-min') win over # the computed local executor parameters. - slurm_args = {**local_executor_parameters(time_minutes), **slurm_args} + slurm_args = {**local_executor_parameters(time_minutes, gpu=False), **slurm_args} executor.update_parameters(**slurm_args) click.echo("Submitting jobs...") @@ -1124,11 +1005,12 @@ def track( position_dirpath=Path(position_dirpath), output_dirpath=output_dirpath, tracking_config=tracking_cfg, - label_name=settings.target_channel, + label_name=settings.output_label, input_images=settings.input_images, blank_frames_path=settings.blank_frames_path, scale=track_scale, - cellpose_config=cellpose_cfg, + input_label=settings.labels, + labels_sigma=settings.labels_sigma, output_mode=settings.output_mode, min_sibling_length=settings.min_sibling_length, ) @@ -1183,12 +1065,13 @@ def track_cli( init_only: bool = False, input_images_path: str | None = None, ): - """Track objects in 2D or 3D time-lapse microscopy data using configurable preprocessing. + """Track the instance labels of a label image (`biahub segment`) with ultrack. - Writes into each -i position: the labels as an OME-NGFF label image - (labels/), the tracks as tracks.geff, and a tracks CSV. 2D tracking - reads (and writes into) the projected Z=1 plate `biahub project` makes; the Ultrack - databases are temporary unless -o names a directory to keep them in. + Reads labels/ from each -i position and writes into it: the tracked labels + as an OME-NGFF label image (labels/_tracked), the tracks as tracks.geff, + and a tracks CSV. 2D tracking reads (and writes into) the projected Z=1 plate + `biahub project` makes; the Ultrack databases are temporary unless -o names a + directory to keep them in. Track does not segment or project. \b Full SLURM fan-out: diff --git a/biahub/utils/ngff.py b/biahub/utils/ngff.py index 6309cd90..022d8b13 100644 --- a/biahub/utils/ngff.py +++ b/biahub/utils/ngff.py @@ -1,9 +1,13 @@ -"""OME-Zarr store helpers: output paths, version, provenance, channels.""" +"""OME-Zarr store helpers: output paths, version, provenance, channels, label images.""" from pathlib import Path -from typing import Literal +from typing import Any, Literal + +import numpy as np +import zarr from iohub.ngff import open_ome_zarr +from iohub.ngff.utils import create_empty_plate from tqdm import tqdm #: ``fnmatch`` patterns for the per-position zattrs keys that steps carry @@ -129,3 +133,84 @@ def append_channels(input_data_path: Path, target_data_path: Path) -> None: position["0"][:, num_channels + i + 1] = appending_dataset[str(name)][0][:, i] dataset.print_tree() appending_dataset.close() + + +# Label chunks, TZYX. 16 x 1 x 128 x 128 uint32 is 1 MB, the DCA v0.2 SHOULD (512 KB +# is the MUST) and its "2D+time" example; growing T keeps per-plane reads cheap. +LABEL_CHUNKS = (16, 1, 128, 128) +# Shards span one time chunk (16 timepoints, the DCA SHOULD) and the whole YX plane, +# capped at 16 chunks (2048 px, the DCA spatial-shard SHOULD). +MAX_LABEL_SHARD_CHUNKS = 16 + + +def label_chunks_and_shards( + shape: tuple[int, int, int, int], +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Chunks and shards ratio for a (T, Z, Y, X) label array of ``shape``.""" + chunks = tuple(min(c, s) for c, s in zip(LABEL_CHUNKS, shape, strict=True)) + shards_ratio = (1, 1) + tuple( + min(-(-s // c), MAX_LABEL_SHARD_CHUNKS) + for c, s in zip(chunks[2:], shape[2:], strict=True) + ) + return chunks, shards_ratio + + +def labels_exist(position_dirpaths: list[Path], label_name: str) -> bool: + """Whether every position has the label image ``labels/``.""" + return all( + (Path(p) / "labels" / label_name / "zarr.json").exists() + or (Path(p) / "labels" / label_name / ".zattrs").exists() + for p in position_dirpaths + ) + + +def create_label_images( + position_dirpaths: list[Path], + label_names: list[str], + provenance_key: str, + provenance: dict[str, Any], +) -> tuple[int, int, int, int, int]: + """Create empty ``labels/`` images in every position of one plate. + + The labels are uint32 TZYX with the image's shape (its Z: one plane for a + projected plate, the stack for a 3D one), chunked per ``label_chunks_and_shards``, + and take the plate's OME-Zarr version. Existing label images are left untouched, + so ``--init`` and per-position workers can both call this. ``provenance`` is + merged into each label group's attributes under ``provenance_key``, keeping what + a worker recorded there before. + + Returns the image's (T, C, Z, Y, X) shape. + """ + position_dirpaths = [Path(p) for p in position_dirpaths] + with open_ome_zarr(str(position_dirpaths[0]), mode="r") as dataset: + T, C, Z, Y, X = dataset.data.shape + scale = dataset.scale + channel_names = dataset.channel_names + dtype = dataset.data.dtype + version = dataset.version + + plate = position_dirpaths[0].parents[2] + position_keys = [p.parts[-3:] for p in position_dirpaths] + label_chunks, label_shards_ratio = label_chunks_and_shards((T, Z, Y, X)) + create_empty_plate( + store_path=plate, + position_keys=position_keys, + channel_names=channel_names, + shape=(T, C, Z, Y, X), + scale=scale, + version=version, + dtype=dtype, + label_names=label_names, + label_dtype=np.uint32, + label_chunks=label_chunks, + label_shards_ratio=label_shards_ratio if version == "0.5" else None, + ) + + for position_dirpath in position_dirpaths: + for name in label_names: + label_group = zarr.open_group(position_dirpath / "labels" / name, mode="r+") + label_group.attrs[provenance_key] = { + **label_group.attrs.get(provenance_key, {}), + **provenance, + } + return (T, C, Z, Y, X) diff --git a/nextflow/configs/a549/qc_project.yaml b/nextflow/configs/a549/qc_project.yaml index e247912b..3755c7e3 100644 --- a/nextflow/configs/a549/qc_project.yaml +++ b/nextflow/configs/a549/qc_project.yaml @@ -1,8 +1,10 @@ # QC of the PROJECTED plate (`4-assemble/_proj.zarr`, Z=1): the plate # `biahub project` writes beside the assembled one and tracking reads and writes # into. It holds the projected channels and, after tracking, each position's -# `labels/` label image, so this is where the cell counts live: -# `instance_count` reads that label image (imaging-qc#243). The pipeline plans +# `labels/` (segment) and `labels/_tracked` (track) label images, +# so this is where the cell counts live: +# `instance_count` reads the TRACKED label image (imaging-qc#243); segment writes +# `labels/` beside it. The pipeline plans # this config after `track --init` (the label image must exist at plan time) and # computes it after tracking; without tracking it runs after project. # @@ -39,12 +41,12 @@ blank: gate: granularity: channel -# Cells per frame, from the label image tracking writes (`label:`, named after the -# track config's `target_channel`). `channels:` is the ONE anchor channel the labels +# Cells per frame, from the label image tracking writes (`label:`, the track config's +# `labels` plus `_tracked`). `channels:` is the ONE anchor channel the labels # were segmented from: the count rows carry its C. The plate has one Z-plane, so # `slice` and `volume` count the same thing; `slice` keeps per-(T, Z) rows. instance_count: - label: nuclei_prediction + label: nuclei_tracked channels: - nuclei_prediction granularity: slice diff --git a/nextflow/configs/a549/segment.yml b/nextflow/configs/a549/segment.yml index ff4623ac..073494a3 100644 --- a/nextflow/configs/a549/segment.yml +++ b/nextflow/configs/a549/segment.yml @@ -1,14 +1,25 @@ -# biahub segment: cellpose 4 nuclei labels from the virtual-stain nuclei channel of the -# PROJECTED plate (-assemble/_proj.zarr, Z=1; mantis v2, 0.1133 um/px). Segment -# does not project: a model without eval_args.do_3D segments in 2D and needs a Z=1 -# input, which `biahub project` (project.yml) makes from the assembled plate. -# Not wired into mantis-v2.nf yet. +# Segmentation config for the mantis-v2 pipeline (`biahub segment`). A549. +# +# Segment reads the PROJECTED plate (-assemble/_proj.zarr, Z=1) and writes +# one label image per model INTO it, `labels/` in every position, which is what +# `biahub track` reads (track.yml: `labels: nuclei`). Segment does not project: a model +# without eval_args.do_3D segments in 2D and needs a Z=1 input, which `biahub project` +# (project.yml) makes from the assembled plate. Channel names are the PRE-rename names +# of the assembled plate, which the projected plate inherits. +# +# The cellpose settings are the ones tracking ran with when it still segmented itself +# (cpsam_v2, diameter 80, min_size 2500): validated on the 2026_04_28 annotated FOVs and +# reproduced position by position on the 2026_08_11 v8 debug reconstruction. Cellpose 4 +# resizes each frame by 30 / diameter before inference, so diameter 80 runs at 37.5 % +# of native resolution, which shrinks small rounded mitotic nuclei; `diameter: null` +# with `min_size: 300` keeps them, at the cost of more false positives, and has not been +# validated for tracking. models: nuclei: pretrained_model: cpsam_v2 # cellpose 4: cpsam_v2, cpsam, cpdino, cpdino-vitb, or a path channels: [nuclei_prediction] # projected-plate channel name(s), at most 3 eval_args: - diameter: null # native resolution (cellpose 4 resizes by 30 / diameter) + diameter: 80 cellprob_threshold: 0.0 flow_threshold: 0.4 - min_size: 300 # px; small mitotic nuclei are ~300-2500 px + min_size: 2500 # px at native resolution diff --git a/nextflow/configs/a549/track.yml b/nextflow/configs/a549/track.yml index 83490dbb..1e937157 100644 --- a/nextflow/configs/a549/track.yml +++ b/nextflow/configs/a549/track.yml @@ -1,24 +1,32 @@ # Tracking config for the mantis-v2 pipeline (`biahub track`). A549. # -# Track reads the PROJECTED plate (-assemble/_proj.zarr, Z=1), so target_channel -# and cellpose_config.input_channel must name channels that `biahub project` wrote -# (project.yml's `channels`). If you rename channels before tracking, update these. +# Track reads the PROJECTED plate (-assemble/_proj.zarr, Z=1) and the label +# image `biahub segment` (segment.yml) wrote into it, `labels/`; it writes the +# tracked labels next to it as `labels/_tracked`, plus tracks.geff and a tracks +# CSV, into every position. Track neither segments nor projects. # # CURRENT SCHEMA. TrackingSettings forbids extra keys; the old fields `mode`, -# `z_range`, `focus_config` and `z_slicing` are rejected. Z-plane selection and -# projection moved to `biahub project` (project.yml): the 60-plane focus window and -# the mean projection that used to be here live there now. Mapping from the old schema: +# `z_range`, `focus_config`, `z_slicing` and `cellpose_config` are rejected. Mapping +# from the old schema: # mode -> output_mode # z_slicing / z_range / focus_config -> project.yml: channels..z, focus # input_images np.mean step -> project.yml: channels..method: mean +# cellpose_config (model, diameter, thresholds, min_size, input_channel) +# -> segment.yml: models. +# cellpose_config.labels_sigma -> labels_sigma +# target_channel -> labels (the segment.yml model name); the +# output is named _tracked -# 2D tracks the Z=1 projected plate and writes a Z=1 label plate. Track does not +# 2D tracks the Z=1 projected plate and writes a Z=1 label image. Track does not # project; a 3D input with output_mode 2D is an error. output_mode: "2D" fov: "*/*/*" -target_channel: nuclei_prediction -segmentation_method: cellpose +# The label image to track: segment.yml's model name. Tracks go to labels/nuclei_tracked. +labels: nuclei +# Smoothing (px) of the labels' contours before ultrack builds its own segmentation +# hypotheses from them. +labels_sigma: 5.0 # Sibling filter (ultrack filter_short_sibling_tracks, same meaning): at a division # where exactly one daughter lives 3 frames or fewer and does not divide again, drop @@ -29,21 +37,6 @@ segmentation_method: cellpose # and from 0.58 to 0.63-0.64 with this config on the annotated 2026_04_28 FOVs. # Remove the line to keep every division. min_sibling_length: 3 -cellpose_config: - pretrained_model: cpsam_v2 - diameter: 80 - cellprob_threshold: 0.0 - flow_threshold: 0.4 - gpu: true - min_size: 2500 - input_channel: nuclei_prediction - labels_sigma: 5.0 - -# The projected plate holds the nuclei already reduced to one plane; nothing to process. -input_images: - - path: null - channels: - nuclei_prediction: [] tracking_config: segmentation_config: diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf index 8c637818..7645007c 100644 --- a/nextflow/mantis-v2.nf +++ b/nextflow/mantis-v2.nf @@ -23,16 +23,18 @@ nextflow.enable.dsl = 2 // workflow body for why that is possible and what it buys. // // Flat-field → deskew → reconstruct → virtual-stain → assemble → project → -// track → QC is the full chain: assemble concatenates the +// segment → track → QC is the full chain: assemble concatenates the // deskew/reconstruct/virtual-stain channels into one plate, project reduces -// that 3D plate to the 2D (Z=1) plate 2D tracking needs, track reads it as its -// single input, and QC reads the finished stores. Follow the chaining below for -// the pattern. +// that 3D plate to the 2D (Z=1) plate 2D segmentation needs, segment writes +// cellpose labels into it, track links those labels and writes the tracked +// labels beside them, and QC reads the finished stores. Follow the chaining +// below for the pattern. // -// The last four are OPTIONAL and selected by whether their config is given, so -// a run performs the prefix it asks for: A549 wants assemble + project + track -// + QC, a neuromast run wants assemble + QC and no tracking (issue #306). A -// skipped step never renumbers the directories of the ones around it. +// The last five are OPTIONAL and selected by whether their config is given, so +// a run performs the prefix it asks for: A549 wants assemble + project + +// segment + track + QC, a neuromast run wants assemble + QC and no tracking +// (issue #306). A skipped step never renumbers the directories of the ones +// around it. // --------------------------------------------------------------------------- params.input = null // raw source — may not be a zarr store @@ -42,6 +44,7 @@ params.flat_field_config = null params.reconstruct_config = null params.virtual_stain_config = null params.project_config = null +params.segment_config = null params.track_config = null params.concatenate_config = null params.max_positions = 0 @@ -49,7 +52,8 @@ params.max_positions = 0 // either, both, or neither. Both QC'd stores become tabs of ONE report. // --qc_config nextflow/configs//qc.yaml (assembled store) // --qc_project_config nextflow/configs//qc_project.yaml (projected plate, -// which holds the labels tracking writes; counts them) +// which holds the labels segment and track write; counts +// the tracked ones) // One self-contained file per store, beside that family's other step configs. params.qc_config = null params.qc_project_config = null @@ -62,6 +66,7 @@ include { reconstruct_init_wf; reconstruct_tf_wf; reconstruct_run_wf } from './m include { virtual_stain_init_wf; virtual_stain_run_wf } from './modules/virtual_stain' include { assemble_init_wf; assemble_run_wf } from './modules/assembly' include { project_init_wf; project_run_wf } from './modules/project' +include { segment_init_wf; segment_run_wf } from './modules/segment' include { track_init_wf; track_run_wf } from './modules/tracking' include { qc_plan_wf; qc_compute_wf; qc_report_wf; qc_report_spec } from './modules/qc' include { notify_step; notify_run_start; notify_run_end } from './modules/notify' @@ -102,8 +107,8 @@ def step_directories(performed) { virtual_stain : 'virtual-stain', assemble : 'assemble', // project writes beside the assembled plate (/_proj.zarr) - // and track writes into the plate it reads; neither has a directory of - // its own or takes a number. + // and segment and track write into the plate they read; none has a + // directory of its own or takes a number. ] def layout = [:] order.each { key, name -> @@ -135,14 +140,17 @@ workflow { // end at `4-assemble`. // // Project turns the assembled 3D plate into the 2D (Z=1) plate that 2D - // tracking reads: `biahub track` no longer selects or projects Z-planes, so - // an A549 run names --project_config as well as --track_config. Tracking - // without it reads the assembled plate itself, which suits 3D tracking. - // The projected plate is a companion of the assembled one and lives beside - // it as `-assemble/_proj.zarr`; track writes its labels and tracks - // into the plate it reads. Neither has a step directory. + // segmentation and tracking read: neither `biahub segment` nor `biahub track` + // selects or projects Z-planes, so an A549 run names --project_config as + // well as --segment_config and --track_config. Without it they read the + // assembled plate itself, which suits 3D work. The projected plate is a + // companion of the assembled one and lives beside it as + // `-assemble/_proj.zarr`; segment writes its label images into the + // plate it reads and track writes the tracked labels and tracks beside + // them. None of the three has a step directory. def assemble_on = params.concatenate_config as boolean def project_on = params.project_config as boolean + def segment_on = params.segment_config as boolean def track_on = params.track_config as boolean def qc_image_on = params.qc_config as boolean def qc_project_on = params.qc_project_config as boolean @@ -154,9 +162,15 @@ workflow { if (project_on && !assemble_on) { error "--project_config needs --concatenate_config: project reads the assembled plate." } + if (segment_on && !assemble_on) { + error "--segment_config needs --concatenate_config: segment reads the assembled (or projected) plate." + } if (track_on && !assemble_on) { error "--track_config needs --concatenate_config: tracking reads the assembled (or projected) plate." } + // Track's `labels:` must name a label image segment wrote; `biahub track + // --init` refuses a plate without it, so a run naming --track_config but + // not --segment_config fails in the init phase with that message. if (qc_image_on && !assemble_on) { error "--qc_config needs --concatenate_config: it QCs the assembled store." } @@ -214,9 +228,9 @@ workflow { // Beside the assembled plate, with a `_proj` suffix: the 2D companion of // that store, not a step of its own. project_output = project_on ? "${out}/${layout.assemble}/${ds}_proj.zarr" : null - // What tracking reads AND writes into (labels/, tracks.geff, the tracks CSV): - // the projected (Z=1) plate when project runs, else the assembled plate. - // Both the plate structure and the pixels come from it. + // What segment and track read AND write into (labels/, tracks.geff, the + // tracks CSV): the projected (Z=1) plate when project runs, else the + // assembled plate. Both the plate structure and the pixels come from it. track_input = project_on ? project_output : assemble_output // ======================================================================== @@ -289,16 +303,30 @@ workflow { init_signals << pj_init.done } + // The plate segment and track work in is scaffolded by project's init when + // project runs, else by assemble's. + def plate_scaffolded = project_on ? pj_init.done : as_init.done + + if (segment_on) { + // Segment reads the projected plate (or the assembled one): its init + // checks that the models' channels exist and that their dimensionality + // matches the plate's Z, creates the empty label image per model in every + // position, and warms the shared cellpose weights cache, which is better + // done here than with N GPU workers racing for it. + sg_init = segment_init_wf(track_input, params.segment_config, plate_scaffolded) + init_signals << sg_init.done + } + if (track_on) { - // Track reads ONE plate for both of its inputs (plate structure and image - // data): the projected plate when project runs, else the assembled plate - // (concatenate preserves channel names, so the track config's channel - // names resolve unchanged). Its init checks that the plate's Z matches - // output_mode, creates the empty label image tracking fills in every - // position of that plate, and warms the shared cellpose weights cache, - // which is better done here than with N GPU workers racing for it. + // Track reads ONE plate for both of its inputs (plate structure and the + // label image to track; image data for a legacy foreground+contour + // config): the projected plate when project runs, else the assembled + // plate. Its init checks that the plate's Z matches output_mode and that + // the label image `labels:` names exists — segment's init created it, + // which is why this init follows that one — and creates the empty + // tracked label image in every position. tk_init = track_init_wf(track_input, params.track_config, - project_on ? pj_init.done : as_init.done) + segment_on ? sg_init.done : plate_scaffolded) init_signals << tk_init.done } @@ -306,10 +334,11 @@ workflow { // inits do — `plan-stage` and `estimate-resources` read the store's // structure and validate the config, and read no pixels. Each store's plan // is gated on the init that scaffolded what it reads: for the projected - // plate with tracking on that is track's init, because a config counting - // instances in `labels/` is refused at plan time until the label image - // exists. The tab label is the step directory, plus the store name for the - // projected plate, which shares the assemble directory. + // plate that is the last of project's, segment's and track's inits that + // runs, because a config counting instances in `labels/` is refused at plan + // time until the label image exists. The tab label is the step directory, + // plus the store name for the projected plate, which shares the assemble + // directory. def qc_stores = [] if (qc_on) { if (qc_image_on) { @@ -317,9 +346,9 @@ workflow { config: params.qc_config, scaffolded: as_init.done] } if (qc_project_on) { + def proj_scaffolded = track_on ? tk_init.done : (segment_on ? sg_init.done : pj_init.done) qc_stores << [label: "${layout.assemble}/${ds}_proj", zarr: project_output, - config: params.qc_project_config, - scaffolded: track_on ? tk_init.done : pj_init.done] + config: params.qc_project_config, scaffolded: proj_scaffolded] } // The trigger is mapped, NOT combined. `combine` concatenates the two @@ -397,32 +426,41 @@ workflow { } // Projection reads the assembled plate, one task per position, and writes the - // 2D plate tracking reads. + // 2D plate segmentation and tracking read. if (project_on) { project_done = project_run_wf(all_positions, assemble_output, project_output, params.project_config, pj_init.resources, assemble_done.done) } + def plate_written = project_on ? project_done.done : assemble_done.done + + // Segmentation (cellpose, GPU) writes its label images into the plate, one + // task per position, once that plate holds data. + if (segment_on) { + segment_done = segment_run_wf(all_positions, track_input, params.segment_config, + sg_init.resources, plate_written) + } - // Tracking runs AFTER assemble (and project) rather than in parallel with it — - // the tradeoff for reading one store instead of two is that the whole plate - // must be assembled first. Two consequences: any Z/Y/X crop or time_indices - // subset in the concatenate config is what tracking sees, and the - // intermediate stores are no longer needed once assemble is verified. To go - // back to the parallel wiring, point the first input at reconstruct_output, - // the second at virtual_stain_output, and gate on virtual_stain_done. + // Tracking runs AFTER assemble, project and segment rather than in parallel + // with them — the tradeoff for reading one store instead of two is that the + // whole plate must be assembled first. Two consequences: any Z/Y/X crop or + // time_indices subset in the concatenate config is what tracking sees, and + // the intermediate stores are no longer needed once assemble is verified. To + // go back to the parallel wiring, point the first input at + // reconstruct_output, the second at virtual_stain_output, and gate on + // virtual_stain_done. if (track_on) { track_done = track_run_wf(all_positions, track_input, track_input, params.track_config, tk_init.resources, - project_on ? project_done.done : assemble_done.done) + segment_on ? segment_done.done : plate_written) } // ----- QC compute ------------------------------------------------------- // Planned in the init phase; the compute it planned is gated on the last // step that writes each store: the ASSEMBLED plate on assemble_done (so its - // pixel QC runs concurrently with projection and tracking), and the - // PROJECTED plate on track_done when tracking is on (its labels are counted - // there), else project_done. + // pixel QC runs concurrently with projection, segmentation and tracking), + // and the PROJECTED plate on the last of project, segment and track that + // runs (the tracked labels are counted there). // // A QC verdict cannot fail the pipeline: `imaging-qc gate` exits 0 whether // positions pass or fail, recording the verdict in each store's own @@ -445,7 +483,7 @@ workflow { assemble_done.done.map { tuple(assemble_output, params.qc_config) } ) } if (qc_project_on) { - def proj_written = track_on ? track_done.done : project_done.done + def proj_written = track_on ? track_done.done : (segment_on ? segment_done.done : project_done.done) qc_compute_ready = qc_compute_ready.mix( proj_written.map { tuple(project_output, params.qc_project_config) } ) } @@ -458,14 +496,15 @@ workflow { // Once the LAST step has finished, delete what only existed to feed the // assembled plate: the flat-field, deskew, reconstruct and virtual-stain step // directories (store, slurm_output/ and .iohub-progress/ alike). The - // assembled and projected plates and the tracking store are DELIVERABLES and + // assembled and projected plates (with the labels and tracks inside) are DELIVERABLES and // are never targets. Beside them, delete only what served the run: // slurm_output/ (submitit placeholders and a pointer README; the real logs // are in nextflow/slurm_output/) and the assembled and projected stores' // .iohub-progress/ resume markers, which left behind would make concatenate // or project skip every write unit of a later run into this directory, - // keeping data built from the old intermediates. Tracking writes no resume - // markers. + // keeping data built from the old intermediates. Segment's resume markers + // (for its label images) live under the projected plate's entry in the same + // .iohub-progress/; tracking writes none. // // What was removed is appended to nextflow/intermediates_cleaned.txt rather // than logged. Nextflow's `cleanup`, switched on above, empties the work @@ -488,6 +527,7 @@ workflow { def final_signals = [assemble_done.done] if (project_on) final_signals << project_done.done + if (segment_on) final_signals << segment_done.done if (track_on) final_signals << track_done.done if (qc_on) final_signals << qc_report.done final_gate = channel.empty() @@ -534,6 +574,7 @@ workflow { ] if (assemble_on) step_events << [label: 'assemble', done: assemble_done.done, output: assemble_output] if (project_on) step_events << [label: 'project', done: project_done.done, output: project_output] + if (segment_on) step_events << [label: 'segment', done: segment_done.done, output: track_input] if (track_on) step_events << [label: 'track', done: track_done.done, output: track_input] if (qc_on) step_events << [label: 'QC', done: qc_report.done, output: qc_report_dir] if (cleanup_on) step_events << [label: 'cleanup intermediates', done: cleanup_run.done, output: cleanup_record] diff --git a/nextflow/modules/segment.nf b/nextflow/modules/segment.nf index f7ed7fdf..91922b11 100644 --- a/nextflow/modules/segment.nf +++ b/nextflow/modules/segment.nf @@ -1,17 +1,25 @@ -// Segmentation subworkflows: init (scaffold + resources + weights cache) and run -// (fan-out × N positions), the same split as tracking. +// Segmentation subworkflows: init (label scaffold + resources + weights cache) and +// run (fan-out × N positions), the same split as tracking. // -// This module is PATH-AGNOSTIC. Callers pass the input zarr, output zarr and -// config explicitly. For 2D models the input is the PROJECTED (Z=1) plate that +// This module is PATH-AGNOSTIC. Callers pass the input zarr and config +// explicitly. For 2D models the input is the PROJECTED (Z=1) plate that // `biahub project` writes from the assembled plate; segment does not project. // +// Segment writes INTO the plate it reads: one OME-NGFF label image per model, +// `labels/`, in every position, which is what `biahub track` reads +// (track.yml `labels: `). There is no step directory and no `-o`; the +// SLURM logs are in nextflow/slurm_output/segment/, like every step's, and +// submitit's debug-mode placeholders land in the task's work directory. +// // `biahub segment --init` reads only the input plate's channels, shape and scale, // so it runs as soon as the input plate is scaffolded: a config naming a channel // the plate lacks, 2D models on a 3D plate (or 3D models on a Z=1 plate), or an -// unknown cellpose model fails before any GPU job. It also warms the shared cellpose weights -// cache, better done once here than by N GPU workers racing for it. +// unknown cellpose model fails before any GPU job. It creates the empty label +// images (so `track --init`, which needs them to exist, can follow it in the +// init phase) and warms the shared cellpose weights cache, better done once here +// than by N GPU workers racing for it. -include { parse_resources; slurm_logs; slurm_log_dir; slurm_output_readme; retry_time } from './common' +include { parse_resources; slurm_logs; slurm_log_dir; retry_time } from './common' process init_segment { @@ -19,7 +27,6 @@ process init_segment { input: val input_zarr - val output_zarr val config path config_file // staged only for the task hash: see common.nf, #397 val trigger @@ -30,10 +37,8 @@ process init_segment { script: """ mkdir -p "${slurm_log_dir('segment')}" - ${slurm_output_readme('segment', output_zarr)} biahub segment --init \ -i "${input_zarr}"/*/*/* \ - -o "${output_zarr}" \ -c "${config}" """ } @@ -42,9 +47,9 @@ process run_segment { tag "${position}" // Preemptable, like virtual staining: `--resume` (below) makes a reclaimed task // cost at most the timepoint it was segmenting, so it runs on the `preempted` - // partition (any GPU node) rather than `gpu`, which is reserved for tracking - // because ultrack cannot resume. It inherits nextflow.config's errorStrategy, - // which retries preemptions; no SLURM --requeue (see nextflow.config). + // partition (any GPU node) rather than `gpu`. It inherits nextflow.config's + // errorStrategy, which retries preemptions; no SLURM --requeue (see + // nextflow.config). label 'gpu_preempted' clusterOptions { "--gres=gpu:1 " + slurm_logs('segment') } cpus { meta.cpus } @@ -54,7 +59,6 @@ process run_segment { input: tuple val(position), val(meta) val input_zarr - val output_zarr val config path config_file // staged only for the task hash: see common.nf, #397 @@ -67,32 +71,29 @@ process run_segment { """ biahub segment --cluster debug --resume \ -i "${input_zarr}/${position}" \ - -o "${output_zarr}" \ -c "${config}" """ } -// Validate the config, scaffold the label plate, warm the cellpose weights. +// Validate the config, create the empty label images, warm the cellpose weights. // Metadata-only and cheap, so it belongs in the pipeline's up-front init phase. // // take: -// input_zarr path to the input plate.zarr -// output_zarr path to the segmentation output plate.zarr +// input_zarr path to the input plate.zarr (receives the label images) // config path to the segment settings YAML // trigger gating channel — init starts once this emits // emit: // resources the RESOURCES payload sizing one position's task -// done fires once the output plate exists +// done fires once every position has its empty label images workflow segment_init_wf { take: input_zarr - output_zarr config trigger main: - init_out = init_segment(input_zarr, output_zarr, config, file(config), trigger.collect().map { 'done' }) + init_out = init_segment(input_zarr, config, file(config), trigger.collect().map { 'done' }) emit: resources = init_out.map { stdout_text -> parse_resources(stdout_text) } @@ -104,8 +105,7 @@ workflow segment_init_wf { // // take: // positions collected channel of position keys -// input_zarr input plate store -// output_zarr path to the segmentation output plate.zarr +// input_zarr input plate store (receives the label images) // config path to the segment settings YAML // resources RESOURCES payload from segment_init_wf // prev_done gating channel — the input store holds data @@ -113,7 +113,6 @@ workflow segment_run_wf { take: positions input_zarr - output_zarr config resources prev_done @@ -127,7 +126,7 @@ workflow segment_run_wf { .combine(prev_done.map { 'done' }) .map { pos, meta, _gate -> [pos, meta] } - sg_done = run_segment(pos_meta, input_zarr, output_zarr, config, file(config)) | collect + sg_done = run_segment(pos_meta, input_zarr, config, file(config)) | collect emit: done = sg_done diff --git a/nextflow/modules/tracking.nf b/nextflow/modules/tracking.nf index 77a2090d..49040adf 100644 --- a/nextflow/modules/tracking.nf +++ b/nextflow/modules/tracking.nf @@ -1,27 +1,31 @@ -// Tracking subworkflows: init (label scaffold + resources + weights cache) and -// run (fan-out × N positions). +// Tracking subworkflows: init (label scaffold + resources) and run (fan-out × N +// positions). // -// This module is PATH-AGNOSTIC. Callers pass the input zarr (plate structure), -// input images zarr (image data for tracking), and config explicitly. +// This module is PATH-AGNOSTIC. Callers pass the input zarr (plate structure and +// the label image to track), input images zarr (image data, for the legacy +// foreground+contour path), and config explicitly. // -// Tracking is a 2-input step: `input_zarr` supplies the plate structure and -// `input_images_zarr` the pixels. mantis-v2.nf passes the assembled plate for -// both. The outputs go INTO `input_zarr`'s positions — `labels/`, -// `tracks.geff` and a tracks CSV — and nothing else is kept: each worker builds -// its Ultrack database in the node's $TMPDIR and deletes it after export (no -// -o), so tracking has no step directory. Its real logs are in -// nextflow/slurm_output/track/, like every step's; submitit's debug-mode -// placeholders land in the task's work directory. +// Tracking reads the label image `biahub segment` wrote into `input_zarr` +// (`labels/`) and writes INTO the same positions — the tracked labels as +// `labels/_tracked`, `tracks.geff` and a tracks CSV — and nothing else +// is kept: each worker builds its Ultrack database in the node's $TMPDIR and +// deletes it after export (no -o), so tracking has no step directory. Its real +// logs are in nextflow/slurm_output/track/, like every step's; submitit's +// debug-mode placeholders land in the task's work directory. `input_images_zarr` +// only matters to a foreground+contour config; mantis-v2.nf passes the same +// plate for both. // // INIT AND RUN ARE SEPARATE SUBWORKFLOWS, and hoisting init is worth more here // than anywhere else: the tracking config is the LAST one a run would otherwise // parse, so a typo in it used to surface after every reconstruction step and // the assembly had finished. `biahub track --init` reads only the input plate's -// shape and scale (z-slice resolution is data-free) and creates an empty label -// image in each position, so it runs against the assembled plate as soon as -// `concatenate --init` has scaffolded it. It also -// warms the shared cellpose weights cache, which is strictly better done before -// any GPU worker exists to race for it. +// shape and scale, checks that the label image it will track exists (segment's +// init creates it, so it runs after that one) and creates the empty tracked +// label image in each position. +// +// Tracking is CPU work: ultrack re-segments from the labels' contours and solves +// an ILP; the GPU segmentation is `biahub segment`'s. It cannot resume, so it +// runs on the non-preemptible `cpu` partition. include { parse_resources; slurm_logs; slurm_log_dir } from './common' @@ -49,8 +53,8 @@ process init_track { process run_track { tag "${position}" - label 'gpu' - clusterOptions { "--gres=gpu:1 " + slurm_logs('track') } + label 'cpu_reserved' + clusterOptions { slurm_logs('track') } cpus { meta.cpus } memory { "${meta.mem_gb} GB" } time '2h' @@ -77,11 +81,12 @@ process run_track { } -// Validate the config, create the empty label images, warm the cellpose weights. +// Validate the config and create the empty tracked label images. // Metadata-only and cheap, so it belongs in the pipeline's up-front init phase. // // take: -// input_zarr path to the input plate.zarr (plate structure; receives the labels) +// input_zarr path to the input plate.zarr (plate structure and the label image +// to track; receives the tracked labels) // config path to the track settings YAML // trigger gating channel — init starts once this emits // emit: diff --git a/nextflow/nextflow.config b/nextflow/nextflow.config index fdaf3fd4..015129c8 100644 --- a/nextflow/nextflow.config +++ b/nextflow/nextflow.config @@ -172,7 +172,7 @@ process { errorStrategy = 'ignore' } - // NOTE: labels 'cpu', 'gpu' and 'gpu_preempted' intentionally have NO + // NOTE: labels 'cpu', 'cpu_reserved', 'gpu' and 'gpu_preempted' intentionally have NO // withLabel body here. Each module sets cpus/memory/time in its own process // body (dynamically from the init step's RESOURCES payload). A withLabel // selector in config overrides process-body directives, so defining @@ -218,9 +218,13 @@ profiles { // preemption costs at most the one it was predicting. Runs // on any of its GPU nodes, the A40 workstations included // (see run_virtual_stain). + // cpu_reserved -> `cpu` partition (not preempted), for tracking. Ultrack + // is CPU work (segmentation runs on the GPU in `biahub + // segment`) and cannot resume, so a preemption would restart + // the position from scratch. // gpu -> `gpu` partition (high PriorityTier, effectively not - // preempted), for tracking. Ultrack cannot resume, so a - // preemption would restart the position from scratch. + // preempted), for GPU work that cannot resume. Unused by the + // mantis-v2 steps since segmentation gained --resume. slurm { executor { name = 'slurm' @@ -265,6 +269,9 @@ profiles { withLabel: 'cpu' { queue = 'preempted' } + withLabel: 'cpu_reserved' { + queue = 'cpu' + } withLabel: 'gpu_preempted' { queue = 'preempted' } diff --git a/pyproject.toml b/pyproject.toml index ddf44327..ee5442aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,11 +139,14 @@ imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipelin # royerlab/ultrack main once #290 merges, and at a release once one contains them. ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5f84ccda3bdf072f5289f939c" } -# Pin iohub to main at the czbiohub-sf/iohub#486 merge, for -# create_empty_plate(label_names=...), with which track --init scaffolds the empty -# label images (#360), and czbiohub-sf/iohub#485, which puts the labels list on -# /labels/zarr.json so ome-zarr-models and dca-validate can open a field -# that has labels. Repoint at a PyPI release once one contains them. +# Pin iohub to the czbiohub-sf/iohub#489 branch (on main at the #486 merge): +# process_single_position(output_label=...), with which `biahub segment` writes its +# labels into the plate's label images with --resume. Main's #486 adds +# create_empty_plate(label_names=...), with which segment --init and track --init +# scaffold the empty label images (#360), and #485 puts the labels list on +# /labels/zarr.json so ome-zarr-models and dca-validate can open a field that +# has labels. Repoint at main once #489 merges, and at a PyPI release once one +# contains them. # # Main's own reasons: #476 syncs a write unit's shards before recording it as # done (without it a preemption between the write and the marker leaves a `.done` @@ -151,7 +154,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "c9832a7675e0fb997a27db2796623d19b7dc9f1a" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "8c85656" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/settings/example_segmentation_settings.yml b/settings/example_segmentation_settings.yml index 5eaaf0b5..641c3801 100644 --- a/settings/example_segmentation_settings.yml +++ b/settings/example_segmentation_settings.yml @@ -1,5 +1,7 @@ # Segmentation settings for cellpose 4 (`biahub segment`). -# Model names become the output channel names (e.g. membrane -> membrane_labels). +# Model names become the label image names in each -i position (e.g. membrane -> +# labels/membrane), which `biahub track` reads; with -o they become the channels of a +# separate label plate (membrane -> membrane_labels). # A model segments in 2D unless eval_args.do_3D (or a stitch_threshold) is set. Segment # does not project: 2D models need a Z=1 input plate (make one from a stack with # `biahub project`), 3D models a stack; a mismatch is an error. All models in one config diff --git a/settings/example_track_settings.yml b/settings/example_track_settings.yml index 92155c7e..8673d9f2 100644 --- a/settings/example_track_settings.yml +++ b/settings/example_track_settings.yml @@ -1,21 +1,27 @@ # The dimensionality of the tracking. "2D" tracks a Z=1 input plate and writes a Z=1 -# label plate; "3D" tracks the stack and writes labels of the same Z. Track does NOT +# label image; "3D" tracks the stack and writes labels of the same Z. Track does NOT # project or select Z-planes: a 3D input with output_mode 2D is an error. To track a # projection, a window or one plane of a stack, make the Z=1 plate first with # `biahub project` (settings/example_project_settings.yml) and track that. output_mode: "2D" fov: "*/*/*" - -# Name of the target channel used to name the tracking output label channel. -target_channel: nuclei_prediction blank_frames_path: blank_frames.csv -# Segmentation method used to feed ultrack: "cellpose" (the default; needs -# cellpose_config) or "foreground_contour" (the legacy path). -# - cellpose: run a Cellpose model per frame to produce instance labels. -# - foreground_contour: derive a foreground mask + contour map (ultrack segments). -segmentation_method: cellpose +# What feeds ultrack: "labels" (the default) tracks the instance labels of a label +# image of the -i plate, written by `biahub segment` +# (settings/example_segmentation_settings.yml); "foreground_contour" is the legacy +# path, which derives a foreground mask and contour map from input_images. Track does +# not segment: the cellpose settings live in the segment config. +segmentation_method: labels + +# The label image to track, `labels/` in every -i position: the segment +# config's model name. The tracked labels are written next to it as +# `labels/_tracked`, with tracks.geff and a tracks CSV. +labels: nuclei +# Smoothing (pixels) ultrack applies to the labels' contours before it builds its own +# segmentation hypotheses from them. +labels_sigma: 5.0 # Optional post-processing (ultrack filter_short_sibling_tracks): at a division where # exactly one daughter lives N frames or fewer and does not divide again, drop that @@ -23,69 +29,51 @@ segmentation_method: cellpose # both that short is kept. Removes short false divisions. Off by default. # min_sibling_length: 3 -# Cellpose model parameters (used only when segmentation_method: cellpose). -cellpose_config: - pretrained_model: cpsam_v2 # cellpose 4: cpsam_v2, cpsam, cpdino, cpdino-vitb, or a path - diameter: 80 # frames are resized by 30/diameter before inference; null = native size - cellprob_threshold: 0.0 - flow_threshold: 0.4 - gpu: true - min_size: 2500 - input_channel: nuclei_prediction # channel (from input_images) to segment - labels_sigma: 5.0 # smoothing applied to labels before ultrack - -# Input image configurations. Each entry defines one dataset (path) and the channels -# to load/process from it. +# --------------------------------------------------------------------------- +# Alternative: segmentation_method: foreground_contour (legacy) +# Replace `labels` / `labels_sigma` above with `target_channel` (the name of the label +# image the tracks are written to) and an input_images block whose processing functions +# produce the foreground mask and contour map: +# +# segmentation_method: foreground_contour +# target_channel: nuclei_prediction # -# Two input sources feed tracking: +# Input image configurations. Each entry defines one dataset (path) and the channels +# to load/process from it. Two input sources feed tracking: # * `-i` (input_position_dirpaths) defines the plate structure (positions, shape, scale). # * input_images[].path supplies the pixel data loaded per FOV. # The FIRST entry with `path: null` is the primary data source: it resolves to -# `--input-images-path` if given (Nextflow), else to the `-i` input plate. So you can -# leave it null here and point `-i` at the data. Any LATER null entries are -# virtual/computed channels (see the foreground_contour example below). -input_images: - - # Cellpose mode: load the channel named by cellpose_config.input_channel (already a - # single plane in a projected plate; no processing needed). - # path: null -> resolved from -i (or --input-images-path). Use an explicit - # /path/to/projected.zarr if the data lives elsewhere than the -i plate. - - path: null - channels: - nuclei_prediction: [] - - # --------------------------------------------------------------------------- - # Alternative: segmentation_method: foreground_contour - # Replace the input_images block above with the following (first entry is the - # primary data source -> path: null; the second entry's null is a virtual channel): - # - # - path: null - # channels: - # nuclei_prediction: - # - function: ultrack.imgproc.normalize - # kwargs: - # gamma: 0.7 - # membrane_prediction: - # - function: ultrack.imgproc.normalize - # kwargs: - # gamma: 0.7 - # - # # Virtual channel computed dynamically from the prior outputs (path: null). - # - path: null - # channels: - # foreground: - # - function: ultrack.imgproc.detect_foreground - # input_channels: - # - nuclei_prediction - # kwargs: - # sigma: 90 - # contour: - # - function: biahub.cli.track.mem_nuc_contour - # input_channels: - # - nuclei_prediction - # - membrane_prediction - # kwargs: {} - # --------------------------------------------------------------------------- +# `--input-images-path` if given (Nextflow), else to the `-i` input plate. Any LATER +# null entries are virtual/computed channels. +# +# input_images: +# - path: null +# channels: +# nuclei_prediction: +# - function: ultrack.imgproc.normalize +# kwargs: +# gamma: 0.7 +# membrane_prediction: +# - function: ultrack.imgproc.normalize +# kwargs: +# gamma: 0.7 +# +# # Virtual channel computed dynamically from the prior outputs (path: null). +# - path: null +# channels: +# foreground: +# - function: ultrack.imgproc.detect_foreground +# input_channels: +# - nuclei_prediction +# kwargs: +# sigma: 90 +# contour: +# - function: biahub.cli.track.mem_nuc_contour +# input_channels: +# - nuclei_prediction +# - membrane_prediction +# kwargs: {} +# --------------------------------------------------------------------------- # Configuration for Ultrack tracking algorithm. # See: https://royerlab.github.io/ultrack/ diff --git a/tests/test_cellpose_inputs.py b/tests/test_cellpose_inputs.py index 4ccdabdc..5d090b49 100644 --- a/tests/test_cellpose_inputs.py +++ b/tests/test_cellpose_inputs.py @@ -5,8 +5,6 @@ The test extra does not install cellpose, so a recording stand-in replaces it. """ -import warnings - import numpy as np import pytest @@ -52,29 +50,6 @@ def test_segment_passes_only_the_configured_channel(fake_cellpose): assert np.all(seen == 30) -def test_track_cellpose_uses_the_requested_model(fake_cellpose): - from biahub.track import run_cellpose_per_frame - - labels = run_cellpose_per_frame( - np.zeros((2, 8, 8), dtype=np.float32), pretrained_model="cpdino", gpu=False - ) - - assert fake_cellpose.built == ["cpdino"] - assert labels.shape == (2, 8, 8) - - -def test_model_type_in_a_track_config_warns(): - from biahub.settings import CellposeConfig - - with pytest.warns(UserWarning, match="model_type"): - config = CellposeConfig(model_type="nuclei") - assert config.pretrained_model == "cpsam_v2" - - with warnings.catch_warnings(): - warnings.simplefilter("error") - CellposeConfig() - - def test_segment_2d_passes_the_single_plane_without_a_z_axis(fake_cellpose): czyx = np.full((1, 1, 8, 8), 30, dtype=np.float32) diff --git a/tests/test_cli/test_segment_cli.py b/tests/test_cli/test_segment_cli.py index 8a0c1449..6570ea7c 100644 --- a/tests/test_cli/test_segment_cli.py +++ b/tests/test_cli/test_segment_cli.py @@ -46,6 +46,12 @@ def cpu_cellpose(monkeypatch): monkeypatch.setattr("biahub.segment.cellpose_device", lambda gpu: torch.device("cpu")) +@pytest.fixture(autouse=True) +def run_in_tmp_path(monkeypatch, tmp_path): + """Without -o the submitit folder is ./slurm_output: keep it out of the repo.""" + monkeypatch.chdir(tmp_path) + + def _run(*args): return CliRunner().invoke(cli, ["segment", *map(str, args)]) @@ -77,6 +83,144 @@ def test_segment_init_prints_resources_and_stamps_provenance( assert fake_cellpose.evaluated == [] # --init segments nothing +# --------------------------------------------------------------------------- +# No -o: the labels go into the -i plate as label images, labels/ +# --------------------------------------------------------------------------- + + +def test_segment_init_creates_label_images_in_the_input_plate( + fake_cellpose, plate_2d, segment_config, tmp_path +): + plate_path, _ = plate_2d # (T=3, C=3, Z=1, Y=5, X=6) + positions = [plate_path / "A" / "1" / "0", plate_path / "B" / "1" / "0"] + + result = _run("--init", "-i", *positions, "-c", segment_config()) + + assert result.exit_code == 0, result.output + assert "RESOURCES:" in result.output + assert "Created labels/{nuc}" in result.output + for position in positions: + with open_ome_zarr(position, mode="r") as pos: + assert pos.label_names() == ["nuc"] + label = pos.get_label("nuc")["0"] + assert label.shape == (3, 1, 5, 6) # the image's (T, Z, Y, X) + assert label.dtype == np.uint32 + assert not label[:].any() + assert pos.channel_names == ["GFP", "RFP", "Phase3D"] # the image is untouched + attrs = yaml.safe_load((position / "labels" / "nuc" / "zarr.json").read_text()) + # Provenance on the label image: the config as written (names, not indices). + assert attrs["attributes"]["biahub-segment"]["models"]["nuc"]["channels"] == ["GFP"] + assert fake_cellpose.evaluated == [] + + +def test_segment_debug_fills_the_label_image( + fake_cellpose, cpu_cellpose, plate_2d, segment_config, tmp_path +): + plate_path, _ = plate_2d + config = segment_config() + position = plate_path / "B" / "1" / "0" + assert _run("--init", "-i", position, "-c", config).exit_code == 0 + fake_cellpose.built.clear() + + result = _run("--cluster", "debug", "-i", position, "-c", config) + + assert result.exit_code == 0, result.output + assert "Segmentation complete" in result.output + with open_ome_zarr(position, mode="r") as pos: + assert np.all(pos.get_label("nuc")["0"][:] == 1) # the stand-in labels every pixel 1 + assert pos.data.shape == (3, 3, 1, 5, 6) + assert fake_cellpose.built == ["cpsam_v2"] + assert len(fake_cellpose.evaluated) == 3 # one per timepoint + + +def test_segment_resume_skips_finished_label_images( + fake_cellpose, cpu_cellpose, create_custom_plate, segment_config, tmp_path +): + """Resume works on the label image as on an image (iohub output_label).""" + from iohub.ngff.utils import create_empty_plate + + # A Zarr v3 (OME-Zarr 0.5) plate: resume needs it. + plate_path = tmp_path / "v05.zarr" + create_empty_plate( + store_path=plate_path, + position_keys=[("A", "1", "0")], + channel_names=["GFP"], + shape=(3, 1, 1, 5, 6), + version="0.5", + ) + with open_ome_zarr(plate_path / "A/1/0", mode="r+") as pos: + pos["0"][:] = np.random.randint(1, 100, size=(3, 1, 1, 5, 6)).astype(np.float32) + config = segment_config() + position = plate_path / "A" / "1" / "0" + assert _run("--init", "-i", position, "-c", config).exit_code == 0 + assert _run("--cluster", "debug", "-i", position, "-c", config).exit_code == 0 + n_first = len(fake_cellpose.evaluated) + assert n_first == 3 + + result = _run("--cluster", "debug", "--resume", "-i", position, "-c", config) + + assert result.exit_code == 0, result.output + assert len(fake_cellpose.evaluated) == n_first + + +def test_segment_writes_one_label_image_per_model( + fake_cellpose, cpu_cellpose, plate_2d, tmp_path +): + plate_path, _ = plate_2d + position = plate_path / "B" / "1" / "0" + config = _two_model_config(tmp_path / "two.yml", ["nuc", "mem"]) + + assert _run("--cluster", "debug", "-i", position, "-c", config).exit_code == 0 + + with open_ome_zarr(position, mode="r") as pos: + assert pos.label_names() == ["mem", "nuc"] + for name in ("nuc", "mem"): + assert np.all(pos.get_label(name)["0"][:] == 1) + assert len(fake_cellpose.evaluated) == 6 # 3 timepoints x 2 models + + +def test_segment_refuses_a_label_image_of_another_shape( + fake_cellpose, plate_2d, example_plate, tmp_path +): + """A model renamed from 3D to 2D would otherwise write into the old stack.""" + model = {"pretrained_model": "cpsam_v2", "channels": ["GFP"]} + three_d = tmp_path / "3d.yml" + three_d.write_text( + yaml.safe_dump({"models": {"nuc": {**model, "eval_args": {"do_3D": True}}}}) + ) + stack = example_plate[0] / "B" / "1" / "0" # Z=4 + assert _run("--init", "-i", stack, "-c", three_d).exit_code == 0 + with open_ome_zarr(stack, mode="r+") as pos: + pos.get_label("nuc")["0"].resize((3, 2, 5, 6)) # pretend the image changed + + result = _run("--init", "-i", stack, "-c", three_d) + + assert result.exit_code != 0 + assert "already has a label image 'nuc'" in str(result.exception) + + +def test_segment_resume_token_names_the_model_in_the_input_plate(tmp_path): + from biahub.segment import _resume_token + from biahub.settings import SegmentationSettings + from biahub.utils.config import yaml_to_model + + a = yaml_to_model( + _two_model_config(tmp_path / "a.yml", ["nuc", "mem"]), SegmentationSettings + ) + b = yaml_to_model( + _two_model_config(tmp_path / "b.yml", ["mem", "nuc"]), SegmentationSettings + ) + + # Each model has its own label image, so the model order does not matter there. + assert _resume_token(a, "nuc") == _resume_token(b, "nuc") + assert _resume_token(a, "nuc") != _resume_token(a, "mem") + + +# --------------------------------------------------------------------------- +# -o: a separate label plate, one channel per model +# --------------------------------------------------------------------------- + + def test_segment_init_rejects_an_unknown_model( fake_cellpose, plate_2d, segment_config, tmp_path ): diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index f975c096..c9a6d8a3 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -1,6 +1,8 @@ import os import tempfile +from pathlib import Path + import numpy as np import pandas as pd import pytest @@ -11,16 +13,22 @@ from biahub.cli.main import cli from biahub.settings import TrackingSettings -from biahub.track import _init_labels, label_chunks_and_shards, track, write_tracking_labels +from biahub.track import _init_labels, track, write_tracking_labels +from biahub.utils.ngff import create_label_images, label_chunks_and_shards -LABEL = "nuclei_prediction" +# The label image `biahub segment` would write (its model name) and the one track writes. +INPUT_LABEL = "nuclei" +LABEL = "nuclei_tracked" @pytest.fixture(scope="function") def example_tracking_plate(tmp_path): """ - Create a test plate with nuclei and membrane prediction channels for tracking + Create a test plate with nuclei and membrane prediction channels for tracking, and + the instance labels of the nuclei spots in ``labels/nuclei`` as segment would write. """ + from skimage.measure import label as connected_components + plate_path = tmp_path / "tracking_plate.zarr" position_list = ( @@ -64,6 +72,16 @@ def example_tracking_plate(tmp_path): position["0"] = data + plate_dataset.close() + positions = [plate_path / Path(*key) for key in position_list] + create_label_images(positions, [INPUT_LABEL], "biahub-segment", {"models": {}}) + for position_path in positions: + with open_ome_zarr(str(position_path), mode="r+") as position: + nuclei = np.asarray(position["0"][:, 0]) # (T, Z, Y, X) + labels = np.stack([connected_components(frame > 0.5) for frame in nuclei]) + position.get_label(INPUT_LABEL)["0"][:] = labels.astype(np.uint32) + plate_dataset = open_ome_zarr(plate_path, mode="r") + yield plate_path, plate_dataset @@ -117,14 +135,50 @@ def _assert_tracking_outputs(position_path): assert "working_dir" not in ultrack_config["data"] +_ULTRACK_CONFIG = { + "segmentation_config": { + "min_area": 100, + "max_area": 80000, + "n_workers": 1, + "min_frontier": 0.4, + "max_noise": 0.05, + }, + "linking_config": { + "n_workers": 1, + "max_distance": 15, + "distance_weight": -0.0001, + "max_neighbors": 3, + }, + "tracking_config": { + "n_threads": 1, + "disappear_weight": -0.0001, + "appear_weight": -0.001, + "division_weight": -0.0001, + }, +} + + def _make_tracking_config(plate_path, tmp_path): - """Create a minimal tracking config pointing at the test plate.""" + """A minimal config tracking the plate's `labels/nuclei` image (the default path).""" config_path = tmp_path / "track_config.yml" config = { "output_mode": "2D", "fov": "*/*/*", - "target_channel": "nuclei_prediction", - # The foreground/contour path needs no cellpose weights or GPU in the tests. + "labels": INPUT_LABEL, + "labels_sigma": 1.0, + "tracking_config": _ULTRACK_CONFIG, + } + config_path.write_text(yaml.dump(config)) + return config_path + + +def _make_foreground_contour_config(plate_path, tmp_path): + """A config on the legacy foreground+contour path, which processes image channels.""" + config_path = tmp_path / "track_config.yml" + config = { + "output_mode": "2D", + "fov": "*/*/*", + "target_channel": LABEL, "segmentation_method": "foreground_contour", "input_images": [ { @@ -156,27 +210,7 @@ def _make_tracking_config(plate_path, tmp_path): }, }, ], - "tracking_config": { - "segmentation_config": { - "min_area": 100, - "max_area": 80000, - "n_workers": 1, - "min_frontier": 0.4, - "max_noise": 0.05, - }, - "linking_config": { - "n_workers": 1, - "max_distance": 15, - "distance_weight": -0.0001, - "max_neighbors": 3, - }, - "tracking_config": { - "n_threads": 1, - "disappear_weight": -0.0001, - "appear_weight": -0.001, - "division_weight": -0.0001, - }, - }, + "tracking_config": _ULTRACK_CONFIG, } config_path.write_text(yaml.dump(config)) return config_path @@ -215,6 +249,36 @@ def test_track_cli_local( for position in ["A/1/0", "B/1/0", "B/2/0"]: _assert_tracking_outputs(plate_path / position) assert (output_path / position.replace("/", "_")).is_dir() # Ultrack database + with open_ome_zarr(str(plate_path / position), mode="r") as pos: + # The segmentation labels are left as they were, beside the tracked ones. + assert pos.label_names() == [INPUT_LABEL, LABEL] + assert pos.get_label(INPUT_LABEL)["0"][:].any() + + +def test_track_cli_foreground_contour_path(tmp_path, example_tracking_plate, monkeypatch): + """The legacy path processes image channels into a mask and contour map.""" + monkeypatch.setenv("ULTRACK_ARRAY_MODULE", "numpy") + plate_path, _ = example_tracking_plate + config_path = _make_foreground_contour_config(plate_path, tmp_path) + + result = CliRunner().invoke( + cli, + [ + "track", + "-i", + str(plate_path / "A" / "1" / "0"), + "-o", + str(tmp_path / "out"), + "-c", + str(config_path), + "--cluster", + "debug", + ], + ) + + assert result.exit_code == 0, result.output + assert "Tracking with foreground + contour" in result.output + _assert_tracking_outputs(plate_path / "A" / "1" / "0") def test_track_cli_with_blank_frames( @@ -312,7 +376,7 @@ def test_track_cli_init_only(tmp_path, example_tracking_plate): for position in ["A/1/0", "B/1/0", "B/2/0"]: with open_ome_zarr(str(plate_path / position), mode="r") as pos: - assert pos.label_names() == [LABEL] + assert pos.label_names() == [INPUT_LABEL, LABEL] assert pos.channel_names == ["nuclei_prediction", "membrane_prediction"] label = pos.get_label(LABEL) assert label["0"].shape == (5, 1, 64, 64) # the plate's Z @@ -322,12 +386,26 @@ def test_track_cli_init_only(tmp_path, example_tracking_plate): label_attrs = yaml.safe_load( (plate_path / position / "labels" / LABEL / "zarr.json").read_text() )["attributes"] - assert label_attrs["biahub-track"]["target_channel"] == LABEL + assert label_attrs["biahub-track"]["labels"] == INPUT_LABEL # The labels listing lives on the labels group (OME-NGFF 0.5). labels_attrs = yaml.safe_load( (plate_path / position / "labels" / "zarr.json").read_text() )["attributes"] - assert labels_attrs["ome"]["labels"] == [LABEL] + assert labels_attrs["ome"]["labels"] == [INPUT_LABEL, LABEL] + + +def test_track_init_requires_the_input_label_image(tmp_path, create_custom_plate): + """A plate segment has not run on fails at --init, naming `biahub segment`.""" + plate_path, _ = create_custom_plate(tmp_path, z_size=1) + config_path = _make_tracking_config(plate_path, tmp_path) + + result = CliRunner().invoke( + cli, ["track", "--init", "-i", str(plate_path / "A/1/0"), "-c", str(config_path)] + ) + + assert result.exit_code != 0 + assert "biahub segment" in str(result.exception) + assert not (plate_path / "A/1/0" / "labels").exists() def test_track_cli_debug_single_position(tmp_path, example_tracking_plate, monkeypatch): @@ -378,8 +456,9 @@ def run(): def _minimal_settings(**overrides): + """Legacy-path settings: they need no input label image, only the plate.""" base = { - "target_channel": "nuclei_prediction", + "target_channel": LABEL, "segmentation_method": "foreground_contour", "input_images": [{"path": None, "channels": {"nuclei_prediction": []}}], } @@ -387,15 +466,30 @@ def _minimal_settings(**overrides): return TrackingSettings(**base) -def test_cellpose_is_the_default_and_needs_its_config(): - base = { - "target_channel": "nuclei_prediction", - "input_images": [{"path": None, "channels": {"nuclei_prediction": []}}], - } - with pytest.raises(ValueError, match="cellpose_config is not set"): - TrackingSettings(**base) - settings = TrackingSettings(**base, cellpose_config={"input_channel": "nuclei_prediction"}) - assert settings.segmentation_method == "cellpose" +def test_labels_is_the_default_and_names_the_output(): + with pytest.raises(ValueError, match="`labels` is not set"): + TrackingSettings() + settings = TrackingSettings(labels="nuclei") + assert settings.segmentation_method == "labels" + assert settings.output_label == "nuclei_tracked" + assert settings.input_images == [] + with pytest.raises(ValueError, match="target_channel is for"): + TrackingSettings(labels="nuclei", target_channel="x") + + +def test_foreground_contour_needs_a_target_channel_and_input_images(): + with pytest.raises(ValueError, match="needs target_channel"): + TrackingSettings(segmentation_method="foreground_contour") + with pytest.raises(ValueError, match="needs input_images"): + TrackingSettings(segmentation_method="foreground_contour", target_channel="x") + assert _minimal_settings().output_label == LABEL + + +def test_cellpose_config_is_rejected_with_a_pointer_to_segment(): + with pytest.raises(ValueError, match="biahub segment.*labels"): + TrackingSettings(labels="nuclei", cellpose_config={"diameter": 80}) + with pytest.raises(ValueError, match="biahub segment"): + TrackingSettings(segmentation_method="cellpose") def test_init_labels_take_the_input_z(tmp_path, example_tracking_plate, create_custom_plate): @@ -513,7 +607,7 @@ def test_input_images_path_override(tmp_path, example_tracking_plate, monkeypatc plate_path, _ = example_tracking_plate output_path = tmp_path / "override_output" - config_path = _make_tracking_config(plate_path, tmp_path) + config_path = _make_foreground_contour_config(plate_path, tmp_path) runner = CliRunner() result = runner.invoke( @@ -537,53 +631,23 @@ def test_input_images_path_override(tmp_path, example_tracking_plate, monkeypatc _assert_tracking_outputs(plate_path / "A" / "1" / "0") -def test_track_init_rejects_an_unknown_cellpose_model( - tmp_path, example_tracking_plate, monkeypatch -): - """A bad model name must fail once at --init, not later in every worker.""" - import sys - import types - - class Model: - def __init__(self, gpu=False, pretrained_model="cpsam_v2", device=None): - # Cellpose 4 would silently substitute cpsam_v2 here. - self.pretrained_model = str(tmp_path / "cpsam_v2") - - models = types.ModuleType("cellpose.models") - models.CellposeModel = Model - models.MODEL_NAMES = ["cpsam_v2", "cpsam", "cpdino", "cpdino-vitb"] - models.get_user_models = lambda: [] - package = types.ModuleType("cellpose") - package.models = models - monkeypatch.setitem(sys.modules, "cellpose", package) - monkeypatch.setitem(sys.modules, "cellpose.models", models) +def test_load_labels_fills_blank_frames(tmp_path, example_tracking_plate): + from biahub.track import load_labels plate_path, _ = example_tracking_plate - config_path = _make_tracking_config(plate_path, tmp_path) - config = yaml.safe_load(config_path.read_text()) - config["segmentation_method"] = "cellpose" - config["cellpose_config"] = { - "pretrained_model": "nuclei", - "input_channel": "nuclei_prediction", - } - config_path.write_text(yaml.safe_dump(config)) + position_path = plate_path / "A" / "1" / "0" + with open_ome_zarr(str(position_path), mode="r+") as pos: + pos.get_label(INPUT_LABEL)["0"][2] = 0 # a dropped timepoint: nothing segmented + csv_path = tmp_path / "blank.csv" + pd.DataFrame({"FOV": ["A/1/0"], "t": ["[2]"]}).to_csv(csv_path, index=False) - result = CliRunner().invoke( - cli, - [ - "track", - "-i", - str(plate_path / "A" / "1" / "0"), - "-o", - str(tmp_path / "out"), - "-c", - str(config_path), - "--init", - ], - ) + labels = load_labels(position_path, INPUT_LABEL, "2D", csv_path) - assert result.exit_code != 0 - assert "Unknown cellpose model 'nuclei'" in str(result.exception) + result.output + assert labels.shape == (5, 64, 64) + assert np.array_equal(labels[2], labels[1]) # filled from the previous frame + assert not load_labels(position_path, INPUT_LABEL, "2D")[2].any() + with pytest.raises(ValueError, match="biahub segment"): + load_labels(position_path, "missing", "2D") def test_filter_short_siblings_merges_a_one_frame_daughter(): @@ -744,8 +808,8 @@ def test_track_debug_run_applies_the_sibling_filter(tmp_path, example_tracking_p assert "Sibling filter (min_length=3)" in result.output -def test_track_gives_local_jobs_a_gpu(tmp_path, example_tracking_plate, monkeypatch): - """Same as segment: submitit's local executor hides GPUs without gpus_per_node.""" +def test_track_jobs_ask_for_no_gpu(tmp_path, example_tracking_plate, monkeypatch): + """Ultrack is CPU work (segmentation is `biahub segment`'s): the cpu partition.""" import contextlib recorded = {} @@ -787,7 +851,8 @@ def submit(self, *args, **kwargs): assert result.exit_code == 0, result.output assert recorded["cluster"] == "local" - assert recorded.get("gpus_per_node") == 1 + assert "gpus_per_node" not in recorded + assert recorded["slurm_partition"] == "cpu" assert ( recorded.get("timeout_min") == recorded["slurm_time"] ) # local default kills at 2 min @@ -842,7 +907,6 @@ def submit(self, *args, **kwargs): assert result.exit_code == 0, result.output assert recorded["timeout_min"] == 5 - assert recorded["gpus_per_node"] == 1 def test_track_local_timeout_ignores_a_slurm_time_string( diff --git a/tests/test_example_settings.py b/tests/test_example_settings.py index 94feb59d..0036784d 100644 --- a/tests/test_example_settings.py +++ b/tests/test_example_settings.py @@ -213,10 +213,14 @@ def test_nextflow_a549_track_config_validates(): with open("nextflow/configs/a549/track.yml") as f: settings = TrackingSettings(**yaml.safe_load(f)) - # Track reads the projected plate: 2D, no processing step on the nuclei channel. + # Track reads the projected plate's `labels/nuclei` (segment.yml's model) and writes + # `labels/nuclei_tracked`; no image channel is processed. assert settings.output_mode == "2D" - assert settings.input_images[0].channels == {"nuclei_prediction": []} - assert settings.cellpose_config.input_channel == "nuclei_prediction" + assert settings.segmentation_method == "labels" + assert settings.labels == "nuclei" + assert settings.output_label == "nuclei_tracked" + assert settings.input_images == [] + assert settings.min_sibling_length == 3 def test_nextflow_a549_project_config_validates(): diff --git a/tests/test_segment_settings.py b/tests/test_segment_settings.py index 217df4ec..20a7166a 100644 --- a/tests/test_segment_settings.py +++ b/tests/test_segment_settings.py @@ -178,3 +178,25 @@ def test_a_config_needs_at_least_one_model(): def test_eval_args_cannot_set_the_image_or_self(key): with pytest.raises(ValueError, match=key): _settings(eval_args={key: 0}) + + +def test_nextflow_a549_segment_config_carries_the_validated_cellpose_settings(): + """The A549 segment config took over track.yml's validated cellpose values.""" + import yaml + + from biahub.settings import SegmentationSettings + + with open("nextflow/configs/a549/segment.yml") as f: + settings = SegmentationSettings(**yaml.safe_load(f)) + (name,) = settings.models + model = settings.models[name] + assert name == "nuclei" # track.yml's `labels: nuclei`; QC counts labels/nuclei_tracked + assert model.pretrained_model == "cpsam_v2" + assert model.channels == ["nuclei_prediction"] + assert model.eval_args == { + "diameter": 80, + "cellprob_threshold": 0.0, + "flow_threshold": 0.4, + "min_size": 2500, + } + assert not settings.is_3d diff --git a/uv.lock b/uv.lock index 4051eb08..db92c998 100644 --- a/uv.lock +++ b/uv.lock @@ -3,25 +3,25 @@ revision = 5 requires-python = ">=3.12" resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", - "python_full_version == '3.14.*' and sys_platform == 'win32'", "python_full_version >= '3.15' and sys_platform == 'emscripten'", - "python_full_version == '3.14.*' and sys_platform == 'emscripten'", "python_full_version >= '3.15' and sys_platform == 'darwin'", - "python_full_version == '3.14.*' and sys_platform == 'darwin'", "python_full_version >= '3.15' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.14.*' and sys_platform == 'win32'", + "python_full_version == '3.14.*' and sys_platform == 'emscripten'", + "python_full_version == '3.14.*' and sys_platform == 'darwin'", "python_full_version == '3.14.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version < '3.13' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version < '3.13' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform == 'darwin'", - "python_full_version < '3.13' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.13' and sys_platform == 'win32'", + "python_full_version < '3.13' and sys_platform == 'emscripten'", + "python_full_version < '3.13' and sys_platform == 'darwin'", "python_full_version < '3.13' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c9832a7675e0fb997a27db2796623d19b7dc9f1a" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=8c85656" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c9832a7675e0fb997a27db2796623d19b7dc9f1a" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=8c85656" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev18+c9832a7" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=c9832a7675e0fb997a27db2796623d19b7dc9f1a#c9832a7675e0fb997a27db2796623d19b7dc9f1a" } +version = "0.3.12.dev19+8c85656" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=8c85656#8c856562af67f8d8abaf8c41dd7d69cda0d6739a" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 266675865491b243a0e63f65a68f49367ff8ac63 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Tue, 20 Oct 2026 19:48:51 -0700 Subject: [PATCH 34/54] refactor(segment,track): scaffold labels with iohub's create_empty_labels; track on preempted The label images are created by iohub's new `create_empty_labels` (czbiohub-sf/iohub#489), which takes each label's geometry from its position and records the provenance on the label group, so biahub's wrapper around create_empty_plate goes; `label_layout` keeps only the DCA chunk sizing. Track goes back to the `preempted` partition (the Nextflow `cpu` label): ultrack cannot resume, but with segmentation out of track a position tracks in a minute or two, so a preemption costs that much and the retry covers it. The `cpu_reserved` label is dropped again. Co-Authored-By: Claude Fable 5.1 --- README.md | 2 +- biahub/segment.py | 14 ++++---- biahub/track.py | 23 +++++++----- biahub/utils/ngff.py | 61 ++++++++------------------------ nextflow/modules/tracking.nf | 7 ++-- nextflow/nextflow.config | 13 +++---- pyproject.toml | 8 ++--- tests/test_cli/test_track_cli.py | 9 +++-- uv.lock | 8 ++--- 9 files changed, 60 insertions(+), 85 deletions(-) diff --git a/README.md b/README.md index 3ce56ba6..15109804 100644 --- a/README.md +++ b/README.md @@ -224,7 +224,7 @@ With `--biahub_project`, each task runs as `uv run --project biahub ...`; Profiles select *where* work runs: - `-profile local` — everything on the current machine. Good for debugging and single-node runs. -- `-profile slurm` — per-position work goes to the `preempted` partition, virtual-stain prediction and segmentation to the `preempted` partition on GPU nodes (both resume where a reclaimed task stopped), tracking to the `cpu` partition (ultrack cannot resume), and lightweight init steps stay local. Preempted, timed-out, and OOM-killed tasks are retried automatically (up to 5 times, with escalating time/memory); genuine errors fail fast instead of burning retries. +- `-profile slurm` — per-position work goes to the `preempted` partition, virtual-stain prediction and segmentation to the `preempted` partition on GPU nodes (both resume where a reclaimed task stopped), tracking to the `preempted` partition as well (ultrack cannot resume, but a position tracks in minutes), and lightweight init steps stay local. Preempted, timed-out, and OOM-killed tasks are retried automatically (up to 5 times, with escalating time/memory); genuine errors fail fast instead of burning retries. Nextflow's own flags still apply — most usefully `-resume` to reuse completed tasks after a failure or a config tweak, and `-w /fast/scratch` to move the work directory off the output filesystem. diff --git a/biahub/segment.py b/biahub/segment.py index 1b2085c4..02949bb3 100644 --- a/biahub/segment.py +++ b/biahub/segment.py @@ -8,7 +8,7 @@ import submitit from iohub.ngff import open_ome_zarr -from iohub.ngff.utils import create_empty_plate, process_single_position +from iohub.ngff.utils import create_empty_labels, create_empty_plate, process_single_position from biahub.cli.monitor import monitor_jobs from biahub.cli.parsing import ( @@ -40,8 +40,8 @@ from biahub.utils.config import settings_fingerprint, yaml_to_model from biahub.utils.ngff import ( PROVENANCE_METADATA_KEYS, - create_label_images, get_output_paths, + label_layout, resolve_ome_zarr_version, ) @@ -219,11 +219,13 @@ def _init_label_images( "Rename the model or remove the old label image." ) - create_label_images( - input_position_dirpaths, + _, layout = label_layout(input_position_dirpaths[0]) + create_empty_labels( + Path(input_position_dirpaths[0]).parents[2], + [Path(p).parts[-3:] for p in input_position_dirpaths], label_names, - PROVENANCE_KEY, - settings.model_dump(mode="json"), + label_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, + **layout, ) return (T, C, Z, Y, X), models diff --git a/biahub/track.py b/biahub/track.py index 85f78d4e..c49a1b6d 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -15,6 +15,7 @@ import zarr from iohub import open_ome_zarr +from iohub.ngff.utils import create_empty_labels from numpy.typing import ArrayLike from biahub.cli.monitor import monitor_jobs @@ -38,7 +39,7 @@ local_executor_parameters, ) from biahub.utils.config import update_model, yaml_to_model -from biahub.utils.ngff import create_label_images, labels_exist +from biahub.utils.ngff import label_layout, labels_exist logger = logging.getLogger(__name__) @@ -852,11 +853,15 @@ def _init_labels( "`biahub segment` on the plate first (its model name is the label name)." ) - shape = create_label_images( - input_position_dirpaths, + shape, layout = label_layout(input_position_dirpaths[0]) + create_empty_labels( + Path(input_position_dirpaths[0]).parents[2], + [Path(p).parts[-3:] for p in input_position_dirpaths], [settings.output_label], - PROVENANCE_KEY, - settings.model_dump(mode="json"), + # The settings as written; the worker adds the resolved Ultrack config under + # the same key once it has tracked (track_one_position). + label_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, + **layout, ) click.echo( f"Created labels/{settings.output_label} in " @@ -974,10 +979,10 @@ def track( "slurm_cpus_per_task": num_cpus, "slurm_array_parallelism": 100, # process up to 100 positions at a time "slurm_time": time_minutes, - # Ultrack is CPU-bound (segmentation is `biahub segment`'s, on the GPU) and - # cannot resume, so the non-preemptible `cpu` partition rather than - # `preempted`. Override via --sbatch-filepath if a different partition is needed. - "slurm_partition": "cpu", + # Ultrack is CPU-bound (segmentation is `biahub segment`'s, on the GPU); it + # cannot resume, but a position tracks in minutes, so a preemption costs little. + # Override via --sbatch-filepath if a different partition is needed. + "slurm_partition": "preempted", "slurm_use_srun": False, } diff --git a/biahub/utils/ngff.py b/biahub/utils/ngff.py index 022d8b13..306860df 100644 --- a/biahub/utils/ngff.py +++ b/biahub/utils/ngff.py @@ -4,10 +4,8 @@ from typing import Any, Literal import numpy as np -import zarr from iohub.ngff import open_ome_zarr -from iohub.ngff.utils import create_empty_plate from tqdm import tqdm #: ``fnmatch`` patterns for the per-position zattrs keys that steps carry @@ -164,53 +162,22 @@ def labels_exist(position_dirpaths: list[Path], label_name: str) -> bool: ) -def create_label_images( - position_dirpaths: list[Path], - label_names: list[str], - provenance_key: str, - provenance: dict[str, Any], -) -> tuple[int, int, int, int, int]: - """Create empty ``labels/`` images in every position of one plate. +def label_layout( + position_dirpath: Path, +) -> tuple[tuple[int, int, int, int, int], dict[str, Any]]: + """Return the image's (T, C, Z, Y, X) shape and the label-image layout for its plate. - The labels are uint32 TZYX with the image's shape (its Z: one plane for a - projected plate, the stack for a 3D one), chunked per ``label_chunks_and_shards``, - and take the plate's OME-Zarr version. Existing label images are left untouched, - so ``--init`` and per-position workers can both call this. ``provenance`` is - merged into each label group's attributes under ``provenance_key``, keeping what - a worker recorded there before. - - Returns the image's (T, C, Z, Y, X) shape. + The layout is the ``label_chunks`` / ``label_shards_ratio`` keyword arguments of + iohub's ``create_empty_labels`` per :func:`label_chunks_and_shards` (no shards for + an OME-Zarr 0.4 plate), for a uint32 label with the image's TZYX shape. """ - position_dirpaths = [Path(p) for p in position_dirpaths] - with open_ome_zarr(str(position_dirpaths[0]), mode="r") as dataset: + with open_ome_zarr(str(position_dirpath), mode="r") as dataset: T, C, Z, Y, X = dataset.data.shape - scale = dataset.scale - channel_names = dataset.channel_names - dtype = dataset.data.dtype version = dataset.version - - plate = position_dirpaths[0].parents[2] - position_keys = [p.parts[-3:] for p in position_dirpaths] label_chunks, label_shards_ratio = label_chunks_and_shards((T, Z, Y, X)) - create_empty_plate( - store_path=plate, - position_keys=position_keys, - channel_names=channel_names, - shape=(T, C, Z, Y, X), - scale=scale, - version=version, - dtype=dtype, - label_names=label_names, - label_dtype=np.uint32, - label_chunks=label_chunks, - label_shards_ratio=label_shards_ratio if version == "0.5" else None, - ) - - for position_dirpath in position_dirpaths: - for name in label_names: - label_group = zarr.open_group(position_dirpath / "labels" / name, mode="r+") - label_group.attrs[provenance_key] = { - **label_group.attrs.get(provenance_key, {}), - **provenance, - } - return (T, C, Z, Y, X) + layout = { + "label_dtype": np.uint32, + "label_chunks": label_chunks, + "label_shards_ratio": label_shards_ratio if version == "0.5" else None, + } + return (T, C, Z, Y, X), layout diff --git a/nextflow/modules/tracking.nf b/nextflow/modules/tracking.nf index 49040adf..ce2ce37a 100644 --- a/nextflow/modules/tracking.nf +++ b/nextflow/modules/tracking.nf @@ -24,8 +24,9 @@ // label image in each position. // // Tracking is CPU work: ultrack re-segments from the labels' contours and solves -// an ILP; the GPU segmentation is `biahub segment`'s. It cannot resume, so it -// runs on the non-preemptible `cpu` partition. +// an ILP; the GPU segmentation is `biahub segment`'s. It cannot resume, but a +// position now tracks in a minute or two, so it runs on the `preempted` +// partition like the other per-position steps and a reclaimed task is retried. include { parse_resources; slurm_logs; slurm_log_dir } from './common' @@ -53,7 +54,7 @@ process init_track { process run_track { tag "${position}" - label 'cpu_reserved' + label 'cpu' clusterOptions { slurm_logs('track') } cpus { meta.cpus } memory { "${meta.mem_gb} GB" } diff --git a/nextflow/nextflow.config b/nextflow/nextflow.config index 015129c8..6120b5fb 100644 --- a/nextflow/nextflow.config +++ b/nextflow/nextflow.config @@ -172,7 +172,7 @@ process { errorStrategy = 'ignore' } - // NOTE: labels 'cpu', 'cpu_reserved', 'gpu' and 'gpu_preempted' intentionally have NO + // NOTE: labels 'cpu', 'gpu' and 'gpu_preempted' intentionally have NO // withLabel body here. Each module sets cpus/memory/time in its own process // body (dynamically from the init step's RESOURCES payload). A withLabel // selector in config overrides process-body directives, so defining @@ -218,10 +218,10 @@ profiles { // preemption costs at most the one it was predicting. Runs // on any of its GPU nodes, the A40 workstations included // (see run_virtual_stain). - // cpu_reserved -> `cpu` partition (not preempted), for tracking. Ultrack - // is CPU work (segmentation runs on the GPU in `biahub - // segment`) and cannot resume, so a preemption would restart - // the position from scratch. + // Tracking runs here too: ultrack cannot resume, but with + // segmentation moved to `biahub segment` a position tracks + // in a minute or two, so a preemption costs that much and + // the retry covers it. // gpu -> `gpu` partition (high PriorityTier, effectively not // preempted), for GPU work that cannot resume. Unused by the // mantis-v2 steps since segmentation gained --resume. @@ -269,9 +269,6 @@ profiles { withLabel: 'cpu' { queue = 'preempted' } - withLabel: 'cpu_reserved' { - queue = 'cpu' - } withLabel: 'gpu_preempted' { queue = 'preempted' } diff --git a/pyproject.toml b/pyproject.toml index ee5442aa..c0f506f0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,9 +141,9 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # Pin iohub to the czbiohub-sf/iohub#489 branch (on main at the #486 merge): # process_single_position(output_label=...), with which `biahub segment` writes its -# labels into the plate's label images with --resume. Main's #486 adds -# create_empty_plate(label_names=...), with which segment --init and track --init -# scaffold the empty label images (#360), and #485 puts the labels list on +# labels into the plate's label images with --resume, and create_empty_labels, with +# which segment --init and track --init scaffold those label images. Main's #486 adds +# create_empty_plate(label_names=...) (#360), and #485 puts the labels list on # /labels/zarr.json so ome-zarr-models and dca-validate can open a field that # has labels. Repoint at main once #489 merges, and at a PyPI release once one # contains them. @@ -154,7 +154,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "8c85656" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "16cdc4f" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index c9a6d8a3..66a6727f 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -10,11 +10,12 @@ from click.testing import CliRunner from iohub.ngff import open_ome_zarr +from iohub.ngff.utils import create_empty_labels from biahub.cli.main import cli from biahub.settings import TrackingSettings from biahub.track import _init_labels, track, write_tracking_labels -from biahub.utils.ngff import create_label_images, label_chunks_and_shards +from biahub.utils.ngff import label_chunks_and_shards # The label image `biahub segment` would write (its model name) and the one track writes. INPUT_LABEL = "nuclei" @@ -74,7 +75,9 @@ def example_tracking_plate(tmp_path): plate_dataset.close() positions = [plate_path / Path(*key) for key in position_list] - create_label_images(positions, [INPUT_LABEL], "biahub-segment", {"models": {}}) + create_empty_labels( + plate_path, list(position_list), [INPUT_LABEL], label_metadata={"biahub-segment": {}} + ) for position_path in positions: with open_ome_zarr(str(position_path), mode="r+") as position: nuclei = np.asarray(position["0"][:, 0]) # (T, Z, Y, X) @@ -852,7 +855,7 @@ def submit(self, *args, **kwargs): assert result.exit_code == 0, result.output assert recorded["cluster"] == "local" assert "gpus_per_node" not in recorded - assert recorded["slurm_partition"] == "cpu" + assert recorded["slurm_partition"] == "preempted" assert ( recorded.get("timeout_min") == recorded["slurm_time"] ) # local default kills at 2 min diff --git a/uv.lock b/uv.lock index db92c998..6f0c1fe6 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=8c85656" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=16cdc4f" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=8c85656" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=16cdc4f" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev19+8c85656" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=8c85656#8c856562af67f8d8abaf8c41dd7d69cda0d6739a" } +version = "0.3.12.dev20+16cdc4f" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=16cdc4f#16cdc4f6d42fba4476833f1501449c3548226d7d" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 8025a44997be50750dbb2fea2c6b580dd0e7c9de Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 21 Oct 2026 13:03:04 -0700 Subject: [PATCH 35/54] build(deps): iohub pin moves to the #461 branch (pyramid levels keep level-0 shard bytes) Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 18 ++++++++++-------- uv.lock | 8 ++++---- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bc068623..245f056e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,10 +99,10 @@ biahub = "biahub.cli.main:cli" index-strategy = "unsafe-best-match" # waveorder's own [tool.uv.sources] pins iohub to an older main commit # (b7bc648, the czbiohub-sf/iohub#476 merge) than the one biahub needs below -# (the czbiohub-sf/iohub#488 branch). uv refuses two git URLs for one package, +# (the czbiohub-sf/iohub#461 branch). uv refuses two git URLs for one package, # so force the whole graph onto biahub's rev. Drop together with the iohub pin. override-dependencies = [ - "iohub @ git+https://github.com/czbiohub-sf/iohub@d8e649e2fe9de3b126854430916149dde745c05c", + "iohub @ git+https://github.com/czbiohub-sf/iohub@4128ff6142bcc2e872da5e147cfcb22e88eed19c", ] [tool.uv.sources] @@ -139,11 +139,13 @@ imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipelin # royerlab/ultrack main once #290 merges, and at a release once one contains them. ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5f84ccda3bdf072f5289f939c" } -# Pin iohub to the czbiohub-sf/iohub#488 branch (feat/label-pyramid, on main at -# the #486 merge): PositionLabel.compute_pyramid, with which `biahub pyramid` -# gives the label images tracking writes the same levels as the image (#359), -# and a fix for the "mode" downsampler. Repoint at main once #488 merges, and at -# a PyPI release once one contains it. +# Pin iohub to the czbiohub-sf/iohub#461 branch (feat/auto-size-shards, stacked +# on #488 feat/label-pyramid, on main at the #486 merge). #488 adds +# PositionLabel.compute_pyramid, with which `biahub pyramid` gives the label +# images tracking writes the same levels as the image (#359), and fixes the +# "mode" downsampler; #461 sizes each pyramid level's shards to the level-0 +# shard's byte size instead of halving them. Repoint at main once both merge, +# and at a PyPI release once one contains them. # # Main's reasons: #486 adds create_empty_plate(label_names=...), with which # track --init scaffolds the empty label images (#360), and #485 puts the labels @@ -156,7 +158,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "d8e649e2fe9de3b126854430916149dde745c05c" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "4128ff6142bcc2e872da5e147cfcb22e88eed19c" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index 7279c2c7..0fe09e42 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=d8e649e2fe9de3b126854430916149dde745c05c" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=4128ff6142bcc2e872da5e147cfcb22e88eed19c" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=d8e649e2fe9de3b126854430916149dde745c05c" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=4128ff6142bcc2e872da5e147cfcb22e88eed19c" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev22+d8e649e" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=d8e649e2fe9de3b126854430916149dde745c05c#d8e649e2fe9de3b126854430916149dde745c05c" } +version = "0.3.12.dev24+4128ff6" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=4128ff6142bcc2e872da5e147cfcb22e88eed19c#4128ff6142bcc2e872da5e147cfcb22e88eed19c" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From f1df4fc1a050e43c66d0b654d5dfe01c39921383 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 21 Oct 2026 13:27:48 -0700 Subject: [PATCH 36/54] feat(settings)!: shards replaces shards_ratio in the concatenate and project configs `shards` takes iohub's forms (czbiohub-sf/iohub#461): a size target such as "5GB", "XY"/"XYZ", or an explicit TCZYX shape in array elements, so a config no longer carries chunk arithmetic that must be redone per output shape. A config that still says shards_ratio is refused with a message naming the new key. Concatenate reads the shard's T extent back from the plate its init created to size its workers. Track scaffolds its label images with create_empty_plate(label_shards=...), an explicit (16, 1, Y, X) shape capped at 2048 px, the same geometry as before. Templates: a549 and zebrafish concatenate `shards: 5GB` (A549 unchanged at 5 timepoints by the full volume; dynatrack goes from 10 to 9 timepoints, under the DCA 5 GB SHOULD); a549 project `shards: 200MB` (21 frames by the full frame instead of 16). Co-Authored-By: Claude Fable 5.1 --- .../references/caveats.md | 34 ++++---- biahub/concatenate.py | 9 ++- biahub/project.py | 4 +- biahub/settings.py | 78 +++++++++++++++---- biahub/track.py | 17 ++-- nextflow/configs/a549/concatenate.yml | 13 ++-- nextflow/configs/a549/project.yml | 11 +-- nextflow/configs/zebrafish/concatenate.yml | 13 ++-- pyproject.toml | 4 +- settings/example_concatenate_settings.yml | 11 +-- settings/example_project_settings.yml | 13 ++-- tests/test_cli/test_project_cli.py | 2 +- tests/test_cli/test_track_cli.py | 7 +- tests/test_concatenate.py | 36 +++++---- tests/test_example_settings.py | 4 +- tests/test_project_settings.py | 16 +++- tests/test_virtual_stain.py | 6 +- uv.lock | 8 +- 18 files changed, 182 insertions(+), 104 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index d8cb5636..175c5775 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -248,22 +248,22 @@ Assemble is one large SLURM job — the single longest step. `biahub concatenate --resume` (passed by the Nextflow task) makes a late kill recoverable per (t, c) unit. -## 7. Shard geometry and `shards_ratio` - -`shards_ratio` is a ratio over the **chunk** shape -(`shards = chunks * shards_ratio`, elementwise TCZYX), not an absolute shard -shape — changing T without recomputing ZYX silently changes the spatial extent. -Larger T shards mean a larger RAM request and a worse blast radius when one -tears; leave the T ratio alone without a measured reason. Channels cannot be -sharded — keep C at 1. The same two keys exist in `project.yml`: the A549 -template uses one-frame `256 x 256` chunks (the assembled and tracking grid) -and `shards_ratio: [16, 1, 1, 7, 5]`, 16-timepoint full-frame shards; its YX -entries are `ceil(Y/256), ceil(X/256)` and must be recomputed for a different -frame size, exactly as for `concatenate.yml`. Auto-sizing is -[iohub#458](https://github.com/czbiohub-sf/iohub/issues/458); most of this -section goes away when it lands. - -**A T ratio above 1 needs [iohub#460](https://github.com/czbiohub-sf/iohub/pull/460), +## 7. Shard geometry and `shards` + +`shards` in `concatenate.yml` and `project.yml` is a **size target** +(`shards: 5GB`, `shards: 200MB`): iohub builds the largest chunk-aligned shard +under it, filling X, Y and Z first and then timepoints, and never spans +channels ([iohub#461](https://github.com/czbiohub-sf/iohub/pull/461)). Nothing +is recomputed per dataset: the same `5GB` gives 5 timepoints by the full volume +on the A549 plate and 9 on the dynatrack one. `"XY"` / `"XYZ"` (one shard per +plane / volume) and an explicit TCZYX shape in array elements are also accepted. +The old `shards_ratio` (a multiplier over the chunk shape) is refused with a +message naming the new key. A larger target means more timepoints per shard, +hence a larger RAM request (concatenate reads the T extent back from the plate it +initialised to size its workers) and a worse blast radius when one tears; leave +the templates' values alone without a measured reason. + +**Shards spanning several timepoints need [iohub#460](https://github.com/czbiohub-sf/iohub/pull/460), which shipped in iohub 0.3.11** — the floor `pyproject.toml` already requires, so a synced checkout has it. @@ -271,7 +271,7 @@ Before that fix, `concatenate` aborted with a message that points nowhere near t cause — `numpy ... Unable to allocate 651. TiB for an array with shape (1048576, 86, 1664, 1193)`, whose first axis is not T, not a shard, not anything. It is **not** evidence that the shard buffer is too large, and -unsetting `shards_ratio` "fixes" it for the wrong reason. What happens: a +unsetting `shards` "fixes" it for the wrong reason. What happens: a shard-aligned batch writes a whole shard's worth of timepoints in one call, iohub drops the blank ones (`Skipping t=1, c=0 due to all zeros or nans` in the step's `.command.out`), and the surviving indices are left with a *gap* — which diff --git a/biahub/concatenate.py b/biahub/concatenate.py index df36e3ed..f1930f5b 100644 --- a/biahub/concatenate.py +++ b/biahub/concatenate.py @@ -651,7 +651,7 @@ def _resolve_concatenate_inputs( output_metadata = { "shape": (len(input_time_indices), len(all_channel_names)) + tuple(cropped_shape_zyx), "chunks": chunk_size, - "shards_ratio": settings.shards_ratio, + "shards": settings.shards, "version": resolve_ome_zarr_version( all_data_paths[0], settings.output_ome_zarr_version ), @@ -793,7 +793,12 @@ def concatenate( # 16-worker tasks took 7-26 min; 0.15 min/volume budgets 60 min here. T_out, C_out, _, _, _ = prep["output_metadata"]["shape"] _, _, Z, Y, X = prep["shape"] - batch_size = settings.shards_ratio[0] if settings.shards_ratio else 1 + # iohub writes one shard's timepoints per unit, so the shard's T extent is the + # batch a worker holds. Read it back from the plate init just created rather than + # predicting it: the config may give the shards as a size target. + with open_ome_zarr(output_dirpath, layout="hcs", mode="r") as output_plate: + _, first_position = next(output_plate.positions()) + batch_size = (first_position.data.shards or first_position.data.chunks)[0] time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( shape=(max(T_out // batch_size, 1), C_out, Z, Y, X), ram_multiplier=8 * batch_size, diff --git a/biahub/project.py b/biahub/project.py index be1989af..1830856b 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -219,7 +219,9 @@ def _init_output_plate( ), dtype=dtype, chunks=(1, *settings.chunks_czyx) if settings.chunks_czyx else None, - shards_ratio=tuple(settings.shards_ratio) if settings.shards_ratio else None, + shards=tuple(settings.shards) + if isinstance(settings.shards, list) + else settings.shards, metadata_sources=input_plate, metadata_keys=PROVENANCE_METADATA_KEYS, extra_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, diff --git a/biahub/settings.py b/biahub/settings.py index 7e6658fb..9b6264e1 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -8,6 +8,7 @@ import torch import yaml +from iohub.core.sharding import parse_shard_size from pydantic import ( BaseModel, BeforeValidator, @@ -27,6 +28,44 @@ class MyBaseModel(BaseModel): model_config = ConfigDict(extra="forbid") +# Shard geometry as iohub's `shards=` takes it: a size target such as "5GB" (the largest +# chunk-aligned shard under it, filled X, Y, Z then T; C is never grown), "XY"/"XYZ" (one +# shard per plane / per volume), or an explicit TCZYX shape in array elements. +ShardsSpec = str | list[int] | None + + +def _validate_shards(v: ShardsSpec) -> ShardsSpec: + if v is None: + return v + if isinstance(v, str): + if v.strip().upper() in ("XY", "XYZ"): + return v + try: + parse_shard_size(v) + except ValueError as err: + raise ValueError( + f"shards: {err}. Use a size such as '5GB', 'XY' / 'XYZ', or a TCZYX shape." + ) from err + return v + if len(v) != 5 or any(i < 1 for i in v): + raise ValueError( + "shards must be 5 positive integers (T, C, Z, Y, X), in array elements" + ) + if v[1] != 1: + raise ValueError("shards: iohub does not shard along C; the C entry must be 1") + return v + + +def _reject_shards_ratio(data): + if isinstance(data, dict) and "shards_ratio" in data: + raise ValueError( + "shards_ratio (a multiplier over the chunk shape) was replaced by shards, which " + "needs no chunk arithmetic: a size such as '5GB', 'XY' / 'XYZ', or a TCZYX shape " + "in array elements." + ) + return data + + class DetectPeaksSettings(MyBaseModel): threshold_abs: float = 110 nms_distance: int = 16 @@ -427,11 +466,11 @@ class ProjectionSettings(MyBaseModel): channels: dict[str, ChannelProjection] = Field(min_length=1) focus: ProjectionFocusSettings | None = None - # Output chunk shape (C, Z, Y, X) and shards ratio (T, C, Z, Y, X), as in concatenate: - # a shard is chunks * ratio, elementwise. None takes iohub's defaults (one-frame - # 256 x 256 chunks, one (t, c) frame per shard). Sharding needs OME-Zarr 0.5. + # Output chunk shape (C, Z, Y, X) and shard geometry (see ShardsSpec), as in + # concatenate. None takes iohub's defaults (one-frame 256 x 256 chunks, one (t, c) + # frame per shard). Sharding needs OME-Zarr 0.5. chunks_czyx: list[int] | None = None - shards_ratio: list[int] | None = None + shards: ShardsSpec = None # When None, preserve the OME-Zarr version of the input store. output_ome_zarr_version: Literal["0.4", "0.5"] | None = None @@ -442,16 +481,15 @@ def check_chunk_size(cls, v): raise ValueError("chunks_czyx must be 4 positive integers (C, Z, Y, X)") return v - @field_validator("shards_ratio") + @field_validator("shards") @classmethod - def check_shards_ratio(cls, v): - if v is not None and (len(v) != 5 or any(i < 1 for i in v)): - raise ValueError("shards_ratio must be 5 positive integers (T, C, Z, Y, X)") - if v is not None and v[1] != 1: - raise ValueError( - "shards_ratio: iohub does not shard along C; the C entry must be 1" - ) - return v + def check_shards(cls, v): + return _validate_shards(v) + + @model_validator(mode="before") + @classmethod + def reject_shards_ratio(cls, data): + return _reject_shards_ratio(data) @model_validator(mode="after") def check_focus(self): @@ -685,7 +723,9 @@ class ConcatenateSettings(MyBaseModel): Y_slice: list | list[list | Literal["all"]] | Literal["all"] = "all" Z_slice: list | list[list | Literal["all"]] | Literal["all"] = "all" chunks_czyx: Literal[None] | list[int] = None - shards_ratio: list[int] | None = None + # Shard geometry of the output (see ShardsSpec); None takes iohub's default, one + # shard per (t, c) volume on OME-Zarr 0.5. + shards: ShardsSpec = None ensure_unique_positions: bool | None = False # Concatenate is the migration path into v0.5 stores, so it defaults to # "0.5". Set to None to preserve the input store's OME-Zarr version, or @@ -809,6 +849,16 @@ def check_chunk_size(cls, v): raise ValueError("chunks_czyx must be a list of 4 integers (C, Z, Y, X)") return v + @field_validator("shards") + @classmethod + def check_shards(cls, v): + return _validate_shards(v) + + @model_validator(mode="before") + @classmethod + def reject_shards_ratio(cls, data): + return _reject_shards_ratio(data) + @model_validator(mode="after") def validate_slice_lengths(self): # Per-source lists (channel_names, X/Y/Z_slice) must be one entry per diff --git a/biahub/track.py b/biahub/track.py index e96afb83..068f635e 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -704,20 +704,17 @@ def cellpose_segmentation( # is the MUST) and its "2D+time" example; growing T keeps per-plane reads cheap. LABEL_CHUNKS = (16, 1, 128, 128) # Shards span one time chunk (16 timepoints, the DCA SHOULD) and the whole YX plane, -# capped at 16 chunks (2048 px, the DCA spatial-shard SHOULD). -MAX_LABEL_SHARD_CHUNKS = 16 +# capped at 2048 px (the DCA spatial-shard SHOULD). +MAX_LABEL_SHARD_PX = 2048 def label_chunks_and_shards( shape: tuple[int, int, int, int], ) -> tuple[tuple[int, ...], tuple[int, ...]]: - """Chunks and shards ratio for a (T, Z, Y, X) label array of ``shape``.""" + """Chunks and shard shape (elements, rounded up to chunks by iohub) for a (T, Z, Y, X) label.""" chunks = tuple(min(c, s) for c, s in zip(LABEL_CHUNKS, shape, strict=True)) - shards_ratio = (1, 1) + tuple( - min(-(-s // c), MAX_LABEL_SHARD_CHUNKS) - for c, s in zip(chunks[2:], shape[2:], strict=True) - ) - return chunks, shards_ratio + shards = chunks[:2] + tuple(min(s, MAX_LABEL_SHARD_PX) for s in shape[2:]) + return chunks, shards def write_tracking_labels(label_array, labels: ArrayLike, output_mode: str) -> None: @@ -940,7 +937,7 @@ def _init_labels( input_plate = Path(input_position_dirpaths[0]).parents[2] position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] - label_chunks, label_shards_ratio = label_chunks_and_shards((T, Z, Y, X)) + label_chunks, label_shards = label_chunks_and_shards((T, Z, Y, X)) create_empty_plate( store_path=input_plate, @@ -953,7 +950,7 @@ def _init_labels( label_names=[settings.target_channel], label_dtype=np.uint32, label_chunks=label_chunks, - label_shards_ratio=label_shards_ratio if version == "0.5" else None, + label_shards=label_shards if version == "0.5" else None, ) provenance = settings.model_dump(mode="json") diff --git a/nextflow/configs/a549/concatenate.yml b/nextflow/configs/a549/concatenate.yml index 5ab9438a..e0df504d 100644 --- a/nextflow/configs/a549/concatenate.yml +++ b/nextflow/configs/a549/concatenate.yml @@ -17,8 +17,11 @@ channel_names: all chunks_czyx: [1, 16, 256, 256] output_ome_zarr_version: "0.5" -# shards_ratio is a ratio over the CHUNK shape (shards = chunks * ratio, -# elementwise TCZYX), not an absolute shard shape. Leave unset unless you have a -# measured reason: raising the T entry multiplies the shard buffer and the RAM -# request, and one torn shard then loses every timepoint in it. -shards_ratio: [5, 1, 6, 7, 5] # ZYX entries recomputed for THIS output shape: ceil(Z/16), ceil(Y/256), ceil(X/256) +# shards is a size target: iohub builds the largest chunk-aligned shard under it, +# filling X, Y, Z first and then timepoints (C is never sharded). 5GB on this +# plate's 86 x 1664 x 1193 float32 volumes gives shards of 5 timepoints by the +# full volume (4.1 GiB), the same geometry a run before this key produced. Do +# not raise it without a measured reason: the T extent multiplies the shard +# buffer and the RAM request, and one torn shard loses every timepoint in it. +# Also accepts "XY" / "XYZ" (one shard per plane / volume) or a TCZYX shape. +shards: 5GB diff --git a/nextflow/configs/a549/project.yml b/nextflow/configs/a549/project.yml index 20c5af85..ad3f4f09 100644 --- a/nextflow/configs/a549/project.yml +++ b/nextflow/configs/a549/project.yml @@ -60,9 +60,10 @@ focus: # plate can follow it. Chunks are one frame, 256 x 256: the same YX grid as the # assembled and tracking stores, so a 256-aligned crop reads one chunk per channel in # every store (the standard's 16-timepoint chunk suits uint8 frames; for float32 it -# would decode 16 frames per single-frame patch). Shards span 16 timepoints and the -# full frame (the standard's floor for the T shard; ~150 MB each here, 30 files per -# position instead of 400). iohub writes one shard's timepoints per unit, so --resume -# is 16 frames coarse and a torn shard loses 16 frames of one channel. +# would decode 16 frames per single-frame patch). Shards are a size target: iohub +# fills the full frame and then as many timepoints as fit, 21 frames of 1792 x 1280 +# float32 (193 MB) here, above the standard's 16-timepoint floor and ~24 files per +# position instead of 400. iohub writes one shard's timepoints per unit, so --resume +# is one shard of frames coarse and a torn shard loses that many frames of one channel. chunks_czyx: [1, 1, 256, 256] -shards_ratio: [16, 1, 1, 7, 5] # YX entries recomputed for THIS output shape: ceil(Y/256), ceil(X/256) +shards: 200MB diff --git a/nextflow/configs/zebrafish/concatenate.yml b/nextflow/configs/zebrafish/concatenate.yml index 3124fe18..5f540c15 100644 --- a/nextflow/configs/zebrafish/concatenate.yml +++ b/nextflow/configs/zebrafish/concatenate.yml @@ -17,8 +17,11 @@ channel_names: all chunks_czyx: [1, 16, 256, 256] output_ome_zarr_version: "0.5" -# shards_ratio is a ratio over the CHUNK shape (shards = chunks * ratio, -# elementwise TCZYX), not an absolute shard shape. Leave unset unless you have a -# measured reason: raising the T entry multiplies the shard buffer and the RAM -# request, and one torn shard then loses every timepoint in it. -shards_ratio: [10, 1, 11, 3, 4] +# shards is a size target: iohub builds the largest chunk-aligned shard under it, +# filling X, Y, Z first and then timepoints (C is never sharded). 5GB on a +# 171 x 768 x 794 float32 volume gives shards of 9 timepoints by the full volume +# (4.98 GB; the old 10-timepoint shards were 5.2 GiB, above the DCA 5 GB SHOULD). +# Do not raise it without a measured reason: the T extent multiplies the shard +# buffer and the RAM request, and one torn shard loses every timepoint in it. +# Also accepts "XY" / "XYZ" (one shard per plane / volume) or a TCZYX shape. +shards: 5GB diff --git a/pyproject.toml b/pyproject.toml index 245f056e..dab3b491 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,7 +102,7 @@ index-strategy = "unsafe-best-match" # (the czbiohub-sf/iohub#461 branch). uv refuses two git URLs for one package, # so force the whole graph onto biahub's rev. Drop together with the iohub pin. override-dependencies = [ - "iohub @ git+https://github.com/czbiohub-sf/iohub@4128ff6142bcc2e872da5e147cfcb22e88eed19c", + "iohub @ git+https://github.com/czbiohub-sf/iohub@1d73e77135f61f9f45fd8a42d656f1550cc4ce32", ] [tool.uv.sources] @@ -158,7 +158,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "4128ff6142bcc2e872da5e147cfcb22e88eed19c" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "1d73e77135f61f9f45fd8a42d656f1550cc4ce32" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/settings/example_concatenate_settings.yml b/settings/example_concatenate_settings.yml index 0fbb61c7..a7cf3c31 100644 --- a/settings/example_concatenate_settings.yml +++ b/settings/example_concatenate_settings.yml @@ -66,11 +66,12 @@ Z_slice: "all" # - [1, 10, 100, 100]: Specify custom chunk sizes chunks_czyx: null -# Number of chunks in a shard for each dimension [T, C, Z, Y, X] -# Options: -# - null: No sharding -# - [1, 1, 4, 8, 8]: Specify custom shards ratio -shards_ratio: null +# Shard geometry (OME-Zarr 0.5 only). Options: +# - null: iohub's default, one shard per (t, c) volume +# - "5GB": the largest chunk-aligned shard under that size, filling X, Y, Z then T +# - "XY" / "XYZ": one shard per plane / per volume +# - [5, 1, 86, 1664, 1193]: an explicit TCZYX shape in array elements (C must be 1) +shards: null # Version of the OME-Zarr format to use for the output # Options: diff --git a/settings/example_project_settings.yml b/settings/example_project_settings.yml index c320ccf0..c8686db4 100644 --- a/settings/example_project_settings.yml +++ b/settings/example_project_settings.yml @@ -57,14 +57,15 @@ focus: na_det: 1.35 lambda_ill: 0.5 -# Output chunk shape [C, Z, Y, X] and shards ratio [T, C, Z, Y, X], as in concatenate: a -# shard is chunks * ratio, elementwise (the C entry must be 1). null takes iohub's -# defaults: one-frame 256 x 256 chunks, one (t, c) frame per shard. For a timelapse -# consider shards spanning 16 timepoints and the full frame, e.g. [16, 1, 1, 7, 5] for -# 1664 x 1193 frames (ceil(Y/256), ceil(X/256)); see nextflow/configs/a549/project.yml. +# Output chunk shape [C, Z, Y, X] and shard geometry, as in concatenate. null takes +# iohub's defaults: one-frame 256 x 256 chunks, one (t, c) frame per shard. shards is a +# size target such as "200MB" (the largest chunk-aligned shard under it, filling the +# frame and then timepoints; C is never sharded), "XY" / "XYZ" (one shard per plane / +# volume), or an explicit TCZYX shape in array elements. For a timelapse a size target +# gives shards of many timepoints by the full frame; see nextflow/configs/a549/project.yml. # Sharding needs OME-Zarr 0.5. chunks_czyx: null -shards_ratio: null +shards: null # When null, keep the OME-Zarr version of the input store. output_ome_zarr_version: null diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py index 2db37c06..694f052d 100644 --- a/tests/test_cli/test_project_cli.py +++ b/tests/test_cli/test_project_cli.py @@ -255,7 +255,7 @@ def test_project_writes_the_configured_chunks_and_shards(plate, tmp_path): {"GFP": {"method": "max"}, "Phase3D": {"method": "slice", "z": {"center": 1}}}, output_ome_zarr_version="0.5", chunks_czyx=[1, 1, 4, 4], - shards_ratio=[2, 1, 1, 2, 2], # (2, 1, 1, 8, 8): two timepoints per shard + shards=[2, 1, 1, 8, 8], # two timepoints per shard ) position = plate_path / "A" / "1" / "0" diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index c0655b92..d7343ce0 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -431,11 +431,12 @@ def test_init_labels_leaves_existing_labels(tmp_path, example_tracking_plate): def test_label_chunks_and_shards(): - # 1 MB uint32 chunks (DCA SHOULD); shards cover YX up to 16 chunks. + # 1 MB uint32 chunks (DCA SHOULD); shards are one time chunk by the YX plane, up to + # 2048 px (iohub rounds the elements up to whole chunks). chunks, shards = label_chunks_and_shards((67, 86, 1664, 1193)) assert chunks == (16, 1, 128, 128) - assert shards == (1, 1, 13, 10) - assert label_chunks_and_shards((67, 86, 4096, 300))[1] == (1, 1, 16, 3) + assert shards == (16, 1, 1664, 1193) + assert label_chunks_and_shards((67, 86, 4096, 300))[1] == (16, 1, 2048, 300) # Clamped to small arrays. assert label_chunks_and_shards((5, 1, 64, 64))[0] == (5, 1, 64, 64) diff --git a/tests/test_concatenate.py b/tests/test_concatenate.py index d1769006..980d9f75 100644 --- a/tests/test_concatenate.py +++ b/tests/test_concatenate.py @@ -490,11 +490,19 @@ def test_concatenate_with_cropping(create_custom_plate, tmp_path, sbatch_file): @pytest.mark.parametrize( - ["version", "shards_ratio_time"], - [["0.4", 1], ["0.5", None], ["0.5", 1], ["0.5", 2], ["0.5", 5]], + ["version", "shards"], + [ + ["0.4", None], + ["0.5", None], + ["0.5", [1, 1, 2, 8, 6]], + ["0.5", [2, 1, 2, 8, 6]], + ["0.5", [5, 1, 2, 8, 6]], # an explicit shape is kept as given, even past T = 3 + ["0.5", "XYZ"], + ["0.5", "1KiB"], # 48 B chunks: the whole volume fits, then timepoints + ], ) def test_concatenate_with_custom_chunks( - create_custom_plate, tmp_path, sbatch_file, version, shards_ratio_time + create_custom_plate, tmp_path, sbatch_file, version, shards ): """ Test concatenating with custom chunk sizes @@ -521,20 +529,20 @@ def test_concatenate_with_custom_chunks( # Define custom chunk sizes chunks = [1, 1, 2, 4, 3] # [C, Z, Y, X] - if version == "0.5": - if shards_ratio_time is None: - shards_ratio = None - else: - shards_ratio = [shards_ratio_time, 1, 1, 2, 2] - elif version == "0.4": - shards_ratio = None + expected_shards = { + None: None, + "XYZ": (1, 1, 4, 8, 6), + "1KiB": (2, 1, 4, 8, 6), # a 384 B volume: two timepoints fit under 1 KiB + }.get(shards if not isinstance(shards, list) else "list") + if isinstance(shards, list): + expected_shards = tuple(shards) # elements, already whole chunks, taken as given settings = ConcatenateSettings( concat_data_paths=[str(plate_1_path) + "/*/*/*", str(plate_2_path) + "/*/*/*"], channel_names=["all", "all"], time_indices="all", chunks_czyx=chunks[1:], - shards_ratio=shards_ratio, + shards=shards, output_ome_zarr_version=version, ) @@ -553,10 +561,8 @@ def test_concatenate_with_custom_chunks( output_plate = open_ome_zarr(output_path) for pos_name, pos in output_plate.positions(): assert pos.data.chunks == tuple(chunks) - if version == "0.5" and shards_ratio is not None: - assert pos.data.shards == tuple( - c * s for c, s in zip(chunks, shards_ratio, strict=True) - ) + if version == "0.5" and shards is not None: + assert pos.data.shards == expected_shards np.testing.assert_array_equal( pos.data.numpy(), np.concatenate( diff --git a/tests/test_example_settings.py b/tests/test_example_settings.py index 94feb59d..cfbacfb5 100644 --- a/tests/test_example_settings.py +++ b/tests/test_example_settings.py @@ -241,9 +241,9 @@ def test_nextflow_a549_project_config_validates(): "membrane_prediction", ] assert settings.channels["Phase3D"].method == "slice" - # One-frame 256 x 256 chunks (the assembled/track grid), 16-timepoint full-frame shards. + # One-frame 256 x 256 chunks (the assembled/track grid), ~200 MB full-frame shards. assert settings.chunks_czyx == [1, 1, 256, 256] - assert settings.shards_ratio == [16, 1, 1, 7, 5] + assert settings.shards == "200MB" assert settings.channels["GFP EX488 EM525-45"].method == "max" assert settings.focus.channel == "Phase3D" and settings.focus.per == "timepoint" diff --git a/tests/test_project_settings.py b/tests/test_project_settings.py index 0409c672..5fc8d46a 100644 --- a/tests/test_project_settings.py +++ b/tests/test_project_settings.py @@ -129,15 +129,23 @@ def test_chunks_and_shards_are_validated(): s = ProjectionSettings( channels={"a": {"method": "max"}}, chunks_czyx=[1, 1, 256, 256], - shards_ratio=[16, 1, 1, 7, 5], + shards=[16, 1, 1, 1664, 1193], ) - assert s.chunks_czyx == [1, 1, 256, 256] and s.shards_ratio == [16, 1, 1, 7, 5] + assert s.chunks_czyx == [1, 1, 256, 256] and s.shards == [16, 1, 1, 1664, 1193] + for form in ("200MB", "1.5GiB", "XY", "xyz"): + assert ( + ProjectionSettings(channels={"a": {"method": "max"}}, shards=form).shards == form + ) with pytest.raises(ValueError, match="4 positive integers"): ProjectionSettings(channels={"a": {"method": "max"}}, chunks_czyx=[1, 256, 256]) with pytest.raises(ValueError, match="5 positive integers"): - ProjectionSettings(channels={"a": {"method": "max"}}, shards_ratio=[16, 1, 7, 5]) + ProjectionSettings(channels={"a": {"method": "max"}}, shards=[16, 1, 7, 5]) with pytest.raises(ValueError, match="C entry must be 1"): - ProjectionSettings(channels={"a": {"method": "max"}}, shards_ratio=[16, 2, 1, 7, 5]) + ProjectionSettings(channels={"a": {"method": "max"}}, shards=[16, 2, 1, 7, 5]) + with pytest.raises(ValueError, match="Use a size such as"): + ProjectionSettings(channels={"a": {"method": "max"}}, shards="big") + with pytest.raises(ValueError, match="replaced by shards"): + ProjectionSettings(channels={"a": {"method": "max"}}, shards_ratio=[16, 1, 1, 7, 5]) def test_unknown_keys_are_refused(): diff --git a/tests/test_virtual_stain.py b/tests/test_virtual_stain.py index af6bb427..de112966 100644 --- a/tests/test_virtual_stain.py +++ b/tests/test_virtual_stain.py @@ -20,14 +20,14 @@ def implementation(request): return request.param -def _output_position(tmp_path, shards_ratio=None): +def _output_position(tmp_path, shards=None): store = tmp_path / "output.zarr" create_empty_plate( store_path=store, position_keys=[("A", "1", "0")], channel_names=["Nuclei", "Membrane"], shape=(T, C, Z, Y, X), - shards_ratio=shards_ratio, + shards=shards, version="0.5", ) return store / "A" / "1" / "0" @@ -91,7 +91,7 @@ def test_write_timepoints_recomputes(tmp_path, implementation, resume, token): def test_write_timepoints_batches_by_time_shard(tmp_path, implementation): """With two timepoints per shard, an interruption inside a shard recomputes it whole.""" - position = _output_position(tmp_path, shards_ratio=(2, 1, 1, 1, 1)) + position = _output_position(tmp_path, shards=(2, 1, Z, Y, X)) # two timepoints per shard with pytest.raises(RuntimeError): _write_timepoints( _predictor([], fail_at=1), position, T, resume=True, resume_token="a" diff --git a/uv.lock b/uv.lock index 0fe09e42..d88100dd 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=4128ff6142bcc2e872da5e147cfcb22e88eed19c" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=1d73e77135f61f9f45fd8a42d656f1550cc4ce32" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=4128ff6142bcc2e872da5e147cfcb22e88eed19c" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=1d73e77135f61f9f45fd8a42d656f1550cc4ce32" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev24+4128ff6" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=4128ff6142bcc2e872da5e147cfcb22e88eed19c#4128ff6142bcc2e872da5e147cfcb22e88eed19c" } +version = "0.3.12.dev25+1d73e77" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=1d73e77135f61f9f45fd8a42d656f1550cc4ce32#1d73e77135f61f9f45fd8a42d656f1550cc4ce32" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 53b8f3607970ec6984579e7a11b5804694f0bc97 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 21 Oct 2026 14:00:01 -0700 Subject: [PATCH 37/54] feat(track): the tracked label image takes the source label's dtype as well label_layout(source_label=...) copies dtype, chunks and shards; track writes its ids in that dtype and refuses, rather than wraps, ids it cannot hold. Co-Authored-By: Claude Fable 5.1 --- biahub/track.py | 13 +++++++++++-- biahub/utils/ngff.py | 16 ++++++++-------- tests/test_cli/test_track_cli.py | 12 ++++++++++-- 3 files changed, 29 insertions(+), 12 deletions(-) diff --git a/biahub/track.py b/biahub/track.py index ce363b03..ad582aad 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -638,9 +638,18 @@ def write_tracking_labels(label_array, labels: ArrayLike, output_mode: str) -> N The label image has the plate's Z: one plane for a projected plate, where 2D tracking writes its (T, Y, X) labels; the full stack for 3D tracking, whose - (T, Z, Y, X) labels fill it. + (T, Z, Y, X) labels fill it. The labels are stored in the label image's dtype + (the source label's, see ``label_layout``); track ids beyond it are refused + rather than wrapped. """ - labels = np.asarray(labels, dtype=np.uint32) + labels = np.asarray(labels) + dtype = np.dtype(label_array.dtype) + if labels.size and labels.max() > np.iinfo(dtype).max: + raise ValueError( + f"Track ids reach {int(labels.max())}, more than the label image's {dtype} holds " + f"({np.iinfo(dtype).max}); segment with a wider label dtype." + ) + labels = labels.astype(dtype, copy=False) if output_mode == "2D": # A single-plane series that kept its Z axis. if labels.ndim == 4 and labels.shape[1] == 1: diff --git a/biahub/utils/ngff.py b/biahub/utils/ngff.py index 526de3d4..88f5547e 100644 --- a/biahub/utils/ngff.py +++ b/biahub/utils/ngff.py @@ -163,21 +163,21 @@ def label_layout( """Return the image's (T, C, Z, Y, X) shape and the label-image layout for its plate. The layout is the ``label_dtype`` / ``label_chunks`` / ``label_shards`` keyword - arguments of iohub's ``create_empty_labels``. With ``source_label``, the chunks and - shards are those of the position's ``labels/`` level 0, so a label - derived from another (track's from segment's) reads like it; otherwise - :data:`LABEL_CHUNKS` and the :data:`LABEL_SHARDS` size target (no shards on an - OME-Zarr 0.4 plate). The dtype is uint32 either way. + arguments of iohub's ``create_empty_labels``. With ``source_label``, the dtype, + chunks and shards are those of the position's ``labels/`` level 0, so + a label derived from another (track's from segment's) is stored like it; otherwise + uint32, :data:`LABEL_CHUNKS` and the :data:`LABEL_SHARDS` size target (no shards on + an OME-Zarr 0.4 plate). """ with open_ome_zarr(str(position_dirpath), mode="r") as dataset: T, C, Z, Y, X = dataset.data.shape version = dataset.version if source_label is not None: source = dataset.get_label(source_label).data - chunks, shards = tuple(source.chunks), source.shards + dtype, chunks, shards = np.dtype(source.dtype), tuple(source.chunks), source.shards shards = tuple(shards) if shards is not None else None else: - chunks = label_chunks((T, Z, Y, X)) + dtype, chunks = np.dtype(np.uint32), label_chunks((T, Z, Y, X)) shards = LABEL_SHARDS if version == "0.5" else None - layout = {"label_dtype": np.uint32, "label_chunks": chunks, "label_shards": shards} + layout = {"label_dtype": dtype, "label_chunks": chunks, "label_shards": shards} return (T, C, Z, Y, X), layout diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 73b230dd..0ded6392 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -556,7 +556,7 @@ def test_label_chunks_and_layout(create_custom_plate, tmp_path): "label_chunks": (3, 1, 5, 6), "label_shards": LABEL_SHARDS, } - # A source label: its level-0 chunks and shards, whatever they are. + # A source label: its level-0 dtype, chunks and shards, whatever they are. with open_ome_zarr(str(position), mode="r+") as dataset: dataset.create_label_zeros( "cells", @@ -567,12 +567,20 @@ def test_label_chunks_and_layout(create_custom_plate, tmp_path): ) _, layout = label_layout(position, source_label="cells") assert layout == { - "label_dtype": np.uint32, + "label_dtype": np.uint16, "label_chunks": (1, 1, 5, 3), "label_shards": (3, 1, 5, 6), } +def test_write_tracking_labels_keeps_the_label_dtype_and_refuses_overflow(): + label_array = np.zeros((2, 1, 4, 4), dtype=np.uint16) + write_tracking_labels(label_array, np.full((2, 4, 4), 65535, dtype=np.int64), "2D") + assert label_array.dtype == np.uint16 and label_array.max() == 65535 + with pytest.raises(ValueError, match="more than the label image's uint16"): + write_tracking_labels(label_array, np.full((2, 4, 4), 65536, dtype=np.int64), "2D") + + def test_write_tracking_labels_2d_fills_the_single_plane(): label_array = np.zeros((2, 1, 4, 4), dtype=np.uint32) From ed7be79717a313edd558893461a8c95c36732450 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 21 Oct 2026 14:01:55 -0700 Subject: [PATCH 38/54] feat(segment,track)!: segment writes only into the plate; track's output_label `biahub segment` loses `-o` and the separate label plate: its labels are the plate's label images, one per model, and segmentation strategies accumulate there as differently named models. A name is only rewritten by the model that made it: `--init` refuses a label image whose recorded `biahub-segment` settings (or shape) differ from the current model's, pointing at a new model name. `output_ome_zarr_version` is rejected, since the labels take the plate's version. `biahub track` gains an optional `output_label`, another name for the tracked label image, for several tracking runs of one segmentation side by side; then the GEFF and CSV are `tracks_.geff` / `tracks__.csv`, while the default keeps `tracks.geff` / `tracks_.csv`, the names the downstream readers expect. Co-Authored-By: Claude Fable 5.1 --- .../references/caveats.md | 6 +- .../templates/check_qc_channels.py | 2 +- README.md | 2 +- biahub/segment.py | 267 +++++-------- biahub/settings.py | 48 ++- biahub/track.py | 26 +- settings/example_segmentation_settings.yml | 4 +- settings/example_track_settings.yml | 4 + tests/test_cellpose_inputs.py | 2 - tests/test_cli/test_segment_cli.py | 360 ++++++------------ tests/test_cli/test_track_cli.py | 56 ++- tests/test_example_settings.py | 2 +- 12 files changed, 332 insertions(+), 447 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 02a6b9ec..da45e448 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -156,8 +156,10 @@ with (cpsam_v2, diameter 80, min_size 2500) — they used to live in `track.yml` **`track.yml` — schema.** `TrackingSettings` forbids extra keys. Current fields: `output_mode`, `fov`, `blank_frames_path`, `segmentation_method` (`labels`, the default, or the legacy `foreground_contour`), `labels`, -`labels_sigma`, `min_sibling_length`, `tracking_config` (plus `target_channel` -and `input_images` on the legacy path). Track neither segments nor selects or +`labels_sigma`, `output_label` (optional; another name for the tracked label +image, which also renames the GEFF and CSV to `tracks_.*`), +`min_sibling_length`, `tracking_config` (plus `target_channel` and +`input_images` on the legacy path). Track neither segments nor selects or projects Z-planes: it reads the label image `biahub segment` wrote into the **projected** plate (`labels: `), writes the tracked labels beside it as `labels/_tracked` with `tracks.geff` and a tracks diff --git a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py index 73ece064..ac12f218 100644 --- a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py +++ b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py @@ -115,7 +115,7 @@ def tracked_label(track_cfg): t = yaml.safe_load(Path(track_cfg).read_text()) if t.get("segmentation_method", "labels") != "labels": return None - return t["labels"], f"{t['labels']}_tracked" + return t["labels"], t.get("output_label") or f"{t['labels']}_tracked" if __name__ == "__main__": diff --git a/README.md b/README.md index 15109804..058acc53 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,7 @@ Some steps include an `estimate-*` command that writes a YAML of parameters you | `concatenate` | Concatenate datasets channel-wise, with optional cropping | | | | | `project` | Project a 3D plate to 2D: max, mean, sum or one plane per channel, optionally around the in-focus plane | -| `segment` | Segment positions with cellpose 4; writes label images into the plate (or a separate label plate with `-o`) | +| `segment` | Segment positions with cellpose 4; writes one label image per model into the plate | | `virtual-stain` | Run VisCy/cytoland virtual staining | | `track` | Track the instance labels of a label image (ultrack); writes tracked labels and tracks.geff into the plate | | | | diff --git a/biahub/segment.py b/biahub/segment.py index 02949bb3..2308ab47 100644 --- a/biahub/segment.py +++ b/biahub/segment.py @@ -8,7 +8,7 @@ import submitit from iohub.ngff import open_ome_zarr -from iohub.ngff.utils import create_empty_labels, create_empty_plate, process_single_position +from iohub.ngff.utils import create_empty_labels, process_single_position from biahub.cli.monitor import monitor_jobs from biahub.cli.parsing import ( @@ -17,7 +17,6 @@ init_only, input_position_dirpaths, monitor, - output_dirpath, resume, sbatch_filepath, sbatch_to_submitit, @@ -38,17 +37,12 @@ local_executor_parameters, ) from biahub.utils.config import settings_fingerprint, yaml_to_model -from biahub.utils.ngff import ( - PROVENANCE_METADATA_KEYS, - get_output_paths, - label_layout, - resolve_ome_zarr_version, -) +from biahub.utils.ngff import label_layout if TYPE_CHECKING: import torch -# Provenance key: on each label image (no -o) or on each output position (-o). +# Provenance key on each label image: the segment settings as written. PROVENANCE_KEY = "biahub-segment" @@ -184,17 +178,48 @@ def _segment_czyx( return np.stack(czyx_segmentation, axis=0) +def _check_existing_label( + position, name: str, model_settings: dict, shape_tzyx: tuple +) -> None: + """Refuse to write model ``name`` into a label image made with other settings. + + Segmentation strategies accumulate in a plate as differently named label images, so + a label image is only ever rewritten by the model that made it. A different shape + (a model renamed from 3D to 2D, say) or different recorded model settings mean + another strategy owns the name. + """ + if name not in position.label_names(): + return + label = position.get_label(name) + existing_shape = tuple(label["0"].shape) + if existing_shape != shape_tzyx: + raise ValueError( + f"{position.zgroup.path} already has a label image {name!r} of shape " + f"{existing_shape}, but the image is {shape_tzyx} (T, Z, Y, X). Rename the " + "model or remove the old label image." + ) + recorded = dict(label.zattrs).get(PROVENANCE_KEY, {}).get("models", {}).get(name) + if recorded is not None and recorded != model_settings: + raise ValueError( + f"{position.zgroup.path} already has a label image {name!r} segmented with " + f"other settings:\n recorded: {recorded}\n this config: {model_settings}\n" + "Give this strategy its own model name to keep both, or remove the old " + "label image to replace it." + ) + + def _init_label_images( input_position_dirpaths: list[Path], settings: SegmentationSettings, ) -> tuple[tuple[int, int, int, int, int], list[ResolvedModel]]: - """Create the empty ``labels/`` images in the input positions (no ``-o``). + """Create the empty ``labels/`` images in the input positions. One label image per model, named after it, with the image's shape (its Z: one plane for a projected plate). Existing label images are left untouched, so both - --init and per-position workers can call this; one whose shape differs from the - image's (a model renamed from 3D to 2D, say) is refused. The settings are recorded - on each label image exactly as written (names, not indices). + --init and per-position workers can call this; one made by another strategy (other + shape or settings under the same name, see :func:`_check_existing_label`) is + refused. The settings are recorded on each label image exactly as written (names, + not indices). Returns the input (T, C, Z, Y, X) shape and the resolved models. """ @@ -205,155 +230,65 @@ def _init_label_images( models = resolve_models(settings, channel_names, scale, Z) label_names = [m.name for m in models] + recorded = settings.model_dump(mode="json") for position_dirpath in input_position_dirpaths: with open_ome_zarr(str(position_dirpath), mode="r") as position: for name in label_names: - if name not in position.label_names(): - continue - existing_shape = tuple(position.get_label(name)["0"].shape) - if existing_shape != (T, Z, Y, X): - raise ValueError( - f"{position_dirpath} already has a label image {name!r} of shape " - f"{existing_shape}, but the image is {(T, Z, Y, X)} (T, Z, Y, X). " - "Rename the model or remove the old label image." - ) + _check_existing_label(position, name, recorded["models"][name], (T, Z, Y, X)) _, layout = label_layout(input_position_dirpaths[0]) create_empty_labels( Path(input_position_dirpaths[0]).parents[2], [Path(p).parts[-3:] for p in input_position_dirpaths], label_names, - label_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, + label_metadata={PROVENANCE_KEY: recorded}, **layout, ) return (T, C, Z, Y, X), models -def _init_output_plate( - input_position_dirpaths: list[Path], - output_dirpath: Path, - settings: SegmentationSettings, -) -> tuple[tuple[int, int, int, int, int], list[ResolvedModel]]: - """Create (or extend) a separate label plate (``-o``) and resolve the models. - - create_empty_plate is idempotent: re-running with the same positions is a no-op and - new positions are appended, so both --init and per-position workers can call it. - The settings are recorded as provenance exactly as written (names, not indices). +def _resume_token(settings: SegmentationSettings, model: str) -> str: + """Fingerprint the settings and the model whose label image is being written. - Returns the input (T, C, Z, Y, X) shape and the resolved models. + Each model has its own label image, so the model order does not matter; a changed + setting does, and makes a resumed run recompute. """ - with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as input_dataset: - T, C, Z, Y, X = input_dataset.data.shape - scale = input_dataset.scale - channel_names = input_dataset.channel_names - - models = resolve_models(settings, channel_names, scale, Z) - - # Label channel i belongs to model i. On an existing output, create_empty_plate keeps - # the stored channels and array shape, so a reordered config would write into swapped - # channels. The labels have the input's Z (1 for a 2D plate). - label_channels = [m.name + "_labels" for m in models] - output_shape = (T, len(models), Z, Y, X) - position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] - for key in position_keys: - output_position = Path(output_dirpath, *key) - if not output_position.exists(): - continue - with open_ome_zarr(str(output_position), mode="r") as output_dataset: - existing = output_dataset.channel_names - existing_shape = tuple(output_dataset.data.shape) - if existing != label_channels or existing_shape != output_shape: - raise ValueError( - f"{output_position} has label channels {existing} and shape " - f"{existing_shape}, but this config writes {label_channels} with shape " - f"{output_shape}. Keep the models and their settings, or segment into " - "a new output." - ) - - input_plate = Path(input_position_dirpaths[0]).parents[2] - create_empty_plate( - store_path=output_dirpath, - position_keys=position_keys, - channel_names=label_channels, - shape=output_shape, - dtype=np.uint32, - scale=scale, - version=resolve_ome_zarr_version( - input_position_dirpaths[0], settings.output_ome_zarr_version - ), - metadata_sources=input_plate, - metadata_keys=PROVENANCE_METADATA_KEYS, - extra_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, - ) - return (T, C, Z, Y, X), models - - -def _resume_token(settings: SegmentationSettings, model: str | None = None) -> str: - """Fingerprint the settings, including the order of the models. - - settings_fingerprint sorts keys, but in a separate label plate (``-o``) the model - order decides which label channel each model writes, so finished work from another - order must not be skipped. In the input plate each model has its own label image, - so the token names the model instead and the order does not matter. - """ - if model is not None: - return settings_fingerprint( - {"settings": settings.model_dump(mode="json"), "model": model} - ) - return settings_fingerprint( - {"settings": settings.model_dump(mode="json"), "model_order": list(settings.models)} - ) + return settings_fingerprint({"settings": settings.model_dump(mode="json"), "model": model}) def _segment_position( input_position_path: Path, - output_position_path: Path | None, models: list[ResolvedModel], n_channels: int, settings: SegmentationSettings, resume: bool, ) -> None: - """Segment one position: into its own label images, or into the ``-o`` plate. + """Segment one position into its own label images, one per model. - Without ``-o`` each model gets one ``process_single_position`` call writing the - label image named after it (iohub's ``output_label``); the cellpose model cache - makes the repeated calls free. With ``-o`` one call writes all the models as - channels of the output position. + Each model gets one ``process_single_position`` call writing the label image named + after it (iohub's ``output_label``); the cellpose model cache makes the repeated + calls free. """ - if output_position_path is None: - for model in models: - process_single_position( - _segment_czyx, - input_position_path, - input_position_path, - input_channel_indices=[list(range(n_channels))], - output_label=model.name, - # One process per position: timepoints share the loaded model and - # the GPU, which is the bottleneck. - num_workers=1, - resume=resume, - resume_token=_resume_token(settings, model.name), - models=[model], - ) - return - process_single_position( - _segment_czyx, - input_position_path, - output_position_path, - input_channel_indices=[list(range(n_channels))], - output_channel_indices=[list(range(len(models)))], - num_workers=1, - resume=resume, - resume_token=_resume_token(settings), - models=models, - ) + for model in models: + process_single_position( + _segment_czyx, + input_position_path, + input_position_path, + input_channel_indices=[list(range(n_channels))], + output_label=model.name, + # One process per position: timepoints share the loaded model and + # the GPU, which is the bottleneck. + num_workers=1, + resume=resume, + resume_token=_resume_token(settings, model.name), + models=[model], + ) def segment( input_position_dirpaths: list[Path], config_filepath: Path, - output_dirpath: Path | None = None, sbatch_filepath: str | None = None, cluster: str = "slurm", monitor: bool = True, @@ -366,14 +301,10 @@ def segment( ---------- input_position_dirpaths : list[Path] Paths to input positions, for example: "input.zarr/0/0/0", "input.zarr/0/0/[0-9]", - or "input.zarr/*/*/*". + or "input.zarr/*/*/*". The labels are written into these positions as OME-NGFF + label images, ``labels/``, which is what `biahub track` reads. config_filepath : Path Path to the segmentation YAML configuration file. - output_dirpath : Path, optional - Path to a separate "output.zarr" label plate, with one channel per model - (``_labels``). If None (the default), the labels are written into the - ``-i`` positions as OME-NGFF label images, ``labels/``, which is what - `biahub track` reads, and the SLURM logs go to ``./slurm_output``. sbatch_filepath : str, optional SBATCH filepath that contains slurm parameters to overwrite defaults. cluster : str, optional @@ -382,17 +313,19 @@ def segment( monitor : bool, optional Monitor of submitted SLURM jobs. init_only : bool, optional - Only initialize the output store (and check/warm the models) and exit. + Only create the empty label images (and check/warm the models) and exit. resume : bool, optional - Skip the (time, channel) units a previous attempt already finished; see + Skip the timepoints a previous attempt already finished; see ``iohub.ngff.utils.process_single_position``. + + Notes + ----- + Segmentation strategies accumulate in a plate as differently named label images + (the model names); a name is only rewritten by the model that made it (see + :func:`_check_existing_label`). The SLURM logs go to ``./slurm_output``. """ - if output_dirpath is not None: - output_dirpath = Path(output_dirpath) - slurm_out_path = output_dirpath.parent / "slurm_output" - else: - # Not $TMPDIR: submitit's folder must be visible from the compute nodes. - slurm_out_path = Path("slurm_output").absolute() + # Not $TMPDIR: submitit's folder must be visible from the compute nodes. + slurm_out_path = Path("slurm_output").absolute() resolved_cluster = get_submitit_cluster(cluster=cluster) if not init_only and resolved_cluster == "debug": @@ -407,12 +340,7 @@ def segment( # Fail on a bad model name before scaffolding anything. for name in dict.fromkeys(m.pretrained_model for m in settings.models.values()): check_cellpose_model_name(name) - if output_dirpath is None: - (T, C, Z, Y, X), models = _init_label_images(input_position_dirpaths, settings) - else: - (T, C, Z, Y, X), models = _init_output_plate( - input_position_dirpaths, output_dirpath, settings - ) + (T, C, Z, Y, X), models = _init_label_images(input_position_dirpaths, settings) # Calibrated on a real 2D run (2026_04_28 SEC61B, 67 x 1664 x 1193, 6 input channels, # cpsam_v2 on one GPU): ~5 s per frame (6-10 min per position, incl. weight staging) and @@ -437,23 +365,13 @@ def segment( # than N workers racing for it. Workers must not (staging needs cellpose unimported). for name in dict.fromkeys(m.pretrained_model for m in models): warm_cellpose_weights(name) - if output_dirpath is None: - click.echo( - f"Created labels/{{{', '.join(m.name for m in models)}}} in " - f"{Path(input_position_dirpaths[0]).parents[2]} " - f"({len(input_position_dirpaths)} positions)" - ) - else: - click.echo( - f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)" - ) + click.echo( + f"Created labels/{{{', '.join(m.name for m in models)}}} in " + f"{Path(input_position_dirpaths[0]).parents[2]} " + f"({len(input_position_dirpaths)} positions)" + ) return - if output_dirpath is None: - output_position_paths = [None] * len(input_position_dirpaths) - else: - output_position_paths = get_output_paths(input_position_dirpaths, output_dirpath) - slurm_args = { "slurm_job_name": "segment", "slurm_mem": f"{mem_gb}G", @@ -480,14 +398,11 @@ def segment( click.echo("Submitting jobs...") jobs = [] with submitit.helpers.clean_env(), executor.batch(): - for input_position_path, output_position_path in zip( - input_position_dirpaths, output_position_paths, strict=True - ): + for input_position_path in input_position_dirpaths: jobs.append( executor.submit( _segment_position, input_position_path, - output_position_path, models=models, n_channels=C, settings=settings, @@ -516,13 +431,6 @@ def segment( @click.command("segment") @input_position_dirpaths() @config_filepath() -@output_dirpath( - required=False, - help="Separate label plate to write, one channel per model (_labels). " - "If omitted, the labels are written into the -i positions as label images, " - "labels/, which `biahub track` reads, and the SLURM logs go to " - "./slurm_output.", -) @sbatch_filepath() @cluster() @monitor() @@ -531,7 +439,6 @@ def segment( def segment_cli( input_position_dirpaths: list[Path], config_filepath: Path, - output_dirpath: Path | None = None, sbatch_filepath: str | None = None, cluster: str = "slurm", monitor: bool = False, @@ -541,7 +448,8 @@ def segment_cli( """Segment positions with cellpose 4 models configured in a YAML file. Writes one label image per model into each -i position (labels/), where - `biahub track` reads them; -o writes a separate label plate instead. + `biahub track` reads them. Strategies accumulate as differently named models; a + name is only rewritten by the model that made it. \b SLURM fan-out of positions across a whole plate: @@ -554,15 +462,10 @@ def segment_cli( \b In-process run of a single position (Nextflow per-position worker): >>> biahub segment --cluster debug -i ./input.zarr/B/3/000000 -c ./segment.yml --resume - - \b - Into a separate label plate: - >>> biahub segment -i ./input.zarr/*/*/* -c ./segment.yml -o ./segment.zarr """ # noqa: D301 segment( input_position_dirpaths=input_position_dirpaths, config_filepath=config_filepath, - output_dirpath=output_dirpath, sbatch_filepath=sbatch_filepath, cluster=cluster, monitor=monitor, diff --git a/biahub/settings.py b/biahub/settings.py index b1a17dd6..a9dba1e6 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -87,8 +87,13 @@ class TrackingSettings(MyBaseModel): # Smoothing (pixels) ultrack applies to the labels' contours before it builds its # own segmentation hypotheses from them (ultrack `labels_to_contours`). labels_sigma: float = 5.0 + # segmentation_method: labels -- the name of the tracked label image, when the + # default `_tracked` is not wanted, e.g. to keep several tracking runs of + # one segmentation side by side. Then the tracks go to `tracks_.geff` + # and `tracks__.csv` instead of `tracks.geff` / `tracks_.csv`. + output_label: str | None = None # segmentation_method: foreground_contour only -- the name of the label image the - # tracks are written to. With `labels` the name is derived (_tracked). + # tracks are written to. target_channel: str | None = None # Image channels to load and process per FOV; needed by foreground_contour (its # processing functions compute the mask and contour) and optional otherwise. @@ -144,7 +149,12 @@ def check_segmentation_method(self): raise ValueError( "target_channel is for segmentation_method: foreground_contour; with " f"labels the tracks are written to labels/{self.labels}" - f"{TRACKED_LABEL_SUFFIX}." + f"{TRACKED_LABEL_SUFFIX}, or to labels/ if set." + ) + if self.output_label == self.labels: + raise ValueError( + f"output_label {self.labels!r} is the label image being tracked; name " + "the tracked one differently." ) else: if self.target_channel is None: @@ -152,6 +162,11 @@ def check_segmentation_method(self): "segmentation_method: foreground_contour needs target_channel, the name " "of the label image the tracks are written to." ) + if self.output_label is not None: + raise ValueError( + "output_label is for segmentation_method: labels; foreground_contour " + "names its label image with target_channel." + ) if not self.input_images: raise ValueError( "segmentation_method: foreground_contour needs input_images whose " @@ -160,12 +175,21 @@ def check_segmentation_method(self): return self @property - def output_label(self) -> str: - """Name of the label image tracking writes, ``labels/``.""" + def tracked_label(self) -> str: + """Name of the label image tracking writes, ``labels/``.""" if self.segmentation_method == "labels": - return f"{self.labels}{TRACKED_LABEL_SUFFIX}" + return self.output_label or f"{self.labels}{TRACKED_LABEL_SUFFIX}" return self.target_channel + @property + def tracks_stem(self) -> str: + """Stem of the tracks files: ``tracks`` (``tracks.geff``, ``tracks_.csv``). + + With an explicit ``output_label`` the stem carries it, so several tracking runs + of one plate keep their GEFFs and CSVs apart: ``tracks_``. + """ + return "tracks" if self.output_label is None else f"tracks_{self.output_label}" + class EdgeGraphSettings(BaseModel): method: Literal["knn", "radius", "full"] = "knn" @@ -1000,11 +1024,21 @@ def is_3d(self) -> bool: class SegmentationSettings(BaseModel): + # One label image per model, `labels/`, written into the plate segment reads. + # The name is the strategy's: strategies accumulate as differently named models. models: dict[str, SegmentationModel] = Field(min_length=1) - # When None, preserve the OME-Zarr version of the input store. - output_ome_zarr_version: Literal["0.4", "0.5"] | None = None model_config = {"extra": "forbid", "protected_namespaces": ()} + @model_validator(mode="before") + @classmethod + def reject_output_version(cls, data): + if isinstance(data, dict) and "output_ome_zarr_version" in data: + raise ValueError( + "output_ome_zarr_version was removed: segment writes its label images into " + "the plate it reads, so they take that plate's OME-Zarr version." + ) + return data + @model_validator(mode="after") def check_models_share_dimensionality(self): # One input plate is either 2D (Z=1) or 3D (Z>1), so it suits only one kind. diff --git a/biahub/track.py b/biahub/track.py index c49a1b6d..bf2dd12c 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -701,6 +701,7 @@ def track_one_position( labels_sigma: float = 5.0, output_mode: str = "2D", min_sibling_length: int | None = None, + tracks_stem: str = "tracks", ) -> None: """ Run tracking on a single field of view. @@ -751,14 +752,18 @@ def track_one_position( If set, remove daughter tracks of this many frames or fewer before writing (see :func:`filter_short_siblings`); the labels, the GEFF and the CSV all come from the filtered tracks. Default None keeps every division. + tracks_stem : str, optional + Stem of the GEFF and CSV file names. Default ``"tracks"``. Returns ------- None Outputs are saved into ``position_dirpath``: - ``labels/``: the tracked labels, painted with ``track_id`` - - ``tracks.geff``: the track graph, linked to the labels by ``seg_id`` - - ``tracks___.csv``: the same tracks as a CSV, kept for now + - ``.geff``: the track graph, linked to the labels by ``seg_id`` + - ``___.csv``: the same tracks as a CSV, kept for now + ``tracks_stem`` is ``tracks`` unless the config names the tracked label itself + (``output_label``), when it is ``tracks_``. Notes ----- @@ -797,7 +802,7 @@ def track_one_position( tracking_labels, tracks_df, _ = run_ultrack( tracking_config=tracking_config, database_path=database_path, - geff_path=position_dirpath / "tracks.geff", + geff_path=position_dirpath / f"{tracks_stem}.geff", geff_kwargs=geff_kwargs, min_sibling_length=min_sibling_length, scale=scale, @@ -805,8 +810,8 @@ def track_one_position( **detection_kwargs, ) - # Kept alongside tracks.geff for downstream readers that still expect a CSV. - tracks_df.to_csv(position_dirpath / f"tracks_{fov}.csv", index=False) + # Kept alongside the GEFF for downstream readers that still expect a CSV. + tracks_df.to_csv(position_dirpath / f"{tracks_stem}_{fov}.csv", index=False) with open_ome_zarr(position_dirpath, mode="r+") as position: label_image = position.get_label(label_name) @@ -826,7 +831,7 @@ def _init_labels( input_position_dirpaths: list[Path], settings: TrackingSettings, ) -> tuple[int, int, int, int, int]: - """Create the empty ``labels/`` image in every input position. + """Create the empty ``labels/`` image in every input position. The input's dimensionality must match ``settings.output_mode`` (see :func:`check_input_dimensionality`), and with ``settings.labels`` every position @@ -857,14 +862,14 @@ def _init_labels( create_empty_labels( Path(input_position_dirpaths[0]).parents[2], [Path(p).parts[-3:] for p in input_position_dirpaths], - [settings.output_label], + [settings.tracked_label], # The settings as written; the worker adds the resolved Ultrack config under # the same key once it has tracked (track_one_position). label_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, **layout, ) click.echo( - f"Created labels/{settings.output_label} in " + f"Created labels/{settings.tracked_label} in " f"{Path(input_position_dirpaths[0]).parents[2]} " f"({len(input_position_dirpaths)} positions)" ) @@ -943,7 +948,7 @@ def track( break # Workers skip the scaffold --init already made: they all share the -i plate. - if init_only or not labels_exist(input_position_dirpaths, settings.output_label): + if init_only or not labels_exist(input_position_dirpaths, settings.tracked_label): _init_labels(input_position_dirpaths, settings) with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as dataset: T, C, Z, Y, X = dataset.data.shape @@ -1010,7 +1015,7 @@ def track( position_dirpath=Path(position_dirpath), output_dirpath=output_dirpath, tracking_config=tracking_cfg, - label_name=settings.output_label, + label_name=settings.tracked_label, input_images=settings.input_images, blank_frames_path=settings.blank_frames_path, scale=track_scale, @@ -1018,6 +1023,7 @@ def track( labels_sigma=settings.labels_sigma, output_mode=settings.output_mode, min_sibling_length=settings.min_sibling_length, + tracks_stem=settings.tracks_stem, ) jobs.append(job) diff --git a/settings/example_segmentation_settings.yml b/settings/example_segmentation_settings.yml index 641c3801..a53915ea 100644 --- a/settings/example_segmentation_settings.yml +++ b/settings/example_segmentation_settings.yml @@ -1,7 +1,7 @@ # Segmentation settings for cellpose 4 (`biahub segment`). # Model names become the label image names in each -i position (e.g. membrane -> -# labels/membrane), which `biahub track` reads; with -o they become the channels of a -# separate label plate (membrane -> membrane_labels). +# labels/membrane), which `biahub track` reads. Strategies accumulate as differently +# named models; a name is only rewritten by the model that made it. # A model segments in 2D unless eval_args.do_3D (or a stitch_threshold) is set. Segment # does not project: 2D models need a Z=1 input plate (make one from a stack with # `biahub project`), 3D models a stack; a mismatch is an error. All models in one config diff --git a/settings/example_track_settings.yml b/settings/example_track_settings.yml index 8673d9f2..a6f018d6 100644 --- a/settings/example_track_settings.yml +++ b/settings/example_track_settings.yml @@ -19,6 +19,10 @@ segmentation_method: labels # config's model name. The tracked labels are written next to it as # `labels/_tracked`, with tracks.geff and a tracks CSV. labels: nuclei +# Optional: another name for the tracked label image, to keep several tracking runs of +# one segmentation side by side. The tracks then go to tracks_.geff and +# tracks__.csv rather than tracks.geff / tracks_.csv. +# output_label: nuclei_tracked_sigma3 # Smoothing (pixels) ultrack applies to the labels' contours before it builds its own # segmentation hypotheses from them. labels_sigma: 5.0 diff --git a/tests/test_cellpose_inputs.py b/tests/test_cellpose_inputs.py index 5d090b49..50e885c8 100644 --- a/tests/test_cellpose_inputs.py +++ b/tests/test_cellpose_inputs.py @@ -112,8 +112,6 @@ def test_segment_cli_rejects_a_stack_for_a_2d_model(fake_cellpose, example_plate "segment", "-i", str(plate_path / "A" / "1" / "0"), - "-o", - str(tmp_path / "out.zarr"), "-c", str(config), "--init", diff --git a/tests/test_cli/test_segment_cli.py b/tests/test_cli/test_segment_cli.py index 6570ea7c..aa181468 100644 --- a/tests/test_cli/test_segment_cli.py +++ b/tests/test_cli/test_segment_cli.py @@ -1,4 +1,7 @@ -"""biahub segment: Nextflow-style modes (--init, --cluster debug, --resume).""" +"""biahub segment: Nextflow-style modes (--init, --cluster debug, --resume). + +Segment writes one label image per model into the -i positions (labels/). +""" import json @@ -9,25 +12,24 @@ from click.testing import CliRunner from iohub.ngff import open_ome_zarr +from iohub.ngff.utils import create_empty_plate from biahub.cli.main import cli @pytest.fixture def segment_config(tmp_path): - def _write(pretrained_model="cpsam_v2", version=None): + def _write(pretrained_model="cpsam_v2", diameter=None, name="nuc"): config = { "models": { - "nuc": { + name: { "pretrained_model": pretrained_model, "channels": ["GFP"], - "eval_args": {"diameter": None}, + "eval_args": {"diameter": diameter}, } } } - if version: - config["output_ome_zarr_version"] = version - path = tmp_path / "segment.yml" + path = tmp_path / f"segment_{name}_{diameter}.yml" path.write_text(yaml.safe_dump(config)) return path @@ -48,7 +50,7 @@ def cpu_cellpose(monkeypatch): @pytest.fixture(autouse=True) def run_in_tmp_path(monkeypatch, tmp_path): - """Without -o the submitit folder is ./slurm_output: keep it out of the repo.""" + """The submitit folder is ./slurm_output: keep it out of the repo.""" monkeypatch.chdir(tmp_path) @@ -56,39 +58,21 @@ def _run(*args): return CliRunner().invoke(cli, ["segment", *map(str, args)]) -def test_segment_init_prints_resources_and_stamps_provenance( - fake_cellpose, plate_2d, segment_config, tmp_path -): - plate_path, _ = plate_2d # (T=3, C=3, Z=1, Y=5, X=6) - out = tmp_path / "seg.zarr" +def _resources(output): + line = next(ln for ln in output.splitlines() if ln.startswith("RESOURCES:")) + return json.loads(line.removeprefix("RESOURCES:")) - result = _run( - "--init", "-i", plate_path / "A" / "1" / "0", "-o", out, "-c", segment_config() - ) - assert result.exit_code == 0, result.output - line = next(ln for ln in result.output.splitlines() if ln.startswith("RESOURCES:")) - assert set(json.loads(line.removeprefix("RESOURCES:"))) == { - "cpus", - "mem_gb", - "time_minutes", +def _two_model_config(path, order): + models = { + "nuc": {"pretrained_model": "cpsam_v2", "channels": ["GFP"]}, + "mem": {"pretrained_model": "cpsam_v2", "channels": ["RFP"]}, } - with open_ome_zarr(out / "A" / "1" / "0", mode="r") as ds: - assert ds.channel_names == ["nuc_labels"] - assert ds.data.shape == (3, 1, 1, 5, 6) - assert ds.data.dtype == np.uint32 - recorded = dict(ds.zattrs)["biahub-segment"] # provenance lives on each position - # Provenance is the config as written: channel names, not resolved indices. - assert recorded["models"]["nuc"]["channels"] == ["GFP"] - assert fake_cellpose.evaluated == [] # --init segments nothing - - -# --------------------------------------------------------------------------- -# No -o: the labels go into the -i plate as label images, labels/ -# --------------------------------------------------------------------------- + path.write_text(yaml.safe_dump({"models": {name: models[name] for name in order}})) + return path -def test_segment_init_creates_label_images_in_the_input_plate( +def test_segment_init_creates_label_images_and_prints_resources( fake_cellpose, plate_2d, segment_config, tmp_path ): plate_path, _ = plate_2d # (T=3, C=3, Z=1, Y=5, X=6) @@ -97,7 +81,7 @@ def test_segment_init_creates_label_images_in_the_input_plate( result = _run("--init", "-i", *positions, "-c", segment_config()) assert result.exit_code == 0, result.output - assert "RESOURCES:" in result.output + assert set(_resources(result.output)) == {"cpus", "mem_gb", "time_minutes"} assert "Created labels/{nuc}" in result.output for position in positions: with open_ome_zarr(position, mode="r") as pos: @@ -110,7 +94,20 @@ def test_segment_init_creates_label_images_in_the_input_plate( attrs = yaml.safe_load((position / "labels" / "nuc" / "zarr.json").read_text()) # Provenance on the label image: the config as written (names, not indices). assert attrs["attributes"]["biahub-segment"]["models"]["nuc"]["channels"] == ["GFP"] - assert fake_cellpose.evaluated == [] + assert fake_cellpose.evaluated == [] # --init segments nothing + + +def test_segment_init_rejects_an_unknown_model_before_creating_labels( + fake_cellpose, plate_2d, segment_config, tmp_path +): + plate_path, _ = plate_2d + position = plate_path / "A" / "1" / "0" + + result = _run("--init", "-i", position, "-c", segment_config(pretrained_model="nuclei")) + + assert result.exit_code != 0 + assert "Unknown cellpose model 'nuclei'" in str(result.exception) + result.output + assert not (position / "labels").exists() def test_segment_debug_fills_the_label_image( @@ -120,7 +117,7 @@ def test_segment_debug_fills_the_label_image( config = segment_config() position = plate_path / "B" / "1" / "0" assert _run("--init", "-i", position, "-c", config).exit_code == 0 - fake_cellpose.built.clear() + fake_cellpose.built.clear() # --init warmed the weights; count the worker's loads only result = _run("--cluster", "debug", "-i", position, "-c", config) @@ -129,16 +126,27 @@ def test_segment_debug_fills_the_label_image( with open_ome_zarr(position, mode="r") as pos: assert np.all(pos.get_label("nuc")["0"][:] == 1) # the stand-in labels every pixel 1 assert pos.data.shape == (3, 3, 1, 5, 6) - assert fake_cellpose.built == ["cpsam_v2"] + assert fake_cellpose.built == ["cpsam_v2"] # one model load for the whole position assert len(fake_cellpose.evaluated) == 3 # one per timepoint +def test_segment_worker_creates_the_labels_when_init_did_not_run( + fake_cellpose, cpu_cellpose, plate_2d, segment_config +): + plate_path, _ = plate_2d + position = plate_path / "B" / "1" / "0" + + result = _run("--cluster", "debug", "-i", position, "-c", segment_config()) + + assert result.exit_code == 0, result.output + with open_ome_zarr(position, mode="r") as pos: + assert np.all(pos.get_label("nuc")["0"][:] == 1) + + def test_segment_resume_skips_finished_label_images( - fake_cellpose, cpu_cellpose, create_custom_plate, segment_config, tmp_path + fake_cellpose, cpu_cellpose, segment_config, tmp_path ): """Resume works on the label image as on an image (iohub output_label).""" - from iohub.ngff.utils import create_empty_plate - # A Zarr v3 (OME-Zarr 0.5) plate: resume needs it. plate_path = tmp_path / "v05.zarr" create_empty_plate( @@ -179,6 +187,56 @@ def test_segment_writes_one_label_image_per_model( assert len(fake_cellpose.evaluated) == 6 # 3 timepoints x 2 models +# --------------------------------------------------------------------------- +# Strategies accumulate: a label image is only rewritten by the model that made it +# --------------------------------------------------------------------------- + + +def test_segment_strategies_accumulate_under_different_model_names( + fake_cellpose, cpu_cellpose, plate_2d, segment_config +): + plate_path, _ = plate_2d + position = plate_path / "B" / "1" / "0" + assert ( + _run( + "--cluster", + "debug", + "-i", + position, + "-c", + segment_config(name="nuc_d80", diameter=80), + ).exit_code + == 0 + ) + + result = _run( + "--cluster", "debug", "-i", position, "-c", segment_config(name="nuc_native") + ) + + assert result.exit_code == 0, result.output + with open_ome_zarr(position, mode="r") as pos: + assert pos.label_names() == ["nuc_d80", "nuc_native"] + for name in pos.label_names(): + recorded = dict(pos.get_label(name).zattrs)["biahub-segment"]["models"] + assert list(recorded) == [name] + + +def test_segment_refuses_to_rewrite_a_label_image_with_other_settings( + fake_cellpose, cpu_cellpose, plate_2d, segment_config +): + plate_path, _ = plate_2d + position = plate_path / "B" / "1" / "0" + assert _run("--init", "-i", position, "-c", segment_config(diameter=80)).exit_code == 0 + + result = _run("--init", "-i", position, "-c", segment_config(diameter=None)) + + assert result.exit_code != 0 + assert "segmented with other settings" in str(result.exception) + assert "own model name" in str(result.exception) + # The same settings again are fine: that is a retry or a resume. + assert _run("--init", "-i", position, "-c", segment_config(diameter=80)).exit_code == 0 + + def test_segment_refuses_a_label_image_of_another_shape( fake_cellpose, plate_2d, example_plate, tmp_path ): @@ -199,7 +257,7 @@ def test_segment_refuses_a_label_image_of_another_shape( assert "already has a label image 'nuc'" in str(result.exception) -def test_segment_resume_token_names_the_model_in_the_input_plate(tmp_path): +def test_segment_resume_token_names_the_model(tmp_path): from biahub.segment import _resume_token from biahub.settings import SegmentationSettings from biahub.utils.config import yaml_to_model @@ -211,81 +269,31 @@ def test_segment_resume_token_names_the_model_in_the_input_plate(tmp_path): _two_model_config(tmp_path / "b.yml", ["mem", "nuc"]), SegmentationSettings ) - # Each model has its own label image, so the model order does not matter there. + # Each model has its own label image, so the model order does not matter. assert _resume_token(a, "nuc") == _resume_token(b, "nuc") assert _resume_token(a, "nuc") != _resume_token(a, "mem") -# --------------------------------------------------------------------------- -# -o: a separate label plate, one channel per model -# --------------------------------------------------------------------------- - - -def test_segment_init_rejects_an_unknown_model( - fake_cellpose, plate_2d, segment_config, tmp_path -): +def test_segment_output_flag_and_version_are_gone(plate_2d, segment_config, tmp_path): plate_path, _ = plate_2d result = _run( - "--init", - "-i", - plate_path / "A" / "1" / "0", - "-o", - tmp_path / "seg.zarr", - "-c", - segment_config(pretrained_model="nuclei"), + "-o", tmp_path / "o.zarr", "-i", plate_path / "A" / "1" / "0", "-c", segment_config() ) - assert result.exit_code != 0 - assert "Unknown cellpose model 'nuclei'" in str(result.exception) + result.output - - -def test_segment_debug_runs_one_position_in_process( - fake_cellpose, cpu_cellpose, plate_2d, segment_config, tmp_path -): - plate_path, _ = plate_2d - out, config = tmp_path / "seg.zarr", segment_config() - position = plate_path / "B" / "1" / "0" - assert _run("--init", "-i", position, "-o", out, "-c", config).exit_code == 0 - fake_cellpose.built.clear() # --init warmed the weights; count the worker's loads only - - result = _run("--cluster", "debug", "-i", position, "-o", out, "-c", config) - - assert result.exit_code == 0, result.output - assert "Segmentation complete" in result.output - with open_ome_zarr(out / "B" / "1" / "0", mode="r") as ds: - assert np.all(np.asarray(ds.data) == 1) # the stand-in labels every pixel 1 - assert fake_cellpose.built == ["cpsam_v2"] # one model load for the whole position - assert len(fake_cellpose.evaluated) == 3 # one per timepoint - - -def test_segment_resume_skips_finished_work( - fake_cellpose, cpu_cellpose, plate_2d, segment_config, tmp_path -): - plate_path, _ = plate_2d - out, config = tmp_path / "seg.zarr", segment_config(version="0.5") # resume needs Zarr v3 - position = plate_path / "B" / "1" / "0" - assert _run("--init", "-i", position, "-o", out, "-c", config).exit_code == 0 - assert _run("--cluster", "debug", "-i", position, "-o", out, "-c", config).exit_code == 0 - n_first = len(fake_cellpose.evaluated) + assert result.exit_code == 2 + assert "No such option '-o'" in result.output - result = _run("--cluster", "debug", "--resume", "-i", position, "-o", out, "-c", config) + from biahub.settings import SegmentationSettings - assert result.exit_code == 0, result.output - assert len(fake_cellpose.evaluated) == n_first + with pytest.raises(ValueError, match="output_ome_zarr_version was removed"): + SegmentationSettings( + models={"nuc": {"pretrained_model": "cpsam_v2", "channels": ["GFP"]}}, + output_ome_zarr_version="0.5", + ) -def test_segment_legacy_local_flag_is_gone(plate_2d, segment_config, tmp_path): - plate_path, _ = plate_2d - result = _run( - "--local", - "-i", - plate_path / "A" / "1" / "0", - "-o", - tmp_path / "o.zarr", - "-c", - segment_config(), - ) - assert result.exit_code == 2 - assert "No such option '--local'" in result.output +# --------------------------------------------------------------------------- +# Executors and resources +# --------------------------------------------------------------------------- class _RecordingExecutor: @@ -323,17 +331,9 @@ def test_segment_gives_local_jobs_a_gpu( monkeypatch.setattr("biahub.segment.submitit.AutoExecutor", _RecordingExecutor) monkeypatch.delenv("CI", raising=False) # CI forces cluster=debug plate_path, _ = plate_2d - out = tmp_path / "seg.zarr" result = _run( - "--cluster", - cluster, - "-i", - plate_path / "A" / "1" / "0", - "-o", - out, - "-c", - segment_config(), + "--cluster", cluster, "-i", plate_path / "A" / "1" / "0", "-c", segment_config() ) assert result.exit_code == 0, result.output @@ -360,21 +360,16 @@ def test_segment_worker_stages_weights_before_cellpose_is_imported( lambda: calls.append("import cellpose") or real_valid_args(), ) plate_path, _ = plate_2d - out, config = tmp_path / "seg.zarr", segment_config() + config = segment_config() position = plate_path / "B" / "1" / "0" - assert _run("--init", "-i", position, "-o", out, "-c", config).exit_code == 0 + assert _run("--init", "-i", position, "-c", config).exit_code == 0 calls.clear() - assert _run("--cluster", "debug", "-i", position, "-o", out, "-c", config).exit_code == 0 + assert _run("--cluster", "debug", "-i", position, "-c", config).exit_code == 0 assert calls and calls[0] == "stage", calls -def _resources(output): - line = next(ln for ln in output.splitlines() if ln.startswith("RESOURCES:")) - return json.loads(line.removeprefix("RESOURCES:")) - - def test_segment_requests_less_time_for_2d_than_3d( fake_cellpose, plate_2d, example_plate, tmp_path ): @@ -399,7 +394,7 @@ def test_segment_requests_less_time_for_2d_than_3d( } ) ) - result = _run("--init", "-i", position, "-o", tmp_path / f"{kind}.zarr", "-c", cfg) + result = _run("--init", "-i", position, "-c", cfg) assert result.exit_code == 0, result.output res[kind] = _resources(result.output) assert res["2d"]["cpus"] <= 4 @@ -418,14 +413,7 @@ def test_segment_slurm_head_does_not_stage_weights( plate_path, _ = plate_2d result = _run( - "--cluster", - "slurm", - "-i", - plate_path / "A" / "1" / "0", - "-o", - tmp_path / "s.zarr", - "-c", - segment_config(), + "--cluster", "slurm", "-i", plate_path / "A" / "1" / "0", "-c", segment_config() ) assert result.exit_code == 0, result.output @@ -441,14 +429,7 @@ def test_segment_local_jobs_get_the_step_time_limit( plate_path, _ = plate_2d result = _run( - "--cluster", - "local", - "-i", - plate_path / "A" / "1" / "0", - "-o", - tmp_path / "s.zarr", - "-c", - segment_config(), + "--cluster", "local", "-i", plate_path / "A" / "1" / "0", "-c", segment_config() ) assert result.exit_code == 0, result.output @@ -458,24 +439,6 @@ def test_segment_local_jobs_get_the_step_time_limit( ) -def test_segment_init_checks_models_before_creating_the_plate( - fake_cellpose, plate_2d, segment_config, tmp_path -): - plate_path, _ = plate_2d - out = tmp_path / "seg.zarr" - result = _run( - "--init", - "-i", - plate_path / "A" / "1" / "0", - "-o", - out, - "-c", - segment_config(pretrained_model="nuclei"), - ) - assert result.exit_code != 0 - assert not out.exists() - - def test_segment_local_sbatch_file_overrides_computed_parameters( fake_cellpose, plate_2d, segment_config, tmp_path, monkeypatch ): @@ -494,8 +457,6 @@ def test_segment_local_sbatch_file_overrides_computed_parameters( sbatch, "-i", plate_path / "A" / "1" / "0", - "-o", - tmp_path / "s.zarr", "-c", segment_config(), ) @@ -505,56 +466,6 @@ def test_segment_local_sbatch_file_overrides_computed_parameters( assert _RecordingExecutor.parameters["gpus_per_node"] == 1 -def _two_model_config(path, order): - models = { - "nuc": {"pretrained_model": "cpsam_v2", "channels": ["GFP"]}, - "mem": {"pretrained_model": "cpsam_v2", "channels": ["RFP"]}, - } - config = { - "models": {name: models[name] for name in order}, - "output_ome_zarr_version": "0.5", - } - path.write_text(yaml.safe_dump(config, sort_keys=False)) - return path - - -def test_segment_refuses_an_output_with_a_different_model_order( - fake_cellpose, cpu_cellpose, plate_2d, tmp_path -): - """Reordering the models used to resume into swapped label channels.""" - plate_path, _ = plate_2d - out, position = tmp_path / "seg.zarr", plate_path / "B" / "1" / "0" - first = _two_model_config(tmp_path / "first.yml", ["nuc", "mem"]) - assert _run("--init", "-i", position, "-o", out, "-c", first).exit_code == 0 - assert _run("--cluster", "debug", "-i", position, "-o", out, "-c", first).exit_code == 0 - n_first = len(fake_cellpose.evaluated) - - swapped = _two_model_config(tmp_path / "swapped.yml", ["mem", "nuc"]) - result = _run("--cluster", "debug", "--resume", "-i", position, "-o", out, "-c", swapped) - - assert result.exit_code != 0 - assert "['nuc_labels', 'mem_labels']" in str(result.exception) - assert len(fake_cellpose.evaluated) == n_first # nothing was segmented into it - with open_ome_zarr(out / "B" / "1" / "0", mode="r") as ds: - assert ds.channel_names == ["nuc_labels", "mem_labels"] - - -def test_segment_resume_token_depends_on_model_order(tmp_path): - from biahub.segment import _resume_token - from biahub.settings import SegmentationSettings - from biahub.utils.config import yaml_to_model - - a = yaml_to_model( - _two_model_config(tmp_path / "a.yml", ["nuc", "mem"]), SegmentationSettings - ) - b = yaml_to_model( - _two_model_config(tmp_path / "b.yml", ["mem", "nuc"]), SegmentationSettings - ) - - assert _resume_token(a) != _resume_token(b) - assert _resume_token(a) == _resume_token(a.model_copy(deep=True)) - - def test_segment_mem_per_cpu_override_drops_the_computed_total( fake_cellpose, plate_2d, segment_config, tmp_path, monkeypatch ): @@ -573,8 +484,6 @@ def test_segment_mem_per_cpu_override_drops_the_computed_total( sbatch, "-i", plate_path / "A" / "1" / "0", - "-o", - tmp_path / "s.zarr", "-c", segment_config(), ) @@ -582,26 +491,3 @@ def test_segment_mem_per_cpu_override_drops_the_computed_total( assert result.exit_code == 0, result.output assert _RecordingExecutor.parameters["slurm_mem_per_cpu"] == "16G" assert "slurm_mem" not in _RecordingExecutor.parameters - - -def test_segment_refuses_an_output_with_a_different_shape( - fake_cellpose, plate_2d, example_plate, tmp_path -): - """Switching a same-named model from 3D to 2D kept the old multi-plane array.""" - out = tmp_path / "seg.zarr" - model = {"pretrained_model": "cpsam_v2", "channels": ["GFP"]} - three_d = tmp_path / "3d.yml" - three_d.write_text( - yaml.safe_dump({"models": {"nuc": {**model, "eval_args": {"do_3D": True}}}}) - ) - two_d = tmp_path / "2d.yml" - two_d.write_text(yaml.safe_dump({"models": {"nuc": model}})) - stack = example_plate[0] / "B" / "1" / "0" # Z=4 - assert _run("--init", "-i", stack, "-o", out, "-c", three_d).exit_code == 0 - - result = _run("--init", "-i", plate_2d[0] / "B" / "1" / "0", "-o", out, "-c", two_d) - - assert result.exit_code != 0 - assert "(3, 1, 4, 5, 6)" in str(result.exception) - with open_ome_zarr(out / "B" / "1" / "0", mode="r") as ds: - assert ds.data.shape == (3, 1, 4, 5, 6) diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 66a6727f..7613d813 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -474,7 +474,7 @@ def test_labels_is_the_default_and_names_the_output(): TrackingSettings() settings = TrackingSettings(labels="nuclei") assert settings.segmentation_method == "labels" - assert settings.output_label == "nuclei_tracked" + assert settings.tracked_label == "nuclei_tracked" assert settings.input_images == [] with pytest.raises(ValueError, match="target_channel is for"): TrackingSettings(labels="nuclei", target_channel="x") @@ -485,7 +485,59 @@ def test_foreground_contour_needs_a_target_channel_and_input_images(): TrackingSettings(segmentation_method="foreground_contour") with pytest.raises(ValueError, match="needs input_images"): TrackingSettings(segmentation_method="foreground_contour", target_channel="x") - assert _minimal_settings().output_label == LABEL + assert _minimal_settings().tracked_label == LABEL + + +def test_output_label_names_the_tracked_image_and_the_tracks_files(): + default = TrackingSettings(labels="nuclei") + assert (default.tracked_label, default.tracks_stem) == ("nuclei_tracked", "tracks") + custom = TrackingSettings(labels="nuclei", output_label="nuclei_sigma3") + assert (custom.tracked_label, custom.tracks_stem) == ( + "nuclei_sigma3", + "tracks_nuclei_sigma3", + ) + with pytest.raises(ValueError, match="is the label image being tracked"): + TrackingSettings(labels="nuclei", output_label="nuclei") + with pytest.raises(ValueError, match="output_label is for"): + _minimal_settings(output_label="x") + + +def test_track_debug_run_with_an_output_label(tmp_path, example_tracking_plate, monkeypatch): + """Two tracking runs of one segmentation keep their labels, GEFFs and CSVs apart.""" + monkeypatch.setenv("ULTRACK_ARRAY_MODULE", "numpy") + plate_path, _ = example_tracking_plate + config_path = _make_tracking_config(plate_path, tmp_path) + config = yaml.safe_load(config_path.read_text()) + config["output_label"] = "nuclei_run2" + config_path.write_text(yaml.safe_dump(config)) + position_path = plate_path / "A" / "1" / "0" + + result = CliRunner().invoke( + cli, + [ + "track", + "--cluster", + "debug", + "-i", + str(position_path), + "-o", + str(tmp_path / "o"), + "-c", + str(config_path), + ], + ) + + assert result.exit_code == 0, result.output + with open_ome_zarr(str(position_path), mode="r") as pos: + assert pos.label_names() == [INPUT_LABEL, "nuclei_run2"] + assert pos.get_label("nuclei_run2")["0"][:].any() + assert (position_path / "tracks_nuclei_run2.geff").is_dir() + assert (position_path / "tracks_nuclei_run2_A_1_0.csv").is_file() + assert not (position_path / "tracks.geff").exists() + import geff + + _, metadata = geff.read(str(position_path / "tracks_nuclei_run2.geff"), backend="networkx") + assert metadata.related_objects[0].path == "../labels/nuclei_run2" def test_cellpose_config_is_rejected_with_a_pointer_to_segment(): diff --git a/tests/test_example_settings.py b/tests/test_example_settings.py index 0036784d..0f71e052 100644 --- a/tests/test_example_settings.py +++ b/tests/test_example_settings.py @@ -218,7 +218,7 @@ def test_nextflow_a549_track_config_validates(): assert settings.output_mode == "2D" assert settings.segmentation_method == "labels" assert settings.labels == "nuclei" - assert settings.output_label == "nuclei_tracked" + assert settings.tracked_label == "nuclei_tracked" assert settings.input_images == [] assert settings.min_sibling_length == 3 From bebe62885236e8213af17ca10ce36f9011372edb Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 21 Oct 2026 14:41:35 -0700 Subject: [PATCH 39/54] build(configs): concatenate chunks 1 x 8 x 256 x 256; iohub pin keeps level-0 chunks down the pyramid The templates' float32 chunks go from 4 MiB to 2 MiB (Z 16 -> 8); with shards: 5GB the A549 plate's shards become 6 timepoints by the full volume (4.5 GiB). iohub dev/pyramid-shards now keeps the level-0 chunk shape at every pyramid level, clamped to the level, instead of halving it. Co-Authored-By: Claude Fable 5.1 --- nextflow/configs/a549/concatenate.yml | 6 +++--- nextflow/configs/zebrafish/concatenate.yml | 2 +- pyproject.toml | 4 ++-- uv.lock | 8 ++++---- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/nextflow/configs/a549/concatenate.yml b/nextflow/configs/a549/concatenate.yml index e0df504d..8377a6ab 100644 --- a/nextflow/configs/a549/concatenate.yml +++ b/nextflow/configs/a549/concatenate.yml @@ -14,13 +14,13 @@ # A range past the end of the shorter positions pads them the same way. time_indices: all channel_names: all -chunks_czyx: [1, 16, 256, 256] +chunks_czyx: [1, 8, 256, 256] # 2 MiB float32 chunks; the pyramid levels keep this shape output_ome_zarr_version: "0.5" # shards is a size target: iohub builds the largest chunk-aligned shard under it, # filling X, Y, Z first and then timepoints (C is never sharded). 5GB on this -# plate's 86 x 1664 x 1193 float32 volumes gives shards of 5 timepoints by the -# full volume (4.1 GiB), the same geometry a run before this key produced. Do +# plate's 86 x 1664 x 1193 float32 volumes gives shards of 6 timepoints by the +# full volume (88 x 1792 x 1280 in chunks, 4.5 GiB). Do # not raise it without a measured reason: the T extent multiplies the shard # buffer and the RAM request, and one torn shard loses every timepoint in it. # Also accepts "XY" / "XYZ" (one shard per plane / volume) or a TCZYX shape. diff --git a/nextflow/configs/zebrafish/concatenate.yml b/nextflow/configs/zebrafish/concatenate.yml index 5f540c15..76d32dda 100644 --- a/nextflow/configs/zebrafish/concatenate.yml +++ b/nextflow/configs/zebrafish/concatenate.yml @@ -14,7 +14,7 @@ # A range past the end of the shorter positions pads them the same way. time_indices: all channel_names: all -chunks_czyx: [1, 16, 256, 256] +chunks_czyx: [1, 8, 256, 256] # 2 MiB float32 chunks; the pyramid levels keep this shape output_ome_zarr_version: "0.5" # shards is a size target: iohub builds the largest chunk-aligned shard under it, diff --git a/pyproject.toml b/pyproject.toml index 5a7225db..2853e830 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,7 +102,7 @@ index-strategy = "unsafe-best-match" # (the dev/pyramid-shards branch). uv refuses two git URLs for one package, # so force the whole graph onto biahub's rev. Drop together with the iohub pin. override-dependencies = [ - "iohub @ git+https://github.com/czbiohub-sf/iohub@7d18a5761deaec89580ef9ee15df78d178a55084", + "iohub @ git+https://github.com/czbiohub-sf/iohub@82808fd225051aa638f804af88afd0a7f8c86744", ] [tool.uv.sources] @@ -162,7 +162,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "7d18a5761deaec89580ef9ee15df78d178a55084" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "82808fd225051aa638f804af88afd0a7f8c86744" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index 71d232f0..3c310815 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=7d18a5761deaec89580ef9ee15df78d178a55084" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=82808fd225051aa638f804af88afd0a7f8c86744" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?rev=f5ae656e0e37320ead74f7c1f95ab1ed304c1829" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=7d18a5761deaec89580ef9ee15df78d178a55084" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=82808fd225051aa638f804af88afd0a7f8c86744" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev28+7d18a57" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=7d18a5761deaec89580ef9ee15df78d178a55084#7d18a5761deaec89580ef9ee15df78d178a55084" } +version = "0.3.12.dev30+82808fd" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=82808fd225051aa638f804af88afd0a7f8c86744#82808fd225051aa638f804af88afd0a7f8c86744" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 7a4de6cd6b808a8ebb772559ba4f26f0707d71b7 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 21 Oct 2026 14:42:51 -0700 Subject: [PATCH 40/54] feat(track)!: the tracked label image is labels/tracks_ The default name of the tracked label image gains the `tracks_` prefix instead of a `_tracked` suffix, so `labels/tracks_nuclei` sits beside `tracks.geff` and `tracks_.csv`, and an explicit `output_label` names all three the same way (`labels/`, `tracks_.geff`, `tracks__.csv`). The A549 QC config counts `tracks_nuclei`. Co-Authored-By: Claude Fable 5.1 --- .../skills/reconstruct-with-nextflow/SKILL.md | 6 +++--- .../references/caveats.md | 4 ++-- .../references/datasets.md | 2 +- .../templates/check_qc_channels.py | 10 +++++----- .../templates/run_mantis_v2.sh | 2 +- biahub/settings.py | 17 +++++++++-------- biahub/track.py | 4 ++-- nextflow/configs/a549/qc_project.yaml | 6 +++--- nextflow/configs/a549/track.yml | 6 +++--- nextflow/modules/tracking.nf | 2 +- settings/example_track_settings.yml | 4 ++-- tests/test_cli/test_track_cli.py | 6 +++--- tests/test_example_settings.py | 4 ++-- tests/test_segment_settings.py | 2 +- 14 files changed, 38 insertions(+), 37 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index c78497a6..3f3f56d0 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -227,9 +227,9 @@ label image segment wrote (`labels/`). **`qc_project.yaml` against the projected plate, `segment.yml` and `track.yml` — the coupling that actually breaks.** Segment writes `labels/` into each -position of the plate it reads and tracking writes `labels/_tracked` +position of the plate it reads and tracking writes `labels/tracks_` beside it. `qc_project.yaml`'s `instance_count` names the tracked one, `label: -nuclei_tracked`, and its anchor `channels: [nuclei_prediction]`, the channel the +tracks_nuclei`, and its anchor `channels: [nuclei_prediction]`, the channel the labels were segmented from. The label only resolves if `track.yml` says `labels: nuclei` and `segment.yml` has a model `nuclei`; the anchor must be that model's channel and a channel `project.yml` writes — four files with no shared source of @@ -567,7 +567,7 @@ directory (`_rerun`) over relaunching in place. size on disk. This is the deliverable for every family. When segmentation and tracking ran, their output is inside the projected plate: check a position has `labels/nuclei` (the segment config's model name) and - `labels/nuclei_tracked` (the track config's `labels` plus `_tracked`; both + `labels/tracks_nuclei` (`tracks_` plus the track config's `labels`; both uint32 TZYX at the image's full Z), `tracks.geff` and `tracks___.csv`, and report them beside the image data. 3. If QC ran, report its verdict: the `QC_SUMMARY` line from the diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index da45e448..a7ecd4fc 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -162,7 +162,7 @@ image, which also renames the GEFF and CSV to `tracks_.*`), `input_images` on the legacy path). Track neither segments nor selects or projects Z-planes: it reads the label image `biahub segment` wrote into the **projected** plate (`labels: `), writes the tracked -labels beside it as `labels/_tracked` with `tracks.geff` and a tracks +labels beside it as `labels/tracks_` with `tracks.geff` and a tracks CSV, and a 3D input with `output_mode: 2D` is an error. Old configs under `/hpc/projects/tlg2_mantis/` still carry rejected fields — migrate: @@ -181,7 +181,7 @@ output_ome_zarr_version -> (delete: labels take the projected plat `focus.lambda_ill`; the pixel size comes from the store.) `segment.yml`'s model `channels` must match channels `project.yml` writes, which carry the **pre-rename** names of the **assembled** plate (§2, §5). `qc_project.yaml`'s -`instance_count.label` must be `_tracked` and its anchor +`instance_count.label` must be `tracks_` and its anchor `channels:` entry the segment model's channel — SKILL.md §5a checks these couplings before launch. diff --git a/.claude/skills/reconstruct-with-nextflow/references/datasets.md b/.claude/skills/reconstruct-with-nextflow/references/datasets.md index 1f0a65b8..18485124 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/datasets.md +++ b/.claude/skills/reconstruct-with-nextflow/references/datasets.md @@ -74,7 +74,7 @@ layout is defined once in `nextflow/mantis-v2.nf` (`directory_layout()`): ├── 4-assemble/_proj.zarr # A549 only — the 2D (Z=1) companion plate; with │ # segmentation and tracking, each position also │ # holds labels/ (segment), -│ # labels/_tracked, tracks.geff and +│ # labels/tracks_, tracks.geff and │ # tracks___.csv (track) ├── qc/ │ ├── report/index.html # one report, one tab per QC'd store diff --git a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py index ac12f218..1283c7a2 100644 --- a/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py +++ b/.claude/skills/reconstruct-with-nextflow/templates/check_qc_channels.py @@ -27,11 +27,11 @@ so they are read from the settings model, not the YAML text) + virtual-stain's data.init_args.target_channel labels = (/labels/) - + _tracked (/labels/_tracked) + + tracks_ (/labels/tracks_) Segment and track write their labels INTO the plate they read, so `qc_project.yaml`'s `instance_count.label` only resolves if it is -`_tracked` (or a `segment.yml` model name), and its one anchor +`tracks_` (or a `segment.yml` model name), and its one anchor channel, the channel the labels were segmented from, must be that model's channel. Four files, no shared source of truth. Checked when `qc_project.yaml` has an `instance_count` block and `track.yml` exists. @@ -115,7 +115,7 @@ def tracked_label(track_cfg): t = yaml.safe_load(Path(track_cfg).read_text()) if t.get("segmentation_method", "labels") != "labels": return None - return t["labels"], t.get("output_label") or f"{t['labels']}_tracked" + return t["labels"], t.get("output_label") or f"tracks_{t['labels']}" if __name__ == "__main__": @@ -161,7 +161,7 @@ def tracked_label(track_cfg): bad += not ok print(f" {'ok ' if ok else 'MISS'} models.{name}.channels: {ch!r}") - # track.yml reads labels/ of that plate and writes labels/_tracked. + # track.yml reads labels/ of that plate and writes labels/tracks_. tracked = None if (cfgs / "track.yml").exists(): tracked = tracked_label(cfgs / "track.yml") @@ -173,7 +173,7 @@ def tracked_label(track_cfg): # qc_project.yaml QCs the plate segment and track write into: its `channels:` # against that plate, and its `instance_count` against the label images there - # (the tracked one, labels/_tracked, or a segment model's) with + # (the tracked one, labels/tracks_, or a segment model's) with # the anchor channel being the one the model segmented. if (cfgs / "qc_project.yaml").exists(): print(f"qc_project.yaml: {source} plate will have {track_reads}") diff --git a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh index 86b39dfe..a4c1be97 100755 --- a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh +++ b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh @@ -15,7 +15,7 @@ # 0-flatfield/ 1-deskew/ 2-reconstruct/ 3-virtual-stain/ 4-assemble/ # Project writes its plate beside the assembled one as 4-assemble/_proj.zarr; # segment writes labels/ into each position of that plate and tracking -# labels/_tracked, tracks.geff and a tracks CSV; none has a directory of its own. +# labels/tracks_, tracks.geff and a tracks CSV; none has a directory of its own. # The number is the step's POSITION among the steps this run performs that # write a directory, so a run that skips a step shifts everything after it down. # The work dir defaults to /nextflow/work; override with `-work-dir`. diff --git a/biahub/settings.py b/biahub/settings.py index a9dba1e6..9387b48d 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -62,9 +62,10 @@ def validate_path_not_plate(cls, v): return v -#: Suffix of the label image tracking writes next to the one it reads: -#: ``labels/`` (from `biahub segment`) -> ``labels/_tracked``. -TRACKED_LABEL_SUFFIX = "_tracked" +#: Prefix of the label image tracking writes next to the one it reads: +#: ``labels/`` (from `biahub segment`) -> ``labels/tracks_``, matching +#: the ``tracks.geff`` / ``tracks_.csv`` beside it. +TRACKED_LABEL_PREFIX = "tracks_" class TrackingSettings(MyBaseModel): @@ -82,13 +83,13 @@ class TrackingSettings(MyBaseModel): segmentation_method: Literal["labels", "foreground_contour"] = "labels" # segmentation_method: labels -- the label image, `labels/` in each -i # position, named after the segment config's model. The tracked labels go to - # `labels/_tracked`. + # `labels/tracks_`. labels: str | None = None # Smoothing (pixels) ultrack applies to the labels' contours before it builds its # own segmentation hypotheses from them (ultrack `labels_to_contours`). labels_sigma: float = 5.0 # segmentation_method: labels -- the name of the tracked label image, when the - # default `_tracked` is not wanted, e.g. to keep several tracking runs of + # default `tracks_` is not wanted, e.g. to keep several tracking runs of # one segmentation side by side. Then the tracks go to `tracks_.geff` # and `tracks__.csv` instead of `tracks.geff` / `tracks_.csv`. output_label: str | None = None @@ -148,8 +149,8 @@ def check_segmentation_method(self): if self.target_channel is not None: raise ValueError( "target_channel is for segmentation_method: foreground_contour; with " - f"labels the tracks are written to labels/{self.labels}" - f"{TRACKED_LABEL_SUFFIX}, or to labels/ if set." + f"labels the tracks are written to labels/{TRACKED_LABEL_PREFIX}{self.labels}" + ", or to labels/ if set." ) if self.output_label == self.labels: raise ValueError( @@ -178,7 +179,7 @@ def check_segmentation_method(self): def tracked_label(self) -> str: """Name of the label image tracking writes, ``labels/``.""" if self.segmentation_method == "labels": - return self.output_label or f"{self.labels}{TRACKED_LABEL_SUFFIX}" + return self.output_label or f"{TRACKED_LABEL_PREFIX}{self.labels}" return self.target_channel @property diff --git a/biahub/track.py b/biahub/track.py index bf2dd12c..417a1334 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -914,7 +914,7 @@ def track( Notes ----- - Outputs go into each ``-i`` position: ``labels/_tracked`` (OME-NGFF label + Outputs go into each ``-i`` position: ``labels/tracks_`` (OME-NGFF label image; ``labels/`` on the foreground+contour path), ``tracks.geff`` and ``tracks___.csv``. The ``-i`` plate is the plate track reads, so 2D tracking writes into the projected (Z=1) plate, next to the @@ -1079,7 +1079,7 @@ def track_cli( """Track the instance labels of a label image (`biahub segment`) with ultrack. Reads labels/ from each -i position and writes into it: the tracked labels - as an OME-NGFF label image (labels/_tracked), the tracks as tracks.geff, + as an OME-NGFF label image (labels/tracks_), the tracks as tracks.geff, and a tracks CSV. 2D tracking reads (and writes into) the projected Z=1 plate `biahub project` makes; the Ultrack databases are temporary unless -o names a directory to keep them in. Track does not segment or project. diff --git a/nextflow/configs/a549/qc_project.yaml b/nextflow/configs/a549/qc_project.yaml index 3755c7e3..3a3fd614 100644 --- a/nextflow/configs/a549/qc_project.yaml +++ b/nextflow/configs/a549/qc_project.yaml @@ -1,7 +1,7 @@ # QC of the PROJECTED plate (`4-assemble/_proj.zarr`, Z=1): the plate # `biahub project` writes beside the assembled one and tracking reads and writes # into. It holds the projected channels and, after tracking, each position's -# `labels/` (segment) and `labels/_tracked` (track) label images, +# `labels/` (segment) and `labels/tracks_` (track) label images, # so this is where the cell counts live: # `instance_count` reads the TRACKED label image (imaging-qc#243); segment writes # `labels/` beside it. The pipeline plans @@ -42,11 +42,11 @@ blank: granularity: channel # Cells per frame, from the label image tracking writes (`label:`, the track config's -# `labels` plus `_tracked`). `channels:` is the ONE anchor channel the labels +# `labels` with a `tracks_` prefix). `channels:` is the ONE anchor channel the labels # were segmented from: the count rows carry its C. The plate has one Z-plane, so # `slice` and `volume` count the same thing; `slice` keeps per-(T, Z) rows. instance_count: - label: nuclei_tracked + label: tracks_nuclei channels: - nuclei_prediction granularity: slice diff --git a/nextflow/configs/a549/track.yml b/nextflow/configs/a549/track.yml index 1e937157..d0e38f82 100644 --- a/nextflow/configs/a549/track.yml +++ b/nextflow/configs/a549/track.yml @@ -2,7 +2,7 @@ # # Track reads the PROJECTED plate (-assemble/_proj.zarr, Z=1) and the label # image `biahub segment` (segment.yml) wrote into it, `labels/`; it writes the -# tracked labels next to it as `labels/_tracked`, plus tracks.geff and a tracks +# tracked labels next to it as `labels/tracks_`, plus tracks.geff and a tracks # CSV, into every position. Track neither segments nor projects. # # CURRENT SCHEMA. TrackingSettings forbids extra keys; the old fields `mode`, @@ -15,14 +15,14 @@ # -> segment.yml: models. # cellpose_config.labels_sigma -> labels_sigma # target_channel -> labels (the segment.yml model name); the -# output is named _tracked +# output is named tracks_ # 2D tracks the Z=1 projected plate and writes a Z=1 label image. Track does not # project; a 3D input with output_mode 2D is an error. output_mode: "2D" fov: "*/*/*" -# The label image to track: segment.yml's model name. Tracks go to labels/nuclei_tracked. +# The label image to track: segment.yml's model name. Tracks go to labels/tracks_nuclei. labels: nuclei # Smoothing (px) of the labels' contours before ultrack builds its own segmentation # hypotheses from them. diff --git a/nextflow/modules/tracking.nf b/nextflow/modules/tracking.nf index ce2ce37a..54bd19e7 100644 --- a/nextflow/modules/tracking.nf +++ b/nextflow/modules/tracking.nf @@ -7,7 +7,7 @@ // // Tracking reads the label image `biahub segment` wrote into `input_zarr` // (`labels/`) and writes INTO the same positions — the tracked labels as -// `labels/_tracked`, `tracks.geff` and a tracks CSV — and nothing else +// `labels/tracks_`, `tracks.geff` and a tracks CSV — and nothing else // is kept: each worker builds its Ultrack database in the node's $TMPDIR and // deletes it after export (no -o), so tracking has no step directory. Its real // logs are in nextflow/slurm_output/track/, like every step's; submitit's diff --git a/settings/example_track_settings.yml b/settings/example_track_settings.yml index a6f018d6..6b007759 100644 --- a/settings/example_track_settings.yml +++ b/settings/example_track_settings.yml @@ -17,12 +17,12 @@ segmentation_method: labels # The label image to track, `labels/` in every -i position: the segment # config's model name. The tracked labels are written next to it as -# `labels/_tracked`, with tracks.geff and a tracks CSV. +# `labels/tracks_`, with tracks.geff and a tracks CSV. labels: nuclei # Optional: another name for the tracked label image, to keep several tracking runs of # one segmentation side by side. The tracks then go to tracks_.geff and # tracks__.csv rather than tracks.geff / tracks_.csv. -# output_label: nuclei_tracked_sigma3 +# output_label: tracks_nuclei_sigma3 # Smoothing (pixels) ultrack applies to the labels' contours before it builds its own # segmentation hypotheses from them. labels_sigma: 5.0 diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 7613d813..4dbd2037 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -19,7 +19,7 @@ # The label image `biahub segment` would write (its model name) and the one track writes. INPUT_LABEL = "nuclei" -LABEL = "nuclei_tracked" +LABEL = "tracks_nuclei" @pytest.fixture(scope="function") @@ -474,7 +474,7 @@ def test_labels_is_the_default_and_names_the_output(): TrackingSettings() settings = TrackingSettings(labels="nuclei") assert settings.segmentation_method == "labels" - assert settings.tracked_label == "nuclei_tracked" + assert settings.tracked_label == "tracks_nuclei" assert settings.input_images == [] with pytest.raises(ValueError, match="target_channel is for"): TrackingSettings(labels="nuclei", target_channel="x") @@ -490,7 +490,7 @@ def test_foreground_contour_needs_a_target_channel_and_input_images(): def test_output_label_names_the_tracked_image_and_the_tracks_files(): default = TrackingSettings(labels="nuclei") - assert (default.tracked_label, default.tracks_stem) == ("nuclei_tracked", "tracks") + assert (default.tracked_label, default.tracks_stem) == ("tracks_nuclei", "tracks") custom = TrackingSettings(labels="nuclei", output_label="nuclei_sigma3") assert (custom.tracked_label, custom.tracks_stem) == ( "nuclei_sigma3", diff --git a/tests/test_example_settings.py b/tests/test_example_settings.py index 0f71e052..b0004544 100644 --- a/tests/test_example_settings.py +++ b/tests/test_example_settings.py @@ -214,11 +214,11 @@ def test_nextflow_a549_track_config_validates(): with open("nextflow/configs/a549/track.yml") as f: settings = TrackingSettings(**yaml.safe_load(f)) # Track reads the projected plate's `labels/nuclei` (segment.yml's model) and writes - # `labels/nuclei_tracked`; no image channel is processed. + # `labels/tracks_nuclei`; no image channel is processed. assert settings.output_mode == "2D" assert settings.segmentation_method == "labels" assert settings.labels == "nuclei" - assert settings.tracked_label == "nuclei_tracked" + assert settings.tracked_label == "tracks_nuclei" assert settings.input_images == [] assert settings.min_sibling_length == 3 diff --git a/tests/test_segment_settings.py b/tests/test_segment_settings.py index 20a7166a..1e9e01e8 100644 --- a/tests/test_segment_settings.py +++ b/tests/test_segment_settings.py @@ -190,7 +190,7 @@ def test_nextflow_a549_segment_config_carries_the_validated_cellpose_settings(): settings = SegmentationSettings(**yaml.safe_load(f)) (name,) = settings.models model = settings.models[name] - assert name == "nuclei" # track.yml's `labels: nuclei`; QC counts labels/nuclei_tracked + assert name == "nuclei" # track.yml's `labels: nuclei`; QC counts labels/tracks_nuclei assert model.pretrained_model == "cpsam_v2" assert model.channels == ["nuclei_prediction"] assert model.eval_args == { From 3d70068d1c6c008237a3d9515e8ca111b1c4532c Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 21 Oct 2026 15:49:36 -0700 Subject: [PATCH 41/54] fix(pyramid): size the task from the level-0 shard, measured on an A549 position iohub writes a pyramid level one target shard at a time and reads the source region feeding it, so RAM follows the shard, not the volume: one assembled A549 position (4.1 GiB level-0 shards, levels sized to the same bytes) peaked at 44 GB and took 41.5 min against the previous 16 GB / 30 min request. The estimate now asks 12x the shard bytes (52 GB there, the 16 GB floor on the projected plate) and 0.15 min per (t, c) volume (70 min). Co-Authored-By: Claude Fable 5.1 --- biahub/pyramid.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/biahub/pyramid.py b/biahub/pyramid.py index 7925986c..c79ab242 100644 --- a/biahub/pyramid.py +++ b/biahub/pyramid.py @@ -8,9 +8,12 @@ position; this module fans it out and records what it did. """ +import math + from pathlib import Path import click +import numpy as np import submitit from iohub.ngff import open_ome_zarr @@ -121,16 +124,23 @@ def pyramid( with open_ome_zarr(input_position_dirpaths[0], mode="r") as position: T, C, Z, Y, X = position.data.shape dtype = position.data.dtype - - # iohub downsamples one target shard at a time from the matching source region - # (16 timepoints by the full frame on a projected plate), single-threaded, so the - # working set is a few volumes, not the position. Time per (t, c) volume is well - # under the 0.05 min requested here. + shards = position.data.shards or position.data.chunks + + # RAM follows the SHARD, not the volume: iohub writes a pyramid level one target + # shard at a time and reads the source region feeding it, 8x the target in 3D. + # Measured on one A549 position (67 x 6 x 86 x 1664 x 1193 float32, level-0 + # shards 5 x 1 x 96 x 1792 x 1280 = 4.1 GiB, levels sized to the same bytes): + # 44 GB peak RSS and 41.5 min single-threaded, i.e. ~10x the shard and ~0.1 min + # per (t, c) volume; the projected plate's 0.14 GiB shards peaked at 1.7 GB in 27 s. + # 12x the shard and 0.15 min/volume leave a margin on both. + itemsize = np.dtype(dtype).itemsize + shard_gb = math.prod(shards) * itemsize / 2**30 + volume_gb = Z * Y * X * itemsize / 2**30 time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( shape=(T, C, Z, Y, X), dtype=dtype, - ram_multiplier=4, - time_multiplier=0.05, + ram_multiplier=max(4.0, 12 * shard_gb / volume_gb / 4), # per CPU, 4 CPUs share it + time_multiplier=0.15, max_num_cpus=4, ) mem_gb = num_cpus * gb_ram_per_cpu From f1966b2c3a6d49f233d5b9b29d11ed14d4f5ab6f Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Thu, 22 Oct 2026 11:40:14 -0700 Subject: [PATCH 42/54] fix(pyramid): double a retried task's memory after an OOM kill run_pyramid requested the init estimate verbatim on every attempt, so an estimate that was too low failed all five attempts of every position with the same 137 and aborted the run. The v9 debug run of 2026_08_11_A549_SEC61B_DENV asked 16 GB 25 times and was OOM-killed 25 times. Use retry_memory like the QC and virtual-stain processes: the next attempt doubles after a 137 and keeps the same request after a preemption. Co-Authored-By: Claude Opus 5.5 --- nextflow/modules/pyramid.nf | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/nextflow/modules/pyramid.nf b/nextflow/modules/pyramid.nf index ee10bdbf..dea8c962 100644 --- a/nextflow/modules/pyramid.nf +++ b/nextflow/modules/pyramid.nf @@ -22,7 +22,7 @@ // work in-process; Nextflow owns the per-position fan-out. See modules/project.nf, // the template this follows. -include { parse_resources; slurm_logs; slurm_log_dir; retry_time } from './common' +include { parse_resources; slurm_logs; slurm_log_dir; retry_time; retry_memory } from './common' process init_pyramid { @@ -52,7 +52,11 @@ process run_pyramid { label 'cpu' clusterOptions { slurm_logs('pyramid') } cpus { meta.cpus } - memory { "${meta.mem_gb} GB" } + // retry_memory, not a fixed request: an OOM kill (137) doubles the next + // attempt's memory. With the bare init estimate every retry asked for the same + // amount, so one low estimate failed all five attempts on every position and + // aborted the run (v9 of 2026_08_11_A549_SEC61B_DENV: 16 GB, 137 each time). + memory { retry_memory(meta.mem_gb, task) } time { retry_time(meta.time_minutes, task) } input: From a84a48004fb2b01bdbfea47c1abef681b6a7e686 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Thu, 22 Oct 2026 11:51:14 -0700 Subject: [PATCH 43/54] fix(project): double a retried task's memory after an OOM kill run_project asked for its init estimate verbatim on every retry, so a low estimate failed all attempts with the same 137. Use retry_memory like the other SLURM steps (fix/nextflow-retry-memory-cap carries the 244 GB cap). Co-Authored-By: Claude Opus 5.5 --- nextflow/modules/project.nf | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nextflow/modules/project.nf b/nextflow/modules/project.nf index c1994ff5..6730801b 100644 --- a/nextflow/modules/project.nf +++ b/nextflow/modules/project.nf @@ -16,7 +16,7 @@ // scheduling, so the CLI must NOT submit its own SLURM jobs. See // modules/flat_field.nf, the template this follows. -include { parse_resources; slurm_logs; slurm_log_dir; slurm_output_readme; retry_time } from './common' +include { parse_resources; slurm_logs; slurm_log_dir; slurm_output_readme; retry_time; retry_memory } from './common' process init_project { @@ -48,7 +48,7 @@ process run_project { label 'cpu' clusterOptions { slurm_logs('project') } cpus { meta.cpus } - memory { "${meta.mem_gb} GB" } + memory { retry_memory(meta.mem_gb, task) } time { retry_time(meta.time_minutes, task) } input: From 466d72766b51da52180d2d45fee42846324162a7 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 13:43:44 -0700 Subject: [PATCH 44/54] fix(pyramid): request 14x the shard and 0.3 min per volume, from v9's 54 positions The v9 debug run of 2026_08_11_A549_SEC61B_DENV measured every position of the assembled plate (level-0 shards 4.51 GiB): 46.5 GB peak RSS each (10.3x the shard) and up to 67.3 min, the slowest on the 120 GB cpu-c nodes. The old 12x / 0.15 asked 56 GB / 70 min, 4% above the slowest position. 14x the shard and 0.3 min per volume ask 64 GB / 130 min there; the projected plate keeps the 16 GB floor. Co-Authored-By: Claude Opus 5.5 --- biahub/pyramid.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/biahub/pyramid.py b/biahub/pyramid.py index c79ab242..31b814eb 100644 --- a/biahub/pyramid.py +++ b/biahub/pyramid.py @@ -132,15 +132,20 @@ def pyramid( # shards 5 x 1 x 96 x 1792 x 1280 = 4.1 GiB, levels sized to the same bytes): # 44 GB peak RSS and 41.5 min single-threaded, i.e. ~10x the shard and ~0.1 min # per (t, c) volume; the projected plate's 0.14 GiB shards peaked at 1.7 GB in 27 s. - # 12x the shard and 0.15 min/volume leave a margin on both. + # All 54 positions of that plate in the v9 debug run (shards 6 x 1 x 88 x 1792 x + # 1280 = 4.51 GiB): 46.5 GB peak RSS on every position (10.3x the shard), and + # 45.6 / 54.5 / 67.3 min min / median / max, i.e. up to 0.17 min per volume, the + # slowest on the 120 GB cpu-c nodes. The old 12x / 0.15 asked 56 GB / 70 min, + # 4% above the slowest position. 14x the shard (~36% over the peak) and + # 0.3 min/volume (~1.8x the slowest) leave room for slower nodes. itemsize = np.dtype(dtype).itemsize shard_gb = math.prod(shards) * itemsize / 2**30 volume_gb = Z * Y * X * itemsize / 2**30 time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( shape=(T, C, Z, Y, X), dtype=dtype, - ram_multiplier=max(4.0, 12 * shard_gb / volume_gb / 4), # per CPU, 4 CPUs share it - time_multiplier=0.15, + ram_multiplier=max(4.0, 14 * shard_gb / volume_gb / 4), # per CPU, 4 CPUs share it + time_multiplier=0.3, max_num_cpus=4, ) mem_gb = num_cpus * gb_ram_per_cpu From c407462e3c550aa62e9edc41b869f3a3c3d13d1e Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 13:59:00 -0700 Subject: [PATCH 45/54] fix(project): max keeps the input dtype; only mean/sum make the plate float32 max picks an existing voxel, so promoting a max-only plate to float32 doubled its storage for nothing and would truncate uint32 labels above 2^24. A store has one dtype, so a plate with any mean/sum channel is still float32 throughout and its max/slice channels are widened on write. Co-Authored-By: Claude Fable 5.1 --- biahub/project.py | 8 +++++--- settings/example_project_settings.yml | 4 ++-- tests/test_cli/test_project_cli.py | 2 +- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/biahub/project.py b/biahub/project.py index be1989af..977596fb 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -49,9 +49,11 @@ PROVENANCE_KEY = "biahub-project" -# Methods that compute a new value; a plate with any of them is float32 (the common -# case is float32 in and out). Only an all-`slice` plate keeps the input dtype. -_FLOAT_METHODS = ("max", "mean", "sum") +# Methods whose result is not one of the input values, so they need float32. `max` and +# `slice` pick an existing voxel and keep the input dtype. A store has one dtype, so a +# plate with any mean/sum channel is float32 throughout and its max/slice channels are +# widened on write (lossless for the integer camera dtypes). +_FLOAT_METHODS = ("mean", "sum") def resolve_z_indices( diff --git a/settings/example_project_settings.yml b/settings/example_project_settings.yml index c320ccf0..1d353a49 100644 --- a/settings/example_project_settings.yml +++ b/settings/example_project_settings.yml @@ -3,8 +3,8 @@ # (the command prints which). # # method: max | mean | sum | slice -# max/mean/sum reduce the selected planes; slice keeps one plane. The output plate is -# float32 unless every channel is a slice, in which case it keeps the input dtype. +# max/mean/sum reduce the selected planes; slice keeps one plane. The output plate +# keeps the input dtype unless a channel uses mean or sum, which make it float32. # z: which planes feed the projection. Omit it for max/mean/sum over the whole stack. # start, stop: absolute planes, half-open [start, stop) like time_indices in concatenate # center: a plane index, or "focus" for the in-focus plane found on focus.channel. diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py index 2db37c06..2ce68724 100644 --- a/tests/test_cli/test_project_cli.py +++ b/tests/test_cli/test_project_cli.py @@ -66,7 +66,7 @@ def test_project_init_scaffolds_a_z1_plate_and_names_dropped_channels(plate, tmp data, channel_names, zattrs = _read(out / "A" / "1" / "0") assert channel_names == ["RFP", "GFP"] # config order, not input order assert data.shape == (3, 2, 1, 5, 6) - assert data.dtype == np.float32 # any max/mean/sum channel makes the plate float32 + assert data.dtype == np.uint16 # max and slice keep the input dtype assert zattrs["biahub-project"]["channels"]["GFP"]["z"]["center"] == 2 assert np.all(data == 0) # --init projects nothing From 48e95a2f1e5824b3478e8e40a600bd51ed2f8881 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 15:41:41 -0700 Subject: [PATCH 46/54] build(deps): pin iohub to main at the #491 merge (80b20a6) One pin for the whole chain: main now carries #489 (output_label, create_empty_labels), #488, #461 and #491, so the 16cdc4f branch pin and the dev/pyramid-shards combination branch retire. Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 9 +++++---- uv.lock | 8 ++++---- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a97af898..a2d31811 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,7 +102,7 @@ index-strategy = "unsafe-best-match" # (d1f9e89, which adds #479). uv refuses two git URLs for one package, so # force the whole graph onto biahub's rev. Drop together with the iohub pin. override-dependencies = [ - "iohub @ git+https://github.com/czbiohub-sf/iohub@d1f9e89294529d365719cf87a9b29622ca77a992", + "iohub @ git+https://github.com/czbiohub-sf/iohub@80b20a693d7dd11d5a14d667001bcf1972a2be23", ] [tool.uv.sources] @@ -141,13 +141,14 @@ imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipelin # royerlab/ultrack main once #290 merges, and at a release once one contains them. ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5f84ccda3bdf072f5289f939c" } -# Pin iohub to the czbiohub-sf/iohub#489 branch (on main at the #486 merge): +# Pin iohub to main at the czbiohub-sf/iohub#491 merge, which carries #489: # process_single_position(output_label=...), with which `biahub segment` writes its # labels into the plate's label images with --resume, and create_empty_labels, with # which segment --init and track --init scaffold those label images. Main's #486 adds # create_empty_plate(label_names=...) (#360), and #485 puts the labels list on # /labels/zarr.json so ome-zarr-models and dca-validate can open a field that -# has labels. Repoint at main once #489 merges, and at a PyPI release once one +# has labels. #488/#461/#491 (label pyramids, shard size targets, resumable +# pyramids) ride along for the pyramid step. Drop to a PyPI release once one # contains them. # # Main's own reasons: #476 syncs a write unit's shards before recording it as @@ -156,7 +157,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 # creates a plate's positions in one batch instead of rewriting the well's # metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "16cdc4f" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "80b20a693d7dd11d5a14d667001bcf1972a2be23" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index ce820bc0..a294cf29 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=16cdc4f" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=80b20a693d7dd11d5a14d667001bcf1972a2be23" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?tag=v0.6.0" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=16cdc4f" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=80b20a693d7dd11d5a14d667001bcf1972a2be23" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev20+16cdc4f" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=16cdc4f#16cdc4f6d42fba4476833f1501449c3548226d7d" } +version = "0.3.12.dev27+80b20a6" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=80b20a693d7dd11d5a14d667001bcf1972a2be23#80b20a693d7dd11d5a14d667001bcf1972a2be23" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From d3cb4699c6223516b5da486f07c45f3441281539 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 16:24:07 -0700 Subject: [PATCH 47/54] feat(segment,track)!: label images are (T, 1, Z, Y, X) A label image now has the image's axes with one channel rather than TZYX: the DCA array standard wants every array 5D, and OME-NGFF's own layout note reads a label's dimensions against the image's. iohub creates and writes them that way (czbiohub-sf/iohub, label channel axis); here the DCA label chunks become (16, 1, 1, 128, 128), segment's shape check and track's label reads and writes index the channel axis, and the tests and skill notes follow. Co-Authored-By: Claude Fable 5.1 --- .../skills/reconstruct-with-nextflow/SKILL.md | 2 +- .../references/caveats.md | 2 +- biahub/segment.py | 10 +++--- biahub/track.py | 29 +++++++++-------- biahub/utils/ngff.py | 18 ++++++----- tests/test_cli/test_segment_cli.py | 4 +-- tests/test_cli/test_track_cli.py | 32 ++++++++++--------- 7 files changed, 52 insertions(+), 45 deletions(-) diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index c5b88d0c..0a96ad72 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -651,7 +651,7 @@ directory (`_rerun`) over relaunching in place. and tracking ran, their output is inside the projected plate: check a position has `labels/nuclei` (the segment config's model name) and `labels/tracks_nuclei` (`tracks_` plus the track config's `labels`; both - uint32 TZYX at the image's full Z), `tracks.geff` and + uint32 (T, 1, Z, Y, X) at the image's full Z), `tracks.geff` and `tracks___.csv`, and report them beside the image data. 3. If QC ran, report its verdict: the `QC_SUMMARY` line from the pipeline log (`pass=`/`fail=`/`gates_fail=`), and point at the report at diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 193f5673..16312670 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -244,7 +244,7 @@ transfer. So simply do not run it: ## 5. Track reads the assembled plate, not the intermediates `track_wf` takes the assembled `.zarr` for *both* inputs, and writes -its output back INTO it: per position, `labels/` (uint32, TZYX +its output back INTO it: per position, `labels/` (uint32, (T, 1, Z, Y, X) at the image's full Z; 2D-tracking masks sit in the middle plane of the tracked z-window), `tracks.geff`, and `tracks___.csv`. Nothing else is kept: the pipeline passes no `-o`, so each worker builds its Ultrack database in diff --git a/biahub/segment.py b/biahub/segment.py index 2308ab47..0a7bc780 100644 --- a/biahub/segment.py +++ b/biahub/segment.py @@ -179,7 +179,7 @@ def _segment_czyx( def _check_existing_label( - position, name: str, model_settings: dict, shape_tzyx: tuple + position, name: str, model_settings: dict, shape_tczyx: tuple ) -> None: """Refuse to write model ``name`` into a label image made with other settings. @@ -192,10 +192,10 @@ def _check_existing_label( return label = position.get_label(name) existing_shape = tuple(label["0"].shape) - if existing_shape != shape_tzyx: + if existing_shape != shape_tczyx: raise ValueError( f"{position.zgroup.path} already has a label image {name!r} of shape " - f"{existing_shape}, but the image is {shape_tzyx} (T, Z, Y, X). Rename the " + f"{existing_shape}, but a label of this image is {shape_tczyx} (T, 1, Z, Y, X). Rename the " "model or remove the old label image." ) recorded = dict(label.zattrs).get(PROVENANCE_KEY, {}).get("models", {}).get(name) @@ -235,7 +235,9 @@ def _init_label_images( for position_dirpath in input_position_dirpaths: with open_ome_zarr(str(position_dirpath), mode="r") as position: for name in label_names: - _check_existing_label(position, name, recorded["models"][name], (T, Z, Y, X)) + _check_existing_label( + position, name, recorded["models"][name], (T, 1, Z, Y, X) + ) _, layout = label_layout(input_position_dirpaths[0]) create_empty_labels( diff --git a/biahub/track.py b/biahub/track.py index 417a1334..8a3de596 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -605,11 +605,12 @@ def load_labels( ) -> np.ndarray: """Read the instance labels ``labels/`` of one position. - Returns (T, Y, X) labels for 2D tracking (the single plane of a projected plate) - and (T, Z, Y, X) for 3D. Frames listed as blank for this FOV in - ``blank_frames_path`` are filled with the nearest non-blank frame's labels, as - the images used to be before segmentation moved out of track, so a dropped - timepoint does not break every track. + The label image is (T, 1, Z, Y, X), the image's axes with one channel. Returns + (T, Y, X) labels for 2D tracking (the single plane of a projected plate) and + (T, Z, Y, X) for 3D. Frames listed as blank for this FOV in ``blank_frames_path`` + are filled with the nearest non-blank frame's labels, as the images used to be + before segmentation moved out of track, so a dropped timepoint does not break + every track. """ with open_ome_zarr(str(position_dirpath), mode="r") as position: if label_name not in position.label_names(): @@ -617,7 +618,7 @@ def load_labels( f"{position_dirpath} has no label image {label_name!r} (labels: " f"{position.label_names()}); segment the plate with `biahub segment` first." ) - labels = np.asarray(position.get_label(label_name)["0"][:]) + labels = np.asarray(position.get_label(label_name)["0"][:, 0]) # (T, Z, Y, X) if output_mode == "2D": labels = labels[:, 0] if blank_frames_path: @@ -634,11 +635,11 @@ def load_labels( def write_tracking_labels(label_array, labels: ArrayLike, output_mode: str) -> None: - """Write tracked labels into the (T, Z, Y, X) label array of the plate track read. + """Write tracked labels into the (T, 1, Z, Y, X) label array of the plate track read. - The label image has the plate's Z: one plane for a projected plate, where 2D - tracking writes its (T, Y, X) labels; the full stack for 3D tracking, whose - (T, Z, Y, X) labels fill it. + The label image has the image's axes with one channel and the plate's Z: one + plane for a projected plate, where 2D tracking writes its (T, Y, X) labels; the + full stack for 3D tracking, whose (T, Z, Y, X) labels fill it. """ labels = np.asarray(labels, dtype=np.uint32) if output_mode == "2D": @@ -650,14 +651,14 @@ def write_tracking_labels(label_array, labels: ArrayLike, output_mode: str) -> N f"output_mode='2D' expects (T, Y, X) labels but tracking produced " f"shape {labels.shape}." ) - label_array[:, 0] = labels + label_array[:, 0, 0] = labels return - if labels.ndim != 4 or labels.shape[1] != label_array.shape[1]: + if labels.ndim != 4 or labels.shape[1] != label_array.shape[2]: raise ValueError( - f"output_mode='3D' expects (T, {label_array.shape[1]}, Y, X) labels but " + f"output_mode='3D' expects (T, {label_array.shape[2]}, Y, X) labels but " f"tracking produced shape {labels.shape}." ) - label_array[:] = labels + label_array[:, 0] = labels @contextmanager diff --git a/biahub/utils/ngff.py b/biahub/utils/ngff.py index 306860df..854e923d 100644 --- a/biahub/utils/ngff.py +++ b/biahub/utils/ngff.py @@ -133,22 +133,23 @@ def append_channels(input_data_path: Path, target_data_path: Path) -> None: appending_dataset.close() -# Label chunks, TZYX. 16 x 1 x 128 x 128 uint32 is 1 MB, the DCA v0.2 SHOULD (512 KB +# Label chunks, TCZYX: a label image has the image's axes with one channel, +# (T, 1, Z, Y, X). 16 x 1 x 1 x 128 x 128 uint32 is 1 MB, the DCA v0.2 SHOULD (512 KB # is the MUST) and its "2D+time" example; growing T keeps per-plane reads cheap. -LABEL_CHUNKS = (16, 1, 128, 128) +LABEL_CHUNKS = (16, 1, 1, 128, 128) # Shards span one time chunk (16 timepoints, the DCA SHOULD) and the whole YX plane, # capped at 16 chunks (2048 px, the DCA spatial-shard SHOULD). MAX_LABEL_SHARD_CHUNKS = 16 def label_chunks_and_shards( - shape: tuple[int, int, int, int], + shape: tuple[int, int, int, int, int], ) -> tuple[tuple[int, ...], tuple[int, ...]]: - """Chunks and shards ratio for a (T, Z, Y, X) label array of ``shape``.""" + """Chunks and shards ratio for a (T, 1, Z, Y, X) label array of ``shape``.""" chunks = tuple(min(c, s) for c, s in zip(LABEL_CHUNKS, shape, strict=True)) - shards_ratio = (1, 1) + tuple( + shards_ratio = (1, 1, 1) + tuple( min(-(-s // c), MAX_LABEL_SHARD_CHUNKS) - for c, s in zip(chunks[2:], shape[2:], strict=True) + for c, s in zip(chunks[3:], shape[3:], strict=True) ) return chunks, shards_ratio @@ -169,12 +170,13 @@ def label_layout( The layout is the ``label_chunks`` / ``label_shards_ratio`` keyword arguments of iohub's ``create_empty_labels`` per :func:`label_chunks_and_shards` (no shards for - an OME-Zarr 0.4 plate), for a uint32 label with the image's TZYX shape. + an OME-Zarr 0.4 plate), for a uint32 label of shape (T, 1, Z, Y, X): the image's + axes with one channel. """ with open_ome_zarr(str(position_dirpath), mode="r") as dataset: T, C, Z, Y, X = dataset.data.shape version = dataset.version - label_chunks, label_shards_ratio = label_chunks_and_shards((T, Z, Y, X)) + label_chunks, label_shards_ratio = label_chunks_and_shards((T, 1, Z, Y, X)) layout = { "label_dtype": np.uint32, "label_chunks": label_chunks, diff --git a/tests/test_cli/test_segment_cli.py b/tests/test_cli/test_segment_cli.py index aa181468..fd194750 100644 --- a/tests/test_cli/test_segment_cli.py +++ b/tests/test_cli/test_segment_cli.py @@ -87,7 +87,7 @@ def test_segment_init_creates_label_images_and_prints_resources( with open_ome_zarr(position, mode="r") as pos: assert pos.label_names() == ["nuc"] label = pos.get_label("nuc")["0"] - assert label.shape == (3, 1, 5, 6) # the image's (T, Z, Y, X) + assert label.shape == (3, 1, 1, 5, 6) # the image's (T, 1, Z, Y, X) assert label.dtype == np.uint32 assert not label[:].any() assert pos.channel_names == ["GFP", "RFP", "Phase3D"] # the image is untouched @@ -249,7 +249,7 @@ def test_segment_refuses_a_label_image_of_another_shape( stack = example_plate[0] / "B" / "1" / "0" # Z=4 assert _run("--init", "-i", stack, "-c", three_d).exit_code == 0 with open_ome_zarr(stack, mode="r+") as pos: - pos.get_label("nuc")["0"].resize((3, 2, 5, 6)) # pretend the image changed + pos.get_label("nuc")["0"].resize((3, 1, 2, 5, 6)) # pretend the image changed result = _run("--init", "-i", stack, "-c", three_d) diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 4dbd2037..883e4f31 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -82,7 +82,9 @@ def example_tracking_plate(tmp_path): with open_ome_zarr(str(position_path), mode="r+") as position: nuclei = np.asarray(position["0"][:, 0]) # (T, Z, Y, X) labels = np.stack([connected_components(frame > 0.5) for frame in nuclei]) - position.get_label(INPUT_LABEL)["0"][:] = labels.astype(np.uint32) + position.get_label(INPUT_LABEL)["0"][:] = labels[:, None].astype( + np.uint32 + ) # (T, 1, Z, Y, X) plate_dataset = open_ome_zarr(plate_path, mode="r") yield plate_path, plate_dataset @@ -117,7 +119,7 @@ def _assert_tracking_outputs(position_path): with open_ome_zarr(str(position_path), mode="r") as pos: labels = pos.get_label(LABEL)["0"][:] assert labels.dtype == np.uint32 - assert labels.shape[1] == 1 # the projected (Z=1) plate's Z + assert labels.shape[1:3] == (1, 1) # one channel; the projected (Z=1) plate's Z assert labels.any() geff_path = position_path / "tracks.geff" @@ -382,10 +384,10 @@ def test_track_cli_init_only(tmp_path, example_tracking_plate): assert pos.label_names() == [INPUT_LABEL, LABEL] assert pos.channel_names == ["nuclei_prediction", "membrane_prediction"] label = pos.get_label(LABEL) - assert label["0"].shape == (5, 1, 64, 64) # the plate's Z + assert label["0"].shape == (5, 1, 1, 64, 64) # one channel, the plate's Z assert label["0"].dtype == np.uint32 assert not label["0"][:].any() - assert [ax.name.lower() for ax in label.axes] == ["t", "z", "y", "x"] + assert [ax.name.lower() for ax in label.axes] == ["t", "c", "z", "y", "x"] label_attrs = yaml.safe_load( (plate_path / position / "labels" / LABEL / "zarr.json").read_text() )["attributes"] @@ -554,14 +556,14 @@ def test_init_labels_take_the_input_z(tmp_path, example_tracking_plate, create_c shape = _init_labels([position_path], _minimal_settings()) assert shape[2] == 1 with open_ome_zarr(str(position_path), mode="r") as pos: - assert pos.get_label(LABEL)["0"].shape == (5, 1, 64, 64) + assert pos.get_label(LABEL)["0"].shape == (5, 1, 1, 64, 64) stack_path, _ = create_custom_plate(tmp_path, z_size=3) position_path = stack_path / "A" / "1" / "0" shape = _init_labels([position_path], _minimal_settings(output_mode="3D")) assert shape[2] == 3 with open_ome_zarr(str(position_path), mode="r") as pos: - assert pos.get_label(LABEL)["0"].shape == (3, 3, 5, 6) + assert pos.get_label(LABEL)["0"].shape == (3, 1, 3, 5, 6) @pytest.mark.parametrize( @@ -585,26 +587,26 @@ def test_init_labels_leaves_existing_labels(tmp_path, example_tracking_plate): _init_labels([position_path], settings) with open_ome_zarr(str(position_path), mode="r+") as pos: - pos.get_label(LABEL)["0"][0, 0, 0, 0] = 7 + pos.get_label(LABEL)["0"][0, 0, 0, 0, 0] = 7 image = pos["0"][:] _init_labels([position_path], settings) with open_ome_zarr(str(position_path), mode="r") as pos: - assert pos.get_label(LABEL)["0"][0, 0, 0, 0] == 7 + assert pos.get_label(LABEL)["0"][0, 0, 0, 0, 0] == 7 assert np.array_equal(pos["0"][:], image) def test_label_chunks_and_shards(): # 1 MB uint32 chunks (DCA SHOULD); shards cover YX up to 16 chunks. - chunks, shards = label_chunks_and_shards((67, 86, 1664, 1193)) - assert chunks == (16, 1, 128, 128) - assert shards == (1, 1, 13, 10) - assert label_chunks_and_shards((67, 86, 4096, 300))[1] == (1, 1, 16, 3) + chunks, shards = label_chunks_and_shards((67, 1, 86, 1664, 1193)) + assert chunks == (16, 1, 1, 128, 128) + assert shards == (1, 1, 1, 13, 10) + assert label_chunks_and_shards((67, 1, 86, 4096, 300))[1] == (1, 1, 1, 16, 3) # Clamped to small arrays. - assert label_chunks_and_shards((5, 1, 64, 64))[0] == (5, 1, 64, 64) + assert label_chunks_and_shards((5, 1, 1, 64, 64))[0] == (5, 1, 1, 64, 64) def test_write_tracking_labels_2d_fills_the_single_plane(): - label_array = np.zeros((2, 1, 4, 4), dtype=np.uint32) + label_array = np.zeros((2, 1, 1, 4, 4), dtype=np.uint32) write_tracking_labels(label_array, np.ones((2, 4, 4), dtype=np.int64), "2D") assert label_array.all() @@ -616,7 +618,7 @@ def test_write_tracking_labels_2d_fills_the_single_plane(): def test_write_tracking_labels_3d_fills_the_stack(): - label_array = np.zeros((2, 6, 4, 4), dtype=np.uint32) + label_array = np.zeros((2, 1, 6, 4, 4), dtype=np.uint32) write_tracking_labels(label_array, np.ones((2, 6, 4, 4), dtype=np.int64), "3D") assert label_array.all() From ee6f1a17b665b40ce33b426fdcb4a4fcdb06fe41 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 16:29:05 -0700 Subject: [PATCH 48/54] build(deps): pin iohub to czbiohub-sf/iohub#496 (labels are (T, 1, Z, Y, X)) Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 28 ++++++++++++---------------- uv.lock | 8 ++++---- 2 files changed, 16 insertions(+), 20 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a2d31811..b3fe1207 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,23 +141,19 @@ imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipelin # royerlab/ultrack main once #290 merges, and at a release once one contains them. ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5f84ccda3bdf072f5289f939c" } -# Pin iohub to main at the czbiohub-sf/iohub#491 merge, which carries #489: -# process_single_position(output_label=...), with which `biahub segment` writes its -# labels into the plate's label images with --resume, and create_empty_labels, with -# which segment --init and track --init scaffold those label images. Main's #486 adds -# create_empty_plate(label_names=...) (#360), and #485 puts the labels list on -# /labels/zarr.json so ome-zarr-models and dca-validate can open a field that -# has labels. #488/#461/#491 (label pyramids, shard size targets, resumable -# pyramids) ride along for the pyramid step. Drop to a PyPI release once one -# contains them. +# Pin iohub to the czbiohub-sf/iohub#496 branch (on main at the #491 merge, 80b20a6): +# label images are (T, 1, Z, Y, X), the image's axes with one channel (the DCA array +# standard wants every array 5D), and a label with more than one channel is refused. +# Main carries #489 (process_single_position(output_label=...), with which `biahub +# segment` writes its labels into the plate's label images with --resume, and +# create_empty_labels, with which segment --init and track --init scaffold those +# label images), #486 (create_empty_plate(label_names=...), #360), #485 (the labels +# list on /labels/zarr.json, so ome-zarr-models and dca-validate can open a +# field that has labels) and #488/#461/#491 (label pyramids, shard size targets, +# resumable pyramids) for the pyramid step. Repoint at main once #496 merges, and at +# a PyPI release once one contains them. # -# Main's own reasons: #476 syncs a write unit's shards before recording it as -# done (without it a preemption between the write and the marker leaves a `.done` -# beside a shard whose bytes never landed, which silently cost several runs on -# 2026_08_14_dynatrack); #472 adds create_empty_plate(extra_metadata=...); #479 -# creates a plate's positions in one batch instead of rewriting the well's -# metadata per position (hours for a ~40k-FOV plate's init). -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "80b20a693d7dd11d5a14d667001bcf1972a2be23" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index a294cf29..c406d774 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=80b20a693d7dd11d5a14d667001bcf1972a2be23" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?tag=v0.6.0" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=80b20a693d7dd11d5a14d667001bcf1972a2be23" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev27+80b20a6" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=80b20a693d7dd11d5a14d667001bcf1972a2be23#80b20a693d7dd11d5a14d667001bcf1972a2be23" } +version = "0.3.12.dev29+c8a5ab6" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef#c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From ca5fdee8597bb9755bd19cd513e42a0fdc03b754 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 16:52:05 -0700 Subject: [PATCH 49/54] build(deps): iohub pin follows czbiohub-sf/iohub#496's review fixes (9f21979) Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 2 +- uv.lock | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b3fe1207..34c96cfe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -153,7 +153,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # resumable pyramids) for the pyramid step. Repoint at main once #496 merges, and at # a PyPI release once one contains them. # -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "9f21979" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index c406d774..8446c58a 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=9f21979" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?tag=v0.6.0" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=9f21979" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev29+c8a5ab6" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef#c8a5ab60d90bae6db65e4d47fe4f4c91c29676ef" } +version = "0.3.12.dev30+9f21979" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=9f21979#9f2197975110e210069aa8fd85ebf7ab813cf355" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From fe09c0d70bf8d667930ecb7250fde3866096db0b Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 17:02:47 -0700 Subject: [PATCH 50/54] build(deps): iohub pin follows czbiohub-sf/iohub#496 (eee7567) Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 2 +- uv.lock | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 34c96cfe..8628a6a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -153,7 +153,7 @@ ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5 # resumable pyramids) for the pyramid step. Repoint at main once #496 merges, and at # a PyPI release once one contains them. # -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "9f21979" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "eee7567" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index 8446c58a..2d26e7b0 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=9f21979" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=eee7567" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?tag=v0.6.0" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=9f21979" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=eee7567" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev30+9f21979" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=9f21979#9f2197975110e210069aa8fd85ebf7ab813cf355" } +version = "0.3.12.dev31+eee7567" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=eee7567#eee75676bda4e6bbafd1d64e98f56d33d96332be" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From d570380e5c69f68bf73811863d965f9ef5560dda Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 17:10:40 -0700 Subject: [PATCH 51/54] build(deps): pin iohub to main at the #496 merge (0cb3e35) Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 14 +++++++------- uv.lock | 8 ++++---- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8628a6a5..5aa3130c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,19 +141,19 @@ imaging_qc_pipeline = { git = "https://github.com/czbiohub-sf/imaging-qc-pipelin # royerlab/ultrack main once #290 merges, and at a release once one contains them. ultrack = { git = "https://github.com/ieivanov/ultrack", rev = "69e7589a050b0aa5f84ccda3bdf072f5289f939c" } -# Pin iohub to the czbiohub-sf/iohub#496 branch (on main at the #491 merge, 80b20a6): -# label images are (T, 1, Z, Y, X), the image's axes with one channel (the DCA array -# standard wants every array 5D), and a label with more than one channel is refused. -# Main carries #489 (process_single_position(output_label=...), with which `biahub +# Pin iohub to main at the czbiohub-sf/iohub#496 merge (0cb3e35): label images are +# (T, 1, Z, Y, X), the image's axes with one channel (the DCA array standard wants +# every array 5D), and a label with more than one channel is refused. Main also +# carries #489 (process_single_position(output_label=...), with which `biahub # segment` writes its labels into the plate's label images with --resume, and # create_empty_labels, with which segment --init and track --init scaffold those # label images), #486 (create_empty_plate(label_names=...), #360), #485 (the labels # list on /labels/zarr.json, so ome-zarr-models and dca-validate can open a # field that has labels) and #488/#461/#491 (label pyramids, shard size targets, -# resumable pyramids) for the pyramid step. Repoint at main once #496 merges, and at -# a PyPI release once one contains them. +# resumable pyramids) for the pyramid step. Drop to a PyPI release once one +# contains them. # -iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "eee7567" } +iohub = { git = "https://github.com/czbiohub-sf/iohub", rev = "0cb3e35b860b6f774e7125e41d457192911a178e" } # Pin waveorder to the main commit that merged mehta-lab/waveorder#583, which # runs apply-inv-tf through iohub's process_single_position with resume (on top diff --git a/uv.lock b/uv.lock index 2d26e7b0..8b41a9c5 100644 --- a/uv.lock +++ b/uv.lock @@ -21,7 +21,7 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=eee7567" }] +overrides = [{ name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=0cb3e35b860b6f774e7125e41d457192911a178e" }] [[package]] name = "absl-py" @@ -423,7 +423,7 @@ requires-dist = [ { name = "humanize" }, { name = "imageio-ffmpeg" }, { name = "imaging-qc-pipeline", extras = ["interactive"], marker = "extra == 'qc'", git = "https://github.com/czbiohub-sf/imaging-qc-pipeline?tag=v0.6.0" }, - { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=eee7567" }, + { name = "iohub", git = "https://github.com/czbiohub-sf/iohub?rev=0cb3e35b860b6f774e7125e41d457192911a178e" }, { name = "largestinteriorrectangle" }, { name = "llvmlite", specifier = ">=0.41.0" }, { name = "markdown" }, @@ -2153,8 +2153,8 @@ wheels = [ [[package]] name = "iohub" -version = "0.3.12.dev31+eee7567" -source = { git = "https://github.com/czbiohub-sf/iohub?rev=eee7567#eee75676bda4e6bbafd1d64e98f56d33d96332be" } +version = "0.3.12.dev28+0cb3e35" +source = { git = "https://github.com/czbiohub-sf/iohub?rev=0cb3e35b860b6f774e7125e41d457192911a178e#0cb3e35b860b6f774e7125e41d457192911a178e" } dependencies = [ { name = "blosc2" }, { name = "dask", extra = ["array"] }, From 0bc9f45890a5243947ea6445afc852f8b8c28268 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 17:21:34 -0700 Subject: [PATCH 52/54] refactor(labels): size label shards with iohub's 200 MB target `label_layout` passes iohub `label_shards="200MB"` (czbiohub-sf/iohub#461: the largest chunk-aligned shard under the target, filled X, Y, Z then T, C never grown) instead of a computed `label_shards_ratio`, which iohub has deprecated. On the A549 projected plate the result is the same geometry as before, one 16-timepoint chunk by the full frame; #430's segment labels use the same constant, so the two branches agree. Co-Authored-By: Claude Fable 5.1 --- biahub/utils/ngff.py | 38 ++++++++++++++------------------ tests/test_cli/test_track_cli.py | 35 +++++++++++++++++++++-------- 2 files changed, 42 insertions(+), 31 deletions(-) diff --git a/biahub/utils/ngff.py b/biahub/utils/ngff.py index 854e923d..038f11bf 100644 --- a/biahub/utils/ngff.py +++ b/biahub/utils/ngff.py @@ -137,21 +137,16 @@ def append_channels(input_data_path: Path, target_data_path: Path) -> None: # (T, 1, Z, Y, X). 16 x 1 x 1 x 128 x 128 uint32 is 1 MB, the DCA v0.2 SHOULD (512 KB # is the MUST) and its "2D+time" example; growing T keeps per-plane reads cheap. LABEL_CHUNKS = (16, 1, 1, 128, 128) -# Shards span one time chunk (16 timepoints, the DCA SHOULD) and the whole YX plane, -# capped at 16 chunks (2048 px, the DCA spatial-shard SHOULD). -MAX_LABEL_SHARD_CHUNKS = 16 - - -def label_chunks_and_shards( - shape: tuple[int, int, int, int, int], -) -> tuple[tuple[int, ...], tuple[int, ...]]: - """Chunks and shards ratio for a (T, 1, Z, Y, X) label array of ``shape``.""" - chunks = tuple(min(c, s) for c, s in zip(LABEL_CHUNKS, shape, strict=True)) - shards_ratio = (1, 1, 1) + tuple( - min(-(-s // c), MAX_LABEL_SHARD_CHUNKS) - for c, s in zip(chunks[3:], shape[3:], strict=True) - ) - return chunks, shards_ratio +# Label shards as a size target for iohub's sizer (czbiohub-sf/iohub#461): the largest +# chunk-aligned shard under 200 MB, filled X, Y, Z first and then T, C never grown. On +# the A549 projected plate (1664 x 1193 uint32) that is one 16-timepoint chunk by the +# full frame, 127 MB: the DCA 16-timepoint SHOULD, well under its 5 GB cap. +LABEL_SHARDS = "200MB" + + +def label_chunks(shape: tuple[int, int, int, int, int]) -> tuple[int, ...]: + """``LABEL_CHUNKS`` clamped to a (T, 1, Z, Y, X) label array of ``shape``.""" + return tuple(min(c, s) for c, s in zip(LABEL_CHUNKS, shape, strict=True)) def labels_exist(position_dirpaths: list[Path], label_name: str) -> bool: @@ -168,18 +163,17 @@ def label_layout( ) -> tuple[tuple[int, int, int, int, int], dict[str, Any]]: """Return the image's (T, C, Z, Y, X) shape and the label-image layout for its plate. - The layout is the ``label_chunks`` / ``label_shards_ratio`` keyword arguments of - iohub's ``create_empty_labels`` per :func:`label_chunks_and_shards` (no shards for - an OME-Zarr 0.4 plate), for a uint32 label of shape (T, 1, Z, Y, X): the image's - axes with one channel. + The layout is the ``label_dtype`` / ``label_chunks`` / ``label_shards`` keyword + arguments of iohub's ``create_empty_labels``: uint32, :data:`LABEL_CHUNKS` and the + :data:`LABEL_SHARDS` size target (no shards on an OME-Zarr 0.4 plate), for a label + of shape (T, 1, Z, Y, X): the image's axes with one channel. """ with open_ome_zarr(str(position_dirpath), mode="r") as dataset: T, C, Z, Y, X = dataset.data.shape version = dataset.version - label_chunks, label_shards_ratio = label_chunks_and_shards((T, 1, Z, Y, X)) layout = { "label_dtype": np.uint32, - "label_chunks": label_chunks, - "label_shards_ratio": label_shards_ratio if version == "0.5" else None, + "label_chunks": label_chunks((T, 1, Z, Y, X)), + "label_shards": LABEL_SHARDS if version == "0.5" else None, } return (T, C, Z, Y, X), layout diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 883e4f31..7d67ca84 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -15,7 +15,7 @@ from biahub.cli.main import cli from biahub.settings import TrackingSettings from biahub.track import _init_labels, track, write_tracking_labels -from biahub.utils.ngff import label_chunks_and_shards +from biahub.utils.ngff import LABEL_SHARDS, label_chunks, label_layout # The label image `biahub segment` would write (its model name) and the one track writes. INPUT_LABEL = "nuclei" @@ -595,14 +595,31 @@ def test_init_labels_leaves_existing_labels(tmp_path, example_tracking_plate): assert np.array_equal(pos["0"][:], image) -def test_label_chunks_and_shards(): - # 1 MB uint32 chunks (DCA SHOULD); shards cover YX up to 16 chunks. - chunks, shards = label_chunks_and_shards((67, 1, 86, 1664, 1193)) - assert chunks == (16, 1, 1, 128, 128) - assert shards == (1, 1, 1, 13, 10) - assert label_chunks_and_shards((67, 1, 86, 4096, 300))[1] == (1, 1, 1, 16, 3) - # Clamped to small arrays. - assert label_chunks_and_shards((5, 1, 1, 64, 64))[0] == (5, 1, 1, 64, 64) +def test_label_layout(tmp_path, example_tracking_plate): + """1 MB uint32 chunks (DCA SHOULD) and a 200 MB shard target, as iohub's sizer takes it.""" + assert label_chunks((67, 1, 86, 1664, 1193)) == (16, 1, 1, 128, 128) + assert label_chunks((5, 1, 1, 64, 64)) == (5, 1, 1, 64, 64) # clamped to small arrays + + plate_path, _ = example_tracking_plate + shape, layout = label_layout(plate_path / "A" / "1" / "0") + assert shape == (5, 2, 1, 64, 64) + assert layout == { + "label_dtype": np.uint32, + "label_chunks": (5, 1, 1, 64, 64), + "label_shards": LABEL_SHARDS, + } + # What the sizer makes of it on the A549 projected plate: the 16-timepoint chunk by + # the full frame, rounded up to whole 128 px chunks (136 MB); a 32-timepoint shard + # would exceed 200 MB. + from iohub.core.sharding import resolve_shards + + assert resolve_shards( + LABEL_SHARDS, + shape=(67, 1, 1, 1664, 1193), + chunks=label_chunks((67, 1, 1, 1664, 1193)), + dtype=np.uint32, + dimension_names=["t", "c", "z", "y", "x"], + ) == (16, 1, 1, 1664, 1280) def test_write_tracking_labels_2d_fills_the_single_plane(): From 6d9d19ae3e9e2d04e2d6a18669ad5d91ac62aabe Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 17:25:59 -0700 Subject: [PATCH 53/54] build(deps): the iohub override follows the pin (0cb3e35) Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 5aa3130c..a3f260ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,7 +102,7 @@ index-strategy = "unsafe-best-match" # (d1f9e89, which adds #479). uv refuses two git URLs for one package, so # force the whole graph onto biahub's rev. Drop together with the iohub pin. override-dependencies = [ - "iohub @ git+https://github.com/czbiohub-sf/iohub@80b20a693d7dd11d5a14d667001bcf1972a2be23", + "iohub @ git+https://github.com/czbiohub-sf/iohub@0cb3e35b860b6f774e7125e41d457192911a178e", ] [tool.uv.sources] From d2b0992dcac7e1d001aa1393639bc034f7b5b243 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Fri, 23 Oct 2026 17:34:33 -0700 Subject: [PATCH 54/54] style: ruff format Co-Authored-By: Claude Fable 5.1 --- tests/test_cli/test_pyramid_cli.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_cli/test_pyramid_cli.py b/tests/test_cli/test_pyramid_cli.py index c659e4f9..02c518bb 100644 --- a/tests/test_cli/test_pyramid_cli.py +++ b/tests/test_cli/test_pyramid_cli.py @@ -61,7 +61,11 @@ def test_pyramid_writes_image_and_label_levels(plate): "1": (3, 3, 2, 8, 6), "2": (3, 3, 1, 4, 3), } - assert label_levels == {"0": (3, 1, 4, 16, 12), "1": (3, 1, 2, 8, 6), "2": (3, 1, 1, 4, 3)} + assert label_levels == { + "0": (3, 1, 4, 16, 12), + "1": (3, 1, 2, 8, 6), + "2": (3, 1, 1, 4, 3), + } assert record == { "levels": 3, "method": "mean",