merge dev into main

Adds the method-agnostic dilate_labels post-segmentation step (API,
Snakemake config wiring, docs).

Author: lguerard

docs/api/postprocess.md

@@ -0,0 +1,3 @@
+# dilate_labels
+
+::: patchworks.dilate_labels

docs/examples/dog.md

@@ -81,6 +81,24 @@ result = tile_process(IMAGE, fn, tile_shape=(1, 1024, 1024), overlap=32)
     so edge tiles keep enough context (a plain intensity/threshold halo is
     too thin).
 
+## Growing the labels afterwards
+
+DoG spots/threads are often thin — grow each label by a few pixels with
+[`dilate_labels`](../api/postprocess.md):
+
+```python
+from patchworks import tile_process, dilate_labels
+from patchworks.plugins.dog import dog_label_fn
+
+fn = dog_label_fn(low_sigma=1.0, high_sigma=3.0, threshold=0.02)
+fn = dilate_labels(fn, iterations=2)
+tile_process(IMAGE, fn, tile_shape=(1, 1024, 1024), overlap=8, write_to=OUTPUT)
+```
+
+On the cluster, set `dilate: 2` in the YAML config instead — it applies to
+`method: "custom"` (this plugin) the same way it does for `cellpose`/
+`threshold`, see [Growing labels after segmentation](../guide/snakemake.md#growing-labels-afterwards-dilation).
+
 ## Using it in the Snakemake workflow
 
 No dedicated wiring needed — `patchworks.plugins.dog` exposes a `segment(tile, **kwargs)`

docs/guide/snakemake.md

@@ -62,8 +62,9 @@ skip_empty: true               # skip background tiles
 empty_threshold: null          # null → Otsu
 
 # segmentation
-method: "cellpose"             # "cellpose" (GPU) or "threshold" (no GPU)
+method: "cellpose"             # "cellpose" (GPU), "threshold" (no GPU), "custom"
 label_name: "cellpose"         # name under image.zarr/labels/
+dilate: 0                      # optional: pixels to grow labels by, any method
 cellpose:
   model: "cyto3"
   diameter: 30
@@ -77,6 +78,18 @@ pyramid_downscale: 2
 sequential_labels: true        # renumber labels to a contiguous 1..N
 ```
 
+!!! tip "Growing labels after segmentation"
+    `dilate: N` grows every label by `N` pixels once segmentation finishes,
+    regardless of `method` (`cellpose`, `threshold`, or `custom`). It runs
+    per-tile, before the overlap halo is trimmed and tiles are merged, so
+    dilated labels still stitch correctly across tile boundaries — just make
+    sure `overlap` covers the dilation amount plus the usual object-diameter
+    halo. `0` (default) disables it. Under the hood this wraps whatever
+    segmentation function `method` builds with
+    [`patchworks.dilate_labels`](../api/postprocess.md); see [Custom
+    segmentation function](#custom-segmentation-function) below for using it
+    directly from Python instead of via config.
+
 !!! tip "Tile size vs runtime"
     `tile_shape: "auto"` sizes each tile to your GPU's VRAM. Smaller tiles =
     more (faster) jobs; very large 3-D tiles are slow. Keep `do_3D: false` (2-D
@@ -433,6 +446,29 @@ custom:
     sigma: 1.5
 ```
 
+### Growing labels afterwards (dilation)
+
+To grow every label by a few pixels after segmentation — any method, not
+just `custom` — set `dilate: N` in the config (see the tip above), or wrap
+your function directly with
+[`patchworks.dilate_labels`](../api/postprocess.md) when calling the API
+yourself:
+
+```python
+from patchworks import tile_process, dilate_labels
+from patchworks.plugins.dog import dog_label_fn
+
+fn = dog_label_fn(low_sigma=1.0, high_sigma=3.0, threshold=0.02)
+fn = dilate_labels(fn, iterations=2)   # grow each label by 2 px, then run
+result = tile_process("image.zarr", fn, tile_shape=(1, 2048, 2048),
+                       overlap=8, write_to="labels.zarr")
+```
+
+`dilate_labels` wraps any `(tile) -> labels` function — the same contract
+described above — so it works with `dog_label_fn`, `cellpose_fn`, or your
+own `segment`. It dilates each tile's labels before the halo is trimmed and
+tiles are merged, so `overlap` must still cover the dilation amount.
+
 ### Real example: StarDist 3-D, with model caching
 
 Heavy models must be loaded **once**, not per tile. On SLURM each tile is its

mkdocs.yml

@@ -54,6 +54,7 @@ nav:
 - API Reference:
   - tile_process: api/tile_process.md
   - merge_tile_labels: api/merge_tile_labels.md
+  - dilate_labels: api/postprocess.md
   - Tile sizing: api/chunks.md
   - I/O helpers: api/io.md
   - Relabelling: api/relabel.md

src/patchworks/__init__.py

@@ -35,6 +35,7 @@
 from ._distributed import create_stage, spatial_tiles, stage_tile
 from ._io import estimate_empty_tiles, load_ome_zarr
 from ._merge import merge_tile_labels
+from ._postprocess import dilate_labels
 from ._relabel import relabel_sequential_array, relabel_sequential_zarr
 from ._relations import label_relations
 
@@ -57,4 +58,5 @@
     "spatial_tiles",
     "create_stage",
     "stage_tile",
+    "dilate_labels",
 ]

src/patchworks/_postprocess.py

@@ -0,0 +1,101 @@
+"""Generic post-segmentation wrappers for patchworks.
+
+These wrap any ``fn(tile) -> labels`` callable (a plugin, a custom function,
+whatever ``method`` in the Snakemake workflow builds) so the same
+post-processing applies regardless of which segmentation method produced the
+labels.
+
+Usage
+-----
+>>> from patchworks import tile_process, dilate_labels
+>>> from patchworks.plugins.dog import dog_label_fn
+>>>
+>>> fn = dog_label_fn(low_sigma=1.0, high_sigma=3.0, threshold=0.02)
+>>> fn = dilate_labels(fn, iterations=2)
+>>> result = tile_process("image.zarr", fn, tile_shape=(1, 2048, 2048),
+...                       overlap=8, write_to="labels.zarr")
+"""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable
+
+import numpy as np
+
+
+def dilate_labels(
+    fn: Callable[[np.ndarray], np.ndarray],
+    iterations: int = 1,
+    *,
+    use_gpu: bool = False,
+) -> Callable[[np.ndarray], np.ndarray]:
+    """Wrap a segmentation callable to grow its labels after each tile.
+
+    Applies a single-pass grey dilation to whatever ``fn`` returns, before
+    ``tile_process``/``stage_tile`` trim the overlap halo and merge across
+    tile boundaries — so dilated labels still stitch correctly at tile
+    edges.
+
+    Parameters
+    ----------
+    fn : Callable[[np.ndarray], np.ndarray]
+        Any segmentation function with the ``tile_process``/``stage_tile``
+        contract (one tile in, integer label array out).
+    iterations : int, optional
+        Pixels to grow each label by (grey-dilation footprint size
+        ``2 * iterations + 1``, single pass). Default 1. Values ``<= 0``
+        disable dilation — ``fn`` is returned unwrapped.
+    use_gpu : bool, optional
+        Dilate via cupyx instead of scipy. Independent of whatever backend
+        ``fn`` itself uses internally.
+
+    Returns
+    -------
+    Callable[[np.ndarray], np.ndarray]
+        Picklable function ready for ``tile_process``/``stage_tile``. If
+        ``iterations <= 0``, this is ``fn`` itself.
+    """
+    if iterations <= 0:
+        return fn
+    return partial(_run, fn=fn, iterations=iterations, use_gpu=use_gpu)
+
+
+def _run(
+    block: np.ndarray,
+    fn: Callable[[np.ndarray], np.ndarray],
+    iterations: int,
+    use_gpu: bool,
+) -> np.ndarray:
+    """Run ``fn`` on ``block``, then grow the resulting labels.
+
+    Parameters
+    ----------
+    block : np.ndarray
+        One image tile.
+    fn : Callable[[np.ndarray], np.ndarray]
+        Segmentation function to run first.
+    iterations : int
+        Pixels to grow each label by.
+    use_gpu : bool
+        Dilate via cupyx instead of scipy.
+
+    Returns
+    -------
+    np.ndarray
+        Dilated integer label array, same shape as ``fn``'s output.
+    """
+    labels = fn(block)
+    size = 2 * iterations + 1
+
+    if use_gpu:
+        import cupy as cp
+        from cupyx.scipy.ndimage import grey_dilation
+
+        labels = cp.asnumpy(grey_dilation(cp.asarray(labels), size=size))
+    else:
+        from scipy.ndimage import grey_dilation
+
+        labels = grey_dilation(labels, size=size)
+
+    return labels

tests/test_postprocess.py

@@ -0,0 +1,41 @@
+"""Self-contained tests for the dilate_labels post-processing wrapper."""
+
+import pickle
+
+import numpy as np
+
+
+def _make_blob_labels(shape=(1, 64, 64)):
+    labels = np.zeros(shape, dtype="int32")
+    labels[0, 28:36, 28:36] = 1
+    return labels
+
+
+def test_dilate_labels_grows_mask():
+    from patchworks import dilate_labels
+
+    fn = lambda tile: _make_blob_labels(tile.shape)  # noqa: E731
+    plain = fn(np.zeros((1, 64, 64)))
+    dilated = dilate_labels(fn, iterations=2)(np.zeros((1, 64, 64)))
+
+    assert (dilated > 0).sum() > (plain > 0).sum()
+
+
+def test_dilate_labels_zero_iterations_is_noop():
+    from patchworks import dilate_labels
+
+    fn = lambda tile: _make_blob_labels(tile.shape)  # noqa: E731
+
+    assert dilate_labels(fn, iterations=0) is fn
+
+
+def test_dilate_labels_picklable():
+    from patchworks.plugins.dog import dog_label_fn
+
+    from patchworks import dilate_labels
+
+    fn = dilate_labels(
+        dog_label_fn(low_sigma=1.0, high_sigma=4.0, threshold=0.01),
+        iterations=2,
+    )
+    pickle.loads(pickle.dumps(fn))

workflow/config/config.yaml

@@ -30,6 +30,7 @@ empty_threshold: null # null → Otsu; or a number
 
 # ---- segmentation -----------------------------------------------------------
 method: "cellpose" # "cellpose" (GPU), "threshold" (no GPU; testing), "custom"
+# dilate: 2 # optional: pixels to grow labels by after segmentation, any method
 # Also namespaces this run's intermediate files under work_dir/<label_name>/,
 # so a second segmentation (different label_name, e.g. nuclei vs cytoplasm)
 # can safely target the same work_dir — see docs/guide/snakemake.md

workflow/config/config_cilia.yaml

@@ -25,6 +25,7 @@ skip_empty: true
 empty_threshold: null
 
 method: "custom"
+# dilate: 2 # optional: pixels to grow labels by after segmentation
 label_name: "cilia_labels"
 custom:
   module: "patchworks.plugins.dog"

workflow/scripts/_pw.py

@@ -133,7 +133,34 @@ def build_fn(cfg):
     cfg : dict
         Snakemake config. ``method`` selects ``"cellpose"`` (default), a simple
         ``"threshold"`` (testing / no-GPU), or ``"custom"`` to import your own
-        function (``cfg["custom"] = {module, function, kwargs}``).
+        function (``cfg["custom"] = {module, function, kwargs}``). Optional
+        ``cfg["dilate"]``: int, pixels to grow labels by after segmentation
+        (via ``patchworks.dilate_labels``), applied regardless of ``method``.
+        Omitted/0 disables dilation.
+
+    Returns
+    -------
+    callable
+        ``(ndarray) -> ndarray`` returning integer labels.
+    """
+    fn = _build_method_fn(cfg)
+
+    dilate = cfg.get("dilate")
+    if dilate:
+        from patchworks import dilate_labels
+
+        fn = dilate_labels(fn, iterations=dilate)
+
+    return fn
+
+
+def _build_method_fn(cfg):
+    """Build the per-tile segmentation function for ``cfg["method"]``.
+
+    Parameters
+    ----------
+    cfg : dict
+        Snakemake config, see :func:`build_fn`.
 
     Returns
     -------