Merge branch dev: optional nuclei channel for Cellpose

Adds a per-config nuclei_channel, guards napari's label auto-load, and
notes the codebase's provenance in the README.

Author: lguerard

.gitignore

@@ -9,3 +9,6 @@ site/
 *.log
 .snakemake/
 .pixi/
+
+# graft's local graph cache — regenerable, not committed (run `graft build`).
+graft/

README.md

@@ -25,6 +25,13 @@ patchworks splits a large image into tiles, runs **any callable** on each
 tile in parallel, and merges the results into a globally consistent label array.
 It handles terabyte-scale images without loading them into memory.
 
+> [!NOTE]
+> **On how this was written.** Large parts of patchworks were vibe coded —
+> written with heavy LLM assistance rather than line by line. It is covered by
+> a test suite and has been run on real data, so it is not untested, but the
+> usual caveats apply: read the code before you trust it with anything
+> irreplaceable, and please open an issue if something looks off.
+
 ---
 
 ## Installation

docs/guide/snakemake.md

@@ -53,7 +53,8 @@ convert_chunks: null           # null → bounded auto chunks; or [c,z,y,x]
 shard: false                   # true → pack chunks into shards (fewer files)
 
 # tiling
-channel: 0                     # channel to segment (null = keep all)
+channel: 0                     # channel to segment, 0-based (null = keep all)
+nuclei_channel: null           # optional 2nd channel for Cellpose (see below)
 level: 0                       # pyramid level (0 = full resolution)
 tile_shape: "auto"             # "auto", or e.g. [16, 1024, 1024] (zyx)
 gpu_memory_gb: null            # for "auto" on SLURM: your segment GPU's VRAM
@@ -299,6 +300,7 @@ cellpose:
 # config/config_cyto.yaml — only the differences
 label_name: "cyto_labels"
 channel: 0                # cytoplasm/membrane channel
+nuclei_channel: 1         # optional: nuclear stain, as Cellpose's 2nd input
 overlap: [4, 30, 30]
 method: "cellpose"
 cellpose:
@@ -307,6 +309,29 @@ cellpose:
   do_3D: true
 ```
 
+### Giving Cellpose a nuclei channel
+
+`nuclei_channel` hands Cellpose a second channel — the nuclear stain — which
+usually improves cytoplasm segmentation. Both indices are 0-based, like
+`channel`.
+
+Only the `segment` step reads it. The pair is stacked on a leading axis that
+is *carried* into each tile rather than tiled, so the tile geometry, the
+occupancy map and the staged labels are byte-for-byte what a single-channel
+run produces, and `merge` and `label_relations` need no changes. Two things
+follow from that:
+
+- A tile holds twice the bytes, so a hand-set `tile_shape` sized to fill a GPU
+  may need halving. `tile_shape: "auto"` sizes from the single-channel array
+  and does not yet know about the pair.
+- The translation is version-specific. Cellpose 3 gets `channels: [1, 2]`
+  (1-based into the channel axis, `0` = grayscale); Cellpose 4 (cpsam) dropped
+  `channels` entirely and simply reads both. Either is overridable by setting
+  `channels:` or `channel_axis:` in the `cellpose:` block.
+
+`nuclei_channel` applies to the SLURM/Snakemake path. The single-process
+`tile_process` API still takes one `channel`.
+
 Run them as two independent SLURM submissions — they touch disjoint files, so
 they can run concurrently. Give each its own `--directory`, because
 Snakemake's lock lives in the working directory, not in the config:

src/patchworks/_distributed.py

@@ -166,6 +166,7 @@ def stage_tile(
     tile_shape: tuple[int, ...],
     overlap: Overlap = 0,
     component: str = "staged",
+    channel_axis: int | None = None,
 ) -> int:
     """Run *fn* on a single tile and write it into the shared stage store.
 
@@ -192,6 +193,13 @@ def stage_tile(
         :func:`normalize_overlap`).
     component : str, optional
         Array name inside the stage store.
+    channel_axis : int or None, optional
+        Axis of *image* holding channels, which is **not** tiled: it is read
+        whole and handed to *fn* alongside the tile's voxels. ``tile_shape``,
+        ``overlap`` and the stage store stay purely spatial, so *fn* still
+        returns one label per voxel with no channel axis (e.g. Cellpose fed a
+        cytoplasm + nuclei pair returns a single label volume). ``None`` (the
+        default) means *image* is already single-channel.
 
     Returns
     -------
@@ -202,18 +210,30 @@ def stage_tile(
         by a cumulative sum, instead of rewriting the whole store to make the
         ids unique.
     """
-    shape = image.shape
-    sl = spatial_tiles(shape, tile_shape)[index]
+    shape = tuple(image.shape)
+    # The channel axis is carried, not tiled: geometry (tiles, halo, the stage
+    # store) stays spatial, so nothing downstream of fn learns about channels.
+    if channel_axis is None:
+        spatial_shape = shape
+    else:
+        channel_axis %= len(shape)
+        spatial_shape = shape[:channel_axis] + shape[channel_axis + 1 :]
+    sl = spatial_tiles(spatial_shape, tile_shape)[index]
     halo = normalize_overlap(overlap, len(sl), tile_shape=tile_shape)
     expanded, trims = [], []
-    for s, dim, ov in zip(sl, shape, halo):
+    for s, dim, ov in zip(sl, spatial_shape, halo):
         lo = max(0, s.start - ov)
         hi = min(dim, s.stop + ov)
         expanded.append(slice(lo, hi))
         trims.append((s.start - lo, hi - s.stop))
-    block = np.asarray(image[tuple(expanded)])
+    read = list(expanded)
+    if channel_axis is not None:
+        read.insert(channel_axis, slice(None))
+    block = np.asarray(image[tuple(read)])
+    # What fn owes us back: one label per voxel, channel axis consumed.
+    block_spatial = tuple(e.stop - e.start for e in expanded)
     out = np.asarray(fn(block))
-    if out.shape != block.shape:
+    if out.shape != block_spatial:
         # Caught here rather than 6 frames deep in zarr's codec pipeline as
         # "could not broadcast input array from shape (13,1020,1020) into
         # shape (14,1024,1024)", which says nothing about which function is
@@ -224,7 +244,7 @@ def stage_tile(
         name = getattr(fn, "__name__", type(fn).__name__)
         raise ValueError(
             f"segmentation function {name!r} returned shape {out.shape} for "
-            f"a tile of shape {block.shape} (tile {index}). It must return "
+            f"a tile of shape {block_spatial} (tile {index}). It must return "
             "one label per input voxel. Some deconvolution backends crop "
             "their output -- pad or centre it back to the input shape before "
             "returning."

src/patchworks/plugins/cellpose.py

@@ -161,9 +161,13 @@ def _make_config(
     gpu : bool
         Run on the GPU.
     channels : list of int or None
-        Cellpose-3 ``[cyto, nucleus]`` channels; defaults to ``[0, 0]``.
+        *Cellpose 3 only.* ``[cyto, nucleus]``, 1-based into the channel axis
+        (0 = grayscale). ``None`` resolves per tile: ``[1, 2]`` when the tile
+        carries two channels, else ``[0, 0]``. Cellpose 4 dropped this
+        argument, so it is ignored there.
     channel_axis : int or None
-        Cellpose-4 channel axis.
+        Axis of the tile holding channels, forwarded to ``eval`` for both
+        Cellpose 3 and 4. ``None`` means single-channel tiles.
     diameter : float or None
         Expected cell diameter in pixels.
     do_3D : bool
@@ -179,7 +183,9 @@ def _make_config(
     return {
         "model": model,
         "gpu": gpu,
-        "channels": channels if channels is not None else [0, 0],
+        # Left as None ("auto") rather than [0, 0]: _run only knows how many
+        # channels a tile actually carries once it has one in hand.
+        "channels": channels,
         "channel_axis": channel_axis,
         "diameter": diameter,
         "do_3D": do_3D,
@@ -284,37 +290,50 @@ def _run(block: np.ndarray, cellpose_dict: dict[str, Any]) -> np.ndarray:
         Integer (``int32``) label array of the same spatial shape.
     """
     do_3D = cellpose_dict["do_3D"]
+    channel_axis = cellpose_dict.get("channel_axis")
+    n_channels = block.shape[channel_axis] if channel_axis is not None else 1
+
+    kwargs: dict[str, Any] = dict(
+        channel_axis=channel_axis,
+        diameter=cellpose_dict["diameter"],
+        do_3D=do_3D,
+        **cellpose_dict.get("cellpose_kwargs", {}),
+    )
+    if not _CELLPOSE_V4:
+        # Cellpose 4 (cpsam) dropped `channels` and reads whatever channels
+        # the array carries; Cellpose 3 needs the cyto/nucleus pairing named.
+        channels = cellpose_dict.get("channels")
+        if channels is None:
+            channels = [1, 2] if n_channels >= 2 else [0, 0]
+        kwargs["channels"] = channels
 
-    if _CELLPOSE_V4:
-        kwargs: dict[str, Any] = dict(
-            channel_axis=cellpose_dict.get("channel_axis"),
-            diameter=cellpose_dict["diameter"],
-            do_3D=do_3D,
-            **cellpose_dict.get("cellpose_kwargs", {}),
-        )
-    else:
-        kwargs = dict(
-            channels=cellpose_dict["channels"],
-            diameter=cellpose_dict["diameter"],
-            do_3D=do_3D,
-            **cellpose_dict.get("cellpose_kwargs", {}),
-        )
+    # Where z sits once the channel axis is accounted for.
+    z_axis = 1 if channel_axis == 0 else 0
 
     if do_3D:
-        kwargs["z_axis"] = 0
+        kwargs["z_axis"] = z_axis
         masks = _eval_with_oom_fallback(block, kwargs, cellpose_dict)
         return masks.astype("int32")
     else:
         # Squeeze singleton z so Cellpose gets a clean 2-D image
-        squeeze = block.ndim == 3 and block.shape[0] == 1
-        if block.ndim == 3 and not squeeze:
+        spatial = list(block.shape)
+        if channel_axis is not None:
+            spatial.pop(channel_axis)
+        squeeze = len(spatial) == 3 and spatial[0] == 1
+        if len(spatial) == 3 and not squeeze:
             raise ValueError(
-                f"do_3D is False but this tile has {block.shape[0]} z-planes. "
+                f"do_3D is False but this tile has {spatial[0]} z-planes. "
                 "Cellpose would receive the stack with no z_axis and treat "
                 "the leading axis as channels. Set do_3D: true, or tile with "
                 "z=1 to segment plane by plane."
             )
-        img = block[0] if squeeze else block
+        if squeeze:
+            img = block[(slice(None),) * z_axis + (0,)]
+            # Dropping z shifts any channel axis that sat behind it.
+            if channel_axis is not None and channel_axis > z_axis:
+                kwargs["channel_axis"] = channel_axis - 1
+        else:
+            img = block
         masks = _eval_with_oom_fallback(img, kwargs, cellpose_dict)
         masks = masks.astype("int32")
         return masks[np.newaxis] if squeeze else masks

src/patchworks/plugins/napari.py

@@ -406,21 +406,57 @@ def view_in_napari(
         # unwrapped to a single array) even for one level, so napari always
         # treats it as multiscale — required for 3D resolution switching, see
         # https://napari.org/stable/gallery/add_multiscale_volume.html
-        for name in _inner_label_names(image):
-            store = f"{image}/labels/{name}"
-            levels = _multiscale_levels(store, None)
-            lab = [lvl.astype("int32") for lvl in levels]
-            lab_scale, lab_units = _pyramid_calibration(store, lab[0].ndim)
-            viewer.add_labels(
-                lab,
-                name=name,
-                multiscale=True,
-                scale=lab_scale,
-                units=lab_units,
-                metadata=_label_hint(store),
-                **label_kwargs,
+        names = _inner_label_names(image)
+        if not names:
+            logger.warning(
+                "%s has no label images under labels/, so nothing was "
+                "overlaid. Pass labels=<path> explicitly if they live "
+                "somewhere else.",
+                image,
+            )
+        else:
+            logger.info(
+                "auto-loading %d label image(s) from %s/labels: %s",
+                len(names),
+                image,
+                ", ".join(names),
             )
+        loaded = 0
+        for name in names:
+            store = f"{image}/labels/{name}"
+            # Guarded per label: without this, one bad label group raised
+            # *after* the image had been added, so the viewer opened showing
+            # the image alone and every remaining label was skipped -- which
+            # looks exactly like "there were no labels".
+            try:
+                levels = _multiscale_levels(store, None)
+                lab = [lvl.astype("int32") for lvl in levels]
+                lab_scale, lab_units = _pyramid_calibration(store, lab[0].ndim)
+                viewer.add_labels(
+                    lab,
+                    name=name,
+                    multiscale=True,
+                    scale=lab_scale,
+                    units=lab_units,
+                    metadata=_label_hint(store),
+                    **label_kwargs,
+                )
+            except Exception:
+                logger.exception(
+                    "could not add labels/%s as a layer; skipping it and "
+                    "continuing with the rest.",
+                    name,
+                )
+                continue
+            loaded += 1
             logger.info("auto-loaded labels/%s from %s", name, image)
+        if names and not loaded:
+            logger.error(
+                "found %d label image(s) in %s/labels but none could be "
+                "added -- see the errors above.",
+                len(names),
+                image,
+            )
 
     if show:
         napari.run()