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.
- {patchworks-0.10.0 → patchworks-0.11.0}/PKG-INFO +4 -1
- patchworks-0.11.0/docs/guide/snakemake.md +213 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/mkdocs.yml +1 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/pyproject.toml +2 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/__init__.py +4 -0
- patchworks-0.11.0/src/patchworks/_distributed.py +140 -0
- patchworks-0.11.0/tests/test_distributed.py +45 -0
- patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2ltYWdlLnphcnI= +1 -0
- patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2xhYmVscy5kb25l +1 -0
- patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3N0YWdlLnphcnI= +1 -0
- patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy80LmRvbmU= +1 -0
- patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8wLmRvbmU= +1 -0
- patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8xLmRvbmU= +1 -0
- patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8zLmRvbmU= +1 -0
- patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3RpbGVzLmpzb24= +1 -0
- patchworks-0.11.0/workflow/README.md +95 -0
- patchworks-0.11.0/workflow/Snakefile +24 -0
- patchworks-0.11.0/workflow/config/config.yaml +42 -0
- patchworks-0.11.0/workflow/profile/slurm/config.yaml +40 -0
- patchworks-0.11.0/workflow/rules/common.smk +13 -0
- patchworks-0.11.0/workflow/rules/convert.smk +7 -0
- patchworks-0.11.0/workflow/rules/merge.smk +9 -0
- patchworks-0.11.0/workflow/rules/segment.smk +22 -0
- patchworks-0.11.0/workflow/scripts/_pw.py +113 -0
- patchworks-0.11.0/workflow/scripts/convert.py +15 -0
- patchworks-0.11.0/workflow/scripts/merge.py +40 -0
- patchworks-0.11.0/workflow/scripts/prepare_tiles.py +56 -0
- patchworks-0.11.0/workflow/scripts/segment_tile.py +28 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/.github/workflows/docs.yml +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/.github/workflows/lint.yml +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/.github/workflows/release.yml +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/.gitignore +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/.markdownlint-cli2.yaml +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/README.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/cliff.toml +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/chunks.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/cluster.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/io.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/merge_tile_labels.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/plugins/cellpose.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/plugins/napari.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/plugins/ome_zarr.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/relabel.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/api/tile_process.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/cellpose_2d.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/cellpose_2d.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/cellpose_3d.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/cellpose_3d.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/custom.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/custom_method.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/standalone_merge.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/stardist.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/examples/stardist_2d.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/getting_started.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/gpu_distributed.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/merging.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/ome_zarr_napari.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/performance.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/pitfalls.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/skip_empty.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/guide/tiling.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/docs/index.md +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_chunks.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_cluster.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_core.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_io.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_merge.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/_relabel.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/plugins/__init__.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/plugins/cellpose.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/plugins/napari.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/src/patchworks/plugins/ome_zarr.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/tests/test_core.py +0 -0
- {patchworks-0.10.0 → patchworks-0.11.0}/tests/test_napari.py +0 -0
- {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.
|
|
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,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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|