patchworks 2.5.2__tar.gz → 2.6.1__tar.gz

{patchworks-2.5.2 → patchworks-2.6.1}/.gitignore

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

{patchworks-2.5.2 → patchworks-2.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchworks
-Version: 2.5.2
+Version: 2.6.1
 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
@@ -115,6 +115,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

{patchworks-2.5.2 → patchworks-2.6.1}/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

{patchworks-2.5.2 → patchworks-2.6.1}/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,30 @@ 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. `tile_shape: "auto"` accounts for this — it is
+  told the tile carries two channels and shrinks each spatial side by ~1/√2,
+  so the tile still fits the same VRAM budget. A **hand-set** `tile_shape`
+  sized to fill a GPU has no such protection and needs halving yourself.
+- 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:

{patchworks-2.5.2 → patchworks-2.6.1}/src/patchworks/_chunks.py

@@ -264,6 +264,7 @@ def auto_tile_shape(
     gpu_memory: int | None = None,
     available_memory: int | None = None,
     n_workers: int | None = None,
+    n_channels: int = 1,
     verbose: bool = False,
 ) -> tuple[int, ...]:
     """Balanced tile shape for general-purpose 3-D processing.
@@ -287,6 +288,10 @@ def auto_tile_shape(
         Available host RAM in bytes; auto-queried when None.
     n_workers:
         Number of parallel workers (divides the RAM budget).
+    n_channels:
+        Channels each tile carries (default 1). Above 1 the per-voxel cost
+        scales with it, so the tile shrinks accordingly -- e.g. the workflow's
+        ``nuclei_channel`` hands Cellpose a cyto+nuclei pair.
     verbose:
         Log the chosen shape and estimated tile size.
 
@@ -302,7 +307,12 @@ def auto_tile_shape(
     (128, 512, 512)
     """
     n_workers = n_workers or cpu_allocation()
-    itemsize = np.dtype(dtype).itemsize
+    # A tile holds n_channels planes per voxel (e.g. Cellpose's
+    # cyto+nuclei pair), so the per-voxel cost -- and every budget
+    # derived from it below -- scales with them.
+    if n_channels < 1:
+        raise ValueError(f"n_channels must be >= 1; got {n_channels!r}")
+    itemsize = np.dtype(dtype).itemsize * n_channels
     n_spatial = min(3, len(shape))
 
     if use_gpu:
@@ -359,6 +369,7 @@ def auto_tile_shape_cellpose(
     n_workers: int | None = None,
     model_memory_bytes: int = 2 * 1024**3,
     cellpose_memory_factor: int = 20,
+    n_channels: int = 1,
     verbose: bool = False,
 ) -> tuple[int, ...]:
     """Cellpose-optimised tile shape.
@@ -391,6 +402,10 @@ def auto_tile_shape_cellpose(
         Memory consumed by the Cellpose model weights (default 2 GiB).
     cellpose_memory_factor:
         Cellpose allocates roughly this multiple of raw input bytes (default 20×).
+    n_channels:
+        Channels each tile carries (default 1). Above 1 the per-voxel cost
+        scales with it, so the tile shrinks accordingly -- e.g. the workflow's
+        ``nuclei_channel`` hands Cellpose a cyto+nuclei pair.
     verbose:
         Log the chosen shape and memory estimates.
 
@@ -406,7 +421,12 @@ def auto_tile_shape_cellpose(
     (1, 2048, 2048)
     """
     n_workers = n_workers or cpu_allocation()
-    itemsize = np.dtype(dtype).itemsize
+    # A tile holds n_channels planes per voxel (e.g. Cellpose's
+    # cyto+nuclei pair), so the per-voxel cost -- and every budget
+    # derived from it below -- scales with them.
+    if n_channels < 1:
+        raise ValueError(f"n_channels must be >= 1; got {n_channels!r}")
+    itemsize = np.dtype(dtype).itemsize * n_channels
 
     if use_gpu:
         total_mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()

{patchworks-2.5.2 → patchworks-2.6.1}/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."

{patchworks-2.5.2 → patchworks-2.6.1}/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

{patchworks-2.5.2 → patchworks-2.6.1}/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()

{patchworks-2.5.2 → patchworks-2.6.1}/tests/test_core.py

@@ -389,3 +389,32 @@ def test_channel_selection_respects_the_stores_axes(tmp_path):
     store4 = str(tmp_path / "czyx.zarr")
     to_ome_zarr(vol4, store4, axes="czyx", n_levels=1, progress=False)
     assert load_ome_zarr(store4, channel=1, level=0).shape == (8, 64, 64)
+
+
+def test_auto_tile_shape_charges_for_extra_channels():
+    """A 2-channel tile must fit the same byte budget, not twice it.
+
+    `nuclei_channel` doubles what a tile holds while the tile geometry stays
+    single-channel, so a sizer blind to it hands the GPU a tile needing twice
+    the VRAM it budgeted for.
+    """
+    import pytest
+
+    from patchworks import auto_tile_shape, auto_tile_shape_cellpose
+
+    shape, dtype = (128, 2048, 2048), "uint16"
+
+    one = auto_tile_shape(shape, dtype)
+    two = auto_tile_shape(shape, dtype, n_channels=2)
+    # Same bytes overall: 2 channels of roughly half the area each.
+    assert np.prod(two) * 2 <= np.prod(one)
+    assert np.prod(two) * 2 >= np.prod(one) * 0.9
+
+    kw = dict(diameter=30, do_3D=True, use_gpu=True, gpu_memory=24 * 1024**3)
+    cp_one = auto_tile_shape_cellpose(shape, dtype, **kw)
+    cp_two = auto_tile_shape_cellpose(shape, dtype, n_channels=2, **kw)
+    assert np.prod(cp_two) * 2 <= np.prod(cp_one)
+    assert cp_two[0] == cp_one[0]  # do_3D still pins z to the full extent
+
+    with pytest.raises(ValueError, match="n_channels"):
+        auto_tile_shape(shape, dtype, n_channels=0)

{patchworks-2.5.2 → patchworks-2.6.1}/tests/test_distributed.py

@@ -399,3 +399,96 @@ def test_separate_objects_keep_distinct_labels(tmp_path):
     ids = np.unique(merged[merged > 0])
     assert ids.size == 4, f"expected 4 distinct objects, got {ids.size}"
     assert set(ids.tolist()) == {1, 2, 3, 4}  # contiguous after relabel
+
+
+def test_stage_tile_carries_a_channel_axis(tmp_path):
+    """A channel axis reaches fn but is never tiled.
+
+    Cellpose's cyto+nucleus pair: tile geometry, the halo and the stage store
+    all stay purely spatial, so fn sees ``(2, z, y, x)`` and owes back
+    ``(z, y, x)``. Without this, a second channel would either be tiled like a
+    spatial axis or trip stage_tile's shape check.
+    """
+    import numpy as np
+    import zarr
+
+    from patchworks import create_stage, stage_tile
+
+    # Channel 0 carries the signal, channel 1 is blank -- so a swapped or
+    # collapsed channel axis fails rather than quietly segmenting the wrong one.
+    img = np.stack(
+        [
+            np.arange(4 * 8 * 8).reshape(4, 8, 8) % 5,
+            np.zeros((4, 8, 8), dtype=int),
+        ]
+    )
+    assert img.shape == (2, 4, 8, 8)
+
+    tile_shape = (4, 4, 4)
+    stage = create_stage(tmp_path / "s.zarr", (4, 8, 8), tile_shape)
+
+    seen = {}
+
+    def fn(block):
+        seen["shape"] = block.shape
+        assert not block[1].any(), "channel 1 should be the blank one"
+        return (block[0] > 2).astype("int32")
+
+    n = stage_tile(
+        img,
+        fn,
+        stage,
+        0,
+        tile_shape=tile_shape,
+        overlap=1,
+        channel_axis=0,
+    )
+
+    # z is already the full extent, y/x grow by the 1-voxel halo.
+    assert seen["shape"] == (2, 4, 5, 5)
+    assert n == 1
+    staged = zarr.open_group(str(stage))["staged"][:]
+    assert staged.shape == (4, 8, 8)  # store never learned about channels
+    assert staged[:, :4, :4].any()  # the tile actually landed
+
+
+def test_cellpose_run_pairs_channels_and_shifts_z(monkeypatch):
+    """A 2-channel tile must set z_axis past the channel axis.
+
+    With channels stacked on axis 0, z moves to axis 1; leaving z_axis=0 would
+    tell Cellpose the channel axis is z. Cellpose 3 also needs the pair named
+    as ``channels=[1, 2]`` (1-based, 0 = grayscale), which Cellpose 4 dropped.
+    """
+    import numpy as np
+
+    from patchworks.plugins import cellpose as cp
+
+    calls = {}
+
+    class _FakeModel:
+        def eval(self, img, **kwargs):
+            calls["img_shape"] = img.shape
+            calls["kwargs"] = kwargs
+            return np.zeros(img.shape[1:], dtype="int32"), None
+
+    monkeypatch.setattr(cp, "_get_model", lambda _cfg: _FakeModel())
+
+    cfg = cp._make_config("cyto3", diameter=30, do_3D=True, channel_axis=0)
+    block = np.zeros((2, 4, 8, 8), dtype="uint16")
+    out = cp._run(block, cfg)
+
+    assert out.shape == (4, 8, 8)  # labels come back spatial
+    assert calls["kwargs"]["z_axis"] == 1
+    assert calls["kwargs"]["channel_axis"] == 0
+    if cp._CELLPOSE_V4:
+        assert "channels" not in calls["kwargs"]
+    else:
+        assert calls["kwargs"]["channels"] == [1, 2]
+
+    # Single-channel tiles keep the old geometry and grayscale pairing.
+    cfg1 = cp._make_config("cyto3", diameter=30, do_3D=True)
+    cp._run(np.zeros((4, 8, 8), dtype="uint16"), cfg1)
+    assert calls["kwargs"]["z_axis"] == 0
+    assert calls["kwargs"]["channel_axis"] is None
+    if not cp._CELLPOSE_V4:
+        assert calls["kwargs"]["channels"] == [0, 0]

patchworks-2.6.1/workflow/config/config_cyto.yaml

@@ -0,0 +1,39 @@
+# Example segmentation config: cytoplasm/membrane channel.
+#
+# Only what differs from config/common.yaml — the input, work_dir, tiling and
+# conversion settings all live there and are merged in ahead of this file.
+# Run it via config/multi.yaml, or on its own with both files:
+#
+#   snakemake -s Snakefile --configfile config/common.yaml config/config_cyto.yaml
+
+channel: 0 # cytoplasm/membrane channel
+
+# Optional second channel handed to Cellpose as the nucleus input, which
+# usually improves cytoplasm segmentation. Channel indices are 0-based, same
+# as `channel`. Leave it out (or null) for a single-channel run.
+#
+# Only the segment step reads it: the tile geometry, the occupancy map and the
+# staged labels all stay exactly as they are for a single channel -- the pair
+# is stacked on a leading axis that is carried into each tile, not tiled.
+#
+# It does double the bytes a tile holds. `tile_shape: "auto"` is told about
+# that and shrinks each spatial side by ~1/sqrt(2) to fit the same VRAM
+# budget; the hand-set tile_shape in common.yaml is not, so halve it yourself.
+#
+# Cellpose 3 turns this into channels: [1, 2]; Cellpose 4 (cpsam) dropped
+# `channels` and just reads both. Override either by setting `channels:` or
+# `channel_axis:` in the cellpose: block below.
+nuclei_channel: 1
+
+# Per-axis halo [z, y, x]. A scalar 30 would expand a [16, 1024, 1024] tile to
+# 76 x 1084 x 1084 = 5.3x the voxels it keeps, nearly all of it wasted z. The
+# z-halo only needs to cover one cell in z, not one cell in x/y.
+overlap: [4, 30, 30]
+
+method: "cellpose"
+label_name: "cyto_labels"
+cellpose:
+  model: "cyto3"
+  diameter: 30
+  do_3D: true
+  gpu: true

{patchworks-2.5.2 → patchworks-2.6.1}/workflow/scripts/_pw.py

@@ -71,7 +71,7 @@ def start_log(path, *, append=True):
     return handle
 
 
-def open_image(work_dir, channel, level):
+def open_image(work_dir, channel, level, nuclei_channel=None):
     """Open the converted image for segmentation.
 
     Parameters
@@ -82,14 +82,38 @@ def open_image(work_dir, channel, level):
         Channel to select.
     level : int
         Pyramid level to read.
+    nuclei_channel : int or None, optional
+        Second channel to stack in front of *channel*, giving a
+        ``(2, z, y, x)`` array whose leading axis is ``[channel,
+        nuclei_channel]``. This is Cellpose's cyto+nucleus pair (see
+        ``config/config_cyto.yaml``). ``None`` (default) returns the plain
+        ``(z, y, x)`` single-channel array.
 
     Returns
     -------
     da.Array
         The (lazy) image array.
     """
-    return load_ome_zarr(
-        str(Path(work_dir) / "image.zarr"), channel=channel, level=level
+    path = str(Path(work_dir) / "image.zarr")
+    if nuclei_channel is None:
+        return load_ome_zarr(path, channel=channel, level=level)
+
+    import dask.array as da
+
+    if nuclei_channel == channel:
+        raise ValueError(
+            f"nuclei_channel ({nuclei_channel}) is the same as channel "
+            f"({channel}); Cellpose would get the same image twice. Channel "
+            "indices are 0-based."
+        )
+    # Two lazy reads of the same spatial box, stacked. The tiling stays
+    # spatial -- stage_tile carries this axis through rather than tiling it.
+    return da.stack(
+        [
+            load_ome_zarr(path, channel=c, level=level)
+            for c in (channel, nuclei_channel)
+        ],
+        axis=0,
     )
 
 
@@ -226,6 +250,24 @@ def validate_config(cfg) -> None:
     elif ts is not None and not all(isinstance(v, int) and v > 0 for v in ts):
         problems.append(f"tile_shape must be positive integers; got {ts!r}")
 
+    nuc = cfg.get("nuclei_channel")
+    if nuc is not None:
+        if isinstance(nuc, bool) or not isinstance(nuc, int):
+            problems.append(
+                "nuclei_channel must be null or a 0-based channel index "
+                f"(e.g. 1); got {nuc!r}"
+            )
+        elif nuc == cfg.get("channel"):
+            problems.append(
+                f"nuclei_channel ({nuc}) is the same as channel; Cellpose "
+                "would get the same image twice. Indices are 0-based."
+            )
+        if cfg.get("method", "cellpose") == "threshold":
+            problems.append(
+                'nuclei_channel has no meaning for method: "threshold", '
+                "which segments a single channel"
+            )
+
     method = cfg.get("method", "cellpose")
     if method not in KNOWN_METHODS:
         listed = ", ".join(f'"{m}"' for m in KNOWN_METHODS)
@@ -378,6 +420,11 @@ def _build_method_fn(cfg):
             for k, v in cp.items()
             if k not in ("model", "diameter", "do_3D", "gpu")
         }
+        if cfg.get("nuclei_channel") is not None:
+            # open_image stacks [channel, nuclei_channel] on axis 0; tell
+            # Cellpose where they are. setdefault so an explicit
+            # cellpose.channel_axis in the config still wins.
+            extra.setdefault("channel_axis", 0)
         return cellpose_fn(
             cp.get("model", "cyto3"),
             gpu=cp.get("gpu", True),

{patchworks-2.5.2 → patchworks-2.6.1}/workflow/scripts/prepare_tiles.py

@@ -44,6 +44,10 @@ if ts == "auto":
     # "GPU memory query failed" fallback). None => the built-in 8 GiB default.
     gpu_gb = cfg.get("gpu_memory_gb")
     gpu_bytes = int(gpu_gb * 1024**3) if gpu_gb else None
+    # `image` is single-channel here (the geometry is spatial), but segment
+    # reads nuclei_channel alongside it, so a tile costs twice the bytes.
+    # Without this the sizer would hand the GPU a tile it cannot hold.
+    n_channels = 2 if cfg.get("nuclei_channel") is not None else 1
     if method == "cellpose":
         cp = cfg["cellpose"]
         sizer = partial(
@@ -52,6 +56,7 @@ if ts == "auto":
             use_gpu=cp.get("gpu", True),
             diameter=cp.get("diameter"),
             gpu_memory=gpu_bytes,
+            n_channels=n_channels,
         )
     else:
         # cfg["cellpose"] used to be read unconditionally here, so a DoG or
@@ -59,7 +64,10 @@ if ts == "auto":
         # 'cellpose'. The Cellpose estimator's memory model wouldn't apply to
         # them anyway.
         sizer = partial(
-            auto_tile_shape, use_gpu=gpu_bytes is not None, gpu_memory=gpu_bytes
+            auto_tile_shape,
+            use_gpu=gpu_bytes is not None,
+            gpu_memory=gpu_bytes,
+            n_channels=n_channels,
         )
     tile_shape = tuple(sizer(image.shape, image.dtype))
 else:

{patchworks-2.5.2 → patchworks-2.6.1}/workflow/scripts/segment_tile.py

@@ -21,7 +21,11 @@ work_dir = cfg["work_dir"]
 label_name = cfg.get("label_name", "labels")
 
 manifest = load_tiles_json(snakemake.input.tiles)  # noqa: F821
-image = open_image(work_dir, cfg["channel"], cfg["level"])
+# Optional second channel (Cellpose's nucleus input): open_image stacks it on
+# a leading axis that stage_tile carries through without tiling it, so the
+# tile geometry and the staged labels stay exactly as single-channel runs.
+nuclei_channel = cfg.get("nuclei_channel")
+image = open_image(work_dir, cfg["channel"], cfg["level"], nuclei_channel)
 indices = manifest["batches"][batch]
 
 # Built once for the whole batch: this is what makes the model load amortize.
@@ -46,6 +50,7 @@ for n, index in enumerate(indices, 1):
         # Scalar (older manifests) or per-axis list; stage_tile normalizes both.
         overlap=manifest["overlap"],
         component=component,
+        channel_axis=0 if nuclei_channel is not None else None,
     )
     # The per-tile time is what `tiles_per_job` has to be sized from: a job's
     # wall time is roughly N x this, and it must stay inside the QOS ceiling.

patchworks-2.5.2/workflow/config/config_cyto.yaml

@@ -1,22 +0,0 @@
-# Example segmentation config: cytoplasm/membrane channel.
-#
-# Only what differs from config/common.yaml — the input, work_dir, tiling and
-# conversion settings all live there and are merged in ahead of this file.
-# Run it via config/multi.yaml, or on its own with both files:
-#
-#   snakemake -s Snakefile --configfile config/common.yaml config/config_cyto.yaml
-
-channel: 0 # cytoplasm/membrane channel
-
-# Per-axis halo [z, y, x]. A scalar 30 would expand a [16, 1024, 1024] tile to
-# 76 x 1084 x 1084 = 5.3x the voxels it keeps, nearly all of it wasted z. The
-# z-halo only needs to cover one cell in z, not one cell in x/y.
-overlap: [4, 30, 30]
-
-method: "cellpose"
-label_name: "cyto_labels"
-cellpose:
-  model: "cyto3"
-  diameter: 30
-  do_3D: true
-  gpu: true