feat(workflow): custom per-tile function via method: custom

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/<i>.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.

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

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,15 +125,31 @@ 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
     -------
     callable
         ``(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):