feat(core): ✨ persist object count for the napari plugin to skip its scan

merge_tile_labels(sequential_labels=True) already computes the exact
object count while renumbering, but discarded it. return_count=True now
surfaces it; write_labels/register_labels persist it as n_objects/
sequential_labels zarr attrs; view_in_napari propagates it into each
Labels layer's metadata via a new _label_hint() helper.

Lets a consumer (napari-dask-ndmeasure) skip a full O(voxels) scan of the
label array to find its id set — the id set is range(1, n_objects+1) by
construction when labels are sequential.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: lguerard

src/patchworks/_merge.py

@@ -459,7 +459,8 @@ def merge_tile_labels(
     stage_dir: Union[str, Path, None] = None,
     keep_stage: bool = False,
     progress: bool = False,
-) -> "da.Array":
+    return_count: bool = False,
+) -> Union["da.Array", tuple["da.Array", Union[int, None]]]:
     """Merge per-tile labels into a globally consistent label array.
 
     Standalone merge step — use this when you already have per-tile labels
@@ -503,11 +504,21 @@ def merge_tile_labels(
         Keep the temp stage zarr after merging. Default False.
     progress:
         Show a progress bar during the relabel step.
+    return_count:
+        Also return the exact object count. Only meaningful (non-``None``)
+        when ``sequential_labels=True``, which already computes it for free
+        while renumbering to ``1..N`` — otherwise no step here knows the
+        final count without an extra full scan, so the second element is
+        ``None``. Useful to persist alongside the labels (e.g.
+        ``write_labels(..., n_objects=...)``) so a downstream consumer with
+        the count can skip re-deriving the id set from the array itself.
 
     Returns
     -------
     da.Array
-        Merged label array (int32) backed by ``write_to``.
+        Merged label array (int32) backed by ``write_to``. Or, when
+        ``return_count=True``, a ``(labels, n_objects)`` tuple —
+        ``n_objects`` is ``None`` unless ``sequential_labels=True``.
 
     Examples
     --------
@@ -600,9 +611,10 @@ def merge_tile_labels(
         show_progress=progress,
     )
 
+    n_objects = None
     if sequential_labels:
         logger.info("Relabelling to contiguous ids…")
-        relabel_sequential_zarr(effective_out, output_component)
+        n_objects = relabel_sequential_zarr(effective_out, output_component)
 
     # -- Cleanup temp stage (only when we created it) --
     if not isinstance(labeled, (str, Path)) and not keep_stage:
@@ -611,4 +623,5 @@ def merge_tile_labels(
         shutil.rmtree(stage_path, ignore_errors=True)
         logger.info("Removed stage store %s", stage_path)
 
-    return da.from_zarr(effective_out, component=output_component)
+    result = da.from_zarr(effective_out, component=output_component)
+    return (result, n_objects) if return_count else result

src/patchworks/plugins/napari.py

@@ -179,6 +179,41 @@ def _pyramid_calibration(
     return scale, units
 
 
+def _label_hint(path: Union[str, Path]) -> dict[str, Any]:
+    """Read the known-object-count hint from a label group's zarr attrs.
+
+    ``write_labels(..., n_objects=...)`` persists this when the labels were
+    renumbered to a contiguous ``1..N`` range (``sequential_labels=True``
+    during the merge) — the exact id set is then ``range(1, n_objects +
+    1)`` by construction, with no scan needed. Passed through as a Labels
+    layer's ``metadata`` so a downstream consumer (e.g.
+    napari-dask-ndmeasure) can use it instead of re-deriving the id set
+    from the array itself.
+
+    Parameters
+    ----------
+    path : str or Path
+        Label group path (e.g. ``f"{image}/labels/{name}"``).
+
+    Returns
+    -------
+    dict
+        ``{"n_objects": int, "sequential_labels": True}`` if the group has
+        the hint, else ``{}`` — safe to splat straight into
+        ``metadata=``/merge into a bigger dict either way.
+    """
+    try:
+        attrs = zarr.open_group(str(path), mode="r").attrs
+    except Exception:
+        return {}
+    if "n_objects" not in attrs:
+        return {}
+    return {
+        "n_objects": attrs["n_objects"],
+        "sequential_labels": attrs.get("sequential_labels", False),
+    }
+
+
 def _inner_label_names(store: Union[str, Path]) -> list[str]:
     """List label images registered under an OME-ZARR's ``labels/`` group.
 
@@ -351,12 +386,14 @@ def view_in_napari(
             if _is_zarr(labels)
             else (None, None)
         )
+        metadata = _label_hint(labels) if _is_zarr(labels) else {}
         viewer.add_labels(
             lab,
             name=labels_name,
             multiscale=isinstance(lab, list),
             scale=lab_scale,
             units=lab_units,
+            metadata=metadata,
             **label_kwargs,
         )
     elif _is_zarr(image):
@@ -377,6 +414,7 @@ def view_in_napari(
                 multiscale=True,
                 scale=lab_scale,
                 units=lab_units,
+                metadata=_label_hint(store),
                 **label_kwargs,
             )
             logger.info("auto-loaded labels/%s from %s", name, image)

src/patchworks/plugins/ome_zarr.py

@@ -1045,6 +1045,7 @@ def register_labels(
     chunks: Union[tuple[int, ...], None] = None,
     shard: ShardSpec = False,
     progress: bool = True,
+    n_objects: Union[int, None] = None,
 ) -> str:
     """Pyramidalise and register an existing ``labels/<name>/0`` base level.
 
@@ -1076,6 +1077,14 @@ def register_labels(
         Sharding request (see :func:`to_ome_zarr`'s *shard*).
     progress : bool, optional
         Show a per-level dask progress bar (default ``True``).
+    n_objects : int or None, optional
+        Exact non-background object count, if known (e.g. from
+        :func:`patchworks.merge_tile_labels`'s ``return_count=True`` after
+        ``sequential_labels=True``, which means ``ids == range(1, n_objects
+        + 1)`` by construction). When given, written into the label group's
+        attrs as ``n_objects``/``sequential_labels`` so a downstream reader
+        (e.g. napari-dask-ndmeasure) can use the known id set instead of
+        re-deriving it with a full-volume scan of its own.
 
     Returns
     -------
@@ -1106,6 +1115,9 @@ def register_labels(
     )
     grp = zarr.open_group(group, mode="a")
     grp.attrs["image-label"] = {"version": _NGFF_VERSION}
+    if n_objects is not None:
+        grp.attrs["n_objects"] = int(n_objects)
+        grp.attrs["sequential_labels"] = True
 
     labels_grp = zarr.open_group(f"{store}/labels", mode="a")
     registered = list(labels_grp.attrs.get("labels", []))
@@ -1128,6 +1140,7 @@ def write_labels(
     shard: ShardSpec = False,
     progress: bool = True,
     overwrite: bool = False,
+    n_objects: Union[int, None] = None,
 ) -> str:
     """Store *labels* inside *image_store* under the NGFF ``labels/`` group.
 
@@ -1166,6 +1179,9 @@ def write_labels(
     overwrite : bool, optional
         Replace an existing label image of the same *name* (default
         ``False``).
+    n_objects : int or None, optional
+        Exact non-background object count, if known — forwarded to
+        :func:`register_labels`; see its docstring for what this enables.
 
     Returns
     -------
@@ -1175,10 +1191,16 @@ def write_labels(
     Examples
     --------
     >>> from patchworks import merge_tile_labels
-    >>> merged = merge_tile_labels(
-    ...     "stage.zarr", input_component="staged", write_to="merged.zarr"
+    >>> merged, n = merge_tile_labels(
+    ...     "stage.zarr",
+    ...     input_component="staged",
+    ...     write_to="merged.zarr",
+    ...     sequential_labels=True,
+    ...     return_count=True,
+    ... )  # doctest: +SKIP
+    >>> write_labels(
+    ...     "scan.zarr", merged, name="cells", n_objects=n
     ... )  # doctest: +SKIP
-    >>> write_labels("scan.zarr", merged, name="cells")  # doctest: +SKIP
     'scan.zarr/labels/cells'
     """
     arr = labels if isinstance(labels, da.Array) else da.asarray(labels)
@@ -1209,4 +1231,5 @@ def write_labels(
         chunks=chunks,
         shard=shard,
         progress=progress,
+        n_objects=n_objects,
     )

tests/test_core.py

@@ -148,6 +148,62 @@ def fn(tile):
     assert ids.size == 1, f"object split into {ids.size} labels, expected 1"
 
 
+def test_merge_tile_labels_return_count(tmp_path):
+    # sequential_labels=True already computes the exact object count while
+    # renumbering to 1..N — return_count=True surfaces it instead of
+    # discarding it, so a caller can persist it (e.g. write_labels'
+    # n_objects=) for a downstream consumer to skip re-deriving the id set.
+    import dask.array as da
+
+    from patchworks import merge_tile_labels
+
+    data = np.zeros((1, 16, 32), dtype="uint16")
+    data[0, 2:6, 2:6] = 1
+    data[0, 2:6, 10:14] = 1  # same tile-local label, different object
+    image = da.from_array(data, chunks=(1, 16, 16))
+
+    def fn(tile):
+        from skimage.measure import label
+
+        return label(tile > 0).astype("int32")
+
+    labeled = image.map_blocks(
+        fn, dtype="int32", meta=np.empty((0,) * image.ndim, dtype="int32")
+    )
+    out = str(tmp_path / "merged.zarr")
+    merged, n_objects = merge_tile_labels(
+        labeled, write_to=out, sequential_labels=True, return_count=True
+    )
+    arr = merged.compute()
+    ids = np.unique(arr[arr > 0])
+    assert n_objects == ids.size
+    assert n_objects == 2
+
+
+def test_merge_tile_labels_return_count_none_without_sequential(tmp_path):
+    import dask.array as da
+
+    from patchworks import merge_tile_labels
+
+    data = np.zeros((1, 16, 32), dtype="uint16")
+    data[0, 4:12, 8:24] = 1
+    image = da.from_array(data, chunks=(1, 16, 16))
+
+    def fn(tile):
+        from skimage.measure import label
+
+        return label(tile > 0).astype("int32")
+
+    labeled = image.map_blocks(
+        fn, dtype="int32", meta=np.empty((0,) * image.ndim, dtype="int32")
+    )
+    out = str(tmp_path / "merged.zarr")
+    merged, n_objects = merge_tile_labels(
+        labeled, write_to=out, sequential_labels=False, return_count=True
+    )
+    assert n_objects is None
+
+
 def test_merge_transitive_three_tiles(tmp_path):
     # A cell that spans 3 tiles (A→B→C) must be merged into one label even
     # though A and C never directly touch. Transitivity via connected_components.

tests/test_napari.py

@@ -68,3 +68,39 @@ def test_inner_label_discovery_none(tmp_path):
         np.zeros((8, 8, 8), "uint16"), tmp_path / "img.zarr", n_levels=1
     )
     assert nplugin._inner_label_names(store) == []
+
+
+def test_label_hint_present_when_n_objects_written(tmp_path):
+    """write_labels(..., n_objects=...) is readable back via _label_hint."""
+    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=1
+    )
+    write_labels(
+        store,
+        np.ones((8, 8, 8), "int32"),
+        name="cells",
+        n_levels=1,
+        n_objects=17,
+    )
+
+    hint = nplugin._label_hint(f"{store}/labels/cells")
+    assert hint == {"n_objects": 17, "sequential_labels": True}
+
+
+def test_label_hint_empty_without_n_objects(tmp_path):
+    """No n_objects= at write time -> no hint, not a misleading default."""
+    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=1
+    )
+    write_labels(store, np.ones((8, 8, 8), "int32"), name="cells", n_levels=1)
+
+    assert nplugin._label_hint(f"{store}/labels/cells") == {}
+
+
+def test_label_hint_missing_store_returns_empty():
+    """A path that doesn't exist (or isn't a label group) just yields {}."""
+    assert nplugin._label_hint("/no/such/store.zarr") == {}

tests/test_ome_zarr.py

@@ -130,6 +130,34 @@ def test_write_labels_into_store(tmp_path):
     assert lg.attrs["image-label"]["version"]
 
 
+def test_write_labels_n_objects_persisted(tmp_path):
+    """n_objects lands in the label group's attrs for a downstream reader."""
+    store = to_ome_zarr(
+        np.zeros((8, 8, 8), "uint16"), tmp_path / "img.zarr", n_levels=2
+    )
+    labels = np.ones((8, 8, 8), dtype="int32")
+
+    group = write_labels(store, labels, name="cells", n_levels=2, n_objects=42)
+
+    lg = zarr.open_group(group, mode="r")
+    assert lg.attrs["n_objects"] == 42
+    assert lg.attrs["sequential_labels"] is True
+
+
+def test_write_labels_no_n_objects_by_default(tmp_path):
+    """Without n_objects=, no misleading count is written."""
+    store = to_ome_zarr(
+        np.zeros((8, 8, 8), "uint16"), tmp_path / "img.zarr", n_levels=2
+    )
+    labels = np.ones((8, 8, 8), dtype="int32")
+
+    group = write_labels(store, labels, name="cells", n_levels=2)
+
+    lg = zarr.open_group(group, mode="r")
+    assert "n_objects" not in lg.attrs
+    assert "sequential_labels" not in lg.attrs
+
+
 def test_reuse_pyramid_ignored_for_arrays(tmp_path):
     """reuse_pyramid only affects .ims inputs; arrays still rebuild."""
     out = to_ome_zarr(

workflow/scripts/merge.py

@@ -28,13 +28,14 @@
 default_workers = int(
     os.environ.get("SLURM_CPUS_PER_TASK", os.cpu_count() or 4)
 )
-merged = merge_tile_labels(
+merged, n_objects = merge_tile_labels(
     stage_path(work_dir, label_name),
     write_to=merged_store,
     input_component="staged",
     sequential_labels=cfg.get("sequential_labels", True),
     n_workers=cfg.get("merge_workers", default_workers),
     progress=False,
+    return_count=True,
 )
 group = write_labels(
     image_store,
@@ -43,6 +44,7 @@
     n_levels=int(cfg.get("pyramid_levels", 5)),
     downscale=int(cfg.get("pyramid_downscale", 2)),
     overwrite=True,
+    n_objects=n_objects,
 )
 
 shutil.rmtree(merged_store, ignore_errors=True)