Merge pull request #17 from imcf/feat/napari-autolabels-gpu-dashboard

feat: napari auto-loads in-store labels + single-GPU dashboard

Author: lguerard

docs/guide/ome_zarr_napari.md

@@ -50,6 +50,13 @@ to_ome_zarr("scan.czi", "scan.zarr", n_levels=5)   # via bioio
 to_ome_zarr("scan.ims", "scan.zarr")               # Imaris, native HDF5
 ```
 
+!!! note "Imaris pyramids are rebuilt, not reused"
+    `.ims` files carry their own resolution pyramid, but `to_ome_zarr` reads
+    only the **full-resolution** level and **builds a fresh NGFF pyramid** from
+    it. This guarantees a consistent pyramid (XY-only, nearest-neighbour,
+    calibrated) rather than inheriting Imaris's own downsampling scheme. It
+    costs some extra compute, but the build is lazy and OOM-safe.
+
 ### Pixel calibration
 
 The physical voxel size is read from the input — bioio's `physical_pixel_sizes`,
@@ -96,13 +103,17 @@ write_labels("scan.zarr", my_labels, name="nuclei")
 layer in one call. OME-ZARR pyramids are handed to napari as a lazy multi-scale
 list, so even huge stores open instantly and only on-screen data is fetched.
 
+Because `tile_process` writes labels **into** the store by default, you usually
+need no `labels=` argument at all — `view_in_napari` auto-loads every label
+image found under `scan.zarr/labels/`:
+
 ```python
 from patchworks.plugins.napari import view_in_napari
 
-# one store holding both image and labels/<name>:
-view_in_napari("scan.zarr", labels="scan.zarr/labels/labels")
+# auto-loads scan.zarr/labels/* as Labels layers:
+view_in_napari("scan.zarr")
 
-# or a separate plain label store written with write_to=:
+# or point at a separate plain label store written with write_to=:
 view_in_napari("scan.zarr", labels="labels.zarr")
 ```
 
@@ -120,7 +131,7 @@ from patchworks.plugins.napari import view_in_napari
 tile_process("scan.zarr", fn, progress=True)
 
 # 2. inspect image + labels together, straight from the one store
-view_in_napari("scan.zarr", labels="scan.zarr/labels/labels")
+view_in_napari("scan.zarr")  # labels auto-loaded from scan.zarr/labels/
 ```
 
 Plugging in a different segmentation method is just swapping `fn` — any

docs/guide/performance.md

@@ -17,6 +17,22 @@ merge step are sized to the host automatically:
 The RAM figure is read live via `psutil`; without it, a conservative default is
 used instead of guessing high.
 
+## Live progress dashboard (GPU runs)
+
+A single-GPU run still gets a **Dask dashboard**: patchworks spins up a tiny
+1-worker / 1-thread in-process cluster, which keeps GPU evaluations serial (no
+VRAM contention) while exposing the dashboard so you can watch tiles stream
+through. The URL is logged at the start of staging:
+
+```text
+INFO:patchworks._core:Dask dashboard for this run: http://127.0.0.1:8787/status
+```
+
+This needs `distributed` (and `bokeh` for the UI) installed; if they are
+missing, patchworks logs a warning and falls back to the threaded scheduler
+(no dashboard, same result). A cluster you start yourself
+(`make_local_cluster`) is used as-is instead.
+
 ## Overriding the worker count
 
 ```python

src/patchworks/_core.py

@@ -332,7 +332,31 @@ def active_fn(block, block_info=None):
     import dask as _dask
 
     _tile_nbytes = int(np.prod(labeled.chunksize)) * labeled.dtype.itemsize
-    if _active is None:
+    _temp_cluster = None
+    _temp_client = None
+    if _active is None and use_gpu:
+        # Single-GPU runs still get a live Dask dashboard: a 1-worker /
+        # 1-thread in-process cluster keeps GPU evals serial (no VRAM
+        # contention) while exposing the dashboard for progress.
+        try:
+            from dask.distributed import Client, LocalCluster
+
+            _temp_cluster = LocalCluster(
+                n_workers=1, threads_per_worker=1, processes=False
+            )
+            _temp_client = Client(_temp_cluster)
+            logger.info(
+                "Dask dashboard for this run: %s",
+                _temp_client.dashboard_link,
+            )
+        except Exception as exc:  # no distributed/bokeh → threaded fallback
+            logger.warning(
+                "Could not start a dashboard cluster (%s); "
+                "falling back to the threaded scheduler.",
+                exc,
+            )
+
+    if _distributed_client() is None:
         _workers = (
             max_workers
             if max_workers is not None
@@ -363,6 +387,9 @@ def active_fn(block, block_info=None):
     logger.info("Staging tiles to %s …", stage_path)
     with _sched_ctx:
         _stage_to_zarr(labeled, stage_path, "staged", progress)
+    if _temp_client is not None:
+        _temp_client.close()
+        _temp_cluster.close()
     labeled = da.from_zarr(stage_path, component="staged")
 
     # NB: no post-staging skip-count pass here — counting skipped tiles by

src/patchworks/plugins/napari.py

@@ -14,13 +14,12 @@
 Usage
 -----
 >>> from patchworks import tile_process
->>> from patchworks.plugins.ome_zarr import to_ome_zarr
 >>> from patchworks.plugins.napari import view_in_napari
 >>>
->>> tile_process("scan.zarr", fn, write_to="labels.zarr")
->>> to_ome_zarr("scan.zarr", "scan_pyramid.zarr")        # optional, for speed
->>>
->>> view_in_napari("scan_pyramid.zarr", labels="labels.zarr")
+>>> # labels are written into scan.zarr/labels/ by default …
+>>> tile_process("scan.zarr", fn)
+>>> # … so the viewer finds and overlays them with no labels= argument:
+>>> view_in_napari("scan.zarr")
 """
 
 from __future__ import annotations
@@ -86,6 +85,15 @@ def _resolve_image(
     return source
 
 
+def _inner_label_names(store: Union[str, Path]) -> list[str]:
+    """Names registered under an OME-ZARR's NGFF ``labels/`` group, if any."""
+    try:
+        grp = zarr.open_group(f"{store}/labels", mode="r")
+    except Exception:
+        return []
+    return list(grp.attrs.get("labels", []))
+
+
 def _resolve_labels(
     source: Union[da.Array, str, Path], component: str
 ) -> Union[da.Array, list[da.Array]]:
@@ -123,7 +131,10 @@ def view_in_napari(
     labels : da.Array, str, Path or None
         Label array to overlay. A plain ``.zarr`` store written by
         ``tile_process`` is read from its ``labels_component``; an OME-ZARR
-        pyramid is shown multi-scale; ``None`` shows the image only.
+        pyramid is shown multi-scale. ``None`` (default) **auto-loads** every
+        label image stored inside the OME-ZARR under ``labels/<name>/`` — the
+        place ``tile_process`` writes them by default — each as its own Labels
+        layer. (Falls back to image-only if there are none.)
     channel : int or None, optional
         Channel to display from the image (``None`` keeps all channels).
     labels_component : str, optional
@@ -145,7 +156,7 @@ def view_in_napari(
 
     Examples
     --------
-    >>> view_in_napari("scan.zarr", labels="labels.zarr")  # doctest: +SKIP
+    >>> view_in_napari("scan.zarr")  # auto-loads scan.zarr/labels/*  # doctest: +SKIP
     """
     napari = _require_napari()
 
@@ -161,6 +172,15 @@ def view_in_napari(
     if labels is not None:
         lab = _resolve_labels(labels, labels_component)
         viewer.add_labels(lab, name=labels_name)
+    elif _is_zarr(image):
+        # No labels given → auto-overlay every label image stored inside the
+        # OME-ZARR under labels/<name>/ (the default place tile_process writes
+        # them), each as its own multi-scale Labels layer.
+        for name in _inner_label_names(image):
+            levels = _multiscale_levels(f"{image}/labels/{name}", None)
+            lab = [lvl.astype("int32") for lvl in levels]
+            viewer.add_labels(lab if len(lab) > 1 else lab[0], name=name)
+            logger.info("auto-loaded labels/%s from %s", name, image)
 
     if show:
         napari.run()

tests/test_napari.py

@@ -39,3 +39,32 @@ def test_require_napari_message(monkeypatch):
             nplugin._require_napari()
     else:
         assert nplugin._require_napari() is napari
+
+
+def test_inner_label_discovery(tmp_path):
+    """Labels written into a store are discoverable for auto-overlay."""
+    import numpy as np
+
+    from patchworks.plugins.ome_zarr import to_ome_zarr, write_labels
+
+    store = to_ome_zarr(
+        np.zeros((8, 8, 8), "uint16"), tmp_path / "scan.zarr", n_levels=2
+    )
+    write_labels(store, np.ones((8, 8, 8), "int32"), name="cells", n_levels=2)
+
+    assert nplugin._inner_label_names(store) == ["cells"]
+    levels = nplugin._multiscale_levels(f"{store}/labels/cells", None)
+    assert len(levels) == 2
+    assert levels[1].shape == (8, 4, 4)  # Z preserved, XY downsampled
+
+
+def test_inner_label_discovery_none(tmp_path):
+    """A store without labels yields an empty list (image-only view)."""
+    import numpy as np
+
+    from patchworks.plugins.ome_zarr import to_ome_zarr
+
+    store = to_ome_zarr(
+        np.zeros((8, 8, 8), "uint16"), tmp_path / "img.zarr", n_levels=1
+    )
+    assert nplugin._inner_label_names(store) == []