patchworks 0.10.0__tar.gz → 0.11.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (75) hide show
  1. {patchworks-0.10.0 → patchworks-0.11.0}/PKG-INFO +4 -1
  2. patchworks-0.11.0/docs/guide/snakemake.md +213 -0
  3. {patchworks-0.10.0 → patchworks-0.11.0}/mkdocs.yml +1 -0
  4. {patchworks-0.10.0 → patchworks-0.11.0}/pyproject.toml +2 -0
  5. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/__init__.py +4 -0
  6. patchworks-0.11.0/src/patchworks/_distributed.py +140 -0
  7. patchworks-0.11.0/tests/test_distributed.py +45 -0
  8. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2ltYWdlLnphcnI= +1 -0
  9. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2xhYmVscy5kb25l +1 -0
  10. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3N0YWdlLnphcnI= +1 -0
  11. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy80LmRvbmU= +1 -0
  12. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8wLmRvbmU= +1 -0
  13. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8xLmRvbmU= +1 -0
  14. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8zLmRvbmU= +1 -0
  15. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3RpbGVzLmpzb24= +1 -0
  16. patchworks-0.11.0/workflow/README.md +95 -0
  17. patchworks-0.11.0/workflow/Snakefile +24 -0
  18. patchworks-0.11.0/workflow/config/config.yaml +42 -0
  19. patchworks-0.11.0/workflow/profile/slurm/config.yaml +40 -0
  20. patchworks-0.11.0/workflow/rules/common.smk +13 -0
  21. patchworks-0.11.0/workflow/rules/convert.smk +7 -0
  22. patchworks-0.11.0/workflow/rules/merge.smk +9 -0
  23. patchworks-0.11.0/workflow/rules/segment.smk +22 -0
  24. patchworks-0.11.0/workflow/scripts/_pw.py +113 -0
  25. patchworks-0.11.0/workflow/scripts/convert.py +15 -0
  26. patchworks-0.11.0/workflow/scripts/merge.py +40 -0
  27. patchworks-0.11.0/workflow/scripts/prepare_tiles.py +56 -0
  28. patchworks-0.11.0/workflow/scripts/segment_tile.py +28 -0
  29. {patchworks-0.10.0 → patchworks-0.11.0}/.github/workflows/docs.yml +0 -0
  30. {patchworks-0.10.0 → patchworks-0.11.0}/.github/workflows/lint.yml +0 -0
  31. {patchworks-0.10.0 → patchworks-0.11.0}/.github/workflows/release.yml +0 -0
  32. {patchworks-0.10.0 → patchworks-0.11.0}/.gitignore +0 -0
  33. {patchworks-0.10.0 → patchworks-0.11.0}/.markdownlint-cli2.yaml +0 -0
  34. {patchworks-0.10.0 → patchworks-0.11.0}/README.md +0 -0
  35. {patchworks-0.10.0 → patchworks-0.11.0}/cliff.toml +0 -0
  36. {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/chunks.md +0 -0
  37. {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/cluster.md +0 -0
  38. {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/io.md +0 -0
  39. {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/merge_tile_labels.md +0 -0
  40. {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/plugins/cellpose.md +0 -0
  41. {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/plugins/napari.md +0 -0
  42. {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/plugins/ome_zarr.md +0 -0
  43. {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/relabel.md +0 -0
  44. {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/tile_process.md +0 -0
  45. {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/cellpose_2d.md +0 -0
  46. {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/cellpose_2d.py +0 -0
  47. {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/cellpose_3d.md +0 -0
  48. {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/cellpose_3d.py +0 -0
  49. {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/custom.md +0 -0
  50. {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/custom_method.py +0 -0
  51. {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/standalone_merge.md +0 -0
  52. {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/stardist.md +0 -0
  53. {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/stardist_2d.py +0 -0
  54. {patchworks-0.10.0 → patchworks-0.11.0}/docs/getting_started.md +0 -0
  55. {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/gpu_distributed.md +0 -0
  56. {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/merging.md +0 -0
  57. {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/ome_zarr_napari.md +0 -0
  58. {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/performance.md +0 -0
  59. {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/pitfalls.md +0 -0
  60. {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/skip_empty.md +0 -0
  61. {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/tiling.md +0 -0
  62. {patchworks-0.10.0 → patchworks-0.11.0}/docs/index.md +0 -0
  63. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_chunks.py +0 -0
  64. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_cluster.py +0 -0
  65. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_core.py +0 -0
  66. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_io.py +0 -0
  67. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_merge.py +0 -0
  68. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_relabel.py +0 -0
  69. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/plugins/__init__.py +0 -0
  70. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/plugins/cellpose.py +0 -0
  71. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/plugins/napari.py +0 -0
  72. {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/plugins/ome_zarr.py +0 -0
  73. {patchworks-0.10.0 → patchworks-0.11.0}/tests/test_core.py +0 -0
  74. {patchworks-0.10.0 → patchworks-0.11.0}/tests/test_napari.py +0 -0
  75. {patchworks-0.10.0 → patchworks-0.11.0}/tests/test_ome_zarr.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: patchworks
3
- Version: 0.10.0
3
+ Version: 0.11.0
4
4
  Summary: Tiled processing of arbitrarily large images with globally consistent labels
5
5
  Project-URL: Homepage, https://github.com/imcf/patchworks
6
6
  Project-URL: Issues, https://github.com/imcf/patchworks/issues
@@ -63,6 +63,9 @@ Requires-Dist: psutil; extra == 'io'
63
63
  Requires-Dist: tqdm; extra == 'io'
64
64
  Provides-Extra: napari
65
65
  Requires-Dist: napari[all]; extra == 'napari'
66
+ Provides-Extra: workflow
67
+ Requires-Dist: snakemake-executor-plugin-slurm; extra == 'workflow'
68
+ Requires-Dist: snakemake>=8; extra == 'workflow'
66
69
  Description-Content-Type: text/markdown
67
70
 
68
71
  # patchworks
@@ -0,0 +1,213 @@
1
+ # Cluster workflow (Snakemake + SLURM)
2
+
3
+ `tile_process` runs every tile **serially on one GPU**. For a large 3-D image
4
+ that can be days. The bundled Snakemake workflow instead submits **one GPU job
5
+ per tile**, so with *N* GPUs the segmentation is ~*N*× faster. This page walks
6
+ through running it from scratch.
7
+
8
+ ```text
9
+ convert ──▶ prepare (checkpoint) ──▶ segment {tile} ──▶ merge
10
+ one GPU SLURM job per tile
11
+ ```
12
+
13
+ ## 1. Get the workflow
14
+
15
+ The workflow lives in the `workflow/` directory of the patchworks repository
16
+ (it is not shipped inside the pip package — it is a set of Snakemake files you
17
+ run):
18
+
19
+ ```bash
20
+ git clone https://github.com/imcf/patchworks
21
+ cd patchworks/workflow
22
+ ```
23
+
24
+ ## 2. Install the dependencies
25
+
26
+ You need patchworks with the workflow + reader + segmentation extras, in the
27
+ environment Snakemake will use:
28
+
29
+ ```bash
30
+ pip install "patchworks[workflow,cellpose,imaris,bioio]"
31
+ ```
32
+
33
+ - `workflow` → Snakemake + the SLURM executor plugin
34
+ - `cellpose` → the segmentation model
35
+ - `imaris` / `bioio` → read your input format (`.ims`, `.czi`, `.lif`, …)
36
+
37
+ On a cluster, do this inside a conda/venv that the compute nodes can see, or let
38
+ each rule activate a conda env (see *Conda*, below).
39
+
40
+ ## 3. Configure the run
41
+
42
+ Copy and edit `config/config.yaml`. Every field:
43
+
44
+ ```yaml
45
+ # input / output
46
+ input: "/data/scan.ims" # .ims/.czi/.lif/.nd2/ome-tiff/.zarr
47
+ work_dir: "/scratch/results" # everything is written here
48
+
49
+ # conversion (input → pyramidal OME-ZARR)
50
+ reuse_pyramid: true # .ims: copy its own pyramid (fast)
51
+ convert_chunks: null # null → bounded auto chunks; or [c,z,y,x]
52
+ shard: false # true → pack chunks into shards (fewer files)
53
+
54
+ # tiling
55
+ channel: 0 # channel to segment (null = keep all)
56
+ level: 0 # pyramid level (0 = full resolution)
57
+ tile_shape: "auto" # "auto", or e.g. [16, 1024, 1024] (zyx)
58
+ overlap: 30 # halo ≈ one object diameter
59
+ skip_empty: true # skip background tiles
60
+ empty_threshold: null # null → Otsu
61
+
62
+ # segmentation
63
+ method: "cellpose" # "cellpose" (GPU) or "threshold" (no GPU)
64
+ label_name: "cellpose" # name under image.zarr/labels/
65
+ cellpose:
66
+ model: "cyto3"
67
+ diameter: 30
68
+ do_3D: true
69
+ gpu: true
70
+ # extra model.eval() kwargs, e.g. flow_threshold: 0.4
71
+
72
+ # label pyramid
73
+ pyramid_levels: 5
74
+ pyramid_downscale: 2
75
+ sequential_labels: true # renumber labels to a contiguous 1..N
76
+ ```
77
+
78
+ !!! tip "Tile size vs runtime"
79
+ `tile_shape: "auto"` sizes each tile to your GPU's VRAM. Smaller tiles =
80
+ more (faster) jobs; very large 3-D tiles are slow. Keep `do_3D: false` (2-D
81
+ per slice) if your objects segment fine per slice — it is much faster.
82
+
83
+ ## 4. Dry-run (always do this first)
84
+
85
+ Check the plan without running anything:
86
+
87
+ ```bash
88
+ python -m snakemake -s Snakefile --configfile config/config.yaml -n -p
89
+ ```
90
+
91
+ You should see `convert`, `prepare`, and a note that the **checkpoint** will add
92
+ the `segment` jobs after `prepare` runs. (The number of segment jobs is only
93
+ known after `prepare` decides which tiles are non-empty.)
94
+
95
+ ## 5a. Run locally (single machine)
96
+
97
+ ```bash
98
+ python -m snakemake -s Snakefile --configfile config/config.yaml --cores 8
99
+ ```
100
+
101
+ Tiles run on the local machine (one at a time on the GPU). Good for a small
102
+ image or a smoke test.
103
+
104
+ ## 5b. Run on SLURM (one GPU job per tile)
105
+
106
+ Edit `profile/slurm/config.yaml` for **your** cluster — partitions, account,
107
+ and the GPU request:
108
+
109
+ ```yaml
110
+ executor: slurm
111
+ jobs: 64 # max concurrent SLURM jobs ≈ GPUs you can grab
112
+ default-resources:
113
+ slurm_partition: "cpu" # your CPU partition
114
+ # slurm_account: "my_account"
115
+ mem_mb: 16000
116
+ cpus_per_task: 4
117
+ runtime: 60
118
+ set-resources:
119
+ segment: # the GPU step
120
+ slurm_partition: "gpu" # your GPU partition
121
+ slurm_extra: "'--gres=gpu:1'"
122
+ mem_mb: 32000
123
+ runtime: 120
124
+ merge:
125
+ mem_mb: 128000
126
+ runtime: 240
127
+ ```
128
+
129
+ Then launch (from a login node — Snakemake submits and watches the jobs):
130
+
131
+ ```bash
132
+ python -m snakemake --workflow-profile profile/slurm \
133
+ --configfile config/config.yaml
134
+ ```
135
+
136
+ Snakemake submits `convert`, then `prepare`, then **one `segment` job per
137
+ non-empty tile** (up to `jobs:` at once → that many GPUs in parallel), then
138
+ `merge`. Raise `jobs:` to use more GPUs.
139
+
140
+ !!! note "GPU request flag"
141
+ Clusters differ. `--gres=gpu:1` is common; some need `--gpus=1` or a
142
+ specific gres name (`--gres=gpu:a100:1`). Put whatever `sbatch` flag your
143
+ cluster needs in `slurm_extra`.
144
+
145
+ ## 6. Monitor
146
+
147
+ - **Snakemake** prints each job as it submits/finishes and a `X of Y steps`
148
+ counter.
149
+ - **SLURM**: `squeue --me` shows your queued/running jobs (`smk-segment`, …);
150
+ logs land where your profile/cluster sends them.
151
+ - **patchworks** logs (`processing tile k/N`, ETA) are inside each job's stdout.
152
+
153
+ ## 7. Output
154
+
155
+ Everything is under `work_dir`:
156
+
157
+ ```text
158
+ results/
159
+ image.zarr/ # converted, pyramidal OME-ZARR
160
+ image.zarr/labels/<name>/ # the segmentation (multi-scale, calibrated)
161
+ ```
162
+
163
+ The labels live **inside** the image store. View image + labels together:
164
+
165
+ ```python
166
+ from patchworks.plugins.napari import view_in_napari
167
+ view_in_napari("/scratch/results/image.zarr") # auto-loads the labels
168
+ ```
169
+
170
+ ## 8. Re-running and resuming
171
+
172
+ Snakemake is resumable — if jobs fail or you cancel, just relaunch the same
173
+ command and it picks up only the missing tiles. To force a clean rerun, delete
174
+ `work_dir` (or the relevant outputs).
175
+
176
+ ## Conda (optional)
177
+
178
+ To have each rule run in a named conda env instead of the active one, add
179
+ `--use-conda` and point the rules at an env; or activate your env in a SLURM
180
+ prologue. The simplest path is a single shared env that the compute nodes see.
181
+
182
+ ## Troubleshooting
183
+
184
+ | Symptom | Fix |
185
+ |---------|-----|
186
+ | `snakemake: command not found` | use `python -m snakemake` |
187
+ | Segment jobs pend forever | wrong `slurm_partition`/`slurm_extra` GPU flag for your cluster |
188
+ | `cellpose is not installed` in a job | the job's env lacks `patchworks[cellpose]` |
189
+ | Reading the input fails | install the matching reader (`patchworks[imaris]`/`[bioio]` + a `bioio-*`) |
190
+ | Out of GPU memory | smaller `tile_shape`, or `do_3D: false` |
191
+ | Very slow | confirm GPU is used (`nvidia-smi`); try 2-D or a lower `level` |
192
+
193
+ ## How it works (for the curious)
194
+
195
+ The rule scripts are thin wrappers over patchworks' public API, so you can build
196
+ the same per-tile distribution yourself:
197
+
198
+ ```python
199
+ from patchworks import (
200
+ load_ome_zarr, spatial_tiles, create_stage, stage_tile, merge_tile_labels
201
+ )
202
+ from patchworks.plugins.ome_zarr import write_labels
203
+
204
+ img = load_ome_zarr("image.zarr", channel=0)
205
+ tiles = spatial_tiles(img.shape, tile_shape=(16, 1024, 1024))
206
+ create_stage("stage.zarr", img.shape, (16, 1024, 1024))
207
+ # (distribute these across jobs:)
208
+ for i in range(len(tiles)):
209
+ stage_tile(img, my_fn, "stage.zarr", i, tile_shape=(16, 1024, 1024), overlap=30)
210
+ merged = merge_tile_labels("stage.zarr", input_component="staged",
211
+ write_to="merged.zarr", sequential_labels=True)
212
+ write_labels("image.zarr", merged, name="cells")
213
+ ```
@@ -39,6 +39,7 @@ nav:
39
39
  - Empty tile skipping: guide/skip_empty.md
40
40
  - GPU & distributed: guide/gpu_distributed.md
41
41
  - Performance & memory: guide/performance.md
42
+ - Cluster workflow (Snakemake): guide/snakemake.md
42
43
  - OME-ZARR & napari: guide/ome_zarr_napari.md
43
44
  - Pitfalls: guide/pitfalls.md
44
45
  - Examples:
@@ -58,6 +58,8 @@ bioio = [
58
58
  imaris = ["imaris-ims-file-reader"]
59
59
  # napari enables the interactive viewer plugin.
60
60
  napari = ["napari[all]"]
61
+ # workflow runs the Snakemake pipeline (per-tile SLURM jobs across GPUs).
62
+ workflow = ["snakemake>=8", "snakemake-executor-plugin-slurm"]
61
63
  dev = ["pytest", "pytest-cov", "scikit-image", "psutil", "tqdm"]
62
64
  docs = ["mkdocs-material>=9.0", "mkdocstrings[python]>=0.24"]
63
65
  all = [
@@ -32,6 +32,7 @@ from importlib.metadata import version as _pkg_version
32
32
  from ._chunks import auto_overlap, auto_tile_shape, auto_tile_shape_cellpose
33
33
  from ._cluster import make_local_cluster
34
34
  from ._core import tile_process
35
+ from ._distributed import create_stage, spatial_tiles, stage_tile
35
36
  from ._io import estimate_empty_tiles, load_ome_zarr
36
37
  from ._merge import merge_tile_labels
37
38
  from ._relabel import relabel_sequential_array, relabel_sequential_zarr
@@ -51,4 +52,7 @@ __all__ = [
51
52
  "make_local_cluster",
52
53
  "relabel_sequential_array",
53
54
  "relabel_sequential_zarr",
55
+ "spatial_tiles",
56
+ "create_stage",
57
+ "stage_tile",
54
58
  ]
@@ -0,0 +1,140 @@
1
+ """Per-tile building blocks for distributed processing.
2
+
3
+ ``tile_process`` runs every tile and merges in one process. To spread tiles
4
+ across separate jobs (e.g. one SLURM GPU job per tile) you need to process a
5
+ *single* tile independently and merge later. These helpers expose exactly that:
6
+ :func:`spatial_tiles` enumerates the tiles, :func:`create_stage` makes the
7
+ shared output store, and :func:`stage_tile` runs ``fn`` on one tile and writes
8
+ it into that store. Stitch the result with
9
+ :func:`patchworks.merge_tile_labels` (or ``zarr_native_merge``).
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ import itertools
15
+ from pathlib import Path
16
+ from typing import Callable, Union
17
+
18
+ import numpy as np
19
+ import zarr
20
+
21
+
22
+ def spatial_tiles(
23
+ shape: tuple[int, ...], tile_shape: tuple[int, ...]
24
+ ) -> list[tuple[slice, ...]]:
25
+ """Enumerate the tiles covering *shape*, in row-major order.
26
+
27
+ Parameters
28
+ ----------
29
+ shape : tuple of int
30
+ Spatial array shape.
31
+ tile_shape : tuple of int
32
+ Tile shape.
33
+
34
+ Returns
35
+ -------
36
+ list of tuple of slice
37
+ One slice tuple per tile (the same order ``estimate_empty_tiles``'s
38
+ ``occupancy`` grid uses when ravelled).
39
+ """
40
+ grids = [range(0, s, t) for s, t in zip(shape, tile_shape)]
41
+ return [
42
+ tuple(
43
+ slice(o, min(o + t, s))
44
+ for o, t, s in zip(starts, tile_shape, shape)
45
+ )
46
+ for starts in itertools.product(*grids)
47
+ ]
48
+
49
+
50
+ def create_stage(
51
+ stage_path: Union[str, Path],
52
+ shape: tuple[int, ...],
53
+ tile_shape: tuple[int, ...],
54
+ *,
55
+ component: str = "staged",
56
+ dtype=np.int32,
57
+ ) -> str:
58
+ """Create the empty (zero-filled) shared stage store for tiled writes.
59
+
60
+ Parameters
61
+ ----------
62
+ stage_path : str or Path
63
+ Destination ``.zarr`` store.
64
+ shape : tuple of int
65
+ Full (spatial) array shape.
66
+ tile_shape : tuple of int
67
+ Chunk = tile shape (one chunk per tile, so jobs write disjoint files).
68
+ component : str, optional
69
+ Array name inside the store (default ``"staged"``).
70
+ dtype : data-type, optional
71
+ Label dtype (default ``int32``).
72
+
73
+ Returns
74
+ -------
75
+ str
76
+ The stage store path.
77
+ """
78
+ root = zarr.open_group(str(stage_path), mode="w")
79
+ root.create_array(
80
+ name=component, shape=shape, chunks=tile_shape, dtype=dtype
81
+ )
82
+ return str(stage_path)
83
+
84
+
85
+ def stage_tile(
86
+ image,
87
+ fn: Callable[[np.ndarray], np.ndarray],
88
+ stage_path: Union[str, Path],
89
+ index: int,
90
+ *,
91
+ tile_shape: tuple[int, ...],
92
+ overlap: int = 0,
93
+ component: str = "staged",
94
+ ) -> int:
95
+ """Run *fn* on a single tile and write it into the shared stage store.
96
+
97
+ Reads the tile (expanded by *overlap* on every side for boundary context),
98
+ runs *fn*, trims the halo back off, and writes the result to the tile's
99
+ disjoint chunk of ``stage_path/component`` — so many of these can run
100
+ concurrently (one per job) without conflicts.
101
+
102
+ Parameters
103
+ ----------
104
+ image : array-like
105
+ The full image (dask/zarr/NumPy), indexable by slices.
106
+ fn : callable
107
+ ``(ndarray) -> ndarray`` returning integer labels of the same shape.
108
+ stage_path : str or Path
109
+ Stage store created by :func:`create_stage`.
110
+ index : int
111
+ Tile index into :func:`spatial_tiles`.
112
+ tile_shape : tuple of int
113
+ Tile shape (must match the stage store's chunks).
114
+ overlap : int, optional
115
+ Halo added on every side before calling *fn*.
116
+ component : str, optional
117
+ Array name inside the stage store.
118
+
119
+ Returns
120
+ -------
121
+ int
122
+ The processed tile *index*.
123
+ """
124
+ shape = image.shape
125
+ sl = spatial_tiles(shape, tile_shape)[index]
126
+ expanded, trims = [], []
127
+ for s, dim in zip(sl, shape):
128
+ lo = max(0, s.start - overlap)
129
+ hi = min(dim, s.stop + overlap)
130
+ expanded.append(slice(lo, hi))
131
+ trims.append((s.start - lo, hi - s.stop))
132
+ block = np.asarray(image[tuple(expanded)])
133
+ out = np.asarray(fn(block))
134
+ sel = tuple(
135
+ slice(left, out.shape[i] - right)
136
+ for i, (left, right) in enumerate(trims)
137
+ )
138
+ dst = zarr.open_group(str(stage_path), mode="r+")[component]
139
+ dst[sl] = out[sel].astype(dst.dtype)
140
+ return index
@@ -0,0 +1,45 @@
1
+ """Tests for the per-tile distributed building blocks."""
2
+
3
+ import numpy as np
4
+
5
+ from patchworks import (
6
+ create_stage,
7
+ merge_tile_labels,
8
+ spatial_tiles,
9
+ stage_tile,
10
+ )
11
+
12
+
13
+ def _fn(tile):
14
+ from skimage.measure import label
15
+
16
+ return label(tile > 0).astype("int32")
17
+
18
+
19
+ def test_spatial_tiles_cover():
20
+ """Tiles tile the array exactly, in row-major order."""
21
+ tiles = spatial_tiles((4, 5), (2, 2))
22
+ assert len(tiles) == 2 * 3 # ceil(4/2) * ceil(5/2)
23
+ assert tiles[0] == (slice(0, 2), slice(0, 2))
24
+ assert tiles[-1] == (slice(2, 4), slice(4, 5)) # clipped last tile
25
+
26
+
27
+ def test_stage_then_merge_stitches_boundary(tmp_path):
28
+ """Per-tile staging + merge reproduces a cross-boundary single object."""
29
+ img = np.zeros((16, 32), "uint16")
30
+ img[4:12, 8:24] = 500 # block straddling the x=16 tile boundary
31
+
32
+ stage = str(tmp_path / "stage.zarr")
33
+ tile = (16, 16)
34
+ create_stage(stage, img.shape, tile)
35
+ for i in range(len(spatial_tiles(img.shape, tile))):
36
+ stage_tile(img, _fn, stage, i, tile_shape=tile, overlap=4)
37
+
38
+ merged = merge_tile_labels(
39
+ stage,
40
+ write_to=str(tmp_path / "out.zarr"),
41
+ input_component="staged",
42
+ sequential_labels=True,
43
+ ).compute()
44
+ ids = np.unique(merged[merged > 0])
45
+ assert ids.size == 1, f"object split into {ids.size} labels"
@@ -0,0 +1 @@
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": {}}
@@ -0,0 +1 @@
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"}}
@@ -0,0 +1 @@
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": {}}
@@ -0,0 +1 @@
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"}}
@@ -0,0 +1 @@
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"}}
@@ -0,0 +1 @@
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"}}
@@ -0,0 +1 @@
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"}}
@@ -0,0 +1 @@
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": {}}
@@ -0,0 +1,95 @@
1
+ # patchworks Snakemake workflow
2
+
3
+ A SLURM-ready pipeline that segments an arbitrarily large image and spreads the
4
+ expensive Cellpose step across **many GPUs** — one tile per SLURM job.
5
+
6
+ > **Full step-by-step guide:**
7
+ > <https://imcf.one/patchworks/guide/snakemake/> — install, configure every
8
+ > field, dry-run, local vs SLURM, monitoring, outputs and troubleshooting.
9
+
10
+ ```text
11
+ convert ──▶ prepare (checkpoint) ──▶ segment {tile} ──▶ merge
12
+ one GPU job/tile
13
+ ```
14
+
15
+ ## Why
16
+
17
+ `tile_process` runs tiles serially on a single GPU. This workflow instead
18
+ submits **one GPU job per (non-empty) tile**, so N GPUs cut the wall-time ~N×.
19
+ Each job writes a disjoint chunk of a shared stage store; a final CPU job
20
+ stitches labels across tile boundaries and writes them back into the image as a
21
+ calibrated, multi-scale `labels/` group.
22
+
23
+ ## Steps
24
+
25
+ 1. **convert** — input (`.ims`, `.czi`, `.lif`, `.nd2`, OME-TIFF, `.zarr`) → a
26
+ pyramidal OME-ZARR (`image.zarr`).
27
+ 2. **prepare** (checkpoint) — plan tiles, skip background tiles, create the
28
+ empty `stage.zarr`, and list the tiles to segment.
29
+ 3. **segment {tile}** — one GPU job per tile: read tile + halo, run Cellpose,
30
+ trim, write to the stage. Scattered across the cluster.
31
+ 4. **merge** — zarr-native boundary stitch + renumber, written into
32
+ `image.zarr/labels/<name>/`.
33
+
34
+ ## Install
35
+
36
+ ```bash
37
+ pip install "patchworks[workflow,cellpose,imaris,bioio]"
38
+ # workflow → snakemake + the SLURM executor plugin
39
+ ```
40
+
41
+ ## Configure
42
+
43
+ Edit `config/config.yaml` (input, output dir, channel, tile shape, Cellpose
44
+ model/diameter/`do_3D`, …) and `profile/slurm/config.yaml` (partitions,
45
+ account, GPU request).
46
+
47
+ ## Run
48
+
49
+ ```bash
50
+ # locally (single machine)
51
+ snakemake --cores 8 --configfile config/config.yaml
52
+
53
+ # on SLURM — one GPU job per tile, up to `jobs:` in parallel
54
+ snakemake --workflow-profile profile/slurm --configfile config/config.yaml
55
+ ```
56
+
57
+ The GPU request lives in `profile/slurm/config.yaml` under
58
+ `set-resources: segment:` (`--gres=gpu:1`). Raise `jobs:` to use more GPUs at
59
+ once.
60
+
61
+ ## Output
62
+
63
+ `<work_dir>/image.zarr` — the image plus `labels/<label_name>/` (multi-scale,
64
+ calibrated). Open it directly:
65
+
66
+ ```python
67
+ from patchworks.plugins.napari import view_in_napari
68
+ view_in_napari("<work_dir>/image.zarr") # auto-loads the labels
69
+ ```
70
+
71
+ ## Layout
72
+
73
+ ```text
74
+ workflow/
75
+ Snakefile # includes the rule files below
76
+ rules/ # convert.smk, segment.smk, merge.smk, common.smk
77
+ scripts/ # thin wrappers over patchworks' public API
78
+ config/config.yaml
79
+ profile/slurm/config.yaml
80
+ ```
81
+
82
+ The rule scripts are intentionally thin — the work is done by patchworks'
83
+ public API (`spatial_tiles`, `create_stage`, `stage_tile`, `merge_tile_labels`,
84
+ `write_labels`), so the same per-tile distribution is available from your own
85
+ code too.
86
+
87
+ ## Notes
88
+
89
+ - Tiles overlap on read (halo) but write **disjoint** regions, so the per-tile
90
+ jobs are safe to run concurrently.
91
+ - Background tiles are skipped (`skip_empty`), so only occupied tiles become
92
+ jobs.
93
+ - `method:` selects `cellpose` (default) or a simple `threshold` (no GPU —
94
+ handy for testing or quick masks).
95
+ - For very large stores, set `shard: true` in the config to cut the file count.
@@ -0,0 +1,24 @@
1
+ # patchworks Snakemake workflow.
2
+ #
3
+ # convert ──▶ prepare (checkpoint) ──▶ segment {tile} ──▶ merge
4
+ # one GPU job per tile
5
+ #
6
+ # Run locally:
7
+ # snakemake --cores 8 --configfile config/config.yaml
8
+ # Run on SLURM (one GPU job per tile, many GPUs in parallel):
9
+ # snakemake --workflow-profile profile/slurm --configfile config/config.yaml
10
+
11
+ import json
12
+ from pathlib import Path
13
+
14
+ configfile: "config/config.yaml"
15
+
16
+ include: "rules/common.smk"
17
+ include: "rules/convert.smk"
18
+ include: "rules/segment.smk"
19
+ include: "rules/merge.smk"
20
+
21
+
22
+ rule all:
23
+ input:
24
+ f"{WORK}/labels.done",
@@ -0,0 +1,42 @@
1
+ # patchworks Snakemake workflow — configuration.
2
+ # Copy and edit, then run: snakemake --workflow-profile profile/slurm
3
+ #
4
+ # The expensive Cellpose step is split into one SLURM job per tile, so many
5
+ # GPUs run in parallel; the merge is a single CPU job.
6
+
7
+ # ---- input / output ---------------------------------------------------------
8
+ input: "/path/to/scan.ims" # .ims, .czi, .lif, .nd2, ome-tiff, or an .zarr
9
+ work_dir: "/path/to/results" # everything is written under here
10
+ # Final outputs (under work_dir):
11
+ # image.zarr — converted, pyramidal OME-ZARR
12
+ # image.zarr/labels/<name>/ — the segmentation (multi-scale), in the image
13
+
14
+ # ---- conversion -------------------------------------------------------------
15
+ reuse_pyramid: true # for .ims: copy its own pyramid (fast); else rebuild
16
+ convert_chunks: null # null → patchworks' bounded auto chunks; or [c,z,y,x]
17
+ shard: false # true → pack chunks into shards (zarr v3; fewer files)
18
+
19
+ # ---- tiling -----------------------------------------------------------------
20
+ channel: 0 # channel to segment
21
+ level: 0 # pyramid level to segment (0 = full res)
22
+ tile_shape: "auto" # "auto", or e.g. [16, 1024, 1024] (zyx)
23
+ overlap: 30 # halo (≈ one object diameter)
24
+ skip_empty: true # skip background tiles
25
+ empty_threshold: null # null → Otsu; or a number
26
+
27
+ # ---- segmentation -----------------------------------------------------------
28
+ method: "cellpose" # "cellpose" (GPU) or "threshold" (simple, no GPU; testing)
29
+ label_name: "cellpose"
30
+ cellpose:
31
+ model: "cyto3"
32
+ diameter: 30
33
+ do_3D: true
34
+ gpu: true
35
+ # extra kwargs forwarded to model.eval(), e.g.:
36
+ # flow_threshold: 0.4
37
+ # cellprob_threshold: 0.0
38
+
39
+ # ---- pyramid for the labels -------------------------------------------------
40
+ pyramid_levels: 5
41
+ pyramid_downscale: 2
42
+ sequential_labels: true # renumber merged labels to a contiguous 1..N range
@@ -0,0 +1,40 @@
1
+ # Snakemake 8+ SLURM profile for the patchworks workflow.
2
+ # Needs: pip install "patchworks[workflow]" (snakemake + slurm executor plugin)
3
+ # Run: snakemake --workflow-profile profile/slurm --configfile config/config.yaml
4
+ #
5
+ # Edit the partitions / account / GPU request for your cluster.
6
+
7
+ executor: slurm
8
+ jobs: 64 # max concurrent SLURM jobs (≈ how many GPUs you can grab at once)
9
+ latency-wait: 60
10
+ rerun-incomplete: true
11
+ printshellcmds: true
12
+ keep-going: true
13
+
14
+ # Defaults for every rule (CPU jobs).
15
+ default-resources:
16
+ slurm_partition: "cpu"
17
+ # slurm_account: "my_account"
18
+ mem_mb: 16000
19
+ cpus_per_task: 4
20
+ runtime: 60 # minutes
21
+
22
+ # Per-rule overrides.
23
+ set-resources:
24
+ convert:
25
+ mem_mb: 64000
26
+ cpus_per_task: 8
27
+ runtime: 360
28
+ prepare:
29
+ mem_mb: 32000
30
+ runtime: 120
31
+ segment: # one GPU per tile — this is what spreads Cellpose across GPUs
32
+ slurm_partition: "gpu"
33
+ slurm_extra: "'--gres=gpu:1'"
34
+ mem_mb: 32000
35
+ cpus_per_task: 4
36
+ runtime: 120
37
+ merge:
38
+ mem_mb: 128000
39
+ cpus_per_task: 8
40
+ runtime: 240
@@ -0,0 +1,13 @@
1
+ # Shared paths and helpers for the patchworks workflow.
2
+
3
+ WORK = config["work_dir"]
4
+ IMAGE = f"{WORK}/image.zarr"
5
+ TILES = f"{WORK}/tiles.json"
6
+ STAGE = f"{WORK}/stage.zarr"
7
+
8
+
9
+ def occupied_done(wildcards):
10
+ """Per-tile markers for the occupied tiles (resolved after the checkpoint)."""
11
+ tiles = checkpoints.prepare.get().output.tiles
12
+ occupied = json.loads(Path(tiles).read_text())["occupied"]
13
+ return [f"{WORK}/seg/{i}.done" for i in occupied]
@@ -0,0 +1,7 @@
1
+ # Convert the input to a pyramidal OME-ZARR.
2
+
3
+ rule convert:
4
+ output:
5
+ directory(IMAGE),
6
+ script:
7
+ "../scripts/convert.py"
@@ -0,0 +1,9 @@
1
+ # Stitch labels across tile boundaries and write them into the image.
2
+
3
+ rule merge:
4
+ input:
5
+ occupied_done,
6
+ output:
7
+ touch(f"{WORK}/labels.done"),
8
+ script:
9
+ "../scripts/merge.py"
@@ -0,0 +1,22 @@
1
+ # Plan tiles (checkpoint) and segment each tile on a GPU.
2
+
3
+ checkpoint prepare:
4
+ input:
5
+ IMAGE,
6
+ output:
7
+ tiles=TILES,
8
+ stage=directory(STAGE),
9
+ script:
10
+ "../scripts/prepare_tiles.py"
11
+
12
+
13
+ rule segment:
14
+ """Segment one tile on a GPU and write it into the stage store."""
15
+ input:
16
+ tiles=TILES,
17
+ stage=STAGE,
18
+ image=IMAGE,
19
+ output:
20
+ f"{WORK}/seg/{{index}}.done",
21
+ script:
22
+ "../scripts/segment_tile.py"
@@ -0,0 +1,113 @@
1
+ """Config glue for the patchworks Snakemake workflow.
2
+
3
+ The heavy lifting lives in patchworks' public API
4
+ (``spatial_tiles``/``create_stage``/``stage_tile``/``merge_tile_labels``);
5
+ these helpers only turn the Snakemake config into the right arguments.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import json
11
+ from pathlib import Path
12
+
13
+ from patchworks import load_ome_zarr
14
+
15
+
16
+ def open_image(work_dir, channel, level):
17
+ """Open the converted image for segmentation.
18
+
19
+ Parameters
20
+ ----------
21
+ work_dir : str or Path
22
+ Workflow output directory containing ``image.zarr``.
23
+ channel : int or None
24
+ Channel to select.
25
+ level : int
26
+ Pyramid level to read.
27
+
28
+ Returns
29
+ -------
30
+ da.Array
31
+ The (lazy) image array.
32
+ """
33
+ return load_ome_zarr(
34
+ str(Path(work_dir) / "image.zarr"), channel=channel, level=level
35
+ )
36
+
37
+
38
+ def stage_path(work_dir):
39
+ """Path of the staged-labels store.
40
+
41
+ Parameters
42
+ ----------
43
+ work_dir : str or Path
44
+ Workflow output directory.
45
+
46
+ Returns
47
+ -------
48
+ str
49
+ ``<work_dir>/stage.zarr``.
50
+ """
51
+ return str(Path(work_dir) / "stage.zarr")
52
+
53
+
54
+ def load_tiles_json(path):
55
+ """Load the tile manifest written by ``prepare_tiles.py``.
56
+
57
+ Parameters
58
+ ----------
59
+ path : str or Path
60
+ Path to ``tiles.json``.
61
+
62
+ Returns
63
+ -------
64
+ dict
65
+ The manifest (``tile_shape``, ``overlap``, ``occupied`` indices, …).
66
+ """
67
+ return json.loads(Path(path).read_text())
68
+
69
+
70
+ def build_fn(cfg):
71
+ """Build the per-tile segmentation function from the config.
72
+
73
+ Parameters
74
+ ----------
75
+ cfg : dict
76
+ Snakemake config. ``method`` selects ``"cellpose"`` (default) or a
77
+ simple ``"threshold"`` (handy for testing / no-GPU runs).
78
+
79
+ Returns
80
+ -------
81
+ callable
82
+ ``(ndarray) -> ndarray`` returning integer labels.
83
+ """
84
+ method = cfg.get("method", "cellpose")
85
+ if method == "threshold":
86
+
87
+ def fn(tile):
88
+ from skimage.filters import threshold_otsu
89
+ from skimage.measure import label
90
+
91
+ thr = threshold_otsu(tile) if tile.max() > tile.min() else 0
92
+ return label(tile > thr).astype("int32")
93
+
94
+ return fn
95
+
96
+ if method == "cellpose":
97
+ from patchworks.plugins.cellpose import cellpose_fn
98
+
99
+ cp = cfg["cellpose"]
100
+ extra = {
101
+ k: v
102
+ for k, v in cp.items()
103
+ if k not in ("model", "diameter", "do_3D", "gpu")
104
+ }
105
+ return cellpose_fn(
106
+ cp.get("model", "cyto3"),
107
+ gpu=cp.get("gpu", True),
108
+ diameter=cp.get("diameter"),
109
+ do_3D=cp.get("do_3D", False),
110
+ **extra,
111
+ )
112
+
113
+ raise ValueError(f"unknown segmentation method: {method!r}")
@@ -0,0 +1,15 @@
1
+ """Snakemake script: convert the input to a pyramidal OME-ZARR."""
2
+
3
+ from patchworks.plugins.ome_zarr import to_ome_zarr
4
+
5
+ cfg = snakemake.config # noqa: F821 (injected by Snakemake)
6
+ chunks = cfg.get("convert_chunks")
7
+ to_ome_zarr(
8
+ cfg["input"],
9
+ snakemake.output[0], # noqa: F821
10
+ chunks=tuple(chunks) if chunks else None,
11
+ shard=bool(cfg.get("shard", False)),
12
+ reuse_pyramid=bool(cfg.get("reuse_pyramid", False)),
13
+ progress=False,
14
+ overwrite=True,
15
+ )
@@ -0,0 +1,40 @@
1
+ """Snakemake script: merge the staged tiles into one labelled OME-ZARR.
2
+
3
+ Runs patchworks' zarr-native boundary merge (stitches labels across tile
4
+ boundaries, optionally renumbers them) and writes the result back into the
5
+ image store under ``labels/<name>/`` as a calibrated, multi-scale pyramid.
6
+ """
7
+
8
+ import shutil
9
+ from pathlib import Path
10
+
11
+ from patchworks import merge_tile_labels
12
+ from patchworks.plugins.ome_zarr import write_labels
13
+
14
+ from _pw import stage_path
15
+
16
+ cfg = snakemake.config # noqa: F821
17
+ work_dir = cfg["work_dir"]
18
+ image_store = str(Path(work_dir) / "image.zarr")
19
+ merged_store = str(Path(work_dir) / "_merged.zarr")
20
+
21
+ merged = merge_tile_labels(
22
+ stage_path(work_dir),
23
+ write_to=merged_store,
24
+ input_component="staged",
25
+ sequential_labels=cfg.get("sequential_labels", True),
26
+ progress=False,
27
+ )
28
+ group = write_labels(
29
+ image_store,
30
+ merged,
31
+ name=cfg.get("label_name", "labels"),
32
+ n_levels=int(cfg.get("pyramid_levels", 5)),
33
+ downscale=int(cfg.get("pyramid_downscale", 2)),
34
+ overwrite=True,
35
+ )
36
+
37
+ shutil.rmtree(merged_store, ignore_errors=True)
38
+ shutil.rmtree(stage_path(work_dir), ignore_errors=True)
39
+ print(f"[patchworks] labels written to {group}")
40
+ open(snakemake.output[0], "w").close() # noqa: F821
@@ -0,0 +1,56 @@
1
+ """Snakemake script: plan tiles, create the empty stage store, list work."""
2
+
3
+ import json
4
+ from functools import partial
5
+ from pathlib import Path
6
+
7
+ from patchworks import (
8
+ auto_tile_shape_cellpose,
9
+ create_stage,
10
+ estimate_empty_tiles,
11
+ spatial_tiles,
12
+ )
13
+
14
+ from _pw import open_image, stage_path
15
+
16
+ cfg = snakemake.config # noqa: F821
17
+ work_dir = cfg["work_dir"]
18
+ image = open_image(work_dir, cfg["channel"], cfg["level"])
19
+
20
+ ts = cfg.get("tile_shape", "auto")
21
+ if ts == "auto":
22
+ cp = cfg["cellpose"]
23
+ tile_shape = tuple(
24
+ partial(
25
+ auto_tile_shape_cellpose,
26
+ do_3D=cp.get("do_3D", False),
27
+ use_gpu=cp.get("gpu", True),
28
+ diameter=cp.get("diameter"),
29
+ )(image.shape, image.dtype)
30
+ )
31
+ else:
32
+ tile_shape = tuple(ts)
33
+
34
+ tiles = spatial_tiles(image.shape, tile_shape)
35
+ occupied = list(range(len(tiles)))
36
+ if cfg.get("skip_empty", True):
37
+ info = estimate_empty_tiles(
38
+ image, tile_shape, threshold=cfg.get("empty_threshold")
39
+ )
40
+ occ = info["occupancy"].ravel() # row-major, matches spatial_tiles
41
+ occupied = [i for i in range(len(tiles)) if occ[i]]
42
+
43
+ create_stage(stage_path(work_dir), image.shape, tile_shape)
44
+
45
+ Path(work_dir, "tiles.json").write_text(
46
+ json.dumps(
47
+ {
48
+ "tile_shape": list(tile_shape),
49
+ "overlap": int(cfg.get("overlap", 0)),
50
+ "n_tiles": len(tiles),
51
+ "occupied": occupied,
52
+ },
53
+ indent=2,
54
+ )
55
+ )
56
+ print(f"[patchworks] {len(occupied)}/{len(tiles)} tiles to segment")
@@ -0,0 +1,28 @@
1
+ """Snakemake script: segment ONE tile and write it to the shared stage.
2
+
3
+ Scattered over tile indices, so each tile is its own SLURM job and many GPUs
4
+ run in parallel. Each job writes a disjoint chunk of the stage store.
5
+ """
6
+
7
+ from patchworks import stage_tile
8
+
9
+ from _pw import build_fn, load_tiles_json, open_image, stage_path
10
+
11
+ cfg = snakemake.config # noqa: F821
12
+ index = int(snakemake.wildcards.index) # noqa: F821
13
+ work_dir = cfg["work_dir"]
14
+
15
+ manifest = load_tiles_json(snakemake.input.tiles) # noqa: F821
16
+ image = open_image(work_dir, cfg["channel"], cfg["level"])
17
+
18
+ stage_tile(
19
+ image,
20
+ build_fn(cfg),
21
+ stage_path(work_dir),
22
+ index,
23
+ tile_shape=tuple(manifest["tile_shape"]),
24
+ overlap=int(manifest["overlap"]),
25
+ )
26
+
27
+ open(snakemake.output[0], "w").close() # noqa: F821
28
+ print(f"[patchworks] segmented tile {index}")
File without changes
File without changes
File without changes
File without changes
File without changes