merge_tile_labels

patchworks.merge_tile_labels(labeled: Union['da.Array', str, Path], write_to: Union[str, Path, None] = None, *, input_component: str = 'labels', output_component: str = 'labels', overlap: int = 0, sequential_labels: bool = False, n_workers: int | None = None, stage_dir: Union[str, Path, None] = None, keep_stage: bool = False, progress: bool = False, return_count: bool = False, label_counts: 'Mapping[int, int] | Sequence[int] | None' = None, output_chunks: 'Sequence[int] | None' = None, shard_output: bool = False, halo_dir: 'str | Path | None' = None, iou_threshold: float = 0.5, connectivity: int = 1, parts_dir: 'str | Path | None' = None, min_voxels: 'int | None' = None, max_voxels: 'int | None' = None, objects_out: 'dict | None' = None) -> 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 (from your own segmentation pipeline) and just need the boundary stitching.

Accepts either:

• A dask array of per-tile integer labels (e.g. output of dask.array.map_blocks on your own segmentation function).
• A zarr store path whose input_component array already contains per-tile labels written by your own pipeline.

Labels that touch across tile boundaries are merged into a single ID. The merge is zarr-native (boundary scan → scipy connected components → parallel relabel) — no dask task graph, scales to thousands of tiles.

Parameters:

Name	Type	Description	Default
labeled	Union['da.Array', str, Path]	Per-tile label array. Either a dask array or a path to a zarr store that contains per-tile labels in input_component. A store is modified: unless label_counts is given, its per-tile ids are renumbered to a global range in place before the merge (the segmentation is unchanged, the ids are not). Copy it first if the original ids matter.	required
write_to	Union[str, Path, None]	Output zarr store path. When None, an auto-temp store is used.	None
input_component	str	Array name inside a zarr input store (ignored for dask arrays).	'labels'
output_component	str	Array name inside write_to. Default "labels".	'labels'
overlap	int	If labeled is a dask array that was computed with da.overlap, pass the same depth here to trim the halos before merging. Set 0 (default) if the array has no overlap halos.	0
sequential_labels	bool	Renumber the merged labels to a contiguous 1..N range. Folded into the merge's own lookup table, so it costs no extra pass over the volume. Default False.	False
label_counts	'Mapping[int, int] | Sequence[int] | None'	Labels written per tile, when the producer recorded them (see :func:patchworks.stage_tile, which returns the count). Lets the merge derive every tile's global id range arithmetically instead of streaming the whole store to renumber it. None keeps the renumber pass.	None
output_chunks	'Sequence[int] | None'	Chunking for the merged store; must divide the staged chunk shape. Lets the merge write straight into a store you will keep (e.g. an OME-ZARR label group's level 0) instead of a scratch store that then has to be copied. See :func:capped_output_chunks.	None
shard_output	bool	Store the output in shards of one staged tile each (zarr v3; the chunks above are inside them). Each relabel worker writes one tile, so it writes whole shards, and the result is sharded with no second pass over it. Ignored when relabeling in place.	False
halo_dir	'str | Path | None'	IoU stitching: the halo strips :func:patchworks.stage_tile saved (its halo_dir=). Labels are then joined across a tile boundary only when both tiles' predictions of the overlap zone agree (IoU >= iou_threshold), so two cells pressed together at a seam stay two. None (default) joins whatever touches.	None
iou_threshold	float	Minimum IoU for that join (default 0.5).	0.5
connectivity	int	Neighbourhood the tiles were labelled with: 1 (default) faces, 2 also edges, 3 also corners (scipy.ndimage.generate_binary_structure). Labels touching across a tile boundary that way are joined, so the result is that of labelling the whole image at once.	1
parts_dir	'str | Path | None'	Object sums saved per tile by :func:stage_tile (parts_dir=), combined into one row per merged object in objects_out -- the object table without a second read of the labels -- and an optional size filter applied in the same relabel pass. See :func:zarr_native_merge. Needs label_counts.	None
min_voxels	'str | Path | None'	Object sums saved per tile by :func:stage_tile (parts_dir=), combined into one row per merged object in objects_out -- the object table without a second read of the labels -- and an optional size filter applied in the same relabel pass. See :func:zarr_native_merge. Needs label_counts.	None
max_voxels	'str | Path | None'	Object sums saved per tile by :func:stage_tile (parts_dir=), combined into one row per merged object in objects_out -- the object table without a second read of the labels -- and an optional size filter applied in the same relabel pass. See :func:zarr_native_merge. Needs label_counts.	None
objects_out	'str | Path | None'	Object sums saved per tile by :func:stage_tile (parts_dir=), combined into one row per merged object in objects_out -- the object table without a second read of the labels -- and an optional size filter applied in the same relabel pass. See :func:zarr_native_merge. Needs label_counts.	None
n_workers	int | None	Parallel workers for the relabel step. Default min(4, cpu_count).	None
stage_dir	Union[str, Path, None]	Directory for the temp stage zarr when labeled is a dask array. Default: a system temp directory.	None
keep_stage	bool	Keep the temp stage zarr after merging. Default False.	False
progress	bool	Show a progress bar during the relabel step.	False
return_count	bool	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.	False

Returns:

Type	Description
Array	Merged label array backed by write_to: the staged dtype (int32 or wider for a dask input), widened if the global object count outgrows it. Or, when return_count=True, a (labels, n_objects) tuple — n_objects is None unless sequential_labels=True.

Examples:

From a dask array of per-tile labels:

>>> import dask.array as da
>>> from patchworks import merge_tile_labels
>>>
>>> # your own tiling + segmentation
>>> image = da.from_zarr("image.zarr").rechunk((1, 1024, 1024))
>>> labeled = image.map_blocks(my_segment_fn, dtype="int32",
...                            meta=np.empty((0,) * image.ndim, dtype="int32"))
>>>
>>> merged = merge_tile_labels(labeled, write_to="labels.zarr", progress=True)

From a pre-staged zarr store (your pipeline already wrote labels):

>>> merged = merge_tile_labels(
...     "my_staged_labels.zarr",
...     input_component="raw_labels",
...     write_to="merged_labels.zarr",
...     sequential_labels=True,
... )

Trim overlap halos before merging:

>>> # if labeled was computed with da.overlap.overlap(depth=20)
>>> merged = merge_tile_labels(labeled, write_to="labels.zarr", overlap=20)

Source code in src/patchworks/_merge.py

def merge_tile_labels(
    labeled: Union["da.Array", str, Path],
    write_to: Union[str, Path, None] = None,
    *,
    input_component: str = "labels",
    output_component: str = "labels",
    overlap: int = 0,
    sequential_labels: bool = False,
    n_workers: int | None = None,
    stage_dir: Union[str, Path, None] = None,
    keep_stage: bool = False,
    progress: bool = False,
    return_count: bool = False,
    label_counts: "Mapping[int, int] | Sequence[int] | None" = None,
    output_chunks: "Sequence[int] | None" = None,
    shard_output: bool = False,
    halo_dir: "str | Path | None" = None,
    iou_threshold: float = 0.5,
    connectivity: int = 1,
    parts_dir: "str | Path | None" = None,
    min_voxels: "int | None" = None,
    max_voxels: "int | None" = None,
    objects_out: "dict | None" = None,
) -> 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
    (from your own segmentation pipeline) and just need the boundary stitching.

    Accepts either:

    - A **dask array** of per-tile integer labels (e.g. output of
      ``dask.array.map_blocks`` on your own segmentation function).
    - A **zarr store path** whose ``input_component`` array already contains
      per-tile labels written by your own pipeline.

    Labels that **touch** across tile boundaries are merged into a single ID.
    The merge is zarr-native (boundary scan → scipy connected components →
    parallel relabel) — no dask task graph, scales to thousands of tiles.

    Parameters
    ----------
    labeled:
        Per-tile label array. Either a dask array or a path to a zarr store
        that contains per-tile labels in ``input_component``. A store is
        **modified**: unless ``label_counts`` is given, its per-tile ids are
        renumbered to a global range in place before the merge (the
        segmentation is unchanged, the ids are not). Copy it first if the
        original ids matter.
    write_to:
        Output zarr store path. When None, an auto-temp store is used.
    input_component:
        Array name inside a zarr *input* store (ignored for dask arrays).
    output_component:
        Array name inside ``write_to``. Default ``"labels"``.
    overlap:
        If ``labeled`` is a dask array that was computed with ``da.overlap``,
        pass the same depth here to trim the halos before merging.
        Set 0 (default) if the array has no overlap halos.
    sequential_labels:
        Renumber the merged labels to a contiguous ``1..N`` range. Folded into
        the merge's own lookup table, so it costs no extra pass over the
        volume. Default False.
    label_counts:
        Labels written per tile, when the producer recorded them (see
        :func:`patchworks.stage_tile`, which returns the count). Lets the
        merge derive every tile's global id range arithmetically instead of
        streaming the whole store to renumber it. ``None`` keeps the
        renumber pass.
    output_chunks:
        Chunking for the merged store; must divide the staged chunk shape.
        Lets the merge write straight into a store you will keep (e.g. an
        OME-ZARR label group's level 0) instead of a scratch store that then
        has to be copied. See :func:`capped_output_chunks`.
    shard_output:
        Store the output in shards of one staged tile each (zarr v3; the
        chunks above are inside them). Each relabel worker writes one tile,
        so it writes whole shards, and the result is sharded with no second
        pass over it. Ignored when relabeling in place.
    halo_dir:
        IoU stitching: the halo strips :func:`patchworks.stage_tile` saved
        (its ``halo_dir=``). Labels are then joined across a tile boundary
        only when both tiles' predictions of the overlap zone agree
        (IoU >= ``iou_threshold``), so two cells pressed together at a seam
        stay two. ``None`` (default) joins whatever touches.
    iou_threshold:
        Minimum IoU for that join (default 0.5).
    connectivity:
        Neighbourhood the tiles were labelled with: 1 (default) faces, 2
        also edges, 3 also corners (``scipy.ndimage.generate_binary_structure``).
        Labels touching across a tile boundary that way are joined, so the
        result is that of labelling the whole image at once.
    parts_dir, min_voxels, max_voxels, objects_out:
        Object sums saved per tile by :func:`stage_tile` (``parts_dir=``),
        combined into one row per merged object in *objects_out* -- the
        object table without a second read of the labels -- and an optional
        size filter applied in the same relabel pass. See
        :func:`zarr_native_merge`. Needs *label_counts*.
    n_workers:
        Parallel workers for the relabel step. Default ``min(4, cpu_count)``.
    stage_dir:
        Directory for the temp stage zarr when *labeled* is a dask array.
        Default: a system temp directory.
    keep_stage:
        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 backed by ``write_to``: the staged dtype (int32
        or wider for a dask input), widened if the global object count
        outgrows it. Or, when
        ``return_count=True``, a ``(labels, n_objects)`` tuple —
        ``n_objects`` is ``None`` unless ``sequential_labels=True``.

    Examples
    --------
    **From a dask array of per-tile labels:**

    >>> import dask.array as da  # doctest: +SKIP
    >>> from patchworks import merge_tile_labels  # doctest: +SKIP
    >>>
    >>> # your own tiling + segmentation
    >>> image = da.from_zarr("image.zarr").rechunk((1, 1024, 1024))  # doctest: +SKIP
    >>> labeled = image.map_blocks(my_segment_fn, dtype="int32",  # doctest: +SKIP
    ...                            meta=np.empty((0,) * image.ndim, dtype="int32"))
    >>>
    >>> merged = merge_tile_labels(labeled, write_to="labels.zarr", progress=True)  # doctest: +SKIP

    **From a pre-staged zarr store (your pipeline already wrote labels):**

    >>> merged = merge_tile_labels(  # doctest: +SKIP
    ...     "my_staged_labels.zarr",
    ...     input_component="raw_labels",
    ...     write_to="merged_labels.zarr",
    ...     sequential_labels=True,
    ... )

    **Trim overlap halos before merging:**

    >>> # if labeled was computed with da.overlap.overlap(depth=20)
    >>> merged = merge_tile_labels(labeled, write_to="labels.zarr", overlap=20)  # doctest: +SKIP
    """
    import dask.array as da

    nw = n_workers if n_workers is not None else min(4, cpu_allocation())

    # -- Stage dask array to zarr if needed --
    stage_cleanup: str | None = None  # set only when we create the stage
    staged_component = "staged"

    if isinstance(labeled, (str, Path)):
        stage_path = str(labeled)
        staged_component = input_component
    else:
        # labeled is a dask array
        if overlap > 0:
            labeled = da.overlap.trim_overlap(
                labeled, depth=overlap, boundary="none"
            )
        # The staged ids are renumbered to a global range in place, so a
        # narrow per-tile dtype (Cellpose's uint16) would overflow.
        if labeled.dtype.itemsize < 4:
            labeled = labeled.astype(np.int32)

        stage_path, stage_cleanup = _scratch_store(stage_dir, "stage")

        import dask

        from ._progress import dask_progress

    try:
        if stage_cleanup is not None:
            ctx = dask_progress("stage tiles", progress)
            logger.info("Staging per-tile labels to %s …", stage_path)
            with ctx:
                dask.compute(
                    labeled.to_zarr(
                        stage_path,
                        component=staged_component,
                        overwrite=True,
                        compute=False,
                    )
                )

        # -- Resolve output path --
        if write_to is not None:
            effective_out = str(write_to)
        else:
            effective_out, _ = _scratch_store(None, "merge")
            logger.info(
                "write_to not set — merged labels in auto-temp %s",
                effective_out,
            )

        # -- Merge (the sequential renumber rides along inside the LUT) --
        n_objects = zarr_native_merge(
            stage_path,
            staged_component,
            effective_out,
            output_component,
            n_workers=nw,
            show_progress=progress,
            label_counts=label_counts,
            sequential=sequential_labels,
            output_chunks=output_chunks,
            shard_output=shard_output,
            halo_dir=halo_dir,
            iou_threshold=iou_threshold,
            connectivity=connectivity,
            parts_dir=parts_dir,
            min_voxels=min_voxels,
            max_voxels=max_voxels,
            objects_out=objects_out,
        )
    finally:
        # -- Cleanup temp stage (only when we created it), even on failure:
        # a half-written stage is no use to a retry, which restages anyway.
        if stage_cleanup is not None:
            if keep_stage:
                logger.info("Keeping stage store %s", stage_path)
            else:
                _remove_scratch(stage_cleanup)

    result = da.from_zarr(effective_out, component=output_component)
    return (result, n_objects) if return_count else result