bioimage_py.segmentation.relabel

  1"""Block-wise relabeling of a segmentation.
  2
  3Two operations live here:
  4
  5- :func:`relabel` applies an explicit, caller-provided relabeling (a *labeling*) to a segmentation:
  6  a mapping from each old segment id to a new id (e.g. the node-to-label result of a graph-based
  7  segmentation such as multicut, or a merge/split assignment). It is a single-stage, disjoint
  8  per-block point op with no halo, so it accepts ``block_ids`` / ``resume_from`` and may be applied
  9  **in place** (the default when ``output`` is omitted). The labeling is applied per block with
 10  ``bioimage_cpp.utils.take_dict`` (when given as a ``dict``) or with ``numpy.take`` (when given as a
 11  dense 1D array/source where ``labeling[old_id]`` is the new id); a dense array passed as an
 12  in-memory numpy array for a distributed job is persisted to a temporary zarr under the runner's
 13  temp root so worker tasks can reopen it.
 14
 15- :func:`relabel_consecutive` relabels a segmentation so its ids are consecutive. It *derives* the
 16  mapping from the data (a global ``unique`` reduction), then delegates the block-wise write to
 17  :func:`relabel`. Being multi-stage, it re-runs whole (it does not accept ``block_ids`` /
 18  ``resume_from``); like :func:`relabel` it relabels in place when ``output`` is omitted.
 19"""
 20from __future__ import annotations
 21
 22import os
 23import shutil
 24import tempfile
 25from typing import Callable, Dict, Mapping, Optional, Sequence, Tuple, Union
 26
 27import bioimage_cpp as bic
 28import numpy as np
 29
 30from ..runner import get_runner
 31from ..runner.config import RunnerConfig
 32from ..sources import Source, SourceLike, SourceSpec, as_source, from_spec
 33from ..stats.unique import unique
 34from ..util import (BlockDescriptor, ComputeFn, check_rerun_args, full_roi, is_direct, take_mapping,
 35                    to_roi)
 36
 37__all__ = ["relabel", "relabel_consecutive"]
 38
 39
 40def _require_integer(source: Source, message: str) -> None:
 41    """Raise ``ValueError`` unless ``source`` has an integer dtype."""
 42    if not np.issubdtype(np.dtype(source.dtype), np.integer):
 43        raise ValueError(f"{message}, got dtype {source.dtype}.")
 44
 45
 46def _is_inmemory_numpy(source: Source) -> bool:
 47    """Return whether ``source`` wraps a plain in-memory numpy array (local-only, not reopenable)."""
 48    return isinstance(getattr(source, "array", None), np.ndarray)
 49
 50
 51def _make_relabel_block(mapping: Dict[int, int]) -> ComputeFn:
 52    """Build the per-block write function applying the label mapping (picklable dict)."""
 53
 54    def _compute(block: BlockDescriptor, inputs: Sequence[Source], outputs: Sequence[Source],
 55                 mask: Optional[Source]) -> None:
 56        input_, output_ = inputs[0], outputs[0]
 57        roi = to_roi(block)
 58        seg = input_[roi]
 59        if mask is None:
 60            output_[roi] = take_mapping(mapping, seg)
 61            return None
 62        m = mask[roi].astype(bool)
 63        if not m.any():
 64            return None
 65        # Only in-mask voxels are in the mapping; out-of-mask output voxels are left unchanged.
 66        out_block = output_[roi].copy()
 67        out_block[m] = take_mapping(mapping, seg[m])
 68        output_[roi] = out_block
 69        return None
 70
 71    return _compute
 72
 73
 74def _resume_placeholder(block: BlockDescriptor, inputs: Sequence[Source],
 75                        outputs: Sequence[Source], mask: Optional[Source]) -> None:
 76    """No-op compute fn for the resume path: ``run(resume_from=...)`` never calls it."""
 77    return None
 78
 79
 80def _make_cleanup(tmp_dir: str) -> Callable[[str], None]:
 81    """Build a success-path callback that removes the persisted labeling temp dir."""
 82
 83    def _cleanup(_tmp: str) -> None:
 84        shutil.rmtree(tmp_dir, ignore_errors=True)
 85
 86    return _cleanup
 87
 88
 89def _persist_labeling(labeling: np.ndarray, tmp_root: Optional[str]) -> Tuple[SourceSpec, str]:
 90    """Persist a 1D labeling array to a temp zarr so distributed workers can reopen it.
 91
 92    Args:
 93        labeling: The dense 1D labeling array to persist.
 94        tmp_root: The runner temp root to create the temp dir under (shared filesystem for slurm);
 95            ``None`` uses the system default.
 96
 97    Returns:
 98        A ``(spec, tmp_dir)`` tuple: the reopen spec of the persisted array and the temp directory
 99        (removed on the run's success path, preserved on failure so ``resume_from`` can reopen it).
100    """
101    import zarr
102
103    tmp_dir = tempfile.mkdtemp(prefix="bioimage_py_labeling_", dir=tmp_root)
104    path = os.path.join(tmp_dir, "labeling.zarr")
105    array = zarr.open_array(path, mode="w", shape=labeling.shape, dtype=labeling.dtype,
106                            chunks=labeling.shape)
107    array[:] = labeling
108    return as_source(array).to_spec(), tmp_dir
109
110
111def _make_take_array_block(labeling: Optional[np.ndarray],
112                           spec: Optional[SourceSpec]) -> ComputeFn:
113    """Build the per-block write function mapping ids through a dense 1D labeling array.
114
115    Exactly one of ``labeling`` (captured directly, for local execution) and ``spec`` (reopened on
116    the worker, for distributed execution) is set; the labeling is read once and cached per worker.
117    """
118    cache: Dict[str, np.ndarray] = {}
119
120    def _labeling() -> np.ndarray:
121        labels = cache.get("labeling")
122        if labels is None:
123            labels = labeling if labeling is not None else np.asarray(from_spec(spec)[full_roi(1)])
124            cache["labeling"] = labels
125        return labels
126
127    def _compute(block: BlockDescriptor, inputs: Sequence[Source], outputs: Sequence[Source],
128                 mask: Optional[Source]) -> None:
129        input_, output_ = inputs[0], outputs[0]
130        roi = to_roi(block)
131        seg = input_[roi]
132        labels = _labeling()
133        if mask is None:
134            output_[roi] = np.take(labels, seg)
135            return None
136        m = mask[roi].astype(bool)
137        if not m.any():
138            return None
139        # Only in-mask voxels are remapped; out-of-mask output voxels are left unchanged.
140        out_block = output_[roi].copy()
141        out_block[m] = np.take(labels, seg[m])
142        output_[roi] = out_block
143        return None
144
145    return _compute
146
147
148def relabel(
149    input: SourceLike,
150    labeling: Union[SourceLike, Mapping[int, int]],
151    output: Optional[SourceLike] = None,
152    *,
153    block_shape: Optional[Tuple[int, ...]] = None,
154    job_type: str = "local",
155    job_config: Optional[RunnerConfig] = None,
156    num_workers: int = 1,
157    mask: Optional[SourceLike] = None,
158    block_ids: Optional[Sequence[int]] = None,
159    resume_from: Optional[str] = None,
160) -> SourceLike:
161    """Apply a labeling (relabeling map) to a segmentation, block-wise.
162
163    Each block of ``input`` is read, its ids are mapped through ``labeling``, and the result is
164    written to ``output`` (or back to ``input`` when ``output`` is omitted). This is a single-stage,
165    disjoint per-block point operation, so it may be applied in place and supports ``block_ids`` /
166    ``resume_from`` re-runs.
167
168    Args:
169        input: The input segmentation (a numpy/zarr/n5 array or a `Source`); must be integer-typed.
170        labeling: The relabeling to apply. Either a ``dict`` ``{old_id: new_id}`` (applied with
171            ``bioimage_cpp.utils.take_dict``; every id present in ``input`` must be a key) or a dense
172            1D array/source where ``labeling[old_id]`` is the new id (applied with ``numpy.take``; it
173            must be long enough to index every id present in ``input``). A dict must map every id of
174            the (masked) input; a dense array must cover the id range ``0 .. max_id``.
175        output: The output array to write the relabeled segmentation into. Optional -- when omitted
176            the relabeling is applied **in place** to ``input`` (which must then be writable, and
177            file-backed for distributed execution). As an exception, a plain in-memory numpy input
178            (which is local-only) is never mutated: a fresh array is allocated and returned. When
179            given, ``output`` must match the input shape; ids are cast to its dtype on write.
180        block_shape: Shape of the processing blocks. Defaults to the input chunk shape; required for
181            unchunked data.
182        job_type: Execution backend: one of ``"local"``, ``"subprocess"`` or ``"slurm"``.
183        job_config: Backend configuration (a `RunnerConfig` / `SlurmConfig`).
184        num_workers: Number of parallel workers (threads for ``local``, tasks for distributed
185            backends).
186        mask: Optional binary mask; only voxels within the mask are relabeled (out-of-mask output
187            voxels are left unchanged).
188        block_ids: Restrict processing to these block ids (e.g. to re-run previously failed blocks
189            into the existing ``output``). Mutually exclusive with ``resume_from``.
190        resume_from: Distributed only; the preserved temp folder of a failed run to resume (see
191            ``runner.run``); the missing blocks are relabeled using the original run's labeling.
192            Mutually exclusive with ``block_ids``. (A numpy labeling array persisted by the failed
193            run is preserved with that temp folder; after a successful resume it is best-effort left
194            behind since the labeling is small.)
195
196    Returns:
197        The output array: the provided ``output`` if given, else ``input`` itself when relabeling a
198        file-backed source in place, or a freshly allocated array for an in-memory numpy input.
199    """
200    check_rerun_args(job_type, resume_from, block_ids)
201    src = as_source(input)
202    _require_integer(src, "relabel expects an integer label image")
203    ndim = src.ndim
204
205    # Resolve the output. By default relabel in place (a disjoint per-block point op with no halo,
206    # so this is safe). Exception: a plain in-memory numpy input is local-only (distributed rejects
207    # it), and silently mutating a passed-in array is surprising -- so allocate a fresh copy instead.
208    if output is not None:
209        out_array: SourceLike = output
210    elif job_type == "local" and _is_inmemory_numpy(src):
211        out_array = np.array(src.array)
212    else:
213        out_array = input
214    out = as_source(out_array)
215
216    runner = get_runner(job_type, job_config)
217
218    # Resume short-circuits inside run() to the preserved run's payload (its own sources and
219    # labeling), so the function/inputs passed here are placeholders that run() does not use.
220    if resume_from is not None:
221        runner.run(_resume_placeholder, [src], outputs=[out], resume_from=resume_from,
222                   name="relabel")
223        return out_array
224
225    direct = is_direct(job_type, num_workers, block_shape) and mask is None and block_ids is None
226
227    # Dict mode: apply the mapping with take_dict.
228    if isinstance(labeling, Mapping):
229        mapping = dict(labeling)
230        if direct:
231            out[full_roi(ndim)] = bic.utils.take_dict(mapping, src[full_roi(ndim)])
232            return out_array
233        runner.run(_make_relabel_block(mapping), [src], outputs=[out], block_shape=block_shape,
234                   mask=mask, num_workers=num_workers, block_ids=block_ids, name="relabel")
235        return out_array
236
237    # Dense 1D labeling array/source: applied with numpy.take.
238    labeling_src = as_source(labeling)
239    if labeling_src.ndim != 1:
240        raise ValueError(f"Dense labeling must be a 1D array; got shape {labeling_src.shape}.")
241    _require_integer(labeling_src, "labeling must be integer-typed")
242
243    if direct:
244        labels = np.asarray(labeling_src[full_roi(1)])
245        out[full_roi(ndim)] = np.take(labels, src[full_roi(ndim)])
246        return out_array
247
248    # Carry the labeling in the (cloudpickled) closure: the array directly for local execution, or
249    # a reopen spec for distributed workers -- persisting an in-memory numpy array to a temp zarr.
250    labeling_array: Optional[np.ndarray] = None
251    spec: Optional[SourceSpec] = None
252    pre_cleanup: Optional[Callable[[str], None]] = None
253    if job_type == "local":
254        labeling_array = np.asarray(labeling_src[full_roi(1)])
255    else:
256        try:
257            spec = labeling_src.to_spec()
258        except ValueError:  # an in-memory numpy array: persist it so worker tasks can reopen it.
259            spec, tmp_dir = _persist_labeling(np.asarray(labeling_src[full_roi(1)]),
260                                              runner.config.tmp_root)
261            pre_cleanup = _make_cleanup(tmp_dir)
262
263    runner.run(_make_take_array_block(labeling_array, spec), [src], outputs=[out],
264               block_shape=block_shape, mask=mask, num_workers=num_workers, block_ids=block_ids,
265               pre_cleanup=pre_cleanup, name="relabel")
266    return out_array
267
268
269def relabel_consecutive(
270    input: SourceLike,
271    output: Optional[SourceLike] = None,
272    *,
273    start_label: int = 0,
274    keep_zeros: bool = True,
275    block_shape: Optional[Tuple[int, ...]] = None,
276    job_type: str = "local",
277    job_config: Optional[RunnerConfig] = None,
278    num_workers: int = 1,
279    mask: Optional[SourceLike] = None,
280) -> Tuple[SourceLike, int, Dict[int, int]]:
281    """Relabel a segmentation to consecutive ids, block-wise.
282
283    This is multi-stage: a global ``unique`` reduction derives the ``{old_id: new_id}`` mapping, then
284    the block-wise write is delegated to :func:`relabel`. Because of the reduction it does **not**
285    accept ``block_ids`` or ``resume_from``: a failed run is re-run whole (it is idempotent given the
286    same ``output``).
287
288    Args:
289        input: The input label image (a numpy/zarr/n5 array or a `Source`); must be integer-typed.
290        output: The output array to write the relabeled segmentation into. Optional -- when omitted
291            the relabeling is applied **in place** to ``input`` (which must then be writable, and
292            file-backed for distributed execution); a plain in-memory numpy input (local-only) is
293            never mutated -- a fresh array is allocated and returned. When given, it must match the
294            input shape.
295        start_label: The value the smallest unique id is mapped to (subsequent ids follow
296            consecutively).
297        keep_zeros: Whether to always keep ``0`` mapped to ``0`` (background), regardless of
298            ``start_label``.
299        block_shape: Shape of the processing blocks. Defaults to the input chunk shape; required
300            for unchunked data.
301        job_type: Execution backend: one of ``"local"``, ``"subprocess"`` or ``"slurm"``.
302        job_config: Backend configuration (a `RunnerConfig` / `SlurmConfig`).
303        num_workers: Number of parallel workers (threads for ``local``, tasks for distributed
304            backends).
305        mask: Optional binary mask; values outside the mask are excluded from the computation and
306            their output voxels are left unchanged.
307
308    Returns:
309        A ``(output, max_id, mapping)`` tuple: the relabeled output array (``input`` itself when
310        relabeling a file-backed source in place, or a freshly allocated array for a numpy input),
311        the maximum label id after relabeling, and the ``{old_id: new_id}`` mapping that was applied.
312
313    """
314    src = as_source(input)
315    _require_integer(src, "relabel_consecutive expects an integer label image")
316    ndim = src.ndim
317
318    # Pass 1: the global set of unique values.
319    direct = is_direct(job_type, num_workers, block_shape) and mask is None
320    if direct:
321        uniques = np.unique(src[full_roi(ndim)])
322    else:
323        uniques = unique(input, block_shape=block_shape, job_type=job_type, job_config=job_config,
324                         num_workers=num_workers, mask=mask)
325
326    # In-process: build the old -> new mapping (consecutive ids from start_label).
327    mapping: Dict[int, int] = {int(v): i for i, v in enumerate(uniques.tolist(), start_label)}
328    if keep_zeros and 0 in mapping:
329        mapping[0] = 0
330    max_id = max(mapping.values()) if mapping else 0
331
332    # Pass 2: apply the mapping (in place when output is omitted), reusing relabel's dict path.
333    out = relabel(input, mapping, output, block_shape=block_shape, job_type=job_type,
334                  job_config=job_config, num_workers=num_workers, mask=mask)
335    return out, max_id, mapping