patchworks 0.11.6__tar.gz → 0.11.8__tar.gz

{patchworks-0.11.6 → patchworks-0.11.8}/.gitignore

@@ -7,3 +7,5 @@ build/
 site/
 .pytest_cache/
 *.log
+.snakemake/
+.pixi/

{patchworks-0.11.6 → patchworks-0.11.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchworks
-Version: 0.11.6
+Version: 0.11.8
 Summary: Tiled processing of arbitrarily large images with globally consistent labels
 Project-URL: Homepage, https://github.com/imcf/patchworks
 Project-URL: Issues, https://github.com/imcf/patchworks/issues
@@ -30,7 +30,8 @@ Requires-Dist: bioio-nd2; extra == 'all'
 Requires-Dist: bioio-ome-tiff; extra == 'all'
 Requires-Dist: bioio-tifffile; extra == 'all'
 Requires-Dist: imaris-ims-file-reader; extra == 'all'
-Requires-Dist: napari[all]; extra == 'all'
+Requires-Dist: lxml-html-clean; extra == 'all'
+Requires-Dist: napari[all]>=0.5.5; extra == 'all'
 Requires-Dist: nvidia-ml-py; extra == 'all'
 Requires-Dist: psutil; extra == 'all'
 Requires-Dist: scikit-image; extra == 'all'
@@ -62,7 +63,8 @@ Provides-Extra: io
 Requires-Dist: psutil; extra == 'io'
 Requires-Dist: tqdm; extra == 'io'
 Provides-Extra: napari
-Requires-Dist: napari[all]; extra == 'napari'
+Requires-Dist: lxml-html-clean; extra == 'napari'
+Requires-Dist: napari[all]>=0.5.5; extra == 'napari'
 Provides-Extra: workflow
 Requires-Dist: snakemake-executor-plugin-slurm; extra == 'workflow'
 Requires-Dist: snakemake>=8; extra == 'workflow'

{patchworks-0.11.6 → patchworks-0.11.8}/docs/guide/snakemake.md

@@ -194,6 +194,146 @@ also re-runs a step when its **code, params or software environment** change —
 so upgrading patchworks would re-do the conversion and overwrite an existing
 result. Keep `mtime` and reruns happen only when an output is missing or stale.
 
+## Custom segmentation function
+
+Not using Cellpose? Run **your own** per-tile function — no need to edit the
+package. You write one function; patchworks handles everything around it
+(tiling, halos, skipping empty tiles, the zarr-native merge, global relabelling,
+resume, logs).
+
+### The contract
+
+Your function is called **once per tile**:
+
+```python
+labels = segment(tile)          # plus any kwargs you configure
+```
+
+| | What you get / must return |
+| --- | --- |
+| **Input `tile`** | A NumPy array of **one** tile, with the overlap halo already included. The channel and pyramid level from the config are already selected, so it is purely spatial: `(z, y, x)` for a 3-D run, `(y, x)` for 2-D. Dtype is the image's (e.g. `uint16`). |
+| **Return** | An integer **label** array (not a boolean mask), **same shape** as `tile`. `0` = background; each object a distinct positive integer. |
+| **Labels** | Only need to be unique **within the tile**. Don't try to make them globally unique — the merge step stitches objects across tile borders and renumbers everything to a contiguous `1..N` (`sequential_labels: true`). |
+| **Shape** | Must match the input exactly — patchworks trims the halo off your output, so a wrong shape is an error. Don't crop or resize inside the function. |
+
+That is the whole interface. Anything that turns an image tile into a label
+image works: classic image processing, StarDist, a trained model, an external
+binary you shell out to, …
+
+### Minimal example (no GPU, no deps beyond scikit-image)
+
+```python
+# my_seg.py
+import numpy as np
+from skimage.measure import label
+
+def segment(tile: np.ndarray, sigma: float = 2.0) -> np.ndarray:
+    """Threshold + connected components. Returns int32 labels (0 = bg)."""
+    from skimage.filters import gaussian, threshold_otsu
+
+    smooth = gaussian(tile, sigma=sigma, preserve_range=True)
+    thr = threshold_otsu(smooth) if smooth.max() > smooth.min() else np.inf
+    return label(smooth > thr).astype("int32")
+```
+
+```yaml
+method: "custom"
+label_name: "my_labels"
+custom:
+  module: "my_seg"        # import name (see "Make it importable")
+  function: "segment"     # default is "segment"
+  kwargs:                 # optional — forwarded as segment(tile, **kwargs)
+    sigma: 1.5
+```
+
+### Real example: StarDist 3-D, with model caching
+
+Heavy models must be loaded **once**, not per tile. On SLURM each tile is its
+own process so this matters less, but for local runs one process segments many
+tiles — cache the model at module level (or with `functools.lru_cache`):
+
+```python
+# stardist_seg.py
+import numpy as np
+
+_MODEL = None
+
+def _model():
+    global _MODEL
+    if _MODEL is None:                       # loaded once per worker process
+        from stardist.models import StarDist3D
+        _MODEL = StarDist3D.from_pretrained("3D_demo")
+    return _MODEL
+
+def segment(tile: np.ndarray, prob_thresh: float = 0.5) -> np.ndarray:
+    from csbdeep.utils import normalize
+
+    labels, _ = _model().predict_instances(
+        normalize(tile), prob_thresh=prob_thresh
+    )
+    return labels.astype("int32")
+```
+
+```yaml
+method: "custom"
+label_name: "stardist"
+custom:
+  module: "stardist_seg"
+  function: "segment"
+  kwargs:
+    prob_thresh: 0.5
+```
+
+Using a GPU? The segment jobs already hold one (the `gres: "gpu:1"` request),
+so just let your framework see it — nothing extra in the config.
+
+### Test it before you submit
+
+Run your function on one real tile first — it catches shape/dtype bugs in
+seconds instead of after a queue wait. Output must be integer, same shape, `0`
+for background:
+
+```python
+from patchworks import load_ome_zarr
+from my_seg import segment
+
+img = load_ome_zarr("results/image.zarr", channel=0, level=0)
+tile = img[:, :512, :512].compute()      # a small spatial block
+out = segment(tile)
+
+assert out.shape == tile.shape, (out.shape, tile.shape)
+assert out.dtype.kind in "iu"            # integer labels, not a float mask
+print("objects in tile:", int(out.max()))
+```
+
+### Make it importable on the cluster
+
+The segment job runs `import <module>`, so the module must be on the path. Pick
+one:
+
+1. **Drop the file in `workflow/scripts/`** — Snakemake adds the script dir to
+   `sys.path`, so `module: "my_seg"` just works. Simplest for a single file.
+2. **Install it** into the run env (`pip install -e .`, `pixi add --pypi …`),
+   then use its import name. Best for a real package with dependencies.
+3. **Set `PYTHONPATH`** to the file's directory before launching Snakemake.
+
+### Cluster checklist
+
+- **Dependencies:** the env that runs the **segment** jobs must have everything
+  your function imports (`pip`/`pixi add` it). A missing import or any crash
+  shows up in `logs/segment/<index>.log`, not the (empty) SLURM log.
+- **Offline GPU nodes:** the built-in `fetch_model` prefetch covers Cellpose
+  only. If your function downloads weights/data on first use, fetch them once on
+  the **login node** (network access) so they land in shared `$HOME`; otherwise
+  the segment jobs fail with `Network is unreachable`. See *Troubleshooting*.
+- **Memory / walltime:** tune `segment:` in `profile/slurm/config.yaml` for your
+  model (`mem_mb`, `runtime`) just as for Cellpose.
+- Everything else — tiling, halos, empty-tile skipping, the zarr-native merge,
+  resume, and per-tile logs — is identical to a Cellpose run.
+
+For full control (your own tiling/merge loop instead of the bundled rules), call
+the public API directly — see *How it works* below.
+
 ## pixi (instead of conda)
 
 Conda is **not** required — Snakemake runs in whatever environment launches it.
@@ -223,10 +363,12 @@ prologue. The simplest path is a single shared env that the compute nodes see.
 | Symptom | Fix |
 |---------|-----|
 | `snakemake: command not found` | use `python -m snakemake` |
-| Segment jobs pend forever | wrong `slurm_partition`/`slurm_extra` GPU flag for your cluster |
+| Segment jobs pend forever | wrong `slurm_partition`/GPU request; on scicore use `gres: "gpu:1"` |
+| Segment dies, `Network is unreachable` | offline GPU nodes — the `fetch_model` localrule caches the model on the submit host first; if it still fails, your submit host has no network either (pre-download manually) |
 | `cellpose is not installed` in a job | the job's env lacks `patchworks[cellpose]` |
 | Reading the input fails | install the matching reader (`patchworks[imaris]`/`[bioio]` + a `bioio-*`) |
 | Out of GPU memory | smaller `tile_shape`, or `do_3D: false` |
+| A job fails with an empty SLURM log | read `logs/segment/<index>.log` (per tile) or `logs/steps.log` — the real traceback is there |
 | Very slow | confirm GPU is used (`nvidia-smi`); try 2-D or a lower `level` |
 
 ## How it works (for the curious)

{patchworks-0.11.6 → patchworks-0.11.8}/pyproject.toml

@@ -57,7 +57,11 @@ bioio = [
 # imaris reads .ims files natively (HDF5, no JVM) for OME-ZARR conversion.
 imaris = ["imaris-ims-file-reader"]
 # napari enables the interactive viewer plugin.
-napari = ["napari[all]"]
+# - napari >= 0.5: NumPy 2.0 compatible (older 0.4.x uses np.array(copy=False),
+#   which raises under NumPy 2 -> ValidationError on Viewer()).
+# - lxml-html-clean: napari's notebook_display imports lxml.html.clean, split
+#   into a separate package in lxml >= 5.2 (else ImportError on Viewer()).
+napari = ["napari[all]>=0.5.5", "lxml-html-clean"]
 # workflow runs the Snakemake pipeline (per-tile SLURM jobs across GPUs).
 workflow = ["snakemake>=8", "snakemake-executor-plugin-slurm"]
 dev = ["pytest", "pytest-cov", "scikit-image", "psutil", "tqdm"]

{patchworks-0.11.6 → patchworks-0.11.8}/workflow/Snakefile

@@ -19,6 +19,11 @@ include: "rules/segment.smk"
 include: "rules/merge.smk"
 
 
+# Runs on the submit host (has network) — never submitted to an offline GPU node.
+localrules:
+    fetch_model,
+
+
 rule all:
     input:
         f"{WORK}/labels.done",

{patchworks-0.11.6 → patchworks-0.11.8}/workflow/config/config.yaml

@@ -12,7 +12,7 @@ work_dir: "/path/to/results" # everything is written under here
 #   image.zarr/labels/<name>/  — the segmentation (multi-scale), in the image
 
 # ---- conversion -------------------------------------------------------------
-reuse_pyramid: true # for .ims: copy its own pyramid (fast); else rebuild
+reuse_pyramid: false # for .ims: copy its own pyramid (fast); else rebuild
 convert_chunks: null # null → patchworks' bounded auto chunks; or [c,z,y,x]
 shard: false # true → pack chunks into shards (zarr v3; fewer files)
 # Re-running reuses an existing image.zarr automatically (skips conversion).
@@ -29,10 +29,10 @@ skip_empty: true # skip background tiles
 empty_threshold: null # null → Otsu; or a number
 
 # ---- segmentation -----------------------------------------------------------
-method: "cellpose" # "cellpose" (GPU) or "threshold" (simple, no GPU; testing)
-label_name: "cellpose"
+method: "cellpose" # "cellpose" (GPU), "threshold" (no GPU; testing), "custom"
+label_name: "cellpose_labels"
 cellpose:
-  model: "cyto3"
+  model: "nuclei"
   diameter: 30
   do_3D: true
   gpu: true
@@ -40,6 +40,13 @@ cellpose:
   # flow_threshold: 0.4
   # cellprob_threshold: 0.0
 
+# Your own per-tile function (method: "custom"). See the "custom function"
+# section of docs/guide/snakemake.md.
+# custom:
+#   module: "my_seg"     # importable on the cluster (workflow/scripts/, PYTHONPATH, or pip-installed)
+#   function: "segment"  # def segment(tile: np.ndarray) -> np.ndarray (int32 labels)
+#   kwargs: {}           # optional extra keyword args
+
 # ---- pyramid for the labels -------------------------------------------------
 pyramid_levels: 5
 pyramid_downscale: 2

{patchworks-0.11.6 → patchworks-0.11.8}/workflow/pixi.toml

@@ -7,7 +7,7 @@
 # On a cluster, put this project on a shared filesystem the compute nodes can
 # read: the SLURM jobs re-launch snakemake from this env's interpreter.
 
-[project]
+[workspace]
 name = "patchworks-workflow"
 channels = ["conda-forge"]
 platforms = ["linux-64"]

{patchworks-0.11.6 → patchworks-0.11.8}/workflow/profile/slurm/config.yaml

@@ -33,11 +33,14 @@ set-resources:
     mem_mb: 32000
     runtime: 120
   segment:
-    # one GPU per tile — this is what spreads Cellpose across GPUs
+    # one GPU per tile — this is what spreads Cellpose across GPUs.
+    # Use the plugin's native `gres` resource → emits --gres=gpu:1, which is
+    # how scicore binds a CUDA device. (The `gpu:` resource emits --gpus and
+    # leaves the job without a device.)
     slurm_partition: "rtx4090"
-    gpu: 1
+    gres: "gpu:1"
     qos: "rtx4090-6hours" # scicore: <partition>-<duration> QOS
-    mem_mb: 32000 # plenty — a tile used ~1G; the failure was a Python error
+    mem_mb: 32000 # plenty — a tile used ~1G
     cpus_per_task: 4
     runtime: 360 # 6 hours — must match the QOS, NOT 120 (=2h → killed early)
   merge:

{patchworks-0.11.6 → patchworks-0.11.8}/workflow/rules/common.smk

@@ -20,6 +20,10 @@ STAGE_OK = f"{STAGE}.done"
 LOGS = f"{WORK}/logs"
 STEPLOG = f"{LOGS}/steps.log"
 
+# Marker that the segmentation model is cached locally. Produced by a local
+# rule (runs on the networked submit host) so offline GPU nodes never download.
+MODEL_OK = f"{WORK}/model.ready"
+
 
 def occupied_done(wildcards):
     """Per-tile markers for the occupied tiles (resolved after the checkpoint)."""

{patchworks-0.11.6 → patchworks-0.11.8}/workflow/rules/segment.smk

@@ -1,5 +1,20 @@
 # Plan tiles (checkpoint) and segment each tile on a GPU.
 
+
+rule fetch_model:
+    """Cache the segmentation model on the (networked) submit host.
+
+    Declared local (see ``localrules`` in the Snakefile) so it never runs on an
+    offline GPU node — Cellpose downloads its weights here, into shared $HOME.
+    """
+    output:
+        touch(MODEL_OK),
+    log:
+        f"{LOGS}/fetch_model.log",
+    script:
+        "../scripts/fetch_model.py"
+
+
 checkpoint prepare:
     input:
         IMAGE_OK,
@@ -18,6 +33,7 @@ rule segment:
         tiles=TILES,
         stage=STAGE_OK,
         image=IMAGE_OK,
+        model=MODEL_OK,
     output:
         f"{WORK}/seg/{{index}}.done",
     log:

{patchworks-0.11.6 → patchworks-0.11.8}/workflow/scripts/_pw.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 import json
 import logging
 import sys
+from functools import partial
 from pathlib import Path
 
 from patchworks import load_ome_zarr
@@ -124,8 +125,9 @@ def build_fn(cfg):
     Parameters
     ----------
     cfg : dict
-        Snakemake config. ``method`` selects ``"cellpose"`` (default) or a
-        simple ``"threshold"`` (handy for testing / no-GPU runs).
+        Snakemake config. ``method`` selects ``"cellpose"`` (default), a simple
+        ``"threshold"`` (testing / no-GPU), or ``"custom"`` to import your own
+        function (``cfg["custom"] = {module, function, kwargs}``).
 
     Returns
     -------
@@ -133,6 +135,21 @@ def build_fn(cfg):
         ``(ndarray) -> ndarray`` returning integer labels.
     """
     method = cfg.get("method", "cellpose")
+    if method == "custom":
+        # Import a user-provided function, e.g.
+        #   custom: {module: my_seg, function: segment, kwargs: {...}}
+        # The module must be importable on the cluster (a file in
+        # workflow/scripts/, on PYTHONPATH, or an installed package).
+        import importlib
+
+        spec = cfg["custom"]
+        fn = getattr(
+            importlib.import_module(spec["module"]),
+            spec.get("function", "segment"),
+        )
+        kwargs = spec.get("kwargs") or {}
+        return partial(fn, **kwargs) if kwargs else fn
+
     if method == "threshold":
 
         def fn(tile):

patchworks-0.11.8/workflow/scripts/fetch_model.py

@@ -0,0 +1,21 @@
+"""Snakemake (local) script: cache the segmentation model before segmenting.
+
+Runs on the submit host, which has network access, so the offline GPU nodes
+never try to download Cellpose weights at run time (they read the shared
+``$HOME/.cellpose`` cache instead).
+"""
+
+from _pw import start_log
+
+start_log(snakemake.log[0])  # noqa: F821
+cfg = snakemake.config  # noqa: F821
+
+if cfg.get("method", "cellpose") == "cellpose":
+    # _get_model downloads + caches the weights keyed by (model, gpu).
+    from patchworks.plugins.cellpose import _get_model
+
+    model = cfg["cellpose"]["model"]
+    _get_model({"model": model, "gpu": False})
+    print(f"[patchworks] cached segmentation model: {model}")
+else:
+    print(f"[patchworks] method={cfg.get('method')!r}; no model to prefetch")