patchworks 0.9.1__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.9.1 → patchworks-0.11.0}/PKG-INFO +4 -1
  2. patchworks-0.11.0/docs/guide/snakemake.md +213 -0
  3. {patchworks-0.9.1 → patchworks-0.11.0}/mkdocs.yml +1 -0
  4. {patchworks-0.9.1 → patchworks-0.11.0}/pyproject.toml +2 -0
  5. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/__init__.py +4 -0
  6. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/_core.py +148 -4
  7. patchworks-0.11.0/src/patchworks/_distributed.py +140 -0
  8. patchworks-0.11.0/tests/test_distributed.py +45 -0
  9. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2ltYWdlLnphcnI= +1 -0
  10. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L2xhYmVscy5kb25l +1 -0
  11. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3N0YWdlLnphcnI= +1 -0
  12. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy80LmRvbmU= +1 -0
  13. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8wLmRvbmU= +1 -0
  14. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8xLmRvbmU= +1 -0
  15. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3NlZy8zLmRvbmU= +1 -0
  16. patchworks-0.11.0/workflow/.snakemake/metadata/L3RtcC90b3kvb3V0L3RpbGVzLmpzb24= +1 -0
  17. patchworks-0.11.0/workflow/README.md +95 -0
  18. patchworks-0.11.0/workflow/Snakefile +24 -0
  19. patchworks-0.11.0/workflow/config/config.yaml +42 -0
  20. patchworks-0.11.0/workflow/profile/slurm/config.yaml +40 -0
  21. patchworks-0.11.0/workflow/rules/common.smk +13 -0
  22. patchworks-0.11.0/workflow/rules/convert.smk +7 -0
  23. patchworks-0.11.0/workflow/rules/merge.smk +9 -0
  24. patchworks-0.11.0/workflow/rules/segment.smk +22 -0
  25. patchworks-0.11.0/workflow/scripts/_pw.py +113 -0
  26. patchworks-0.11.0/workflow/scripts/convert.py +15 -0
  27. patchworks-0.11.0/workflow/scripts/merge.py +40 -0
  28. patchworks-0.11.0/workflow/scripts/prepare_tiles.py +56 -0
  29. patchworks-0.11.0/workflow/scripts/segment_tile.py +28 -0
  30. {patchworks-0.9.1 → patchworks-0.11.0}/.github/workflows/docs.yml +0 -0
  31. {patchworks-0.9.1 → patchworks-0.11.0}/.github/workflows/lint.yml +0 -0
  32. {patchworks-0.9.1 → patchworks-0.11.0}/.github/workflows/release.yml +0 -0
  33. {patchworks-0.9.1 → patchworks-0.11.0}/.gitignore +0 -0
  34. {patchworks-0.9.1 → patchworks-0.11.0}/.markdownlint-cli2.yaml +0 -0
  35. {patchworks-0.9.1 → patchworks-0.11.0}/README.md +0 -0
  36. {patchworks-0.9.1 → patchworks-0.11.0}/cliff.toml +0 -0
  37. {patchworks-0.9.1 → patchworks-0.11.0}/docs/api/chunks.md +0 -0
  38. {patchworks-0.9.1 → patchworks-0.11.0}/docs/api/cluster.md +0 -0
  39. {patchworks-0.9.1 → patchworks-0.11.0}/docs/api/io.md +0 -0
  40. {patchworks-0.9.1 → patchworks-0.11.0}/docs/api/merge_tile_labels.md +0 -0
  41. {patchworks-0.9.1 → patchworks-0.11.0}/docs/api/plugins/cellpose.md +0 -0
  42. {patchworks-0.9.1 → patchworks-0.11.0}/docs/api/plugins/napari.md +0 -0
  43. {patchworks-0.9.1 → patchworks-0.11.0}/docs/api/plugins/ome_zarr.md +0 -0
  44. {patchworks-0.9.1 → patchworks-0.11.0}/docs/api/relabel.md +0 -0
  45. {patchworks-0.9.1 → patchworks-0.11.0}/docs/api/tile_process.md +0 -0
  46. {patchworks-0.9.1 → patchworks-0.11.0}/docs/examples/cellpose_2d.md +0 -0
  47. {patchworks-0.9.1 → patchworks-0.11.0}/docs/examples/cellpose_2d.py +0 -0
  48. {patchworks-0.9.1 → patchworks-0.11.0}/docs/examples/cellpose_3d.md +0 -0
  49. {patchworks-0.9.1 → patchworks-0.11.0}/docs/examples/cellpose_3d.py +0 -0
  50. {patchworks-0.9.1 → patchworks-0.11.0}/docs/examples/custom.md +0 -0
  51. {patchworks-0.9.1 → patchworks-0.11.0}/docs/examples/custom_method.py +0 -0
  52. {patchworks-0.9.1 → patchworks-0.11.0}/docs/examples/standalone_merge.md +0 -0
  53. {patchworks-0.9.1 → patchworks-0.11.0}/docs/examples/stardist.md +0 -0
  54. {patchworks-0.9.1 → patchworks-0.11.0}/docs/examples/stardist_2d.py +0 -0
  55. {patchworks-0.9.1 → patchworks-0.11.0}/docs/getting_started.md +0 -0
  56. {patchworks-0.9.1 → patchworks-0.11.0}/docs/guide/gpu_distributed.md +0 -0
  57. {patchworks-0.9.1 → patchworks-0.11.0}/docs/guide/merging.md +0 -0
  58. {patchworks-0.9.1 → patchworks-0.11.0}/docs/guide/ome_zarr_napari.md +0 -0
  59. {patchworks-0.9.1 → patchworks-0.11.0}/docs/guide/performance.md +0 -0
  60. {patchworks-0.9.1 → patchworks-0.11.0}/docs/guide/pitfalls.md +0 -0
  61. {patchworks-0.9.1 → patchworks-0.11.0}/docs/guide/skip_empty.md +0 -0
  62. {patchworks-0.9.1 → patchworks-0.11.0}/docs/guide/tiling.md +0 -0
  63. {patchworks-0.9.1 → patchworks-0.11.0}/docs/index.md +0 -0
  64. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/_chunks.py +0 -0
  65. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/_cluster.py +0 -0
  66. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/_io.py +0 -0
  67. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/_merge.py +0 -0
  68. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/_relabel.py +0 -0
  69. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/plugins/__init__.py +0 -0
  70. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/plugins/cellpose.py +0 -0
  71. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/plugins/napari.py +0 -0
  72. {patchworks-0.9.1 → patchworks-0.11.0}/src/patchworks/plugins/ome_zarr.py +0 -0
  73. {patchworks-0.9.1 → patchworks-0.11.0}/tests/test_core.py +0 -0
  74. {patchworks-0.9.1 → patchworks-0.11.0}/tests/test_napari.py +0 -0
  75. {patchworks-0.9.1 → 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.9.1
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
  ]
@@ -20,6 +20,71 @@ from ._relabel import relabel_sequential_zarr
20
20
  logger = logging.getLogger(__name__)
21
21
 
22
22
 
23
+ def _attach_log_file(path: str) -> None:
24
+ """Tee the ``patchworks`` INFO log to *path* (replacing a prior auto file).
25
+
26
+ A previous auto-attached handler is removed first so repeated calls do not
27
+ stack handlers. The package logger level is raised to INFO if needed so the
28
+ messages are actually emitted.
29
+
30
+ Parameters
31
+ ----------
32
+ path : str
33
+ Destination log-file path.
34
+
35
+ Returns
36
+ -------
37
+ None
38
+ """
39
+ pkg = logging.getLogger("patchworks")
40
+ for h in list(pkg.handlers):
41
+ if getattr(h, "_patchworks_auto", False):
42
+ pkg.removeHandler(h)
43
+ h.close()
44
+ handler = logging.FileHandler(path)
45
+ handler._patchworks_auto = True # tag so we can find/replace it later
46
+ handler.setLevel(logging.INFO)
47
+ handler.setFormatter(
48
+ logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")
49
+ )
50
+ pkg.addHandler(handler)
51
+ if pkg.level == logging.NOTSET or pkg.level > logging.INFO:
52
+ pkg.setLevel(logging.INFO)
53
+
54
+
55
+ def _read_amplification(
56
+ native_chunks: tuple[int, ...], tile_shape: tuple[int, ...]
57
+ ) -> float:
58
+ """Estimate how much extra data each tile read pulls off disk.
59
+
60
+ When the store's on-disk chunks are larger than the tile, every tile read
61
+ decodes whole chunks and discards most of them. Returns the ratio of bytes
62
+ read to bytes used (1.0 = no waste).
63
+
64
+ Parameters
65
+ ----------
66
+ native_chunks : tuple of int
67
+ The store's on-disk chunk shape.
68
+ tile_shape : tuple of int
69
+ The processing tile shape.
70
+
71
+ Returns
72
+ -------
73
+ float
74
+ Read-amplification factor (``bytes_read / bytes_used``).
75
+ """
76
+ import math
77
+
78
+ read = used = 1.0
79
+ for n, t in zip(native_chunks, tile_shape):
80
+ if n <= 0 or t <= 0:
81
+ continue
82
+ touched = math.ceil(t / n) # chunks a tile spans on this axis
83
+ read *= touched * n
84
+ used *= t
85
+ return read / used if used else 1.0
86
+
87
+
23
88
  def _stage_to_zarr(
24
89
  arr: da.Array, path: str, component: str, show_progress: bool
25
90
  ) -> None:
@@ -83,6 +148,7 @@ def tile_process(
83
148
  empty_threshold: float | None = None,
84
149
  stage_dir: Union[str, Path, None] = None,
85
150
  keep_stage: bool = False,
151
+ log_file: Union[str, Path, bool, None] = True,
86
152
  verbose: bool = False,
87
153
  ) -> da.Array:
88
154
  """Apply *fn* to every tile of *image* and merge labels globally.
@@ -180,6 +246,11 @@ def tile_process(
180
246
  keep_stage:
181
247
  Keep the temp stage store after merging (default: delete it). Useful
182
248
  for debugging or resuming an interrupted run.
249
+ log_file:
250
+ Where to tee the ``patchworks`` INFO log (including a per-tile
251
+ ``processing tile k/N`` counter). ``True`` (default) auto-writes
252
+ ``patchworks.log`` next to the output; a path writes there; ``False``/
253
+ ``None`` disables the file (logs still go to stderr).
183
254
  verbose:
184
255
  Log each tile's location and shape as it is processed.
185
256
 
@@ -250,10 +321,27 @@ def tile_process(
250
321
 
251
322
  # Load + tile
252
323
  image_source_path = None if isinstance(image, da.Array) else str(image)
324
+
325
+ # Auto log file (default): tee patchworks' INFO logs to a file next to the
326
+ # output, so a long run leaves a tailable record without notebook setup.
327
+ if log_file:
328
+ if log_file is True:
329
+ if write_to is not None:
330
+ _ldir = os.path.dirname(os.path.abspath(str(write_to)))
331
+ elif image_source_path is not None:
332
+ _ldir = os.path.dirname(os.path.abspath(image_source_path))
333
+ else:
334
+ _ldir = os.getcwd()
335
+ log_file = os.path.join(_ldir, "patchworks.log")
336
+ _attach_log_file(str(log_file))
337
+ logger.info("patchworks log → %s", log_file)
338
+
253
339
  _load_chunks: tuple[int, ...] | None = None
340
+ _native_chunks: tuple[int, ...] | None = None
254
341
 
255
342
  if not isinstance(image, da.Array):
256
343
  _peek = load_ome_zarr(image, channel=channel, level=level)
344
+ _native_chunks = _peek.chunksize # on-disk zarr chunk shape
257
345
  if callable(tile_shape):
258
346
  _load_chunks = tuple(tile_shape(_peek.shape, _peek.dtype))
259
347
  elif isinstance(tile_shape, str):
@@ -291,13 +379,30 @@ def tile_process(
291
379
  logger.info("Rechunked to %s", tile_shape)
292
380
 
293
381
  n_tiles = int(np.prod([len(c) for c in image.chunks]))
382
+ _tile = tuple(c[0] for c in image.chunks)
294
383
  logger.info(
295
384
  "Processing %d tiles (per-axis %s, tile shape %s)",
296
385
  n_tiles,
297
386
  tuple(len(c) for c in image.chunks),
298
- tuple(c[0] for c in image.chunks),
387
+ _tile,
299
388
  )
300
389
 
390
+ # Warn when the store's on-disk chunks are much larger than the tile: every
391
+ # tile read then decodes whole chunks and throws most away (slow I/O).
392
+ if _native_chunks is not None:
393
+ _amp = _read_amplification(_native_chunks, _tile)
394
+ if _amp >= 4:
395
+ logger.warning(
396
+ "Input chunks %s are much larger than the tile %s → ~%.0fx "
397
+ "read amplification (each tile decodes whole chunks and "
398
+ "discards most). Re-chunk the store near the tile size "
399
+ "(e.g. to_ome_zarr(..., chunks=...) without shard=) or read "
400
+ "the source file directly to avoid wasted I/O.",
401
+ _native_chunks,
402
+ _tile,
403
+ _amp,
404
+ )
405
+
301
406
  image_for_threshold = image
302
407
 
303
408
  # Overlap — build a per-axis depth dict (clips to fit each axis).
@@ -313,6 +418,23 @@ def tile_process(
313
418
  if skip_empty and _skip_thr is None:
314
419
  _skip_thr = _auto_empty_threshold(image_for_threshold, channel, level)
315
420
 
421
+ # Up-front heads-up for a big 3-D GPU job (z-stack tiles, many of them):
422
+ # an accurate ETA is logged after the first few tiles, but warn early that
423
+ # this is the expensive path and point at the faster alternatives.
424
+ if use_gpu and image.ndim >= 3 and _tile[0] > 4 and n_tiles >= 50:
425
+ logger.warning(
426
+ "Large 3-D GPU job: %d tiles of %s. Per-tile 3-D segmentation is "
427
+ "slow on a single device (a live ETA is logged after the first "
428
+ "tiles). If per-slice results are acceptable, 2-D (z=1 tiles) is "
429
+ "typically ~10x faster, or segment a lower pyramid level.",
430
+ n_tiles,
431
+ _tile,
432
+ )
433
+
434
+ # Tile counters + timing (GIL-safe enough for the threaded / in-process
435
+ # schedulers patchworks uses) so the log shows live "tile k/N + ETA".
436
+ _progress = {"done": 0, "seen": 0, "time": 0.0}
437
+
316
438
  def active_fn(block, block_info=None):
317
439
  """Run *fn* on one tile, or return zeros for an empty tile.
318
440
 
@@ -328,14 +450,36 @@ def tile_process(
328
450
  np.ndarray
329
451
  Integer labels, or an all-zero tile when skipped.
330
452
  """
453
+ import time
454
+
331
455
  loc = block_info[0].get("chunk-location") if block_info else "?"
456
+ _progress["seen"] += 1
332
457
  if skip_empty and block.size and block.max() <= _skip_thr:
333
458
  if verbose:
334
459
  logger.debug("skip empty tile %s (max<=%s)", loc, _skip_thr)
335
460
  return np.zeros(block.shape, dtype=np.int32)
336
- if verbose:
337
- logger.debug("process tile %s shape=%s", loc, block.shape)
338
- return fn(block)
461
+
462
+ t0 = time.perf_counter()
463
+ out = fn(block)
464
+ dt = time.perf_counter() - t0
465
+
466
+ _progress["done"] += 1
467
+ _progress["time"] += dt
468
+ done, seen = _progress["done"], _progress["seen"]
469
+ avg = _progress["time"] / done
470
+ # Extrapolate remaining non-empty tiles from the empty fraction seen.
471
+ remaining_nonempty = max(0, n_tiles - seen) * (done / seen)
472
+ eta_h = avg * remaining_nonempty / 3600
473
+ logger.info(
474
+ "tile %d done in %.1fs (avg %.1fs); %d/%d tiles seen; ETA ~%.1fh",
475
+ done,
476
+ dt,
477
+ avg,
478
+ seen,
479
+ n_tiles,
480
+ eta_h,
481
+ )
482
+ return out
339
483
 
340
484
  _meta = np.empty((0,) * image.ndim, dtype=np.int32)
341
485
  if overlap > 0:
@@ -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"}}