I/O helpers

patchworks.load_ome_zarr(store_path: Union[str, Path], channel: int | None = 0, level: int = 0, chunks: tuple[int, ...] | None = None) -> da.Array

Load one spatial array from an OME-ZARR store.

Parameters:

Name	Type	Description	Default
store_path	Union[str, Path]	Path to the OME-ZARR store (.zarr directory).	required
channel	int | None	Channel index to select (axis is dropped). Pass None to keep it.	0
level	int	Resolution pyramid level (0 = full resolution).	0
chunks	tuple[int, ...] | None	Target chunk shape for the returned dask array.	None

Returns:

Type	Description
Array	Shape (z, y, x) when channel is an int, or (c, z, y, x) when channel is None.

Examples:

>>> arr = load_ome_zarr("image.zarr", channel=0)
>>> arr.shape
(128, 2048, 2048)

Source code in src/patchworks/_io.py

def load_ome_zarr(
    store_path: Union[str, Path],
    channel: int | None = 0,
    level: int = 0,
    chunks: tuple[int, ...] | None = None,
) -> da.Array:
    """Load one spatial array from an OME-ZARR store.

    Parameters
    ----------
    store_path:
        Path to the OME-ZARR store (.zarr directory).
    channel:
        Channel index to select (axis is dropped). Pass ``None`` to keep it.
    level:
        Resolution pyramid level (0 = full resolution).
    chunks:
        Target chunk shape for the returned dask array.

    Returns
    -------
    da.Array
        Shape ``(z, y, x)`` when *channel* is an int, or ``(c, z, y, x)``
        when *channel* is None.

    Examples
    --------
    >>> arr = load_ome_zarr("image.zarr", channel=0)
    >>> arr.shape
    (128, 2048, 2048)
    """
    root = zarr.open_group(str(store_path), mode="r")
    # OME-ZARR 0.5 nests under "ome" key; older stores use "multiscales" directly
    _attrs = dict(root.attrs)
    _ms = _attrs.get("multiscales") or _attrs.get("ome", {}).get("multiscales")
    try:
        path = _ms[0]["datasets"][level]["path"]
    except (KeyError, IndexError, TypeError) as exc:
        raise ValueError(
            f"Cannot read OME-ZARR multiscales metadata at level {level} "
            f"in {store_path!r}"
        ) from exc

    zarr_chunks = chunks
    if chunks is not None and channel is not None:
        zarr_ndim = len(root[path].shape)
        if zarr_ndim > len(chunks):
            zarr_chunks = (1,) * (zarr_ndim - len(chunks)) + tuple(chunks)

    arr = da.from_zarr(str(store_path), component=path, chunks=zarr_chunks)
    if channel is not None:
        arr = _select_channel(arr, channel, _ms[0], store_path)
    return arr

Deciding which tiles hold signal

estimate_empty_tiles is a fast preview — it samples a centred window per tile, so it can miss signal at a tile's edge. build_occupancy_map + tile_occupancy are exact: a brick maximum exceeds the threshold exactly when some voxel in that brick does. Use the latter pair when the result is used as a skip list. See Skipping empty tiles.

patchworks.estimate_empty_tiles(image: Union[da.Array, str, Path], tile_shape: tuple[int, ...], threshold: float | None = None, channel: int | None = 0, level: int = 0, sample_window: tuple[int, ...] = (24, 256, 256)) -> dict[str, Any]

Fast preview of which tiles are background before processing.

For each tile, reads a small centred window (sample_window) and tests whether its max exceeds threshold. Bounded I/O — runs in seconds to minutes on terabyte arrays.

APPROXIMATE: only the tile centre is inspected. The actual tile_process run always tests the whole tile inline. Use this only to pick a threshold and gauge the empty fraction before committing to a full run.

Parameters:

Name	Type	Description	Default
image	Union[Array, str, Path]	Dask array or OME-ZARR path.	required
tile_shape	tuple[int, ...]	Tile shape you plan to use, e.g. (120, 697, 697).	required
threshold	float | None	Empty cutoff (signal <= threshold → empty). None → Otsu on samples.	None
channel	int | None	Used only when image is a path.	0
level	int | None	Used only when image is a path.	0
sample_window	tuple[int, ...]	Size of the centred window read per tile.	(24, 256, 256)

Returns:

Type	Description
dict with keys:	threshold, n_tiles, n_occupied, empty_fraction, occupancy (bool ndarray, one entry per tile in the grid).

Examples:

>>> info = estimate_empty_tiles("image.zarr", (120, 697, 697))
>>> print(f"{info['empty_fraction']:.0%} of tiles are background")
>>> labels = tile_process("image.zarr", fn, tile_shape=(120, 697, 697),
...                       skip_empty=True, empty_threshold=info["threshold"])

Source code in src/patchworks/_io.py

def estimate_empty_tiles(
    image: Union[da.Array, str, Path],
    tile_shape: tuple[int, ...],
    threshold: float | None = None,
    channel: int | None = 0,
    level: int = 0,
    sample_window: tuple[int, ...] = (24, 256, 256),
) -> dict[str, Any]:
    """Fast preview of which tiles are background before processing.

    For each tile, reads a small centred window (``sample_window``) and tests
    whether its max exceeds *threshold*. Bounded I/O — runs in seconds to
    minutes on terabyte arrays.

    APPROXIMATE: only the tile centre is inspected. The actual ``tile_process``
    run always tests the whole tile inline. Use this only to pick a threshold
    and gauge the empty fraction before committing to a full run.

    Parameters
    ----------
    image:
        Dask array or OME-ZARR path.
    tile_shape:
        Tile shape you plan to use, e.g. ``(120, 697, 697)``.
    threshold:
        Empty cutoff (signal <= threshold → empty). None → Otsu on samples.
    channel, level:
        Used only when *image* is a path.
    sample_window:
        Size of the centred window read per tile.

    Returns
    -------
    dict with keys:
        ``threshold``, ``n_tiles``, ``n_occupied``, ``empty_fraction``,
        ``occupancy`` (bool ndarray, one entry per tile in the grid).

    Examples
    --------
    >>> info = estimate_empty_tiles("image.zarr", (120, 697, 697))
    >>> print(f"{info['empty_fraction']:.0%} of tiles are background")
    >>> labels = tile_process("image.zarr", fn, tile_shape=(120, 697, 697),
    ...                       skip_empty=True, empty_threshold=info["threshold"])
    """
    n_spatial = len(tile_shape)

    z_src: Any = None
    if isinstance(image, (str, Path)):
        _root = zarr.open_group(str(image), mode="r")
        _rattr = dict(_root.attrs)
        _rms = _rattr.get("multiscales") or _rattr.get("ome", {}).get(
            "multiscales"
        )
        try:
            _zpath = _rms[0]["datasets"][level]["path"]
        except (KeyError, IndexError, TypeError) as exc:
            raise ValueError(
                f"Cannot read OME-ZARR multiscales metadata at level {level} "
                f"in {image!r}"
            ) from exc
        z_src = _root[_zpath]
        sp_shape = tuple(z_src.shape[-n_spatial:])
    else:
        arr = image
        sp_shape = tuple(arr.shape[-n_spatial:])

    win = [min(w, t, s) for w, t, s in zip(sample_window, tile_shape, sp_shape)]
    grid = [int(np.ceil(s / t)) for s, t in zip(sp_shape, tile_shape)]

    _ch_prefix: tuple = ()
    if z_src is not None:
        n_leading = z_src.ndim - n_spatial
        if channel is not None and n_leading > 0:
            _ch_prefix = (0,) * (n_leading - 1) + (channel,)

    # Streaming single pass: store only per-tile max (a scalar) and a bounded
    # sample list for Otsu. The old approach stored every tile's full block in
    # `blocks` dict — for 2000 tiles × 24×256×256 × 2 bytes = ~6 GB in RAM.
    _MAX_OTSU_SAMPLES = 500
    samples: list[np.ndarray] = []
    tile_maxes: dict[tuple, float] = {}

    for idx in np.ndindex(*grid):
        sl: list[slice] = []
        for i, t, w, s in zip(idx, tile_shape, win, sp_shape):
            # Centre the window in this tile, then clamp it to the tile's own
            # extent -- NOT to the array's. Clamping to ``s - w`` used to drag
            # the last (partial) tile's window backwards into its neighbour,
            # so an edge tile's verdict came partly from the tile before it.
            lo, hi = i * t, min((i + 1) * t, s)
            start = max(lo, min(lo + (hi - lo - w) // 2, hi - w))
            sl.append(slice(start, min(start + w, hi)))

        if z_src is not None:
            block = np.asarray(z_src[_ch_prefix + tuple(sl)])
        else:
            sub = (
                arr[(...,) + tuple(sl)]
                if arr.ndim > n_spatial
                else arr[tuple(sl)]
            )
            block = np.asarray(sub)

        tile_maxes[idx] = float(block.max()) if block.size else 0.0
        if threshold is None and len(samples) < _MAX_OTSU_SAMPLES:
            samples.append(block.ravel())
        # block freed here — not stored

    if threshold is None:
        threshold = _otsu_threshold(
            np.concatenate(samples) if samples else np.zeros(1)
        )

    occupancy = np.zeros(grid, dtype=bool)
    for idx, mx in tile_maxes.items():
        occupancy[idx] = mx > threshold

    n_tiles = int(occupancy.size)
    n_occ = int(occupancy.sum())
    empty_frac = 1.0 - n_occ / n_tiles if n_tiles else 0.0
    logger.info(
        "estimate_empty_tiles: threshold=%.4g  occupied %d/%d tiles  empty=%.0f%%",
        threshold,
        n_occ,
        n_tiles,
        empty_frac * 100,
    )
    return {
        "threshold": float(threshold),
        "n_tiles": n_tiles,
        "n_occupied": n_occ,
        "empty_fraction": empty_frac,
        "occupancy": occupancy,
    }

patchworks.build_occupancy_map(image_store: Union[str, Path], *, level: int = 0, block: tuple[int, ...] = DEFAULT_BLOCK, overwrite: bool = False) -> str

Max-pool every channel of an OME-ZARR level into a small summary array.

Reads the image once and writes <image_store>/occupancy/<level>, an array of shape (n_channels, *ceil(spatial_shape / block)) holding the maximum of each brick. Cheap to keep (~1/16384 of the image by default) and reusable by every config that segments this image.

Idempotent: an existing map is reused unless overwrite is set, so concurrent segmentation runs against one work_dir build it at most once. The map is written to a temporary sibling and moved into place, so a crash mid-build never leaves a partial map behind.

Parameters:

Name	Type	Description	Default
image_store	str or Path	OME-ZARR store to summarise.	required
level	int	Pyramid level to read (default 0, full resolution).	0
block	tuple of int	Brick shape over the spatial axes (default (1, 128, 128): no z reduction, so any z-tiling works).	DEFAULT_BLOCK
overwrite	bool	Rebuild even if a map already exists.	False

Returns:

Type	Description
str	Path of the occupancy array.

Examples:

>>> build_occupancy_map("image.zarr")
'image.zarr/occupancy/0'

Source code in src/patchworks/_occupancy.py

def build_occupancy_map(
    image_store: Union[str, Path],
    *,
    level: int = 0,
    block: tuple[int, ...] = DEFAULT_BLOCK,
    overwrite: bool = False,
) -> str:
    """Max-pool every channel of an OME-ZARR level into a small summary array.

    Reads the image once and writes ``<image_store>/occupancy/<level>``, an
    array of shape ``(n_channels, *ceil(spatial_shape / block))`` holding the
    maximum of each brick. Cheap to keep (~1/16384 of the image by default)
    and reusable by every config that segments this image.

    Idempotent: an existing map is reused unless *overwrite* is set, so
    concurrent segmentation runs against one ``work_dir`` build it at most
    once. The map is written to a temporary sibling and moved into place, so a
    crash mid-build never leaves a partial map behind.

    Parameters
    ----------
    image_store : str or Path
        OME-ZARR store to summarise.
    level : int, optional
        Pyramid level to read (default 0, full resolution).
    block : tuple of int, optional
        Brick shape over the spatial axes (default ``(1, 128, 128)``: no z
        reduction, so any z-tiling works).
    overwrite : bool, optional
        Rebuild even if a map already exists.

    Returns
    -------
    str
        Path of the occupancy array.

    Examples
    --------
    >>> build_occupancy_map("image.zarr")  # doctest: +SKIP
    'image.zarr/occupancy/0'
    """
    store = str(image_store)
    out_path = occupancy_path(store, level)
    if not overwrite and Path(out_path).exists():
        # Reuse only if it was built at the block we want. A map left over
        # from a run with a different tile_shape is coarser (or finer) than
        # this run needs, and silently reusing it would degrade every
        # occupancy answer that follows.
        try:
            existing = tuple(zarr.open_array(out_path, mode="r").attrs["block"])
        except Exception:
            existing = None
        if existing == tuple(block):
            logger.info("occupancy map already present at %s", out_path)
            return out_path
        logger.info(
            "rebuilding occupancy map at %s: it was built with block %s, "
            "this run needs %s",
            out_path,
            existing,
            tuple(block),
        )
        overwrite = True

    root = zarr.open_group(store, mode="r")
    src = _level_array(root, level, store)
    n_spatial = len(block)
    sp_shape = tuple(src.shape[-n_spatial:])
    n_leading = src.ndim - n_spatial
    n_channels = src.shape[n_leading - 1] if n_leading > 0 else 1

    grid = tuple(-(-s // b) for s, b in zip(sp_shape, block))
    itemsize = np.dtype(src.dtype).itemsize
    # Read this many output cells per axis at a time -- cube-ish, sized so one
    # read stays near _READ_TARGET_BYTES.
    voxels_per_block = int(np.prod(block))
    cells = max(1, _READ_TARGET_BYTES // (voxels_per_block * itemsize))
    step = max(1, int(round(cells ** (1.0 / n_spatial))))
    steps = tuple(min(step, g) for g in grid)

    logger.info(
        "building occupancy map: %d channel(s), block=%s, grid=%s (%.1f MB)",
        n_channels,
        block,
        grid,
        n_channels * float(np.prod(grid)) * itemsize / 1024**2,
    )

    tmp_path = f"{out_path}.building.{os.getpid()}"
    shutil.rmtree(tmp_path, ignore_errors=True)
    Path(tmp_path).parent.mkdir(parents=True, exist_ok=True)
    dst = zarr.open_array(
        tmp_path,
        mode="w",
        shape=(n_channels, *grid),
        chunks=(1, *steps),
        dtype=src.dtype,
    )

    ranges = [range(0, g, s) for g, s in zip(grid, steps)]
    regions = list(_iproduct(*ranges))

    def _one(starts: tuple[int, ...]) -> None:
        out_sl = tuple(
            slice(o, min(o + s, g)) for o, s, g in zip(starts, steps, grid)
        )
        src_sl = tuple(
            slice(o.start * b, min(o.stop * b, s))
            for o, b, s in zip(out_sl, block, sp_shape)
        )
        # Read every channel of this region in ONE go. Looping channels on the
        # outside would traverse the whole image once per channel -- three
        # full reads for a three-channel stack, where one does.
        for channel in range(n_channels):
            prefix = _leading_index(src.ndim, n_spatial, channel)
            region = np.asarray(src[prefix + src_sl])
            dst[(channel, *out_sl)] = _block_max(region, block)

    try:
        n_workers = max(1, min(cpu_allocation(), len(regions)))
        if n_workers <= 1:
            for _ in track(
                (_one(starts) for starts in regions),
                "occupancy map",
                len(regions),
            ):
                pass
        else:
            # Reads and decompression release the GIL, so threads are enough
            # and there is no worker payload to pickle.
            with ThreadPoolExecutor(max_workers=n_workers) as pool:
                for _ in track(
                    pool.map(_one, regions), "occupancy map", len(regions)
                ):
                    pass
        dst.attrs["block"] = list(block)
        dst.attrs["level"] = int(level)
        dst.attrs["source_shape"] = list(sp_shape)
    except BaseException:
        shutil.rmtree(tmp_path, ignore_errors=True)
        raise

    if Path(out_path).exists():
        if not overwrite:
            # A concurrent run finished first; theirs is as good as ours.
            shutil.rmtree(tmp_path, ignore_errors=True)
            return out_path
        # Explicit rebuild: clear the old one, or os.replace would fail on a
        # non-empty directory and we would keep serving the stale map.
        shutil.rmtree(out_path, ignore_errors=True)
    try:
        os.replace(tmp_path, out_path)
    except OSError:
        shutil.rmtree(tmp_path, ignore_errors=True)
        if not Path(out_path).exists():
            raise
    logger.info("occupancy map written to %s", out_path)
    return out_path

patchworks.tile_occupancy(image_store: Union[str, Path], tile_shape: tuple[int, ...], *, channel: int = 0, threshold: float, level: int = 0) -> dict[str, Any]

Decide which tiles hold signal, using the whole tile, not a sample.

Reduces the occupancy map over each tile's full footprint and marks the tile occupied when any brick maximum exceeds threshold -- equivalent to testing every voxel, because a brick maximum exceeds the threshold exactly when some voxel in that brick does.

Blocks are only ever over-covered, never under-covered: when a tile edge falls inside a brick, that brick counts for both neighbours. A tile can therefore be occupied because of a neighbour's signal in a shared brick, which costs one extra segmentation job but can never drop a tile. Choosing a block that divides tile_shape (the default 128 divides 1024) avoids even that.

Parameters:

Name	Type	Description	Default
image_store	str or Path	OME-ZARR store holding the occupancy map (see :func:build_occupancy_map).	required
tile_shape	tuple of int	Tile shape, in full-resolution voxels.	required
channel	int	Channel to test.	0
threshold	float	Empty cutoff, derived from raw voxel values (signal <= threshold → empty).	required
level	int	Pyramid level the map was built from.	0

Returns:

Type	Description
dict	threshold, n_tiles, n_occupied, empty_fraction and occupancy (bool array over the tile grid), matching :func:patchworks.estimate_empty_tiles.

Source code in src/patchworks/_occupancy.py

def tile_occupancy(
    image_store: Union[str, Path],
    tile_shape: tuple[int, ...],
    *,
    channel: int = 0,
    threshold: float,
    level: int = 0,
) -> dict[str, Any]:
    """Decide which tiles hold signal, using the whole tile, not a sample.

    Reduces the occupancy map over each tile's full footprint and marks the
    tile occupied when any brick maximum exceeds *threshold* -- equivalent to
    testing every voxel, because a brick maximum exceeds the threshold exactly
    when some voxel in that brick does.

    Blocks are only ever over-covered, never under-covered: when a tile edge
    falls inside a brick, that brick counts for both neighbours. A tile can
    therefore be occupied because of a neighbour's signal in a shared brick,
    which costs one extra segmentation job but can never drop a tile. Choosing
    a *block* that divides *tile_shape* (the default 128 divides 1024) avoids
    even that.

    Parameters
    ----------
    image_store : str or Path
        OME-ZARR store holding the occupancy map (see
        :func:`build_occupancy_map`).
    tile_shape : tuple of int
        Tile shape, in full-resolution voxels.
    channel : int, optional
        Channel to test.
    threshold : float
        Empty cutoff, derived from raw voxel values (signal <= threshold →
        empty).
    level : int, optional
        Pyramid level the map was built from.

    Returns
    -------
    dict
        ``threshold``, ``n_tiles``, ``n_occupied``, ``empty_fraction`` and
        ``occupancy`` (bool array over the tile grid), matching
        :func:`patchworks.estimate_empty_tiles`.
    """
    arr = zarr.open_array(occupancy_path(image_store, level), mode="r")
    block = tuple(arr.attrs["block"])
    sp_shape = tuple(arr.attrs["source_shape"])
    if len(tile_shape) != len(block):
        raise ValueError(
            f"tile_shape is {len(tile_shape)}-D but the occupancy map is "
            f"{len(block)}-D"
        )
    coarse = [
        (ax, b, t) for ax, (b, t) in enumerate(zip(block, tile_shape)) if b >= t
    ]
    if coarse:
        logger.warning(
            "occupancy blocks are as large as the tile on axes %s; every tile "
            "will over-cover the same block and test occupied. Rebuild the "
            "map with block=block_for_tile(tile_shape).",
            [ax for ax, _, _ in coarse],
        )

    tile_grid = tuple(-(-s // t) for s, t in zip(sp_shape, tile_shape))
    occupancy = np.zeros(tile_grid, dtype=bool)
    pooled = np.asarray(arr[channel])
    for idx in np.ndindex(*tile_grid):
        sl = tuple(
            slice(
                (i * t) // b,
                min(-(-min((i + 1) * t, s) // b), g),
            )
            for i, t, b, s, g in zip(
                idx, tile_shape, block, sp_shape, pooled.shape
            )
        )
        window = pooled[sl]
        occupancy[idx] = bool(window.size) and bool(window.max() > threshold)

    n_tiles = int(occupancy.size)
    n_occ = int(occupancy.sum())
    logger.info(
        "tile_occupancy: threshold=%.4g  occupied %d/%d tiles",
        threshold,
        n_occ,
        n_tiles,
    )
    return {
        "threshold": float(threshold),
        "n_tiles": n_tiles,
        "n_occupied": n_occ,
        "empty_fraction": 1.0 - n_occ / n_tiles if n_tiles else 0.0,
        "occupancy": occupancy,
    }

patchworks.auto_empty_threshold(image: da.Array, channel: int | None, level: int) -> float

Pick an empty-tile threshold from a cheap bounded sample (Otsu).

Parameters:

Name	Type	Description	Default
image	Array	Image to sample.	required
channel	int or None	Channel hint (kept for signature symmetry).	required
level	int	Pyramid level hint (kept for signature symmetry).	required

Returns:

Type	Description
float	Otsu threshold over a few small centred windows.

Source code in src/patchworks/_io.py

def auto_empty_threshold(
    image: da.Array, channel: int | None, level: int
) -> float:
    """Pick an empty-tile threshold from a cheap bounded sample (Otsu).

    Parameters
    ----------
    image : da.Array
        Image to sample.
    channel : int or None
        Channel hint (kept for signature symmetry).
    level : int
        Pyramid level hint (kept for signature symmetry).

    Returns
    -------
    float
        Otsu threshold over a few small centred windows.
    """
    n = image.ndim
    win = [min(64 if i >= n - 3 else s, s) for i, s in enumerate(image.shape)]
    win = [min(w, 256) if i >= n - 2 else w for i, w in enumerate(win)]
    samples = []
    for frac in (0.33, 0.5, 0.66):
        sl = tuple(
            slice(
                int(s * frac) - w // 2 if s > w else 0,
                (int(s * frac) - w // 2 if s > w else 0) + w,
            )
            for s, w in zip(image.shape, win)
        )
        samples.append(np.asarray(image[sl]).ravel())
    sample = np.concatenate(samples)
    thr = _otsu_threshold(sample)
    logger.info(
        "Auto empty_threshold=%.3g (Otsu on %d samples)", thr, len(samples)
    )
    return thr