From 2b219ab98d34d511930b7018fd16a3538c88c980 Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Thu, 25 Jun 2026 14:14:59 +0200 Subject: [PATCH] =?UTF-8?q?feat(workflow):=20=E2=9C=A8=20custom=20per-tile?= =?UTF-8?q?=20function=20via=20`method:=20custom`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Let users plug their own segmentation function without editing the package: `method: "custom"` imports `cfg.custom.{module,function}` and forwards optional kwargs. The module just needs to be importable on the cluster (a file in workflow/scripts/, on PYTHONPATH, or pip-installed). Document the function contract, the import options and the cluster checklist (env deps, offline prefetch) in the guide. Co-Authored-By: Claude Opus 4.8 --- docs/guide/snakemake.md | 53 +++++++++++++++++++++++++++++++++++++ workflow/config/config.yaml | 9 ++++++- workflow/scripts/_pw.py | 21 +++++++++++++-- 3 files changed, 80 insertions(+), 3 deletions(-) diff --git a/docs/guide/snakemake.md b/docs/guide/snakemake.md index c633f02..6a01ceb 100644 --- a/docs/guide/snakemake.md +++ b/docs/guide/snakemake.md @@ -194,6 +194,59 @@ 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. It just has to take one tile and return integer labels of the same +spatial shape: + +```python +# my_seg.py +import numpy as np +from skimage.feature import blob_log +from skimage.measure import label + +def segment(tile: np.ndarray) -> np.ndarray: + """One tile in, int32 label image out (0 = background).""" + mask = tile > tile.mean() + 2 * tile.std() + return label(mask).astype("int32") +``` + +Point the config at it: + +```yaml +method: "custom" +label_name: "my_labels" +custom: + module: "my_seg" # import name (see below) + function: "segment" # default is "segment" + kwargs: {} # optional, forwarded as segment(tile, **kwargs) +``` + +### Make it importable on the cluster + +Pick one (the workflow imports `module` in each segment job): + +1. **Drop the file in `workflow/scripts/`** — Snakemake puts the script dir on + `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 wherever the file lives before launching Snakemake. + +### Cluster checklist + +- The env that runs the **segment** jobs must have your function's imports + (`pip`/`pixi add` them). A missing import shows up in `logs/segment/.log`. +- **Offline GPU nodes:** the `fetch_model` prefetch only covers Cellpose. If + your function downloads weights/data at run time, fetch them once on the + **login node** first (they must land in shared `$HOME`), or the segment jobs + hit `Network is unreachable` — see *Troubleshooting*. +- Everything else is unchanged: tiling, halos, the zarr-native merge, resume, + and per-tile logs all work exactly as for Cellpose. + +For full control (your own tiling/merge loop, not 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. diff --git a/workflow/config/config.yaml b/workflow/config/config.yaml index 7f6a58c..95f174e 100644 --- a/workflow/config/config.yaml +++ b/workflow/config/config.yaml @@ -29,7 +29,7 @@ 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) +method: "cellpose" # "cellpose" (GPU), "threshold" (no GPU; testing), "custom" label_name: "cellpose_labels" cellpose: model: "nuclei" @@ -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 diff --git a/workflow/scripts/_pw.py b/workflow/scripts/_pw.py index 08c83a0..6fa4301 100644 --- a/workflow/scripts/_pw.py +++ b/workflow/scripts/_pw.py @@ -10,6 +10,7 @@ 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):