tile_process

patchworks.tile_process(image: Union[da.Array, str, Path], fn: Callable[[np.ndarray], np.ndarray], *, tile_shape: Union[tuple[int, ...], Callable[[tuple, Any], tuple], str, None] = None, overlap: int = 16, channel: int | None = 0, level: int = 0, use_gpu: bool = False, max_workers: int | None = None, progress: bool = True, write_to: Union[str, Path, None] = None, output_component: str = 'labels', pyramid_levels: int = 5, pyramid_downscale: int = 2, sequential_labels: bool = False, skip_empty: bool = False, empty_threshold: float | None = None, stage_dir: Union[str, Path, None] = None, keep_stage: bool = False, log_file: Union[str, Path, bool, None] = None, verbose: bool = False, stitch: str = 'touch', iou_threshold: float = 0.5, resume: bool = False, gpus: 'int | Sequence[int | str] | None' = None, dry_run: bool = False, plan_sample: int = 0) -> Any

Apply fn to every tile of image and merge labels globally.

The core workhorse of patchworks. fn can be any callable that takes a NumPy array and returns an integer label array of the same shape — Cellpose, StarDist, Otsu threshold, your own model, anything.

Parameters:

Name	Type	Description	Default
image	Union[Array, str, Path]	Dask array or path to an OME-ZARR store.	required
fn	Callable[[ndarray], ndarray]	(ndarray) -> ndarray returning integer labels of the same shape. Must be picklable when using distributed schedulers.	required
tile_shape	Union[tuple[int, ...], Callable[[tuple, Any], tuple], str, None]	Controls tiling before calling fn. Accepted values: None : keep existing dask chunks. tuple : use this exact tile shape. "auto" : call auto_tile_shape based on shape and dtype. Callable[[shape, dtype], tuple] : called with the image's shape and dtype; the return value is used. Use this with auto_tile_shape_cellpose: .. code-block:: python from functools import partial from patchworks import auto_tile_shape_cellpose, tile_process tile_fn = partial(auto_tile_shape_cellpose, diameter=30, use_gpu=True) result = tile_process("image.zarr", fn, tile_shape=tile_fn)	None
overlap	int	Voxels of overlap (halo) added to each tile before fn is called, so objects near tile boundaries have enough spatial context to be segmented correctly (Cellpose, StarDist, …). The halo is trimmed off before merging — the output has the original shape. Defaults to 16; set it to roughly one object diameter (see auto_overlap) for best results, or 0 to disable. Merging is always touching-label based: after the halo is trimmed, labels that touch across a tile boundary are merged into one object.	16
channel	int | None	Channel index when image is a path. Ignored for arrays.	0
level	int	Pyramid level when image is a path (0 = full resolution).	0
use_gpu	bool	When tile_shape="auto", size tiles against GPU VRAM instead of RAM. Also forces staging to one tile at a time (no VRAM contention).	False
max_workers	int | None	Cap the worker threads/processes used for staging and merging. None (default) auto-sizes to the machine: bounded by available RAM (tile size) and CPU (leaves one core free) so a run can neither OOM nor pin every core. Ignored when a distributed client is active (it manages its own concurrency).	None
progress	bool	Show progress bars for staging, the label write and the pyramid (default True). Set False to silence them.	True
write_to	Union[str, Path, None]	Explicit output zarr store path. Overrides the default behaviour: the merged labels are written here as a single-resolution array named output_component (no pyramid). When None (default) and image is a .zarr store, labels are written back into that store under the NGFF labels/<output_component>/ group with an auto pyramid, so the image and its segmentation live in one file. When None and image is an array, an auto-temp store is used. Every array written uses the active :func:patchworks.compression codec (zstd level 1 unless changed); wrap the call in with compression("blosc"): to pick another.	None
output_component	str	Label name. The array inside write_to, or the NGFF label image name under labels/ when writing into the input store. Default "labels".	'labels'
pyramid_levels	int	Number of resolution levels for the in-store label pyramid (only when writing into the input .zarr). Default 5.	5
pyramid_downscale	int	Per-level X/Y downsampling factor for that pyramid (Z is kept at full resolution). Default 2.	2
sequential_labels	bool	Renumber merged labels to a contiguous 1..N range. Default False — labels stay globally unique but gappy (block-encoded), which is fine for counting/measurement. Uses a cheap linear post-pass (O(voxels)), not the O(n_chunks²) dask built-in.	False
skip_empty	bool	Skip fn on background tiles. A tile whose max signal is <= empty_threshold returns all-zeros immediately. Biggest speed-up for sparse/mostly-background volumes. Use estimate_empty_tiles() first to pick a threshold.	False
empty_threshold	float | None	Intensity at or below which a tile is empty (skip_empty=True only). None → auto-derive via Otsu on a bounded sample.	None
stage_dir	Union[str, Path, None]	Where to put the temporary stage store. fn is always run once per tile to this store, then the merge reads it back from disk (running fn again is never needed). Default → next to write_to, else next to the input store, else a system temp directory. The store gets a unique _pws_stage_<id>.zarr name, so concurrent runs sharing a directory never overwrite each other's tiles.	None
keep_stage	bool	Keep the temp stage store after merging (default: delete it, also when the run fails). Its path is logged. Useful for debugging.	False
log_file	Union[str, Path, bool, None]	Where to tee the patchworks INFO log (including a per-tile tile k done counter and ETA). None/False (default) writes no file -- configure :mod:logging as usual; True writes patchworks.log next to the output; a path writes there. Asking for a file also raises the patchworks logger to INFO.	None
verbose	bool	Log each tile's location and shape as it is processed.	False
stitch	str	How labels are joined across tile boundaries. "touch" (default) joins any two labels that touch there. "iou" joins them only when both tiles' predictions of the overlap zone agree (IoU >= iou_threshold), so two distinct cells pressed together at a seam stay two; it needs overlap > 0 to have a zone to compare (an axis without one falls back to the IoU of the two boundary slices).	'touch'
iou_threshold	float	Minimum IoU for stitch="iou" (default 0.5).	0.5
resume	bool	Make an interrupted run resumable: tiles are staged into a store named after this run's inputs (image, tiling, overlap, fn and its bound arguments, output), with a record of finished tiles. Rerunning the same call skips those tiles; the store is removed once the run succeeds, and kept when it fails. stitch="iou" and resume stage tile by tile in threads rather than through dask, so an active distributed client is not used for staging.	False
dry_run	bool	Only plan: return a dict describing the run -- tile count and grid, tiles with signal (skip_empty), workers, memory per tile, output size -- without segmenting or writing anything. Check a tiling before a job waits hours in a queue for it.	False
plan_sample	int	With dry_run, also time fn on this many real tiles (spread over the image) and extrapolate estimated_seconds.	0
gpus	'int | Sequence[int | str] | None'	Segment on several GPUs at once: an int (the first N devices, of CUDA_VISIBLE_DEVICES if set) or a list of device ids. Each GPU gets its own worker process, pinned to it before CUDA starts, and tiles are handed out as workers free up. Linux only. Stages tile by tile like stitch="iou"; combine freely with it and resume.	None

Returns:

Type	Description
Array	Globally relabeled array (int32) backed by the output zarr (the input store's labels/<name>/0 by default, write_to when given, else an auto-temp zarr). Never loads the full volume into RAM. Call .compute() yourself only if the result fits in RAM. With dry_run=True, the plan dict instead.

Examples:

Any threshold function:

>>> from skimage.filters import threshold_otsu
>>> from skimage.measure import label
>>>
>>> def my_fn(tile):
...     return label(tile > threshold_otsu(tile)).astype("int32")
>>>
>>> result = tile_process("image.zarr", my_fn, write_to="labels.zarr")

Cellpose (via the plugin):

>>> from patchworks.plugins.cellpose import cellpose_fn
>>>
>>> fn = cellpose_fn("cyto3", gpu=True, diameter=30)
>>> result = tile_process(
...     "image.zarr", fn,
...     tile_shape=(1, 2048, 2048),
...     overlap=20,
...     write_to="labels.zarr",
...     progress=True,
... )

StarDist:

>>> from stardist.models import StarDist2D
>>> model = StarDist2D.from_pretrained("2D_versatile_fluo")
>>>
>>> def stardist_fn(tile):
...     norm = tile.astype("float32") / tile.max()
...     labels, _ = model.predict_instances(norm)
...     return labels.astype("int32")
>>>
>>> result = tile_process("image.zarr", stardist_fn,
...                       tile_shape=(1, 1024, 1024), overlap=32)

Write directly to zarr (no RAM accumulation):

>>> tile_process("image.zarr", fn, write_to="labels.zarr", progress=True)

Source code in src/patchworks/_core.py

def tile_process(
    image: Union[da.Array, str, Path],
    fn: Callable[[np.ndarray], np.ndarray],
    *,
    tile_shape: Union[
        tuple[int, ...], Callable[[tuple, Any], tuple], str, None
    ] = None,
    overlap: int = 16,
    channel: int | None = 0,
    level: int = 0,
    use_gpu: bool = False,
    max_workers: int | None = None,
    progress: bool = True,
    write_to: Union[str, Path, None] = None,
    output_component: str = "labels",
    pyramid_levels: int = 5,
    pyramid_downscale: int = 2,
    sequential_labels: bool = False,
    skip_empty: bool = False,
    empty_threshold: float | None = None,
    stage_dir: Union[str, Path, None] = None,
    keep_stage: bool = False,
    log_file: Union[str, Path, bool, None] = None,
    verbose: bool = False,
    stitch: str = "touch",
    iou_threshold: float = 0.5,
    resume: bool = False,
    gpus: "int | Sequence[int | str] | None" = None,
    dry_run: bool = False,
    plan_sample: int = 0,
) -> Any:
    """Apply *fn* to every tile of *image* and merge labels globally.

    The core workhorse of patchworks. ``fn`` can be any callable that takes a
    NumPy array and returns an integer label array of the same shape — Cellpose,
    StarDist, Otsu threshold, your own model, anything.

    Parameters
    ----------
    image:
        Dask array *or* path to an OME-ZARR store.
    fn:
        ``(ndarray) -> ndarray`` returning integer labels of the same shape.
        Must be picklable when using distributed schedulers.
    tile_shape:
        Controls tiling before calling *fn*. Accepted values:

        - ``None`` : keep existing dask chunks.
        - ``tuple`` : use this exact tile shape.
        - ``"auto"`` : call ``auto_tile_shape`` based on shape and dtype.
        - ``Callable[[shape, dtype], tuple]`` : called with the image's shape
          and dtype; the return value is used. Use this with
          ``auto_tile_shape_cellpose``:

          .. code-block:: python

              from functools import partial
              from patchworks import auto_tile_shape_cellpose, tile_process
              tile_fn = partial(auto_tile_shape_cellpose, diameter=30, use_gpu=True)
              result = tile_process("image.zarr", fn, tile_shape=tile_fn)

    overlap:
        Voxels of overlap (halo) added to each tile before *fn* is called, so
        objects near tile boundaries have enough spatial context to be
        segmented correctly (Cellpose, StarDist, …). The halo is trimmed off
        before merging — the output has the original shape. Defaults to ``16``;
        set it to roughly one object diameter (see ``auto_overlap``) for best
        results, or ``0`` to disable.

        Merging is always **touching-label** based: after the halo is trimmed,
        labels that touch across a tile boundary are merged into one object.
    channel:
        Channel index when *image* is a path. Ignored for arrays.
    level:
        Pyramid level when *image* is a path (0 = full resolution).
    use_gpu:
        When ``tile_shape="auto"``, size tiles against GPU VRAM instead of RAM.
        Also forces staging to one tile at a time (no VRAM contention).
    max_workers:
        Cap the worker threads/processes used for staging and merging. ``None``
        (default) auto-sizes to the machine: bounded by available RAM (tile
        size) and CPU (leaves one core free) so a run can neither OOM nor pin
        every core. Ignored when a distributed client is active (it manages its
        own concurrency).
    progress:
        Show progress bars for staging, the label write and the pyramid
        (default ``True``). Set ``False`` to silence them.
    write_to:
        Explicit output zarr store path. Overrides the default behaviour: the
        merged labels are written here as a single-resolution array named
        ``output_component`` (no pyramid). When None (default) and *image* is a
        ``.zarr`` store, labels are written back into that store under the NGFF
        ``labels/<output_component>/`` group with an auto pyramid, so the image
        and its segmentation live in one file. When None and *image* is an
        array, an auto-temp store is used. Every array written uses the
        active :func:`patchworks.compression` codec (zstd level 1 unless
        changed); wrap the call in ``with compression("blosc"):`` to pick
        another.
    output_component:
        Label name. The array inside ``write_to``, or the NGFF label image name
        under ``labels/`` when writing into the input store. Default
        ``"labels"``.
    pyramid_levels:
        Number of resolution levels for the in-store label pyramid (only when
        writing into the input ``.zarr``). Default 5.
    pyramid_downscale:
        Per-level X/Y downsampling factor for that pyramid (Z is kept at full
        resolution). Default 2.
    sequential_labels:
        Renumber merged labels to a contiguous ``1..N`` range. Default False —
        labels stay globally unique but gappy (block-encoded), which is fine for
        counting/measurement. Uses a cheap linear post-pass (O(voxels)), not the
        O(n_chunks²) dask built-in.
    skip_empty:
        Skip *fn* on background tiles. A tile whose max signal is <=
        ``empty_threshold`` returns all-zeros immediately. Biggest speed-up for
        sparse/mostly-background volumes. Use ``estimate_empty_tiles()`` first
        to pick a threshold.
    empty_threshold:
        Intensity at or below which a tile is empty (``skip_empty=True`` only).
        None → auto-derive via Otsu on a bounded sample.
    stage_dir:
        Where to put the temporary stage store. ``fn`` is always run once per
        tile to this store, then the merge reads it back from disk (running
        ``fn`` again is never needed). Default → next to ``write_to``, else next
        to the input store, else a system temp directory. The store gets a
        unique ``_pws_stage_<id>.zarr`` name, so concurrent runs sharing a
        directory never overwrite each other's tiles.
    keep_stage:
        Keep the temp stage store after merging (default: delete it, also
        when the run fails). Its path is logged. Useful for debugging.
    log_file:
        Where to tee the ``patchworks`` INFO log (including a per-tile
        ``tile k done`` counter and ETA). ``None``/``False`` (default) writes
        no file -- configure :mod:`logging` as usual; ``True`` writes
        ``patchworks.log`` next to the output; a path writes there. Asking
        for a file also raises the ``patchworks`` logger to INFO.
    verbose:
        Log each tile's location and shape as it is processed.
    stitch:
        How labels are joined across tile boundaries. ``"touch"`` (default)
        joins any two labels that touch there. ``"iou"`` joins them only when
        both tiles' predictions of the overlap zone agree (IoU >=
        ``iou_threshold``), so two distinct cells pressed together at a seam
        stay two; it needs ``overlap > 0`` to have a zone to compare (an axis
        without one falls back to the IoU of the two boundary slices).
    iou_threshold:
        Minimum IoU for ``stitch="iou"`` (default 0.5).
    resume:
        Make an interrupted run resumable: tiles are staged into a store
        named after this run's inputs (image, tiling, overlap, ``fn`` and its
        bound arguments, output), with a record of finished tiles. Rerunning
        the same call skips those tiles; the store is removed once the run
        succeeds, and kept when it fails. ``stitch="iou"`` and ``resume``
        stage tile by tile in threads rather than through dask, so an
        active distributed client is not used for staging.
    dry_run:
        Only plan: return a dict describing the run -- tile count and grid,
        tiles with signal (``skip_empty``), workers, memory per tile, output
        size -- without segmenting or writing anything. Check a tiling
        before a job waits hours in a queue for it.
    plan_sample:
        With ``dry_run``, also time *fn* on this many real tiles (spread
        over the image) and extrapolate ``estimated_seconds``.
    gpus:
        Segment on several GPUs at once: an int (the first N devices, of
        ``CUDA_VISIBLE_DEVICES`` if set) or a list of device ids. Each GPU
        gets its own worker process, pinned to it before CUDA starts, and
        tiles are handed out as workers free up. Linux only. Stages tile by
        tile like ``stitch="iou"``; combine freely with it and ``resume``.

    Returns
    -------
    da.Array
        Globally relabeled array (int32) backed by the output zarr (the input
        store's ``labels/<name>/0`` by default, ``write_to`` when given, else an
        auto-temp zarr). Never loads the full volume into RAM. Call
        ``.compute()`` yourself only if the result fits in RAM. With
        ``dry_run=True``, the plan dict instead.

    Examples
    --------
    **Any threshold function:**

    >>> from skimage.filters import threshold_otsu  # doctest: +SKIP
    >>> from skimage.measure import label  # doctest: +SKIP
    >>>
    >>> def my_fn(tile):  # doctest: +SKIP
    ...     return label(tile > threshold_otsu(tile)).astype("int32")
    >>>
    >>> result = tile_process("image.zarr", my_fn, write_to="labels.zarr")  # doctest: +SKIP

    **Cellpose (via the plugin):**

    >>> from patchworks.plugins.cellpose import cellpose_fn  # doctest: +SKIP
    >>>
    >>> fn = cellpose_fn("cyto3", gpu=True, diameter=30)  # doctest: +SKIP
    >>> result = tile_process(  # doctest: +SKIP
    ...     "image.zarr", fn,
    ...     tile_shape=(1, 2048, 2048),
    ...     overlap=20,
    ...     write_to="labels.zarr",
    ...     progress=True,
    ... )

    **StarDist:**

    >>> from stardist.models import StarDist2D  # doctest: +SKIP
    >>> model = StarDist2D.from_pretrained("2D_versatile_fluo")  # doctest: +SKIP
    >>>
    >>> def stardist_fn(tile):  # doctest: +SKIP
    ...     norm = tile.astype("float32") / tile.max()
    ...     labels, _ = model.predict_instances(norm)
    ...     return labels.astype("int32")
    >>>
    >>> result = tile_process("image.zarr", stardist_fn,  # doctest: +SKIP
    ...                       tile_shape=(1, 1024, 1024), overlap=32)

    **Write directly to zarr (no RAM accumulation):**

    >>> tile_process("image.zarr", fn, write_to="labels.zarr", progress=True)  # doctest: +SKIP
    """
    # In-process dask workers break the label merge. A GIL-holding fn starves
    # the worker heartbeat and the P2P barrier drops inputs →
    # "FutureCancelledError: lost dependencies".
    _active = _distributed_client()
    if _active is not None and _client_is_in_process(_active):
        raise RuntimeError(
            "Active Dask client uses an in-process worker (processes=False). "
            "This breaks the label merge when fn holds the GIL. Use a "
            "process-based cluster instead:\n"
            "    from patchworks import make_local_cluster\n"
            "    client, cluster = make_local_cluster(use_gpu=True)\n"
            "or drop the client to use the threaded scheduler "
            "(client.close(); cluster.close())."
        )

    if stitch not in ("touch", "iou"):
        raise ValueError(f"stitch must be 'touch' or 'iou', got {stitch!r}")
    devices = _resolve_gpus(gpus)
    per_tile = stitch == "iou" or resume or bool(devices and len(devices) > 1)

    # Load + tile
    image_source_path = None if isinstance(image, da.Array) else str(image)

    # Auto log file (default): tee patchworks' INFO logs to a file next to the
    # output, so a long run leaves a tailable record without notebook setup.
    if log_file:
        if log_file is True:
            if write_to is not None and not is_remote(write_to):
                _ldir = os.path.dirname(os.path.abspath(str(write_to)))
            elif image_source_path is not None and not is_remote(
                image_source_path
            ):
                _ldir = os.path.dirname(os.path.abspath(image_source_path))
            else:
                _ldir = os.getcwd()
            log_file = os.path.join(_ldir, "patchworks.log")
        _attach_log_file(str(log_file))
        logger.info("patchworks log → %s", log_file)

    _load_chunks: tuple[int, ...] | None = None
    _native_chunks: tuple[int, ...] | None = None

    if not isinstance(image, da.Array):
        _peek = load_ome_zarr(image, channel=channel, level=level)
        _native_chunks = _peek.chunksize  # on-disk zarr chunk shape
        _load_chunks = _resolve_tile_shape(
            tile_shape, _peek.shape, _peek.dtype, use_gpu
        )
        tile_shape = None  # already handled at load time
        if _load_chunks is not None:
            logger.info("Loading zarr with target tiles %s", _load_chunks)
            image = load_ome_zarr(
                image, channel=channel, level=level, chunks=_load_chunks
            )
        else:
            image = _peek

    tile_shape = _resolve_tile_shape(
        tile_shape, image.shape, image.dtype, use_gpu
    )
    if tile_shape is not None:
        image = image.rechunk(tile_shape)
        logger.info("Rechunked to %s", tile_shape)

    n_tiles = int(np.prod([len(c) for c in image.chunks]))
    _tile = tuple(c[0] for c in image.chunks)
    logger.info(
        "Processing %d tiles (per-axis %s, tile shape %s)",
        n_tiles,
        tuple(len(c) for c in image.chunks),
        _tile,
    )

    # Warn when the store's on-disk chunks are much larger than the tile: every
    # tile read then decodes whole chunks and throws most away (slow I/O).
    if _native_chunks is not None:
        _amp = _read_amplification(_native_chunks, _tile)
        if _amp >= 4:
            logger.warning(
                "Input chunks %s are much larger than the tile %s → ~%.0fx "
                "read amplification (each tile decodes whole chunks and "
                "discards most). Re-chunk the store near the tile size "
                "(e.g. to_ome_zarr(..., chunks=...) without shard=) or read "
                "the source file directly to avoid wasted I/O.",
                _native_chunks,
                _tile,
                _amp,
            )

    image_for_threshold = image

    # Overlap — build a per-axis depth dict (clips to fit each axis).
    # An integer depth raises if any axis is smaller than the depth, so we
    # cap per axis. In practice z-axis of size 1 (2-D Cellpose) gets depth=0.
    _depth: dict[int, int] = {
        ax: min(overlap, max(0, sum(c) - 1))
        for ax, c in enumerate(image.chunks)
    }

    # Wrap fn with optional empty-tile skipping
    _skip_thr = empty_threshold
    if skip_empty and _skip_thr is None:
        _skip_thr = auto_empty_threshold(image_for_threshold, channel, level)

    # Up-front heads-up for a big 3-D GPU job (z-stack tiles, many of them):
    # an accurate ETA is logged after the first few tiles, but warn early that
    # this is the expensive path and point at the faster alternatives.
    if use_gpu and image.ndim >= 3 and _tile[0] > 4 and n_tiles >= 50:
        logger.warning(
            "Large 3-D GPU job: %d tiles of %s. Per-tile 3-D segmentation is "
            "slow on a single device (a live ETA is logged after the first "
            "tiles). If per-slice results are acceptable, 2-D (z=1 tiles) is "
            "typically ~10x faster, or segment a lower pyramid level.",
            n_tiles,
            _tile,
        )

    if dry_run:
        return _plan(
            image,
            fn,
            overlap=[_depth[ax] for ax in range(image.ndim)],
            skip_empty=skip_empty,
            threshold=_skip_thr,
            use_gpu=use_gpu,
            max_workers=max_workers,
            devices=devices,
            stitch=stitch,
            sample=plan_sample,
        )

    # Tile counters + timing so the log shows live "tile k/N + ETA". The
    # threaded scheduler runs tiles concurrently and ``+=`` is not atomic, so
    # the counters are updated under a lock.
    _progress = {"done": 0, "seen": 0, "time": 0.0}
    _progress_lock = threading.Lock()

    def active_fn(block, block_info=None):
        """Run *fn* on one tile, or return zeros for an empty tile.

        Parameters
        ----------
        block : np.ndarray
            One image tile.
        block_info : dict or None
            Dask block metadata (used for logging the tile location).

        Returns
        -------
        np.ndarray
            Integer labels, or an all-zero tile when skipped.
        """
        loc = block_info[0].get("chunk-location") if block_info else "?"
        with _progress_lock:
            _progress["seen"] += 1
        if skip_empty and block.size and block.max() <= _skip_thr:
            if verbose:
                logger.debug("skip empty tile %s (max<=%s)", loc, _skip_thr)
            return np.zeros(block.shape, dtype=np.int32)

        t0 = time.perf_counter()
        out = np.asarray(fn(block))
        dt = time.perf_counter() - t0
        if out.shape != block.shape:
            # Otherwise this surfaces deep in dask/zarr as a broadcast error
            # that names neither fn nor the tile.
            name = getattr(fn, "__name__", type(fn).__name__)
            raise ValueError(
                f"segmentation function {name!r} returned shape {out.shape} "
                f"for a tile of shape {block.shape} (tile {loc}). It must "
                "return one label per input voxel."
            )

        with _progress_lock:
            _progress["done"] += 1
            _progress["time"] += dt
            done, seen = _progress["done"], _progress["seen"]
            avg = _progress["time"] / done
        # Extrapolate remaining non-empty tiles from the empty fraction seen.
        remaining_nonempty = max(0, n_tiles - seen) * (done / seen)
        eta_h = avg * remaining_nonempty / 3600
        logger.info(
            "tile %d done in %.1fs (avg %.1fs); %d/%d tiles seen; ETA ~%.1fh",
            done,
            dt,
            avg,
            seen,
            n_tiles,
            eta_h,
        )
        return out

    _meta = np.empty((0,) * image.ndim, dtype=np.int32)
    if overlap > 0:
        # One fused pass: add the halo, run fn, trim it back off. map_overlap
        # materialises only the halos it needs (no separate overlapped array)
        # and keeps the task graph small. boundary="none" + trim recovers the
        # original shape, so the boundary-slab scan reads clean tiles.
        labeled = da.map_overlap(
            active_fn,
            image,
            depth=_depth,
            boundary="none",
            trim=True,
            dtype=np.int32,
            meta=_meta,
        )
    else:
        labeled = image.map_blocks(  # type: ignore[call-arg]
            active_fn, dtype=np.int32, meta=_meta
        )

    _tile_nbytes = int(np.prod(labeled.chunksize)) * labeled.dtype.itemsize

    # Stage: run fn once per tile to a temp zarr, then the zarr-native merge
    # reads concrete data from disk (fn is never re-run). Required because the
    # merge scans the labels directly on disk.
    if stage_dir is not None:
        base: str | None = str(stage_dir)
    elif write_to is not None and not is_remote(write_to):
        base = os.path.dirname(os.path.abspath(str(write_to)))
    elif image_source_path is not None and not is_remote(image_source_path):
        base = os.path.dirname(os.path.abspath(image_source_path))
    else:
        # A fresh system temp dir -- also for a remote input or output:
        # scratch data stays local, and os.path on a URL would invent a
        # directory named after the URL scheme.
        base = None
    halo_dir: str | None = None
    checkpoint: str | None = None
    if resume:
        # Named after the run, so the same call finds it again -- not
        # unique per call like a scratch store, which is the point.
        fingerprint = _run_fingerprint(
            source=image_source_path or image.name,
            shape=image.shape,
            tile=tuple(image.chunksize),
            overlap=_depth,
            fn=_fn_key(fn),
            skip_empty=skip_empty,
            threshold=_skip_thr,
            out=[str(write_to), output_component],
        )
        root = base if base is not None else tempfile.gettempdir()
        os.makedirs(root, exist_ok=True)
        stage_path = os.path.join(root, f"_pws_resume_{fingerprint}.zarr")
        stage_cleanup = stage_path
        checkpoint = os.path.join(stage_path, ".patchworks_done.json")
        logger.info("resumable stage store: %s", stage_path)
    else:
        stage_path, stage_cleanup = _scratch_store(base, "stage")
    if stitch == "iou":
        halo_dir = f"{stage_path}.halo"
    label_counts: dict[int, int] | None = None
    succeeded = False

    # Default: input is a .zarr store and no explicit write_to → labels go back
    # *into* the input store under the NGFF labels/<name>/ group with an auto
    # pyramid, so image + segmentation live in one OME-ZARR.
    _into_input = (
        write_to is None
        and image_source_path is not None
        and _is_zarr_path(image_source_path)
    )
    _merge_cleanup: str | None = None

    # Everything from here on can fail halfway (fn raising, disk full, a
    # killed worker). The dashboard cluster and the scratch stores are torn
    # down whatever happens, so a failed run leaves no process or a stage
    # the size of the whole image behind.
    try:
        if per_tile:
            _workers = (
                max_workers
                if max_workers is not None
                else safe_worker_count(_tile_nbytes, use_gpu=use_gpu)
            )
            _workers = max(1, min(_workers, cpu_allocation()))
            logger.info(
                "Staging tile by tile with %d thread(s) to %s …",
                _workers,
                stage_path,
            )
            label_counts = _stage_tiles(
                image,
                active_fn,
                stage_path,
                tuple(image.chunksize),
                [_depth[ax] for ax in range(image.ndim)],
                _workers,
                halo_dir,
                checkpoint,
                progress,
                devices=devices,
            )
        else:
            _stage_fused(
                labeled,
                stage_path,
                _active,
                use_gpu,
                max_workers,
                _tile_nbytes,
                progress,
            )

        # NB: no post-staging skip-count pass here — counting skipped tiles by
        # re-reading the whole staged store off disk would double the I/O of
        # the entire run just for a log line. Use estimate_empty_tiles() up
        # front for that figure instead.

        # Merge runs in worker processes (each holds one chunk + an mmap'd
        # LUT); size it to RAM/CPU like staging, capped so we don't spawn a
        # process storm.
        _nw = max_workers or max(1, min(safe_worker_count(_tile_nbytes), 8))

        # The merge always writes its result to a concrete store first.
        if write_to is not None:
            _merge_out = str(write_to)
        else:
            _merge_out, _merge_tmp = _scratch_store(None, "merge")
            if _into_input:
                _merge_cleanup = _merge_tmp

        # sequential=True folds the contiguous renumbering into the merge's
        # own LUT, so it costs a np.unique over the object count rather than
        # the extra full read+write (plus a Python set of every id) that a
        # separate relabel_sequential_zarr pass would.
        zarr_native_merge(
            stage_path,
            "staged",
            _merge_out,
            output_component,
            n_workers=_nw,
            show_progress=progress,
            sequential=sequential_labels,
            label_counts=label_counts,
            halo_dir=halo_dir,
            iou_threshold=iou_threshold,
        )
        succeeded = True
    finally:
        if keep_stage:
            logger.info("Keeping stage store %s", stage_path)
        elif resume and not succeeded:
            logger.info(
                "Keeping stage store %s to resume from: rerun the same call",
                stage_path,
            )
        else:
            _remove_scratch(stage_cleanup)
            _remove_scratch(halo_dir)

    record = provenance(
        input=image_source_path or f"<dask array {image.name}>",
        fn=_fn_key(fn),
        tile_shape=tuple(image.chunksize),
        overlap=[_depth[ax] for ax in range(image.ndim)],
        stitch=stitch,
        iou_threshold=iou_threshold if stitch == "iou" else None,
        channel=channel,
        level=level,
        skip_empty=skip_empty,
        empty_threshold=_skip_thr,
        sequential_labels=sequential_labels,
        compression=_COMPRESSION.get(),
    )
    merged = da.from_zarr(_merge_out, component=output_component)
    if not _into_input:
        # write_to: the record rides on the labels array itself.
        write_provenance(
            zarr.open_group(_merge_out, mode="r+")[output_component], record
        )
        # Lazy dask array backed by the merge store. Never loads the full
        # volume into RAM. Caller can .compute() if it fits.
        return merged

    # Stream the merged labels into the input store as an NGFF label pyramid,
    # then drop the temporary merge store. write_labels uses da.to_zarr, so
    # this is chunk-streamed and OOM-safe.
    from .plugins.ome_zarr import write_labels

    try:
        label_group = write_labels(
            image_source_path,
            merged,
            name=output_component,
            n_levels=pyramid_levels,
            downscale=pyramid_downscale,
            progress=progress,
            overwrite=True,
            # Segmented at `level`, so calibrated as that level, not level 0.
            level=level,
            provenance=record,
        )
    finally:
        _remove_scratch(_merge_cleanup)
    logger.info("labels stored in input OME-ZARR under %s", label_group)
    return da.from_zarr(label_group, component="0")