bioimage_py.segmentation.label

  1"""Block-wise connected-component labeling (multi-stage: label -> merge -> relabel).
  2
  3The block-wise path runs three :meth:`~bioimage_py.runner.base.Runner.run` calls plus one
  4in-process reduction, with the labeled volume persisted in the ``output`` source between
  5stages. This mirrors cluster_tools' connected-components workflow.
  6
  7Inter-stage data visibility (distributed/slurm): stage N+1 is only launched after every stage-N
  8task's ``.success`` sentinel is visible, i.e. after all stage-N workers have exited and flushed
  9their writes. Stage N+1 then reads the output's chunks by direct path (zarr/n5 address a chunk by
 10key, with no directory listing), so NFS close-to-open consistency guarantees it sees the fresh
 11data -- the attribute-cache lag that affects sentinel *discovery* does not apply to reads of an
 12already-named chunk file.
 13"""
 14from __future__ import annotations
 15
 16from typing import Dict, Optional, Sequence, Tuple
 17
 18import bioimage_cpp as bic
 19import numpy as np
 20
 21from ..runner import get_runner
 22from ..runner.config import RunnerConfig
 23from ..sources import Source, SourceLike, as_source
 24from ..util import BlockDescriptor, ComputeFn, full_roi, get_blocking, is_direct, to_roi
 25from .relabel import relabel
 26
 27__all__ = ["label"]
 28
 29
 30def _binarize(data: np.ndarray, threshold: Optional[float]) -> np.ndarray:
 31    """Binarize ``data`` by threshold, or interpret it as a boolean foreground mask."""
 32    if threshold is not None:
 33        return data > threshold
 34    return data if data.dtype == np.dtype(bool) else data.astype(bool)
 35
 36
 37def _resolve_block_shape(src: Source, out: Source,
 38                         block_shape: Optional[Sequence[int]]) -> Tuple[int, ...]:
 39    """Resolve the block shape from the explicit value or the input/output chunks."""
 40    if block_shape is not None:
 41        return tuple(int(b) for b in block_shape)
 42    chunks = src.chunks or out.chunks
 43    if chunks is None:
 44        raise ValueError("block_shape is required for block-wise labeling of an unchunked array.")
 45    return tuple(int(c) for c in chunks)
 46
 47
 48# --- per-block stage functions (built as closures, capturing only picklable values) ----
 49
 50def _make_stage1(shape: Tuple[int, ...], block_shape: Tuple[int, ...], connectivity: int,
 51                 threshold: Optional[float], offset_factor: int) -> ComputeFn:
 52    """Build stage 1: label each block independently and apply a globally-unique offset."""
 53
 54    def _compute(block: BlockDescriptor, inputs: Sequence[Source], outputs: Sequence[Source],
 55                 mask: Optional[Source]) -> Optional[np.ndarray]:
 56        input_, output_ = inputs[0], outputs[0]
 57        roi = to_roi(block)
 58        blocking = get_blocking(shape, block_shape)
 59        block_id = blocking.coordinates_to_block_id([int(c) for c in block.begin])
 60
 61        binary = _binarize(input_[roi], threshold)
 62        if mask is not None:
 63            binary = binary & mask[roi].astype(bool)
 64
 65        if not binary.any():
 66            output_[roi] = np.zeros(binary.shape, dtype="uint64")
 67            return None
 68
 69        comp = bic.segmentation.label(binary, connectivity=connectivity).astype("uint64", copy=False)
 70        offset = np.uint64(int(block_id) * int(offset_factor))
 71        comp[comp != 0] += offset
 72        output_[roi] = comp
 73        # Return the block's actual (globally-unique) labels so stage 3 can relabel only
 74        # over labels that exist, not over the sparse offset space.
 75        return np.unique(comp[comp != 0])
 76
 77    return _compute
 78
 79
 80def _make_stage2(shape: Tuple[int, ...], block_shape: Tuple[int, ...]) -> ComputeFn:
 81    """Build stage 2: collect label equivalences across lower block faces."""
 82
 83    def _compute(block: BlockDescriptor, inputs: Sequence[Source], outputs: Sequence[Source],
 84                 mask: Optional[Source]) -> Optional[np.ndarray]:
 85        output_ = inputs[0]  # the labeled volume, passed as a read-only input
 86        ndim = len(shape)
 87        blocking = get_blocking(shape, block_shape)
 88        block_id = blocking.coordinates_to_block_id([int(c) for c in block.begin])
 89
 90        pairs = []
 91        block_roi = to_roi(block)
 92        for axis in range(ndim):
 93            if blocking.get_neighbor_id(block_id, axis, True) == -1:  # no lower neighbor
 94                continue
 95            b0 = int(block.begin[axis])
 96            slab_roi = list(block_roi)
 97            slab_roi[axis] = slice(b0 - 1, b0 + 1)  # 2-thick slab straddling the boundary
 98            slab = output_[tuple(slab_roi)]
 99            lo = tuple(slice(0, 1) if d == axis else slice(None) for d in range(ndim))
100            hi = tuple(slice(1, 2) if d == axis else slice(None) for d in range(ndim))
101            labels_b = np.squeeze(slab[lo], axis=axis)  # neighbor side
102            labels_a = np.squeeze(slab[hi], axis=axis)  # this block side
103            keep = (labels_a != 0) & (labels_b != 0)
104            if keep.any():
105                pairs.append(np.stack([labels_a[keep], labels_b[keep]], axis=1).astype("uint64"))
106
107        if not pairs:
108            return None
109        return np.unique(np.concatenate(pairs, axis=0), axis=0)
110
111    return _compute
112
113
114def label(
115    input: SourceLike,
116    output: Optional[SourceLike] = None,
117    *,
118    threshold: Optional[float] = None,
119    connectivity: Optional[int] = None,
120    block_shape: Optional[Tuple[int, ...]] = None,
121    job_type: str = "local",
122    job_config: Optional[RunnerConfig] = None,
123    num_workers: int = 1,
124    mask: Optional[SourceLike] = None,
125) -> SourceLike:
126    """Label connected components of (optionally thresholded) data, block-wise.
127
128    Unlike the single-pass operations, ``label`` is multi-stage with a global cross-block merge
129    (per-block labeling, then a union-find over touching components across block faces), so it
130    does **not** accept ``block_ids`` or ``resume_from``: a failed run must be re-run whole (it is
131    idempotent given the same ``output``).
132
133    Args:
134        input: The input data (a numpy/zarr/n5 array or a `Source`).
135        output: The ``uint64`` output array to write the labels into. Optional for local
136            execution — a numpy array is allocated and returned if omitted; **required** for
137            distributed execution.
138        threshold: If given, the input is binarized as ``input > threshold``; otherwise the
139            input is treated as a binary foreground mask.
140        connectivity: Neighbour connectivity in ``[1, ndim]`` (``1`` = orthogonal). Defaults
141            to ``1``; values ``> 1`` are only supported for the direct (single-block) path.
142        block_shape: Shape of the processing blocks. Defaults to the input/output chunk shape;
143            required for unchunked data.
144        job_type: Execution backend: one of ``"local"``, ``"subprocess"`` or ``"slurm"``.
145        job_config: Backend configuration (a `RunnerConfig` / `SlurmConfig`).
146        num_workers: Number of parallel workers (threads for ``local``, tasks for distributed
147            backends).
148        mask: Optional binary mask; values outside the mask are excluded from the foreground.
149
150    Returns:
151        The output array (the provided ``output``, or a newly allocated numpy array), labeled
152        with consecutive ids (background stays ``0``).
153    """
154    src = as_source(input)
155    ndim = src.ndim
156    conn = 1 if connectivity is None else int(connectivity)
157    if not 1 <= conn <= ndim:
158        raise ValueError(f"connectivity must be in [1, {ndim}], got {conn}.")
159
160    direct = is_direct(job_type, num_workers, block_shape) and mask is None
161    if conn > 1 and not direct:
162        raise NotImplementedError(
163            "Block-wise labeling only supports connectivity=1 (orthogonal). Use the direct "
164            "path (local, single worker, no block_shape, no mask) for higher connectivity."
165        )
166
167    if output is None:
168        if job_type != "local":
169            raise ValueError(
170                f"'output' is required for distributed execution (job_type={job_type!r}); "
171                "pass a file-backed (zarr/n5) output array."
172            )
173        out_array: SourceLike = np.zeros(tuple(src.shape), dtype="uint64")
174    else:
175        out_array = output
176
177    out = as_source(out_array)
178    if out.dtype != np.dtype("uint64"):
179        raise ValueError(f"output must have dtype uint64, got {out.dtype}.")
180
181    if direct:
182        binary = _binarize(src[full_roi(ndim)], threshold)
183        comp = bic.segmentation.label(binary, connectivity=conn).astype("uint64", copy=False)
184        out[full_roi(ndim)] = comp
185        return out_array
186
187    block_shape = _resolve_block_shape(src, out, block_shape)
188    offset_factor = int(np.prod(block_shape))
189    blocking = get_blocking(src.shape, block_shape)
190    n_blocks = int(blocking.number_of_blocks)
191    if (n_blocks * offset_factor) >= int(np.iinfo(np.uint64).max):
192        raise ValueError(
193            "Label id overflow: number_of_blocks * prod(block_shape) exceeds uint64. "
194            "Reduce the block shape or the volume size."
195        )
196
197    runner = get_runner(job_type, job_config)
198
199    # Stage 1: label each block independently with a globally-unique offset.
200    stage1 = _make_stage1(tuple(src.shape), block_shape, conn, threshold, offset_factor)
201    id_results = runner.run(stage1, [input], outputs=[out_array], block_shape=block_shape,
202                            mask=mask, num_workers=num_workers, has_return_val=True,
203                            name="label-blocks")
204    id_arrays = [a for a in id_results if a is not None and len(a)]
205    real_labels = np.unique(np.concatenate(id_arrays)) if id_arrays else np.zeros((0,), dtype="uint64")
206
207    # Stage 2: collect label equivalences across lower block faces.
208    stage2 = _make_stage2(tuple(src.shape), block_shape)
209    pair_results = runner.run(stage2, [out_array], block_shape=block_shape,
210                              num_workers=num_workers, has_return_val=True, name="merge-faces")
211    pairs = [p for p in pair_results if p is not None]
212    assignments = (np.unique(np.concatenate(pairs, axis=0), axis=0)
213                   if pairs else np.zeros((0, 2), dtype="uint64"))
214
215    # Stage 3 (in process): union-find merge, then relabel the labels that exist to consecutive ids.
216    # The stage-1 offset space is sparse (ids run up to ~voxel count), so the union-find is built over
217    # a dense [0..K) compaction of the labels that actually exist -- sized to the component count K,
218    # not the max offset id -- keeping this in-process step O(components) rather than O(voxels). Every
219    # id in `assignments` exists in `real_labels` (stage 2 only reads what stage 1 wrote), so the
220    # compaction covers all pair ids.
221    mapping: Dict[int, int] = {0: 0}
222    if real_labels.size:
223        n_components = int(real_labels.size)
224        dense = {int(lab): idx for idx, lab in enumerate(real_labels.tolist())}
225        uf = bic.utils.UnionFind(n_components)
226        if len(assignments):
227            pu = bic.utils.take_dict(dense, np.ascontiguousarray(assignments[:, 0].astype("uint64")))
228            pv = bic.utils.take_dict(dense, np.ascontiguousarray(assignments[:, 1].astype("uint64")))
229            uf.merge(np.stack([pu, pv], axis=1).astype("uint64"))
230        roots = np.asarray(uf.find(np.arange(n_components, dtype="uint64")))
231        _, new_ids = np.unique(roots, return_inverse=True)  # consecutive component ranks (0-based)
232        for lab, new_id in zip(real_labels.tolist(), new_ids.tolist()):
233            mapping[int(lab)] = int(new_id) + 1  # reserve 0 for background
234
235    # Stage 4: apply the mapping in place through the canonical node-label writer (relabel).
236    out_array = relabel(out_array, mapping, output=out_array, block_shape=block_shape,
237                        job_type=job_type, job_config=job_config, num_workers=num_workers)
238    return out_array