From bb5812ecedf698951d229c5eb6a5350152bff161 Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Wed, 24 Jun 2026 08:57:13 +0200 Subject: [PATCH 1/3] =?UTF-8?q?feat:=20=E2=9C=A8=20Snakemake=20workflow=20?= =?UTF-8?q?with=20per-tile=20GPU=20SLURM=20jobs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add workflow/ (Snakefile + scripts + SLURM profile + config) that runs the full pipeline — convert → prepare (checkpoint) → segment → merge — and spreads the expensive Cellpose step across many GPUs by submitting one GPU SLURM job per (non-empty) tile. Each job writes a disjoint chunk of a shared stage store; a final CPU job stitches labels across tile boundaries and writes them into the image as a calibrated, multi-scale labels/ group. - convert: any input → pyramidal OME-ZARR (auto chunks, optional shard) - prepare: plan tiles, skip empties, create the empty stage store - segment {tile}: read tile+halo → Cellpose → trim → write to stage (GPU) - merge: zarr-native boundary merge + sequential relabel (default on) - new [workflow] extra: snakemake + slurm executor plugin Inspired by the sopa workflow's patch-scatter pattern. Co-Authored-By: Claude Opus 4.8 --- pyproject.toml | 2 + workflow/.snakemake/iocache/latest.pkl | Bin 0 -> 270 bytes workflow/README.md | 73 ++++++++++++ workflow/Snakefile | 70 +++++++++++ workflow/config/config.yaml | 41 +++++++ workflow/profile/slurm/config.yaml | 40 +++++++ workflow/scripts/_pw.py | 153 +++++++++++++++++++++++++ workflow/scripts/convert.py | 15 +++ workflow/scripts/merge.py | 45 ++++++++ workflow/scripts/prepare_tiles.py | 70 +++++++++++ workflow/scripts/segment_tile.py | 48 ++++++++ 11 files changed, 557 insertions(+) create mode 100644 workflow/.snakemake/iocache/latest.pkl create mode 100644 workflow/README.md create mode 100644 workflow/Snakefile create mode 100644 workflow/config/config.yaml create mode 100644 workflow/profile/slurm/config.yaml create mode 100644 workflow/scripts/_pw.py create mode 100644 workflow/scripts/convert.py create mode 100644 workflow/scripts/merge.py create mode 100644 workflow/scripts/prepare_tiles.py create mode 100644 workflow/scripts/segment_tile.py diff --git a/pyproject.toml b/pyproject.toml index 94b9b12..5887a11 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,8 @@ bioio = [ imaris = ["imaris-ims-file-reader"] # napari enables the interactive viewer plugin. napari = ["napari[all]"] +# workflow runs the Snakemake pipeline (per-tile SLURM jobs across GPUs). +workflow = ["snakemake>=8", "snakemake-executor-plugin-slurm"] dev = ["pytest", "pytest-cov", "scikit-image", "psutil", "tqdm"] docs = ["mkdocs-material>=9.0", "mkdocstrings[python]>=0.24"] all = [ diff --git a/workflow/.snakemake/iocache/latest.pkl b/workflow/.snakemake/iocache/latest.pkl new file mode 100644 index 0000000000000000000000000000000000000000..01ad6b1476db2ba735005f3c071ff73bc2ef68dd GIT binary patch literal 270 zcmYk1!Ab)`42COfE0wN@`Ve~Q9{K`;ii!sxAxwv%3Ei2M&1|iLpf`m;4~>s2bvi{b z5D0&gpYMO)|1S1=xn#t_%$-j?rbgx9{_)CKbh_4y4}H^f8Gw)EU7Do?w8G5wDGLP(Gk7)>A*0`3PBuSxDX)~? MEdMPY&bsON7wnX1n*aa+ literal 0 HcmV?d00001 diff --git a/workflow/README.md b/workflow/README.md new file mode 100644 index 0000000..112e28c --- /dev/null +++ b/workflow/README.md @@ -0,0 +1,73 @@ +# patchworks Snakemake workflow + +A SLURM-ready pipeline that segments an arbitrarily large image and spreads the +expensive Cellpose step across **many GPUs** — one tile per SLURM job. + +```text +convert ──▶ prepare (checkpoint) ──▶ segment {tile} ──▶ merge + one GPU job/tile +``` + +## Why + +`tile_process` runs tiles serially on a single GPU. This workflow instead +submits **one GPU job per (non-empty) tile**, so N GPUs cut the wall-time ~N×. +Each job writes a disjoint chunk of a shared stage store; a final CPU job +stitches labels across tile boundaries and writes them back into the image as a +calibrated, multi-scale `labels/` group. + +## Steps + +1. **convert** — input (`.ims`, `.czi`, `.lif`, `.nd2`, OME-TIFF, `.zarr`) → a + pyramidal OME-ZARR (`image.zarr`). +2. **prepare** (checkpoint) — plan tiles, skip background tiles, create the + empty `stage.zarr`, and list the tiles to segment. +3. **segment {tile}** — one GPU job per tile: read tile + halo, run Cellpose, + trim, write to the stage. Scattered across the cluster. +4. **merge** — zarr-native boundary stitch + renumber, written into + `image.zarr/labels//`. + +## Install + +```bash +pip install "patchworks[workflow,cellpose,imaris,bioio]" +# workflow → snakemake + the SLURM executor plugin +``` + +## Configure + +Edit `config/config.yaml` (input, output dir, channel, tile shape, Cellpose +model/diameter/`do_3D`, …) and `profile/slurm/config.yaml` (partitions, +account, GPU request). + +## Run + +```bash +# locally (single machine) +snakemake --cores 8 --configfile config/config.yaml + +# on SLURM — one GPU job per tile, up to `jobs:` in parallel +snakemake --workflow-profile profile/slurm --configfile config/config.yaml +``` + +The GPU request lives in `profile/slurm/config.yaml` under +`set-resources: segment:` (`--gres=gpu:1`). Raise `jobs:` to use more GPUs at +once. + +## Output + +`/image.zarr` — the image plus `labels//` (multi-scale, +calibrated). Open it directly: + +```python +from patchworks.plugins.napari import view_in_napari +view_in_napari("/image.zarr") # auto-loads the labels +``` + +## Notes + +- Tiles overlap on read (halo) but write **disjoint** regions, so the per-tile + jobs are safe to run concurrently. +- Background tiles are skipped (`skip_empty`), so only occupied tiles become + jobs. +- For very large stores, set `shard: true` in the config to cut the file count. diff --git a/workflow/Snakefile b/workflow/Snakefile new file mode 100644 index 0000000..1c19c8b --- /dev/null +++ b/workflow/Snakefile @@ -0,0 +1,70 @@ +# patchworks Snakemake workflow. +# +# convert ──▶ prepare (checkpoint) ──▶ segment {tile} ──▶ merge +# (one GPU job per tile) +# +# Run locally: +# snakemake --cores 8 --configfile config/config.yaml +# Run on SLURM (one GPU job per tile, many GPUs in parallel): +# snakemake --workflow-profile profile/slurm --configfile config/config.yaml + +import json +from pathlib import Path + +configfile: "config/config.yaml" + +WORK = config["work_dir"] +IMAGE = f"{WORK}/image.zarr" +TILES = f"{WORK}/tiles.json" +STAGE = f"{WORK}/stage.zarr" + + +rule all: + input: + f"{WORK}/labels.done", + + +rule convert: + output: + directory(IMAGE), + script: + "scripts/convert.py" + + +checkpoint prepare: + input: + IMAGE, + output: + tiles=TILES, + stage=directory(STAGE), + script: + "scripts/prepare_tiles.py" + + +def occupied_done(wildcards): + """Expand the per-tile markers for the occupied tiles (post-checkpoint).""" + tiles = checkpoints.prepare.get().output.tiles + occupied = json.loads(Path(tiles).read_text())["occupied"] + return [f"{WORK}/seg/{i}.done" for i in occupied] + + +rule segment: + """Segment one tile on a GPU and write it into the stage store.""" + input: + tiles=TILES, + stage=STAGE, + image=IMAGE, + output: + f"{WORK}/seg/{{index}}.done", + script: + "scripts/segment_tile.py" + + +rule merge: + """Stitch labels across tile boundaries and write them into the image.""" + input: + occupied_done, + output: + touch(f"{WORK}/labels.done"), + script: + "scripts/merge.py" diff --git a/workflow/config/config.yaml b/workflow/config/config.yaml new file mode 100644 index 0000000..c10ebca --- /dev/null +++ b/workflow/config/config.yaml @@ -0,0 +1,41 @@ +# patchworks Snakemake workflow — configuration. +# Copy and edit, then run: snakemake --workflow-profile profile/slurm +# +# The expensive Cellpose step is split into one SLURM job per tile, so many +# GPUs run in parallel; the merge is a single CPU job. + +# ---- input / output --------------------------------------------------------- +input: "/path/to/scan.ims" # .ims, .czi, .lif, .nd2, ome-tiff, or an .zarr +work_dir: "/path/to/results" # everything is written under here +# Final outputs (under work_dir): +# image.zarr — converted, pyramidal OME-ZARR +# image.zarr/labels// — the segmentation (multi-scale), in the image + +# ---- conversion ------------------------------------------------------------- +reuse_pyramid: true # for .ims: copy its own pyramid (fast); else rebuild +convert_chunks: null # null → patchworks' bounded auto chunks; or [c,z,y,x] +shard: false # true → pack chunks into shards (zarr v3; fewer files) + +# ---- tiling ----------------------------------------------------------------- +channel: 0 # channel to segment +level: 0 # pyramid level to segment (0 = full res) +tile_shape: "auto" # "auto", or e.g. [16, 1024, 1024] (zyx) +overlap: 30 # halo (≈ one object diameter) +skip_empty: true # skip background tiles +empty_threshold: null # null → Otsu; or a number + +# ---- segmentation (Cellpose) ------------------------------------------------ +label_name: "cellpose" +cellpose: + model: "cyto3" + diameter: 30 + do_3D: true + gpu: true + # extra kwargs forwarded to model.eval(), e.g.: + # flow_threshold: 0.4 + # cellprob_threshold: 0.0 + +# ---- pyramid for the labels ------------------------------------------------- +pyramid_levels: 5 +pyramid_downscale: 2 +sequential_labels: true # renumber merged labels to a contiguous 1..N range diff --git a/workflow/profile/slurm/config.yaml b/workflow/profile/slurm/config.yaml new file mode 100644 index 0000000..49169c0 --- /dev/null +++ b/workflow/profile/slurm/config.yaml @@ -0,0 +1,40 @@ +# Snakemake 8+ SLURM profile for the patchworks workflow. +# Needs: pip install "patchworks[workflow]" (snakemake + slurm executor plugin) +# Run: snakemake --workflow-profile profile/slurm --configfile config/config.yaml +# +# Edit the partitions / account / GPU request for your cluster. + +executor: slurm +jobs: 64 # max concurrent SLURM jobs (≈ how many GPUs you can grab at once) +latency-wait: 60 +rerun-incomplete: true +printshellcmds: true +keep-going: true + +# Defaults for every rule (CPU jobs). +default-resources: + slurm_partition: "cpu" + # slurm_account: "my_account" + mem_mb: 16000 + cpus_per_task: 4 + runtime: 60 # minutes + +# Per-rule overrides. +set-resources: + convert: + mem_mb: 64000 + cpus_per_task: 8 + runtime: 360 + prepare: + mem_mb: 32000 + runtime: 120 + segment: # one GPU per tile — this is what spreads Cellpose across GPUs + slurm_partition: "gpu" + slurm_extra: "'--gres=gpu:1'" + mem_mb: 32000 + cpus_per_task: 4 + runtime: 120 + merge: + mem_mb: 128000 + cpus_per_task: 8 + runtime: 240 diff --git a/workflow/scripts/_pw.py b/workflow/scripts/_pw.py new file mode 100644 index 0000000..0414ecf --- /dev/null +++ b/workflow/scripts/_pw.py @@ -0,0 +1,153 @@ +"""Shared helpers for the patchworks Snakemake workflow. + +These wrap patchworks' public API so each Snakemake rule can act on a single +tile (for SLURM scatter) or the whole store. +""" + +from __future__ import annotations + +import itertools +import json +from pathlib import Path + +import numpy as np +import zarr + +from patchworks import load_ome_zarr + + +def spatial_tile_slices( + shape: tuple[int, ...], tile_shape: tuple[int, ...] +) -> list[tuple[slice, ...]]: + """Row-major list of per-tile slice tuples covering *shape*. + + Parameters + ---------- + shape : tuple of int + Spatial array shape. + tile_shape : tuple of int + Tile shape. + + Returns + ------- + list of tuple of slice + One slice tuple per tile, in row-major order. + """ + grids = [range(0, s, t) for s, t in zip(shape, tile_shape)] + tiles = [] + for starts in itertools.product(*grids): + tiles.append( + tuple( + slice(o, min(o + t, s)) + for o, t, s in zip(starts, tile_shape, shape) + ) + ) + return tiles + + +def process_one_tile( + image, sl: tuple[slice, ...], overlap: int, fn +) -> np.ndarray: + """Read one tile (with halo), run *fn*, and trim the halo back off. + + Parameters + ---------- + image : array-like + The full image (dask/zarr/numpy), indexable by slices. + sl : tuple of slice + The tile's slice (without halo). + overlap : int + Halo size added on every side before calling *fn*. + fn : callable + ``(ndarray) -> ndarray`` returning integer labels of the same shape. + + Returns + ------- + np.ndarray + Labels for exactly the region *sl* (halo trimmed). + """ + shape = image.shape + expanded, trims = [], [] + for s, dim in zip(sl, shape): + lo = max(0, s.start - overlap) + hi = min(dim, s.stop + overlap) + expanded.append(slice(lo, hi)) + trims.append((s.start - lo, hi - s.stop)) # halo added left / right + block = np.asarray(image[tuple(expanded)]) + out = np.asarray(fn(block)) + sel = tuple( + slice(left, out.shape[i] - right) + for i, (left, right) in enumerate(trims) + ) + return out[sel] + + +def load_tiles_json(path: str | Path) -> dict: + """Load the tile manifest written by ``prepare_tiles.py``. + + Parameters + ---------- + path : str or Path + Path to ``tiles.json``. + + Returns + ------- + dict + The manifest (``tile_shape``, ``overlap``, ``occupied`` indices, …). + """ + return json.loads(Path(path).read_text()) + + +def open_image(work_dir: str | Path, channel, level): + """Open the converted image for segmentation. + + Parameters + ---------- + work_dir : str or Path + Workflow output directory containing ``image.zarr``. + channel : int or None + Channel to select. + level : int + Pyramid level to read. + + Returns + ------- + da.Array + The (lazy) image array. + """ + store = str(Path(work_dir) / "image.zarr") + return load_ome_zarr(store, channel=channel, level=level) + + +def stage_path(work_dir: str | Path) -> str: + """Path of the staged-labels zarr store. + + Parameters + ---------- + work_dir : str or Path + Workflow output directory. + + Returns + ------- + str + ``/stage.zarr``. + """ + return str(Path(work_dir) / "stage.zarr") + + +def open_stage(work_dir: str | Path, mode: str = "r+"): + """Open the staged-labels array. + + Parameters + ---------- + work_dir : str or Path + Workflow output directory. + mode : str + Zarr open mode. + + Returns + ------- + zarr.Array + The ``staged`` array inside ``stage.zarr``. + """ + return zarr.open_group(stage_path(work_dir), mode=mode)["staged"] diff --git a/workflow/scripts/convert.py b/workflow/scripts/convert.py new file mode 100644 index 0000000..99fa118 --- /dev/null +++ b/workflow/scripts/convert.py @@ -0,0 +1,15 @@ +"""Snakemake script: convert the input to a pyramidal OME-ZARR.""" + +from patchworks.plugins.ome_zarr import to_ome_zarr + +cfg = snakemake.config # noqa: F821 (injected by Snakemake) +chunks = cfg.get("convert_chunks") +to_ome_zarr( + cfg["input"], + snakemake.output[0], # noqa: F821 + chunks=tuple(chunks) if chunks else None, + shard=bool(cfg.get("shard", False)), + reuse_pyramid=bool(cfg.get("reuse_pyramid", False)), + progress=False, + overwrite=True, +) diff --git a/workflow/scripts/merge.py b/workflow/scripts/merge.py new file mode 100644 index 0000000..01732c3 --- /dev/null +++ b/workflow/scripts/merge.py @@ -0,0 +1,45 @@ +"""Snakemake script: merge the staged tiles into one labelled OME-ZARR. + +Runs patchworks' zarr-native boundary merge (stitches labels across tile +boundaries), optionally renumbers them, and writes the result back into the +image store under ``labels//`` as a calibrated, multi-scale pyramid. +""" + +import os +import shutil +from pathlib import Path + +import dask.array as da + +from patchworks._merge import zarr_native_merge +from patchworks._relabel import relabel_sequential_zarr +from patchworks.plugins.ome_zarr import write_labels + +from _pw import stage_path + +cfg = snakemake.config # noqa: F821 +work_dir = cfg["work_dir"] +image_store = str(Path(work_dir) / "image.zarr") +merged_store = str(Path(work_dir) / "_merged.zarr") + +n_workers = min(8, os.cpu_count() or 1) +zarr_native_merge( + stage_path(work_dir), "staged", merged_store, "labels", n_workers=n_workers +) +if cfg.get("sequential_labels", True): + relabel_sequential_zarr(merged_store, "labels") + +merged = da.from_zarr(merged_store, component="labels") +group = write_labels( + image_store, + merged, + name=cfg.get("label_name", "labels"), + n_levels=int(cfg.get("pyramid_levels", 5)), + downscale=int(cfg.get("pyramid_downscale", 2)), + overwrite=True, +) + +shutil.rmtree(merged_store, ignore_errors=True) +shutil.rmtree(stage_path(work_dir), ignore_errors=True) +print(f"[patchworks] labels written to {group}") +open(snakemake.output[0], "w").close() # noqa: F821 diff --git a/workflow/scripts/prepare_tiles.py b/workflow/scripts/prepare_tiles.py new file mode 100644 index 0000000..1dddbf4 --- /dev/null +++ b/workflow/scripts/prepare_tiles.py @@ -0,0 +1,70 @@ +"""Snakemake script: plan tiles, create the empty stage store, list work. + +Writes ``tiles.json`` (tile shape, overlap, and the indices of *occupied* +tiles to segment) and an empty ``stage.zarr/staged`` array that the per-tile +segment jobs fill in parallel. +""" + +import json +from functools import partial +from pathlib import Path + +import numpy as np +import zarr + +from patchworks import auto_tile_shape_cellpose, estimate_empty_tiles + +from _pw import open_image, spatial_tile_slices, stage_path + +cfg = snakemake.config # noqa: F821 +work_dir = cfg["work_dir"] +image = open_image(work_dir, cfg["channel"], cfg["level"]) + +# Resolve the tile shape (spatial, matches the loaded image's ndim). +ts = cfg.get("tile_shape", "auto") +if ts == "auto": + cp = cfg["cellpose"] + tile_shape = tuple( + partial( + auto_tile_shape_cellpose, + do_3D=cp.get("do_3D", False), + use_gpu=cp.get("gpu", True), + diameter=cp.get("diameter"), + )(image.shape, image.dtype) + ) +else: + tile_shape = tuple(ts) + +tiles = spatial_tile_slices(image.shape, tile_shape) +n_tiles = len(tiles) + +# Decide which tiles to actually segment (skip background). +occupied = list(range(n_tiles)) +if cfg.get("skip_empty", True): + info = estimate_empty_tiles( + image, tile_shape, threshold=cfg.get("empty_threshold") + ) + occ = info["occupancy"].ravel() # row-major, matches spatial_tile_slices + occupied = [i for i in range(n_tiles) if occ[i]] + +# Create the empty staged-labels array (zeros), one chunk per tile. +root = zarr.open_group(stage_path(work_dir), mode="w") +root.create_array( + name="staged", + shape=image.shape, + chunks=tile_shape, + dtype=np.int32, +) + +Path(work_dir, "tiles.json").write_text( + json.dumps( + { + "tile_shape": list(tile_shape), + "overlap": int(cfg.get("overlap", 0)), + "n_tiles": n_tiles, + "occupied": occupied, + }, + indent=2, + ) +) +print(f"[patchworks] {len(occupied)}/{n_tiles} tiles to segment") diff --git a/workflow/scripts/segment_tile.py b/workflow/scripts/segment_tile.py new file mode 100644 index 0000000..7dd042b --- /dev/null +++ b/workflow/scripts/segment_tile.py @@ -0,0 +1,48 @@ +"""Snakemake script: segment ONE tile on a GPU and write it to the stage. + +Scattered over tile indices by Snakemake, so each tile is its own SLURM job +and many GPUs run in parallel. Each job writes a disjoint chunk of +``stage.zarr/staged`` (safe to run concurrently). +""" + +from patchworks.plugins.cellpose import cellpose_fn + +from _pw import ( + load_tiles_json, + open_image, + open_stage, + process_one_tile, + spatial_tile_slices, +) + +cfg = snakemake.config # noqa: F821 +index = int(snakemake.wildcards.index) # noqa: F821 +work_dir = cfg["work_dir"] + +manifest = load_tiles_json(snakemake.input.tiles) # noqa: F821 +tile_shape = tuple(manifest["tile_shape"]) +overlap = int(manifest["overlap"]) + +image = open_image(work_dir, cfg["channel"], cfg["level"]) +sl = spatial_tile_slices(image.shape, tile_shape)[index] + +cp = cfg["cellpose"] +extra = { + k: v + for k, v in cp.items() + if k not in ("model", "diameter", "do_3D", "gpu") +} +fn = cellpose_fn( + cp.get("model", "cyto3"), + gpu=cp.get("gpu", True), + diameter=cp.get("diameter"), + do_3D=cp.get("do_3D", False), + **extra, +) + +labels = process_one_tile(image, sl, overlap, fn) +open_stage(work_dir, mode="r+")[sl] = labels.astype("int32") + +# touch the per-tile done marker +open(snakemake.output[0], "w").close() # noqa: F821 +print(f"[patchworks] segmented tile {index}") From a361ed26bb35e5003d89223dc9d4aad5413f242f Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Wed, 24 Jun 2026 09:15:11 +0200 Subject: [PATCH 2/3] =?UTF-8?q?refactor:=20=E2=99=BB=EF=B8=8F=20public=20p?= =?UTF-8?q?er-tile=20API=20+=20.smk=20rule=20files;=20e2e-tested?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the 'why scripts if the API was fine' point: tile_process does stage+merge in one process, so there was no public way to run a single tile for distributed scatter. Add that to the public API and make the workflow thin: - new patchworks public API: spatial_tiles, create_stage, stage_tile (run fn on one tile → write a disjoint chunk of a shared stage), with tests - workflow scripts now just glue config → public API (cellpose, or a 'threshold' method for no-GPU testing) - split the workflow into rules/*.smk (convert / segment / merge / common) included from the Snakefile, like sopa - verified the whole pipeline runs end-to-end (snakemake, toy input): convert → prepare → segment ×N → merge, cross-boundary object stitched to a single label written into image.zarr/labels/ Co-Authored-By: Claude Opus 4.8 --- src/patchworks/__init__.py | 4 + src/patchworks/_distributed.py | 140 +++++++++++++++ tests/test_distributed.py | 45 +++++ workflow/.snakemake/iocache/latest.pkl | Bin 270 -> 0 bytes .../metadata/L3RtcC90b3kvb3V0L2ltYWdlLnphcnI= | 1 + .../metadata/L3RtcC90b3kvb3V0L2xhYmVscy5kb25l | 1 + .../metadata/L3RtcC90b3kvb3V0L3N0YWdlLnphcnI= | 1 + .../metadata/L3RtcC90b3kvb3V0L3NlZy80LmRvbmU= | 1 + .../metadata/L3RtcC90b3kvb3V0L3NlZy8wLmRvbmU= | 1 + .../metadata/L3RtcC90b3kvb3V0L3NlZy8xLmRvbmU= | 1 + .../metadata/L3RtcC90b3kvb3V0L3NlZy8zLmRvbmU= | 1 + .../metadata/L3RtcC90b3kvb3V0L3RpbGVzLmpzb24= | 1 + workflow/README.md | 18 ++ workflow/Snakefile | 58 +------ workflow/config/config.yaml | 3 +- workflow/rules/common.smk | 13 ++ workflow/rules/convert.smk | 7 + workflow/rules/merge.smk | 9 + workflow/rules/segment.smk | 22 +++ workflow/scripts/_pw.py | 162 +++++++----------- workflow/scripts/merge.py | 21 +-- workflow/scripts/prepare_tiles.py | 44 ++--- workflow/scripts/segment_tile.py | 44 ++--- 23 files changed, 370 insertions(+), 228 deletions(-) create mode 100644 src/patchworks/_distributed.py create mode 100644 tests/test_distributed.py delete mode 100644 workflow/.snakemake/iocache/latest.pkl create mode 100644 workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2ltYWdlLnphcnI= create mode 100644 workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2xhYmVscy5kb25l create mode 100644 workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3N0YWdlLnphcnI= create mode 100644 workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy80LmRvbmU= create mode 100644 workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8wLmRvbmU= create mode 100644 workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8xLmRvbmU= create mode 100644 workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8zLmRvbmU= create mode 100644 workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3RpbGVzLmpzb24= create mode 100644 workflow/rules/common.smk create mode 100644 workflow/rules/convert.smk create mode 100644 workflow/rules/merge.smk create mode 100644 workflow/rules/segment.smk diff --git a/src/patchworks/__init__.py b/src/patchworks/__init__.py index d8e2b6c..c5358ac 100644 --- a/src/patchworks/__init__.py +++ b/src/patchworks/__init__.py @@ -32,6 +32,7 @@ from ._chunks import auto_overlap, auto_tile_shape, auto_tile_shape_cellpose from ._cluster import make_local_cluster from ._core import tile_process +from ._distributed import create_stage, spatial_tiles, stage_tile from ._io import estimate_empty_tiles, load_ome_zarr from ._merge import merge_tile_labels from ._relabel import relabel_sequential_array, relabel_sequential_zarr @@ -51,4 +52,7 @@ "make_local_cluster", "relabel_sequential_array", "relabel_sequential_zarr", + "spatial_tiles", + "create_stage", + "stage_tile", ] diff --git a/src/patchworks/_distributed.py b/src/patchworks/_distributed.py new file mode 100644 index 0000000..822a182 --- /dev/null +++ b/src/patchworks/_distributed.py @@ -0,0 +1,140 @@ +"""Per-tile building blocks for distributed processing. + +``tile_process`` runs every tile and merges in one process. To spread tiles +across separate jobs (e.g. one SLURM GPU job per tile) you need to process a +*single* tile independently and merge later. These helpers expose exactly that: +:func:`spatial_tiles` enumerates the tiles, :func:`create_stage` makes the +shared output store, and :func:`stage_tile` runs ``fn`` on one tile and writes +it into that store. Stitch the result with +:func:`patchworks.merge_tile_labels` (or ``zarr_native_merge``). +""" + +from __future__ import annotations + +import itertools +from pathlib import Path +from typing import Callable, Union + +import numpy as np +import zarr + + +def spatial_tiles( + shape: tuple[int, ...], tile_shape: tuple[int, ...] +) -> list[tuple[slice, ...]]: + """Enumerate the tiles covering *shape*, in row-major order. + + Parameters + ---------- + shape : tuple of int + Spatial array shape. + tile_shape : tuple of int + Tile shape. + + Returns + ------- + list of tuple of slice + One slice tuple per tile (the same order ``estimate_empty_tiles``'s + ``occupancy`` grid uses when ravelled). + """ + grids = [range(0, s, t) for s, t in zip(shape, tile_shape)] + return [ + tuple( + slice(o, min(o + t, s)) + for o, t, s in zip(starts, tile_shape, shape) + ) + for starts in itertools.product(*grids) + ] + + +def create_stage( + stage_path: Union[str, Path], + shape: tuple[int, ...], + tile_shape: tuple[int, ...], + *, + component: str = "staged", + dtype=np.int32, +) -> str: + """Create the empty (zero-filled) shared stage store for tiled writes. + + Parameters + ---------- + stage_path : str or Path + Destination ``.zarr`` store. + shape : tuple of int + Full (spatial) array shape. + tile_shape : tuple of int + Chunk = tile shape (one chunk per tile, so jobs write disjoint files). + component : str, optional + Array name inside the store (default ``"staged"``). + dtype : data-type, optional + Label dtype (default ``int32``). + + Returns + ------- + str + The stage store path. + """ + root = zarr.open_group(str(stage_path), mode="w") + root.create_array( + name=component, shape=shape, chunks=tile_shape, dtype=dtype + ) + return str(stage_path) + + +def stage_tile( + image, + fn: Callable[[np.ndarray], np.ndarray], + stage_path: Union[str, Path], + index: int, + *, + tile_shape: tuple[int, ...], + overlap: int = 0, + component: str = "staged", +) -> int: + """Run *fn* on a single tile and write it into the shared stage store. + + Reads the tile (expanded by *overlap* on every side for boundary context), + runs *fn*, trims the halo back off, and writes the result to the tile's + disjoint chunk of ``stage_path/component`` — so many of these can run + concurrently (one per job) without conflicts. + + Parameters + ---------- + image : array-like + The full image (dask/zarr/NumPy), indexable by slices. + fn : callable + ``(ndarray) -> ndarray`` returning integer labels of the same shape. + stage_path : str or Path + Stage store created by :func:`create_stage`. + index : int + Tile index into :func:`spatial_tiles`. + tile_shape : tuple of int + Tile shape (must match the stage store's chunks). + overlap : int, optional + Halo added on every side before calling *fn*. + component : str, optional + Array name inside the stage store. + + Returns + ------- + int + The processed tile *index*. + """ + shape = image.shape + sl = spatial_tiles(shape, tile_shape)[index] + expanded, trims = [], [] + for s, dim in zip(sl, shape): + lo = max(0, s.start - overlap) + hi = min(dim, s.stop + overlap) + expanded.append(slice(lo, hi)) + trims.append((s.start - lo, hi - s.stop)) + block = np.asarray(image[tuple(expanded)]) + out = np.asarray(fn(block)) + sel = tuple( + slice(left, out.shape[i] - right) + for i, (left, right) in enumerate(trims) + ) + dst = zarr.open_group(str(stage_path), mode="r+")[component] + dst[sl] = out[sel].astype(dst.dtype) + return index diff --git a/tests/test_distributed.py b/tests/test_distributed.py new file mode 100644 index 0000000..0a0decf --- /dev/null +++ b/tests/test_distributed.py @@ -0,0 +1,45 @@ +"""Tests for the per-tile distributed building blocks.""" + +import numpy as np + +from patchworks import ( + create_stage, + merge_tile_labels, + spatial_tiles, + stage_tile, +) + + +def _fn(tile): + from skimage.measure import label + + return label(tile > 0).astype("int32") + + +def test_spatial_tiles_cover(): + """Tiles tile the array exactly, in row-major order.""" + tiles = spatial_tiles((4, 5), (2, 2)) + assert len(tiles) == 2 * 3 # ceil(4/2) * ceil(5/2) + assert tiles[0] == (slice(0, 2), slice(0, 2)) + assert tiles[-1] == (slice(2, 4), slice(4, 5)) # clipped last tile + + +def test_stage_then_merge_stitches_boundary(tmp_path): + """Per-tile staging + merge reproduces a cross-boundary single object.""" + img = np.zeros((16, 32), "uint16") + img[4:12, 8:24] = 500 # block straddling the x=16 tile boundary + + stage = str(tmp_path / "stage.zarr") + tile = (16, 16) + create_stage(stage, img.shape, tile) + for i in range(len(spatial_tiles(img.shape, tile))): + stage_tile(img, _fn, stage, i, tile_shape=tile, overlap=4) + + merged = merge_tile_labels( + stage, + write_to=str(tmp_path / "out.zarr"), + input_component="staged", + sequential_labels=True, + ).compute() + ids = np.unique(merged[merged > 0]) + assert ids.size == 1, f"object split into {ids.size} labels" diff --git a/workflow/.snakemake/iocache/latest.pkl b/workflow/.snakemake/iocache/latest.pkl deleted file mode 100644 index 01ad6b1476db2ba735005f3c071ff73bc2ef68dd..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 270 zcmYk1!Ab)`42COfE0wN@`Ve~Q9{K`;ii!sxAxwv%3Ei2M&1|iLpf`m;4~>s2bvi{b z5D0&gpYMO)|1S1=xn#t_%$-j?rbgx9{_)CKbh_4y4}H^f8Gw)EU7Do?w8G5wDGLP(Gk7)>A*0`3PBuSxDX)~? MEdMPY&bsON7wnX1n*aa+ diff --git a/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2ltYWdlLnphcnI= b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2ltYWdlLnphcnI= new file mode 100644 index 0000000..7db5ece --- /dev/null +++ b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2ltYWdlLnphcnI= @@ -0,0 +1 @@ +{"rule": "convert", "input": [], "log": [], "shellcmd": null, "params": [], "code": " \"../scripts/convert.py\"\n", "record_format_version": 6, "conda_env": null, "container_img_url": null, "software_stack_hash": "d41d8cd98f00b204e9800998ecf8427e", "job_hash": 7968324357465, "starttime": 1782285191.3091571, "endtime": 1782285191.3090107, "incomplete": false, "external_jobid": null, "input_checksums": {}} \ No newline at end of file diff --git a/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2xhYmVscy5kb25l b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2xhYmVscy5kb25l new file mode 100644 index 0000000..ee12982 --- /dev/null +++ b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2xhYmVscy5kb25l @@ -0,0 +1 @@ +{"rule": "merge", "input": ["/tmp/toy/out/seg/0.done", "/tmp/toy/out/seg/1.done", "/tmp/toy/out/seg/3.done", "/tmp/toy/out/seg/4.done"], "log": [], "shellcmd": null, "params": [], "code": " \"../scripts/merge.py\"\n", "record_format_version": 6, "conda_env": null, "container_img_url": null, "software_stack_hash": "d41d8cd98f00b204e9800998ecf8427e", "job_hash": 7968324177309, "starttime": 1782285198.9120235, "endtime": 1782285198.9118745, "incomplete": false, "external_jobid": null, "input_checksums": {"/tmp/toy/out/seg/0.done": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "/tmp/toy/out/seg/1.done": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "/tmp/toy/out/seg/3.done": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "/tmp/toy/out/seg/4.done": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"}} \ No newline at end of file diff --git a/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3N0YWdlLnphcnI= b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3N0YWdlLnphcnI= new file mode 100644 index 0000000..d495aed --- /dev/null +++ b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3N0YWdlLnphcnI= @@ -0,0 +1 @@ +{"rule": "prepare", "input": ["/tmp/toy/out/image.zarr"], "log": [], "shellcmd": null, "params": [], "code": " \"../scripts/prepare_tiles.py\"\n", "record_format_version": 6, "conda_env": null, "container_img_url": null, "software_stack_hash": "d41d8cd98f00b204e9800998ecf8427e", "job_hash": 7968324357329, "starttime": 1782285193.0925362, "endtime": 1782285193.0922585, "incomplete": false, "external_jobid": null, "input_checksums": {}} \ No newline at end of file diff --git a/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy80LmRvbmU= b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy80LmRvbmU= new file mode 100644 index 0000000..5bb2342 --- /dev/null +++ b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy80LmRvbmU= @@ -0,0 +1 @@ +{"rule": "segment", "input": ["/tmp/toy/out/image.zarr", "/tmp/toy/out/stage.zarr", "/tmp/toy/out/tiles.json"], "log": [], "shellcmd": null, "params": [], "code": " \"../scripts/segment_tile.py\"\n", "record_format_version": 6, "conda_env": null, "container_img_url": null, "software_stack_hash": "d41d8cd98f00b204e9800998ecf8427e", "job_hash": 7968316803128, "starttime": 1782285194.9512916, "endtime": 1782285194.9498596, "incomplete": false, "external_jobid": null, "input_checksums": {"/tmp/toy/out/tiles.json": "sha256:a7ba7de7a100b9eed2f69158d779580251e48da2c4bdb722e28c4d9a79abb411"}} \ No newline at end of file diff --git a/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8wLmRvbmU= b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8wLmRvbmU= new file mode 100644 index 0000000..08672c4 --- /dev/null +++ b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8wLmRvbmU= @@ -0,0 +1 @@ +{"rule": "segment", "input": ["/tmp/toy/out/image.zarr", "/tmp/toy/out/stage.zarr", "/tmp/toy/out/tiles.json"], "log": [], "shellcmd": null, "params": [], "code": " \"../scripts/segment_tile.py\"\n", "record_format_version": 6, "conda_env": null, "container_img_url": null, "software_stack_hash": "d41d8cd98f00b204e9800998ecf8427e", "job_hash": 7968316911209, "starttime": 1782285194.9791708, "endtime": 1782285194.9778593, "incomplete": false, "external_jobid": null, "input_checksums": {"/tmp/toy/out/tiles.json": "sha256:a7ba7de7a100b9eed2f69158d779580251e48da2c4bdb722e28c4d9a79abb411"}} \ No newline at end of file diff --git a/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8xLmRvbmU= b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8xLmRvbmU= new file mode 100644 index 0000000..e2b01f3 --- /dev/null +++ b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8xLmRvbmU= @@ -0,0 +1 @@ +{"rule": "segment", "input": ["/tmp/toy/out/image.zarr", "/tmp/toy/out/stage.zarr", "/tmp/toy/out/tiles.json"], "log": [], "shellcmd": null, "params": [], "code": " \"../scripts/segment_tile.py\"\n", "record_format_version": 6, "conda_env": null, "container_img_url": null, "software_stack_hash": "d41d8cd98f00b204e9800998ecf8427e", "job_hash": 7968316911362, "starttime": 1782285196.8732357, "endtime": 1782285196.871846, "incomplete": false, "external_jobid": null, "input_checksums": {"/tmp/toy/out/tiles.json": "sha256:a7ba7de7a100b9eed2f69158d779580251e48da2c4bdb722e28c4d9a79abb411"}} \ No newline at end of file diff --git a/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8zLmRvbmU= b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8zLmRvbmU= new file mode 100644 index 0000000..fb47b54 --- /dev/null +++ b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8zLmRvbmU= @@ -0,0 +1 @@ +{"rule": "segment", "input": ["/tmp/toy/out/image.zarr", "/tmp/toy/out/stage.zarr", "/tmp/toy/out/tiles.json"], "log": [], "shellcmd": null, "params": [], "code": " \"../scripts/segment_tile.py\"\n", "record_format_version": 6, "conda_env": null, "container_img_url": null, "software_stack_hash": "d41d8cd98f00b204e9800998ecf8427e", "job_hash": 7968316911515, "starttime": 1782285196.9122667, "endtime": 1782285196.9108458, "incomplete": false, "external_jobid": null, "input_checksums": {"/tmp/toy/out/tiles.json": "sha256:a7ba7de7a100b9eed2f69158d779580251e48da2c4bdb722e28c4d9a79abb411"}} \ No newline at end of file diff --git a/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3RpbGVzLmpzb24= b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3RpbGVzLmpzb24= new file mode 100644 index 0000000..a52f756 --- /dev/null +++ b/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3RpbGVzLmpzb24= @@ -0,0 +1 @@ +{"rule": "prepare", "input": ["/tmp/toy/out/image.zarr"], "log": [], "shellcmd": null, "params": [], "code": " \"../scripts/prepare_tiles.py\"\n", "record_format_version": 6, "conda_env": null, "container_img_url": null, "software_stack_hash": "d41d8cd98f00b204e9800998ecf8427e", "job_hash": 7968324357329, "starttime": 1782285193.0925362, "endtime": 1782285193.0908728, "incomplete": false, "external_jobid": null, "input_checksums": {}} \ No newline at end of file diff --git a/workflow/README.md b/workflow/README.md index 112e28c..7e6b98b 100644 --- a/workflow/README.md +++ b/workflow/README.md @@ -64,10 +64,28 @@ from patchworks.plugins.napari import view_in_napari view_in_napari("/image.zarr") # auto-loads the labels ``` +## Layout + +```text +workflow/ + Snakefile # includes the rule files below + rules/ # convert.smk, segment.smk, merge.smk, common.smk + scripts/ # thin wrappers over patchworks' public API + config/config.yaml + profile/slurm/config.yaml +``` + +The rule scripts are intentionally thin — the work is done by patchworks' +public API (`spatial_tiles`, `create_stage`, `stage_tile`, `merge_tile_labels`, +`write_labels`), so the same per-tile distribution is available from your own +code too. + ## Notes - Tiles overlap on read (halo) but write **disjoint** regions, so the per-tile jobs are safe to run concurrently. - Background tiles are skipped (`skip_empty`), so only occupied tiles become jobs. +- `method:` selects `cellpose` (default) or a simple `threshold` (no GPU — + handy for testing or quick masks). - For very large stores, set `shard: true` in the config to cut the file count. diff --git a/workflow/Snakefile b/workflow/Snakefile index 1c19c8b..d59a8ed 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -1,7 +1,7 @@ # patchworks Snakemake workflow. # -# convert ──▶ prepare (checkpoint) ──▶ segment {tile} ──▶ merge -# (one GPU job per tile) +# convert ──▶ prepare (checkpoint) ──▶ segment {tile} ──▶ merge +# one GPU job per tile # # Run locally: # snakemake --cores 8 --configfile config/config.yaml @@ -13,58 +13,12 @@ from pathlib import Path configfile: "config/config.yaml" -WORK = config["work_dir"] -IMAGE = f"{WORK}/image.zarr" -TILES = f"{WORK}/tiles.json" -STAGE = f"{WORK}/stage.zarr" +include: "rules/common.smk" +include: "rules/convert.smk" +include: "rules/segment.smk" +include: "rules/merge.smk" rule all: input: f"{WORK}/labels.done", - - -rule convert: - output: - directory(IMAGE), - script: - "scripts/convert.py" - - -checkpoint prepare: - input: - IMAGE, - output: - tiles=TILES, - stage=directory(STAGE), - script: - "scripts/prepare_tiles.py" - - -def occupied_done(wildcards): - """Expand the per-tile markers for the occupied tiles (post-checkpoint).""" - tiles = checkpoints.prepare.get().output.tiles - occupied = json.loads(Path(tiles).read_text())["occupied"] - return [f"{WORK}/seg/{i}.done" for i in occupied] - - -rule segment: - """Segment one tile on a GPU and write it into the stage store.""" - input: - tiles=TILES, - stage=STAGE, - image=IMAGE, - output: - f"{WORK}/seg/{{index}}.done", - script: - "scripts/segment_tile.py" - - -rule merge: - """Stitch labels across tile boundaries and write them into the image.""" - input: - occupied_done, - output: - touch(f"{WORK}/labels.done"), - script: - "scripts/merge.py" diff --git a/workflow/config/config.yaml b/workflow/config/config.yaml index c10ebca..bfa717c 100644 --- a/workflow/config/config.yaml +++ b/workflow/config/config.yaml @@ -24,7 +24,8 @@ overlap: 30 # halo (≈ one object diameter) skip_empty: true # skip background tiles empty_threshold: null # null → Otsu; or a number -# ---- segmentation (Cellpose) ------------------------------------------------ +# ---- segmentation ----------------------------------------------------------- +method: "cellpose" # "cellpose" (GPU) or "threshold" (simple, no GPU; testing) label_name: "cellpose" cellpose: model: "cyto3" diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk new file mode 100644 index 0000000..7a40024 --- /dev/null +++ b/workflow/rules/common.smk @@ -0,0 +1,13 @@ +# Shared paths and helpers for the patchworks workflow. + +WORK = config["work_dir"] +IMAGE = f"{WORK}/image.zarr" +TILES = f"{WORK}/tiles.json" +STAGE = f"{WORK}/stage.zarr" + + +def occupied_done(wildcards): + """Per-tile markers for the occupied tiles (resolved after the checkpoint).""" + tiles = checkpoints.prepare.get().output.tiles + occupied = json.loads(Path(tiles).read_text())["occupied"] + return [f"{WORK}/seg/{i}.done" for i in occupied] diff --git a/workflow/rules/convert.smk b/workflow/rules/convert.smk new file mode 100644 index 0000000..5c45cb6 --- /dev/null +++ b/workflow/rules/convert.smk @@ -0,0 +1,7 @@ +# Convert the input to a pyramidal OME-ZARR. + +rule convert: + output: + directory(IMAGE), + script: + "../scripts/convert.py" diff --git a/workflow/rules/merge.smk b/workflow/rules/merge.smk new file mode 100644 index 0000000..249e171 --- /dev/null +++ b/workflow/rules/merge.smk @@ -0,0 +1,9 @@ +# Stitch labels across tile boundaries and write them into the image. + +rule merge: + input: + occupied_done, + output: + touch(f"{WORK}/labels.done"), + script: + "../scripts/merge.py" diff --git a/workflow/rules/segment.smk b/workflow/rules/segment.smk new file mode 100644 index 0000000..72330fb --- /dev/null +++ b/workflow/rules/segment.smk @@ -0,0 +1,22 @@ +# Plan tiles (checkpoint) and segment each tile on a GPU. + +checkpoint prepare: + input: + IMAGE, + output: + tiles=TILES, + stage=directory(STAGE), + script: + "../scripts/prepare_tiles.py" + + +rule segment: + """Segment one tile on a GPU and write it into the stage store.""" + input: + tiles=TILES, + stage=STAGE, + image=IMAGE, + output: + f"{WORK}/seg/{{index}}.done", + script: + "../scripts/segment_tile.py" diff --git a/workflow/scripts/_pw.py b/workflow/scripts/_pw.py index 0414ecf..39222bf 100644 --- a/workflow/scripts/_pw.py +++ b/workflow/scripts/_pw.py @@ -1,88 +1,57 @@ -"""Shared helpers for the patchworks Snakemake workflow. +"""Config glue for the patchworks Snakemake workflow. -These wrap patchworks' public API so each Snakemake rule can act on a single -tile (for SLURM scatter) or the whole store. +The heavy lifting lives in patchworks' public API +(``spatial_tiles``/``create_stage``/``stage_tile``/``merge_tile_labels``); +these helpers only turn the Snakemake config into the right arguments. """ from __future__ import annotations -import itertools import json from pathlib import Path -import numpy as np -import zarr - from patchworks import load_ome_zarr -def spatial_tile_slices( - shape: tuple[int, ...], tile_shape: tuple[int, ...] -) -> list[tuple[slice, ...]]: - """Row-major list of per-tile slice tuples covering *shape*. +def open_image(work_dir, channel, level): + """Open the converted image for segmentation. Parameters ---------- - shape : tuple of int - Spatial array shape. - tile_shape : tuple of int - Tile shape. + work_dir : str or Path + Workflow output directory containing ``image.zarr``. + channel : int or None + Channel to select. + level : int + Pyramid level to read. Returns ------- - list of tuple of slice - One slice tuple per tile, in row-major order. + da.Array + The (lazy) image array. """ - grids = [range(0, s, t) for s, t in zip(shape, tile_shape)] - tiles = [] - for starts in itertools.product(*grids): - tiles.append( - tuple( - slice(o, min(o + t, s)) - for o, t, s in zip(starts, tile_shape, shape) - ) - ) - return tiles + return load_ome_zarr( + str(Path(work_dir) / "image.zarr"), channel=channel, level=level + ) -def process_one_tile( - image, sl: tuple[slice, ...], overlap: int, fn -) -> np.ndarray: - """Read one tile (with halo), run *fn*, and trim the halo back off. +def stage_path(work_dir): + """Path of the staged-labels store. Parameters ---------- - image : array-like - The full image (dask/zarr/numpy), indexable by slices. - sl : tuple of slice - The tile's slice (without halo). - overlap : int - Halo size added on every side before calling *fn*. - fn : callable - ``(ndarray) -> ndarray`` returning integer labels of the same shape. + work_dir : str or Path + Workflow output directory. Returns ------- - np.ndarray - Labels for exactly the region *sl* (halo trimmed). + str + ``/stage.zarr``. """ - shape = image.shape - expanded, trims = [], [] - for s, dim in zip(sl, shape): - lo = max(0, s.start - overlap) - hi = min(dim, s.stop + overlap) - expanded.append(slice(lo, hi)) - trims.append((s.start - lo, hi - s.stop)) # halo added left / right - block = np.asarray(image[tuple(expanded)]) - out = np.asarray(fn(block)) - sel = tuple( - slice(left, out.shape[i] - right) - for i, (left, right) in enumerate(trims) - ) - return out[sel] + return str(Path(work_dir) / "stage.zarr") -def load_tiles_json(path: str | Path) -> dict: +def load_tiles_json(path): """Load the tile manifest written by ``prepare_tiles.py``. Parameters @@ -98,56 +67,47 @@ def load_tiles_json(path: str | Path) -> dict: return json.loads(Path(path).read_text()) -def open_image(work_dir: str | Path, channel, level): - """Open the converted image for segmentation. - - Parameters - ---------- - work_dir : str or Path - Workflow output directory containing ``image.zarr``. - channel : int or None - Channel to select. - level : int - Pyramid level to read. - - Returns - ------- - da.Array - The (lazy) image array. - """ - store = str(Path(work_dir) / "image.zarr") - return load_ome_zarr(store, channel=channel, level=level) - - -def stage_path(work_dir: str | Path) -> str: - """Path of the staged-labels zarr store. +def build_fn(cfg): + """Build the per-tile segmentation function from the config. Parameters ---------- - work_dir : str or Path - Workflow output directory. + cfg : dict + Snakemake config. ``method`` selects ``"cellpose"`` (default) or a + simple ``"threshold"`` (handy for testing / no-GPU runs). Returns ------- - str - ``/stage.zarr``. + callable + ``(ndarray) -> ndarray`` returning integer labels. """ - return str(Path(work_dir) / "stage.zarr") - - -def open_stage(work_dir: str | Path, mode: str = "r+"): - """Open the staged-labels array. - - Parameters - ---------- - work_dir : str or Path - Workflow output directory. - mode : str - Zarr open mode. + method = cfg.get("method", "cellpose") + if method == "threshold": + + def fn(tile): + from skimage.filters import threshold_otsu + from skimage.measure import label + + thr = threshold_otsu(tile) if tile.max() > tile.min() else 0 + return label(tile > thr).astype("int32") + + return fn + + if method == "cellpose": + from patchworks.plugins.cellpose import cellpose_fn + + cp = cfg["cellpose"] + extra = { + k: v + for k, v in cp.items() + if k not in ("model", "diameter", "do_3D", "gpu") + } + return cellpose_fn( + cp.get("model", "cyto3"), + gpu=cp.get("gpu", True), + diameter=cp.get("diameter"), + do_3D=cp.get("do_3D", False), + **extra, + ) - Returns - ------- - zarr.Array - The ``staged`` array inside ``stage.zarr``. - """ - return zarr.open_group(stage_path(work_dir), mode=mode)["staged"] + raise ValueError(f"unknown segmentation method: {method!r}") diff --git a/workflow/scripts/merge.py b/workflow/scripts/merge.py index 01732c3..4e3e712 100644 --- a/workflow/scripts/merge.py +++ b/workflow/scripts/merge.py @@ -1,18 +1,14 @@ """Snakemake script: merge the staged tiles into one labelled OME-ZARR. Runs patchworks' zarr-native boundary merge (stitches labels across tile -boundaries), optionally renumbers them, and writes the result back into the +boundaries, optionally renumbers them) and writes the result back into the image store under ``labels//`` as a calibrated, multi-scale pyramid. """ -import os import shutil from pathlib import Path -import dask.array as da - -from patchworks._merge import zarr_native_merge -from patchworks._relabel import relabel_sequential_zarr +from patchworks import merge_tile_labels from patchworks.plugins.ome_zarr import write_labels from _pw import stage_path @@ -22,14 +18,13 @@ image_store = str(Path(work_dir) / "image.zarr") merged_store = str(Path(work_dir) / "_merged.zarr") -n_workers = min(8, os.cpu_count() or 1) -zarr_native_merge( - stage_path(work_dir), "staged", merged_store, "labels", n_workers=n_workers +merged = merge_tile_labels( + stage_path(work_dir), + write_to=merged_store, + input_component="staged", + sequential_labels=cfg.get("sequential_labels", True), + progress=False, ) -if cfg.get("sequential_labels", True): - relabel_sequential_zarr(merged_store, "labels") - -merged = da.from_zarr(merged_store, component="labels") group = write_labels( image_store, merged, diff --git a/workflow/scripts/prepare_tiles.py b/workflow/scripts/prepare_tiles.py index 1dddbf4..c178699 100644 --- a/workflow/scripts/prepare_tiles.py +++ b/workflow/scripts/prepare_tiles.py @@ -1,26 +1,22 @@ -"""Snakemake script: plan tiles, create the empty stage store, list work. - -Writes ``tiles.json`` (tile shape, overlap, and the indices of *occupied* -tiles to segment) and an empty ``stage.zarr/staged`` array that the per-tile -segment jobs fill in parallel. -""" +"""Snakemake script: plan tiles, create the empty stage store, list work.""" import json from functools import partial from pathlib import Path -import numpy as np -import zarr - -from patchworks import auto_tile_shape_cellpose, estimate_empty_tiles +from patchworks import ( + auto_tile_shape_cellpose, + create_stage, + estimate_empty_tiles, + spatial_tiles, +) -from _pw import open_image, spatial_tile_slices, stage_path +from _pw import open_image, stage_path cfg = snakemake.config # noqa: F821 work_dir = cfg["work_dir"] image = open_image(work_dir, cfg["channel"], cfg["level"]) -# Resolve the tile shape (spatial, matches the loaded image's ndim). ts = cfg.get("tile_shape", "auto") if ts == "auto": cp = cfg["cellpose"] @@ -35,36 +31,26 @@ else: tile_shape = tuple(ts) -tiles = spatial_tile_slices(image.shape, tile_shape) -n_tiles = len(tiles) - -# Decide which tiles to actually segment (skip background). -occupied = list(range(n_tiles)) +tiles = spatial_tiles(image.shape, tile_shape) +occupied = list(range(len(tiles))) if cfg.get("skip_empty", True): info = estimate_empty_tiles( image, tile_shape, threshold=cfg.get("empty_threshold") ) - occ = info["occupancy"].ravel() # row-major, matches spatial_tile_slices - occupied = [i for i in range(n_tiles) if occ[i]] + occ = info["occupancy"].ravel() # row-major, matches spatial_tiles + occupied = [i for i in range(len(tiles)) if occ[i]] -# Create the empty staged-labels array (zeros), one chunk per tile. -root = zarr.open_group(stage_path(work_dir), mode="w") -root.create_array( - name="staged", - shape=image.shape, - chunks=tile_shape, - dtype=np.int32, -) +create_stage(stage_path(work_dir), image.shape, tile_shape) Path(work_dir, "tiles.json").write_text( json.dumps( { "tile_shape": list(tile_shape), "overlap": int(cfg.get("overlap", 0)), - "n_tiles": n_tiles, + "n_tiles": len(tiles), "occupied": occupied, }, indent=2, ) ) -print(f"[patchworks] {len(occupied)}/{n_tiles} tiles to segment") +print(f"[patchworks] {len(occupied)}/{len(tiles)} tiles to segment") diff --git a/workflow/scripts/segment_tile.py b/workflow/scripts/segment_tile.py index 7dd042b..0398d52 100644 --- a/workflow/scripts/segment_tile.py +++ b/workflow/scripts/segment_tile.py @@ -1,48 +1,28 @@ -"""Snakemake script: segment ONE tile on a GPU and write it to the stage. +"""Snakemake script: segment ONE tile and write it to the shared stage. -Scattered over tile indices by Snakemake, so each tile is its own SLURM job -and many GPUs run in parallel. Each job writes a disjoint chunk of -``stage.zarr/staged`` (safe to run concurrently). +Scattered over tile indices, so each tile is its own SLURM job and many GPUs +run in parallel. Each job writes a disjoint chunk of the stage store. """ -from patchworks.plugins.cellpose import cellpose_fn +from patchworks import stage_tile -from _pw import ( - load_tiles_json, - open_image, - open_stage, - process_one_tile, - spatial_tile_slices, -) +from _pw import build_fn, load_tiles_json, open_image, stage_path cfg = snakemake.config # noqa: F821 index = int(snakemake.wildcards.index) # noqa: F821 work_dir = cfg["work_dir"] manifest = load_tiles_json(snakemake.input.tiles) # noqa: F821 -tile_shape = tuple(manifest["tile_shape"]) -overlap = int(manifest["overlap"]) - image = open_image(work_dir, cfg["channel"], cfg["level"]) -sl = spatial_tile_slices(image.shape, tile_shape)[index] -cp = cfg["cellpose"] -extra = { - k: v - for k, v in cp.items() - if k not in ("model", "diameter", "do_3D", "gpu") -} -fn = cellpose_fn( - cp.get("model", "cyto3"), - gpu=cp.get("gpu", True), - diameter=cp.get("diameter"), - do_3D=cp.get("do_3D", False), - **extra, +stage_tile( + image, + build_fn(cfg), + stage_path(work_dir), + index, + tile_shape=tuple(manifest["tile_shape"]), + overlap=int(manifest["overlap"]), ) -labels = process_one_tile(image, sl, overlap, fn) -open_stage(work_dir, mode="r+")[sl] = labels.astype("int32") - -# touch the per-tile done marker open(snakemake.output[0], "w").close() # noqa: F821 print(f"[patchworks] segmented tile {index}") From e93292f4eb969903e5666d4506504e73e441ed03 Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Wed, 24 Jun 2026 09:22:35 +0200 Subject: [PATCH 3/3] =?UTF-8?q?docs:=20=F0=9F=93=9D=20add=20a=20step-by-st?= =?UTF-8?q?ep=20cluster-workflow=20guide?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a thorough 'Cluster workflow (Snakemake + SLURM)' guide page (install → configure every field → dry-run → local vs SLURM → monitor → outputs → troubleshooting), wire it into the nav, and link it from the workflow README. Co-Authored-By: Claude Opus 4.8 --- docs/guide/snakemake.md | 213 ++++++++++++++++++++++++++++++++++++++++ mkdocs.yml | 1 + workflow/README.md | 4 + 3 files changed, 218 insertions(+) create mode 100644 docs/guide/snakemake.md diff --git a/docs/guide/snakemake.md b/docs/guide/snakemake.md new file mode 100644 index 0000000..e3a8723 --- /dev/null +++ b/docs/guide/snakemake.md @@ -0,0 +1,213 @@ +# Cluster workflow (Snakemake + SLURM) + +`tile_process` runs every tile **serially on one GPU**. For a large 3-D image +that can be days. The bundled Snakemake workflow instead submits **one GPU job +per tile**, so with *N* GPUs the segmentation is ~*N*× faster. This page walks +through running it from scratch. + +```text +convert ──▶ prepare (checkpoint) ──▶ segment {tile} ──▶ merge + one GPU SLURM job per tile +``` + +## 1. Get the workflow + +The workflow lives in the `workflow/` directory of the patchworks repository +(it is not shipped inside the pip package — it is a set of Snakemake files you +run): + +```bash +git clone https://github.com/imcf/patchworks +cd patchworks/workflow +``` + +## 2. Install the dependencies + +You need patchworks with the workflow + reader + segmentation extras, in the +environment Snakemake will use: + +```bash +pip install "patchworks[workflow,cellpose,imaris,bioio]" +``` + +- `workflow` → Snakemake + the SLURM executor plugin +- `cellpose` → the segmentation model +- `imaris` / `bioio` → read your input format (`.ims`, `.czi`, `.lif`, …) + +On a cluster, do this inside a conda/venv that the compute nodes can see, or let +each rule activate a conda env (see *Conda*, below). + +## 3. Configure the run + +Copy and edit `config/config.yaml`. Every field: + +```yaml +# input / output +input: "/data/scan.ims" # .ims/.czi/.lif/.nd2/ome-tiff/.zarr +work_dir: "/scratch/results" # everything is written here + +# conversion (input → pyramidal OME-ZARR) +reuse_pyramid: true # .ims: copy its own pyramid (fast) +convert_chunks: null # null → bounded auto chunks; or [c,z,y,x] +shard: false # true → pack chunks into shards (fewer files) + +# tiling +channel: 0 # channel to segment (null = keep all) +level: 0 # pyramid level (0 = full resolution) +tile_shape: "auto" # "auto", or e.g. [16, 1024, 1024] (zyx) +overlap: 30 # halo ≈ one object diameter +skip_empty: true # skip background tiles +empty_threshold: null # null → Otsu + +# segmentation +method: "cellpose" # "cellpose" (GPU) or "threshold" (no GPU) +label_name: "cellpose" # name under image.zarr/labels/ +cellpose: + model: "cyto3" + diameter: 30 + do_3D: true + gpu: true + # extra model.eval() kwargs, e.g. flow_threshold: 0.4 + +# label pyramid +pyramid_levels: 5 +pyramid_downscale: 2 +sequential_labels: true # renumber labels to a contiguous 1..N +``` + +!!! tip "Tile size vs runtime" + `tile_shape: "auto"` sizes each tile to your GPU's VRAM. Smaller tiles = + more (faster) jobs; very large 3-D tiles are slow. Keep `do_3D: false` (2-D + per slice) if your objects segment fine per slice — it is much faster. + +## 4. Dry-run (always do this first) + +Check the plan without running anything: + +```bash +python -m snakemake -s Snakefile --configfile config/config.yaml -n -p +``` + +You should see `convert`, `prepare`, and a note that the **checkpoint** will add +the `segment` jobs after `prepare` runs. (The number of segment jobs is only +known after `prepare` decides which tiles are non-empty.) + +## 5a. Run locally (single machine) + +```bash +python -m snakemake -s Snakefile --configfile config/config.yaml --cores 8 +``` + +Tiles run on the local machine (one at a time on the GPU). Good for a small +image or a smoke test. + +## 5b. Run on SLURM (one GPU job per tile) + +Edit `profile/slurm/config.yaml` for **your** cluster — partitions, account, +and the GPU request: + +```yaml +executor: slurm +jobs: 64 # max concurrent SLURM jobs ≈ GPUs you can grab +default-resources: + slurm_partition: "cpu" # your CPU partition + # slurm_account: "my_account" + mem_mb: 16000 + cpus_per_task: 4 + runtime: 60 +set-resources: + segment: # the GPU step + slurm_partition: "gpu" # your GPU partition + slurm_extra: "'--gres=gpu:1'" + mem_mb: 32000 + runtime: 120 + merge: + mem_mb: 128000 + runtime: 240 +``` + +Then launch (from a login node — Snakemake submits and watches the jobs): + +```bash +python -m snakemake --workflow-profile profile/slurm \ + --configfile config/config.yaml +``` + +Snakemake submits `convert`, then `prepare`, then **one `segment` job per +non-empty tile** (up to `jobs:` at once → that many GPUs in parallel), then +`merge`. Raise `jobs:` to use more GPUs. + +!!! note "GPU request flag" + Clusters differ. `--gres=gpu:1` is common; some need `--gpus=1` or a + specific gres name (`--gres=gpu:a100:1`). Put whatever `sbatch` flag your + cluster needs in `slurm_extra`. + +## 6. Monitor + +- **Snakemake** prints each job as it submits/finishes and a `X of Y steps` + counter. +- **SLURM**: `squeue --me` shows your queued/running jobs (`smk-segment`, …); + logs land where your profile/cluster sends them. +- **patchworks** logs (`processing tile k/N`, ETA) are inside each job's stdout. + +## 7. Output + +Everything is under `work_dir`: + +```text +results/ + image.zarr/ # converted, pyramidal OME-ZARR + image.zarr/labels// # the segmentation (multi-scale, calibrated) +``` + +The labels live **inside** the image store. View image + labels together: + +```python +from patchworks.plugins.napari import view_in_napari +view_in_napari("/scratch/results/image.zarr") # auto-loads the labels +``` + +## 8. Re-running and resuming + +Snakemake is resumable — if jobs fail or you cancel, just relaunch the same +command and it picks up only the missing tiles. To force a clean rerun, delete +`work_dir` (or the relevant outputs). + +## Conda (optional) + +To have each rule run in a named conda env instead of the active one, add +`--use-conda` and point the rules at an env; or activate your env in a SLURM +prologue. The simplest path is a single shared env that the compute nodes see. + +## Troubleshooting + +| Symptom | Fix | +|---------|-----| +| `snakemake: command not found` | use `python -m snakemake` | +| Segment jobs pend forever | wrong `slurm_partition`/`slurm_extra` GPU flag for your cluster | +| `cellpose is not installed` in a job | the job's env lacks `patchworks[cellpose]` | +| Reading the input fails | install the matching reader (`patchworks[imaris]`/`[bioio]` + a `bioio-*`) | +| Out of GPU memory | smaller `tile_shape`, or `do_3D: false` | +| Very slow | confirm GPU is used (`nvidia-smi`); try 2-D or a lower `level` | + +## How it works (for the curious) + +The rule scripts are thin wrappers over patchworks' public API, so you can build +the same per-tile distribution yourself: + +```python +from patchworks import ( + load_ome_zarr, spatial_tiles, create_stage, stage_tile, merge_tile_labels +) +from patchworks.plugins.ome_zarr import write_labels + +img = load_ome_zarr("image.zarr", channel=0) +tiles = spatial_tiles(img.shape, tile_shape=(16, 1024, 1024)) +create_stage("stage.zarr", img.shape, (16, 1024, 1024)) +# (distribute these across jobs:) +for i in range(len(tiles)): + stage_tile(img, my_fn, "stage.zarr", i, tile_shape=(16, 1024, 1024), overlap=30) +merged = merge_tile_labels("stage.zarr", input_component="staged", + write_to="merged.zarr", sequential_labels=True) +write_labels("image.zarr", merged, name="cells") +``` diff --git a/mkdocs.yml b/mkdocs.yml index 8567e10..3c975a4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -39,6 +39,7 @@ nav: - Empty tile skipping: guide/skip_empty.md - GPU & distributed: guide/gpu_distributed.md - Performance & memory: guide/performance.md + - Cluster workflow (Snakemake): guide/snakemake.md - OME-ZARR & napari: guide/ome_zarr_napari.md - Pitfalls: guide/pitfalls.md - Examples: diff --git a/workflow/README.md b/workflow/README.md index 7e6b98b..a5d4e07 100644 --- a/workflow/README.md +++ b/workflow/README.md @@ -3,6 +3,10 @@ A SLURM-ready pipeline that segments an arbitrarily large image and spreads the expensive Cellpose step across **many GPUs** — one tile per SLURM job. +> **Full step-by-step guide:** +> — install, configure every +> field, dry-run, local vs SLURM, monitoring, outputs and troubleshooting. + ```text convert ──▶ prepare (checkpoint) ──▶ segment {tile} ──▶ merge one GPU job/tile