patchworks 0.2.0__py3-none-any.whl

patchworks/__init__.py

@@ -0,0 +1,48 @@
+"""patchworks — tiled processing for any image, any function.
+
+Process arbitrarily large images by splitting them into overlapping tiles,
+running any callable on each tile, and stitching the results back into globally
+consistent labels.
+
+📖 **Full documentation, guides and tutorials:**
+<https://imcf.one/patchworks/>
+
+Quick start
+-----------
+>>> from patchworks import tile_process
+>>>
+>>> def my_fn(tile):
+...     from skimage.filters import threshold_otsu
+...     from skimage.measure import label
+...     return label(tile > threshold_otsu(tile)).astype("int32")
+>>>
+>>> result = tile_process("image.zarr", my_fn, write_to="labels.zarr")
+
+With Cellpose:
+
+>>> from patchworks.plugins.cellpose import cellpose_fn
+>>> fn = cellpose_fn("cyto3", gpu=True, diameter=30)
+>>> tile_process("image.zarr", fn, tile_shape=(1, 2048, 2048),
+...              overlap=20, write_to="labels.zarr", progress=True)
+"""
+
+from ._chunks import auto_overlap, auto_tile_shape, auto_tile_shape_cellpose
+from ._cluster import make_local_cluster
+from ._core import tile_process
+from ._io import estimate_empty_tiles, load_ome_zarr
+from ._merge import merge_tile_labels
+from ._relabel import relabel_sequential_array, relabel_sequential_zarr
+
+__version__ = "0.2.0"
+__all__ = [
+    "tile_process",
+    "merge_tile_labels",
+    "auto_overlap",
+    "auto_tile_shape",
+    "auto_tile_shape_cellpose",
+    "load_ome_zarr",
+    "estimate_empty_tiles",
+    "make_local_cluster",
+    "relabel_sequential_array",
+    "relabel_sequential_zarr",
+]

patchworks/_chunks.py

@@ -0,0 +1,258 @@
+"""Auto tile-shape estimation."""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def auto_overlap(diameter: float, safety: float = 1.0) -> int:
+    """Recommended overlap (halo) for a given cell diameter.
+
+    Rule: overlap >= diameter so the segmentation function always sees at
+    least one full cell's worth of context on every tile edge. Cells near
+    tile boundaries are then segmented correctly and only genuinely split
+    cells produce touching labels at the boundary → correct merge.
+
+    Parameters
+    ----------
+    diameter:
+        Expected cell diameter in pixels (same unit as your image).
+    safety:
+        Multiplier on top of diameter. Default 1.0 (= one cell width).
+        Use 1.5–2.0 for elongated or irregularly-shaped cells.
+
+    Returns
+    -------
+    int
+        Overlap depth to pass to ``tile_process(..., overlap=...)``.
+
+    Examples
+    --------
+    >>> from patchworks import auto_overlap, tile_process
+    >>> 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=auto_overlap(30))
+    """
+    return max(1, int(np.ceil(diameter * safety)))
+
+_GPU_MEMORY_FALLBACK = 8 * 1024**3
+
+
+def _get_available_memory() -> int:
+    try:
+        import psutil
+        return int(psutil.virtual_memory().available)
+    except Exception:
+        return 8 * 1024**3
+
+
+def _get_gpu_memory() -> int:
+    """Return free GPU VRAM in bytes. Falls back to 8 GiB default."""
+    try:
+        import pynvml
+        pynvml.nvmlInit()
+        handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+        info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        pynvml.nvmlShutdown()
+        return int(info.free)
+    except Exception:
+        logger.warning(
+            "GPU memory query failed (nvidia-ml-py not installed?); "
+            "using %.0f GiB default.",
+            _GPU_MEMORY_FALLBACK / 1024**3,
+        )
+        return _GPU_MEMORY_FALLBACK
+
+
+def auto_tile_shape(
+    shape: tuple[int, ...],
+    dtype: Any,
+    target_bytes: int = 64 * 1024**2,
+    use_gpu: bool = False,
+    gpu_memory: int | None = None,
+    available_memory: int | None = None,
+    n_workers: int | None = None,
+    verbose: bool = False,
+) -> tuple[int, ...]:
+    """Balanced tile shape for general-purpose 3-D processing.
+
+    Sizes the last three axes (spatial) to stay within the memory budget while
+    keeping the shape as cubic as possible. Leading axes (t, c) are always 1.
+
+    Parameters
+    ----------
+    shape:
+        Full array shape, e.g. ``(z, y, x)`` or ``(t, c, z, y, x)``.
+    dtype:
+        Array dtype.
+    target_bytes:
+        Memory ceiling per tile. Default 64 MiB.
+    use_gpu:
+        Size tiles against GPU VRAM rather than host RAM.
+    gpu_memory:
+        Available GPU VRAM in bytes; auto-queried when None.
+    available_memory:
+        Available host RAM in bytes; auto-queried when None.
+    n_workers:
+        Number of parallel workers (divides the RAM budget).
+    verbose:
+        Log the chosen shape and estimated tile size.
+
+    Returns
+    -------
+    tuple[int, ...]
+        Tile shape with the same number of dimensions as *shape*.
+
+    Examples
+    --------
+    >>> tile = auto_tile_shape((128, 2048, 2048), "uint16")
+    >>> tile
+    (8, 2048, 2048)
+    """
+    n_workers = n_workers or os.cpu_count() or 1
+    itemsize = np.dtype(dtype).itemsize
+    n_spatial = min(3, len(shape))
+
+    if use_gpu:
+        mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()
+        budget = min(target_bytes * 2, mem // 2)
+    else:
+        mem = available_memory or _get_available_memory()
+        budget = min(target_bytes, mem // (n_workers * 4))
+
+    budget = max(32 * 1024**2, budget)
+
+    leading = [1] * (len(shape) - n_spatial)
+    spatial = list(shape[-n_spatial:])
+    target_voxels = budget / itemsize
+    target_side = int(target_voxels ** (1.0 / n_spatial))
+    chunk_spatial = [min(s, target_side) for s in spatial]
+
+    capped = [i for i, (c, s) in enumerate(zip(chunk_spatial, spatial)) if c == s]
+    uncapped = [i for i in range(n_spatial) if i not in capped]
+    if uncapped:
+        used_by_capped = np.prod([chunk_spatial[i] for i in capped]) if capped else 1
+        remaining = target_voxels / max(1, used_by_capped)
+        new_side = int(remaining ** (1.0 / len(uncapped)))
+        for i in uncapped:
+            chunk_spatial[i] = min(spatial[i], new_side)
+
+    result = tuple(leading + chunk_spatial)
+
+    if verbose:
+        mib = np.prod(result) * itemsize / 1024**2
+        logger.info(
+            "auto_tile_shape: shape=%s dtype=%s → tiles=%s (~%.0f MiB/tile)",
+            shape, np.dtype(dtype).name, result, mib,
+        )
+
+    return result
+
+
+def auto_tile_shape_cellpose(
+    shape: tuple[int, ...],
+    dtype: Any,
+    diameter: float | None = None,
+    do_3D: bool = False,
+    use_gpu: bool = False,
+    gpu_memory: int | None = None,
+    available_memory: int | None = None,
+    n_workers: int | None = None,
+    model_memory_bytes: int = 2 * 1024**3,
+    cellpose_memory_factor: int = 20,
+    verbose: bool = False,
+) -> tuple[int, ...]:
+    """Cellpose-optimised tile shape.
+
+    Cellpose is fundamentally 2-D: even in 3-D mode it runs 2-D segmentation
+    on orthogonal planes and takes a consensus.
+
+    **do_3D=False (default)**
+        z is set to 1. Each tile is one 2-D ``(y, x)`` slice.
+
+    **do_3D=True**
+        z is kept at its full extent per tile. y and x are tiled based on the
+        available memory, accounting for the 3× overhead of three plane orientations.
+
+    Parameters
+    ----------
+    shape:
+        Spatial shape, e.g. ``(z, y, x)``.
+    dtype:
+        Array dtype.
+    diameter:
+        Expected cell diameter in pixels. Tile will be at least ``4 × diameter``.
+    do_3D:
+        Whether Cellpose will run in 3-D mode.
+    use_gpu:
+        Size tiles for GPU VRAM.
+    gpu_memory, available_memory, n_workers:
+        Memory parameters (auto-queried when None).
+    model_memory_bytes:
+        Memory consumed by the Cellpose model weights (default 2 GiB).
+    cellpose_memory_factor:
+        Cellpose allocates roughly this multiple of raw input bytes (default 20×).
+    verbose:
+        Log the chosen shape and memory estimates.
+
+    Returns
+    -------
+    tuple[int, ...]
+        Tile shape with the same number of dimensions as *shape*.
+
+    Examples
+    --------
+    >>> tile = auto_tile_shape_cellpose((128, 2048, 2048), "uint16", diameter=30)
+    >>> tile
+    (1, 2048, 2048)
+    """
+    n_workers = n_workers or os.cpu_count() or 1
+    itemsize = np.dtype(dtype).itemsize
+
+    if use_gpu:
+        total_mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()
+    else:
+        total_mem = (available_memory or _get_available_memory()) // n_workers
+
+    usable = max(32 * 1024**2, total_mem - model_memory_bytes)
+    max_raw_bytes = usable // cellpose_memory_factor
+
+    n_spatial = min(3, len(shape))
+    leading = [1] * (len(shape) - n_spatial)
+    min_tile = int(4 * diameter) if diameter is not None else 1
+
+    if n_spatial == 2 or not do_3D:
+        max_pixels_2d = max(1, max_raw_bytes // itemsize)
+        tile_side = max(min_tile, int(max_pixels_2d**0.5))
+        if n_spatial == 2:
+            y, x = shape[-2], shape[-1]
+            chunk_spatial = [min(y, tile_side), min(x, tile_side)]
+        else:
+            z, y, x = shape[-3], shape[-2], shape[-1]
+            chunk_spatial = [1, min(y, tile_side), min(x, tile_side)]
+    else:
+        z, y, x = shape[-3], shape[-2], shape[-1]
+        max_pixels_per_slice = max(1, (max_raw_bytes // 3) // (z * itemsize))
+        tile_side = max(min_tile, int(max_pixels_per_slice**0.5))
+        chunk_spatial = [z, min(y, tile_side), min(x, tile_side)]
+
+    result = tuple(leading + chunk_spatial)
+
+    if verbose:
+        raw_mib = np.prod(result) * itemsize / 1024**2
+        logger.info(
+            "auto_tile_shape_cellpose: shape=%s dtype=%s do_3D=%s "
+            "→ tiles=%s (~%.0f MiB raw, ~%.0f MiB Cellpose estimate)",
+            shape, np.dtype(dtype).name, do_3D, result,
+            raw_mib, raw_mib * cellpose_memory_factor,
+        )
+
+    return result

patchworks/_cluster.py

@@ -0,0 +1,93 @@
+"""Dask cluster helpers."""
+from __future__ import annotations
+
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+
+def _distributed_client():
+    """Return the active dask.distributed Client, or None."""
+    try:
+        from dask.distributed import get_client
+        return get_client()
+    except Exception:
+        return None
+
+
+def _client_is_in_process(client) -> bool:
+    """True if *client* runs its worker in this process (processes=False).
+
+    An in-process worker shares the GIL. A long task that holds the GIL
+    (e.g. a Cellpose/torch eval) starves the worker heartbeat, the scheduler
+    declares it dead, and the P2P merge barrier drops its inputs →
+    "FutureCancelledError: lost dependencies".
+    """
+    try:
+        for addr in client.scheduler_info().get("workers", {}):
+            if str(addr).startswith("inproc://"):
+                return True
+    except Exception:
+        pass
+    return False
+
+
+def make_local_cluster(
+    use_gpu: bool = False,
+    n_workers: int | None = None,
+    threads_per_worker: int = 1,
+    memory_limit: str | None = None,
+    **cluster_kwargs,
+):
+    """Create a process-based Dask cluster for tiled processing.
+
+    Always uses worker subprocesses (``processes=True``). An in-process
+    (threaded) worker breaks the label merge when ``segment_fn`` holds the
+    GIL — see the patchworks docs for details.
+
+    For GPU work defaults to a single worker (one CUDA context, no contention).
+    For CPU scales to available cores.
+
+    Parameters
+    ----------
+    use_gpu:
+        Single-worker cluster for GPU. When False, use multiple CPU workers.
+    n_workers:
+        Override the worker count. Defaults to 1 for GPU, min(8, cpu_count).
+    threads_per_worker:
+        Keep at 1 so a GIL-holding tile function doesn't block heartbeats.
+    memory_limit:
+        Per-worker memory cap (e.g. ``"8GB"``).
+    **cluster_kwargs:
+        Extra arguments forwarded to ``dask.distributed.LocalCluster``.
+
+    Returns
+    -------
+    (client, cluster)
+
+    Examples
+    --------
+    >>> client, cluster = make_local_cluster(use_gpu=True)
+    >>> print("dashboard:", client.dashboard_link)
+    >>> result = tile_process("image.zarr", fn, write_to="labels.zarr")
+    >>> client.close(); cluster.close()
+    """
+    from dask.distributed import Client, LocalCluster
+
+    if n_workers is None:
+        n_workers = 1 if use_gpu else min(8, os.cpu_count() or 1)
+
+    cluster = LocalCluster(
+        processes=True,
+        n_workers=n_workers,
+        threads_per_worker=threads_per_worker,
+        memory_limit=memory_limit,
+        **cluster_kwargs,
+    )
+    client = Client(cluster)
+    logger.info(
+        "Started %d-worker process cluster (use_gpu=%s). Dashboard: %s",
+        n_workers, use_gpu, client.dashboard_link,
+    )
+    return client, cluster