ocdkit 0.0.1__py3-none-any.whl

ocdkit/__init__.py

@@ -0,0 +1,10 @@
+"""
+ocdkit — Obsessively precise utilities for scientific Python.
+
+Array manipulation, GPU dispatch, image I/O, morphology, and plotting
+tools shared across projects.
+"""
+
+from .load import enable_submodules
+
+enable_submodules(__name__)

ocdkit/array/__init__.py

@@ -0,0 +1,3 @@
+from ..load import enable_submodules
+
+enable_submodules(__name__)

ocdkit/array/convert.py

@@ -0,0 +1,121 @@
+"""Type detection, conversion, and rescaling utilities."""
+
+from .imports import *
+
+
+def get_module(x):
+    """Return ``np`` or ``torch`` depending on the type of *x*."""
+    if isinstance(x, da.Array):
+        return np
+    if isinstance(x, (np.ndarray, tuple, int, float)) or np.isscalar(x):
+        return np
+    if torch.is_tensor(x):
+        return torch
+    raise ValueError(
+        "Input must be a numpy array, a tuple, a torch tensor, "
+        "an integer, or a float"
+    )
+
+
+def safe_divide(num, den, cutoff=0):
+    """Division ignoring zeros and NaNs in the denominator."""
+    module = get_module(num)
+    valid_den = (den > cutoff) & module.isfinite(den)
+
+    if isinstance(num, da.Array) or isinstance(den, da.Array):
+        return da.where(valid_den, num / den, 0)
+    elif module == np:
+        r = num.astype(np.float32, copy=False)
+        r = np.divide(r, den, out=np.zeros_like(r), where=valid_den)
+    elif module == torch:
+        r = num.float()
+        den = den.float()
+        safe_den = torch.where(valid_den, den, torch.ones_like(den))
+        r = torch.where(valid_den, torch.div(r, safe_den), torch.zeros_like(r))
+    else:
+        raise TypeError("num must be a numpy array or a PyTorch tensor")
+
+    return r
+
+
+def rescale(T, floor=None, ceiling=None, exclude_dims=None):
+    """Min-max rescale to [0, 1].
+
+    Works on numpy arrays and torch tensors.  When *exclude_dims* is
+    given, normalization is applied independently along those axes.
+    """
+    module = get_module(T)
+    if exclude_dims is not None:
+        if isinstance(exclude_dims, int):
+            exclude_dims = (exclude_dims,)
+        axes = tuple(i for i in range(T.ndim) if i not in exclude_dims)
+        newshape = [T.shape[i] if i in exclude_dims else 1 for i in range(T.ndim)]
+    else:
+        axes = None
+        newshape = T.shape
+
+    if ceiling is None:
+        ceiling = module.amax(T, axis=axes)
+        if exclude_dims is not None:
+            ceiling = ceiling.reshape(*newshape)
+    if floor is None:
+        floor = module.amin(T, axis=axes)
+        if exclude_dims is not None:
+            floor = floor.reshape(*newshape)
+
+    T = safe_divide(T - floor, ceiling - floor)
+
+    return T
+
+
+def to_16_bit(im):
+    """Rescale image to [0, 2**16 - 1] and cast to uint16."""
+    return np.uint16(rescale(im) * (2 ** 16 - 1))
+
+
+def to_8_bit(im):
+    """Rescale image to [0, 2**8 - 1] and cast to uint8."""
+    return np.uint8(rescale(im) * (2 ** 8 - 1))
+
+
+def is_integer(var):
+    """Check whether *var* is an integer or integer-typed array/tensor."""
+    if isinstance(var, int):
+        return True
+    if isinstance(var, np.integer):
+        return True
+    if isinstance(var, (np.ndarray, np.memmap)) and np.issubdtype(var.dtype, np.integer):
+        return True
+    if isinstance(var, da.Array) and np.issubdtype(var.dtype, np.integer):
+        return True
+    if isinstance(var, torch.Tensor) and not var.is_floating_point():
+        return True
+    return False
+
+
+def move_axis(img, axis=-1, pos="last"):
+    """Move ndarray axis to a new location, preserving order of other axes."""
+    if axis == -1:
+        axis = img.ndim - 1
+    axis = min(img.ndim - 1, axis)
+    if pos in ("first", 0):
+        pos = 0
+    elif pos in ("last", -1):
+        pos = img.ndim - 1
+    perm = list(range(img.ndim))
+    perm.pop(axis)
+    perm.insert(pos, axis)
+    return np.transpose(img, perm)
+
+
+def move_min_dim(img, force=False):
+    """Move the minimum-sized dimension last (as channels) if < 10."""
+    if len(img.shape) > 2:
+        min_dim = min(img.shape)
+        if min_dim < 10 or force:
+            if img.shape[-1] == min_dim:
+                channel_axis = -1
+            else:
+                channel_axis = (img.shape).index(min_dim)
+            img = move_axis(img, axis=channel_axis, pos="last")
+    return img

ocdkit/array/filters.py

@@ -0,0 +1,56 @@
+"""Spatial filtering utilities for labeled arrays."""
+
+from .imports import *
+from numba import njit
+
+from .spatial import kernel_setup, get_neighbors
+
+
+@njit
+def _most_frequent(neighbor_masks):
+    """Column-wise mode: for each column, return the most common value."""
+    return np.array([np.bincount(row).argmax() for row in neighbor_masks.T])
+
+
+def mode_filter(masks):
+    """Replace each nonzero pixel with the most frequent label in its neighborhood.
+
+    Uses ocdkit's spatial neighbor primitives and a numba-JIT inner loop
+    for the per-pixel bincount. Background (0) results are replaced with
+    the original label to avoid erosion.
+
+    Parameters
+    ----------
+    masks : ndarray
+        Integer label array (2D or ND).
+
+    Returns
+    -------
+    ndarray
+        Filtered label array, same shape as input.
+    """
+    pad = 1
+    masks = np.pad(masks, pad).astype(int)
+    d = masks.ndim
+    shape = masks.shape
+    coords = np.nonzero(masks)
+
+    if coords[0].size == 0:
+        unpad = tuple([slice(pad, -pad)] * d)
+        return masks[unpad]
+
+    steps, inds, idx, fact, sign = kernel_setup(d)
+    subinds = np.concatenate(inds)
+    substeps = steps[subinds]
+    neighbors = get_neighbors(coords, substeps, d, shape)
+
+    neighbor_masks = masks[tuple(neighbors)]
+
+    mask_filt = np.zeros_like(masks)
+    most_f = _most_frequent(neighbor_masks)
+    z = most_f == 0
+    most_f[z] = masks[coords][z]
+    mask_filt[coords] = most_f
+
+    unpad = tuple([slice(pad, -pad)] * d)
+    return mask_filt[unpad]

ocdkit/array/imports.py

@@ -0,0 +1,4 @@
+"""Common imports for ocdkit.array subpackage."""
+from ..imports import *
+
+__all__ = ['np', 'torch', 'da', 'fastremap']

ocdkit/array/index.py

@@ -0,0 +1,242 @@
+"""Indexing, slicing, partitioning, and coordinate generation."""
+
+from collections.abc import Iterable
+
+from .imports import *
+
+
+def ravel_index(b, shp):
+    """Row-major flat index from multi-dim indices *b* and array shape *shp*.
+
+    *b* has shape ``(ndim, npts)``; returns a 1D array of length ``npts``.
+    """
+    return np.concatenate((np.asarray(shp[1:])[::-1].cumprod()[::-1], [1])).dot(b)
+
+
+def unravel_index(index, shape):
+    """Multi-dim indices (row-major) from a flat *index* and array *shape*.
+
+    Returns a tuple of per-axis index arrays.
+    """
+    out = []
+    for dim in reversed(shape):
+        out.append(index % dim)
+        index = index // dim
+    return tuple(reversed(out))
+
+
+def border_indices(tyx):
+    """Flat indices of the border pixels of an ND array with shape *tyx*.
+
+    Use as ``A.flat[border_indices(A.shape)]``.
+    """
+    dim_indices = [np.arange(dim_size) for dim_size in tyx]
+    dim_indices = np.meshgrid(*dim_indices, indexing='ij')
+    dim_indices = [indices.ravel() for indices in dim_indices]
+
+    indices = []
+    for i in range(len(tyx)):
+        for j in [0, tyx[i] - 1]:
+            mask = (dim_indices[i] == j)
+            indices.append(np.where(mask)[0])
+    return np.concatenate(indices)
+
+
+def split_array(array, parts, axes=None):
+    """Split an ndarray into *parts* along specified *axes*.
+
+    Parameters
+    ----------
+    array : ndarray
+        The array to split.
+    parts : int or tuple of int
+        Number of parts to split along each axis. If an integer, applies to
+        all specified axes.
+    axes : int, tuple of int, or None
+        The axes to split. If None, splits all axes.
+
+    Returns
+    -------
+    nested list of ndarray
+        Nested sub-arrays after splitting. The nesting depth equals
+        ``len(axes)``.
+    """
+    if isinstance(parts, int):
+        parts = (parts,) * array.ndim
+
+    if axes is None:
+        axes = tuple(range(array.ndim))
+    elif isinstance(axes, int):
+        axes = (axes,)
+
+    if len(parts) != len(axes):
+        raise ValueError("Length of 'parts' must match the number of axes specified.")
+
+    splits = []
+    warnings = []
+
+    for ax, num_parts in zip(axes, parts):
+        dim_size = array.shape[ax]
+        chunk_sizes = [
+            dim_size // num_parts + (1 if i < dim_size % num_parts else 0)
+            for i in range(num_parts)
+        ]
+        if dim_size % num_parts != 0:
+            warnings.append(f"Axis {ax} ({dim_size}) is not evenly divisible by {num_parts}.")
+        split_indices = np.cumsum(chunk_sizes[:-1])
+        splits.append(np.split(np.arange(dim_size), split_indices))
+
+    for warning in warnings:
+        print("Warning:", warning)
+
+    def _recursive_split(array, splits, axes):
+        if not splits:
+            return array
+        ax = axes[0]
+        subarrays = []
+        for idxs in splits[0]:
+            sliced = np.take(array, idxs, axis=ax)
+            subarrays.append(_recursive_split(sliced, splits[1:], axes[1:]))
+        return subarrays
+
+    return _recursive_split(array, splits, axes)
+
+
+def reconstruct_array(nested_list, axes=None):
+    """Reconstruct an ndarray from a nested list produced by :func:`split_array`.
+
+    Parameters
+    ----------
+    nested_list : list of ndarray
+        Nested sub-arrays to concatenate back together.
+    axes : int, tuple of int, or None
+        The axes used for splitting. If None, infers from the outer list's
+        nesting depth.
+    """
+    if axes is None:
+        axes = tuple(range(
+            len(nested_list[0].shape) if isinstance(nested_list[0], np.ndarray)
+            else len(nested_list)
+        ))
+    elif isinstance(axes, int):
+        axes = (axes,)
+
+    def _recursive_reconstruct(nested, level):
+        if isinstance(nested, np.ndarray):
+            return nested
+        if level == len(axes):
+            return np.array(nested)
+        return np.concatenate(
+            [_recursive_reconstruct(sub, level + 1) for sub in nested],
+            axis=axes[level],
+        )
+
+    return _recursive_reconstruct(nested_list, 0)
+
+
+def meshgrid(shape):
+    """Generate a tuple of coordinate grids for an ND array of given *shape*.
+
+    Parameters
+    ----------
+    shape : tuple of int
+        Shape of the ND array (e.g., ``(Y, X)`` for 2D, ``(Z, Y, X)`` for 3D).
+
+    Returns
+    -------
+    tuple of ndarray
+        ``N`` coordinate arrays, one per dimension, in ``ij`` indexing.
+    """
+    ranges = [np.arange(dim) for dim in shape]
+    return np.meshgrid(*ranges, indexing='ij')
+
+
+def generate_flat_coordinates(shape):
+    """Generate flat coordinate arrays for an ND array.
+
+    Parameters
+    ----------
+    shape : tuple of int
+        Shape of the array (e.g., ``(Y, X)`` for 2D).
+
+    Returns
+    -------
+    tuple of ndarray
+        ``N`` 1D arrays containing the coordinates of every element of an
+        array with ``shape``, in ``ij`` order.
+    """
+    return tuple(grid.ravel() for grid in meshgrid(shape))
+
+
+def get_slice_tuple(start, stop, shape, axis=None):
+    """Build a tuple of slice objects for ND array indexing.
+
+    Parameters
+    ----------
+    start : int or iterable of int
+        Starting index(es).
+    stop : int or iterable of int
+        Stopping index(es).
+    shape : tuple of int
+        Shape of the array to slice.
+    axis : int, iterable of int, or None
+        Axis or axes to apply slices to. Default: all axes if *start*/*stop*
+        are iterable, else axis ``0``.
+
+    Returns
+    -------
+    tuple of slice
+        Length ``len(shape)``. Axes not addressed by *axis* get ``slice(None)``.
+    """
+    ndim = len(shape)
+    slices = [slice(None)] * ndim
+
+    if isinstance(start, Iterable) and isinstance(stop, Iterable):
+        if axis is None:
+            axis = list(range(ndim))
+
+        if len(start) != len(stop):
+            raise ValueError("start and stop must be the same length")
+
+        if isinstance(axis, Iterable):
+            if len(axis) != len(start):
+                raise ValueError("axis must be the same length as start and stop")
+        else:
+            axis = [axis] * len(start)
+
+        for a, s, e in zip(axis, start, stop):
+            slices[a] = slice(s, e, None)
+    else:
+        if axis is None:
+            axis = 0
+        slices[axis] = slice(start, stop, None)
+
+    return tuple(slices)
+
+
+def intersect_slices(s1, s2):
+    """Return a slice that is the intersection of *s1* and *s2*.
+
+    ``None`` boundaries are treated as unbounded (``-inf`` / ``+inf``).
+    Steps are unified: the first non-None step wins, defaulting to 1.
+    """
+    import math
+
+    lo1 = -math.inf if s1.start is None else s1.start
+    hi1 = math.inf if s1.stop is None else s1.stop
+    lo2 = -math.inf if s2.start is None else s2.start
+    hi2 = math.inf if s2.stop is None else s2.stop
+
+    new_lo = max(lo1, lo2)
+    new_hi = min(hi1, hi2)
+    if new_hi < new_lo:
+        new_hi = new_lo
+
+    step = s1.step if s1.step is not None else s2.step
+    if step is None:
+        step = 1
+
+    start = None if math.isinf(new_lo) else int(new_lo)
+    stop = None if math.isinf(new_hi) else int(new_hi)
+
+    return slice(start, stop, step)

ocdkit/array/morphology.py

@@ -0,0 +1,194 @@
+"""Morphology utilities — boundaries, skeletonization, outlines, thresholding."""
+
+from .imports import *
+import skimage.morphology
+
+from .spatial import kernel_setup
+
+
+def find_boundaries(labels, connectivity=1, use_symmetry=False):
+    """Compute boundaries of labeled instances in an N-dimensional array.
+
+    Replicates ``skimage.segmentation.find_boundaries`` with
+    ``mode='inner'``, but is much faster.
+    """
+    boundaries = np.zeros_like(labels, dtype=bool)
+    ndim = labels.ndim
+    shape = labels.shape
+
+    steps, inds, idx, fact, sign = kernel_setup(ndim)
+
+    if use_symmetry:
+        allowed_inds = []
+        for i in range(1, 1 + connectivity):
+            j = inds[i][:len(inds[i]) // 2]
+            allowed_inds.append(j)
+        allowed_inds = np.concatenate(allowed_inds)
+    else:
+        allowed_inds = np.concatenate(inds[1:1 + connectivity])
+
+    shifts = steps[allowed_inds]
+
+    if use_symmetry:
+        for shift in shifts:
+            slices_main = tuple(
+                slice(max(-s, 0), min(shape[d] - s, shape[d]))
+                for d, s in enumerate(shift)
+            )
+            slices_shifted = tuple(
+                slice(max(s, 0), min(shape[d] + s, shape[d]))
+                for d, s in enumerate(shift)
+            )
+            boundary_main = (
+                (labels[slices_main] != labels[slices_shifted])
+                & (labels[slices_main] != 0)
+            )
+            boundary_shifted = (
+                (labels[slices_shifted] != labels[slices_main])
+                & (labels[slices_shifted] != 0)
+            )
+            boundaries[slices_main] |= boundary_main
+            boundaries[slices_shifted] |= boundary_shifted
+    else:
+        for shift in shifts:
+            slices_main = tuple(
+                slice(max(-s, 0), min(shape[d] - s, shape[d]))
+                for d, s in enumerate(shift)
+            )
+            slices_shifted = tuple(
+                slice(max(s, 0), min(shape[d] + s, shape[d]))
+                for d, s in enumerate(shift)
+            )
+            boundaries[slices_main] |= (
+                (labels[slices_main] != labels[slices_shifted])
+                & (labels[slices_main] != 0)
+            )
+
+    return boundaries.astype(np.uint8)
+
+
+def skeletonize(labels, dt_thresh=1, dt=None, method='zhang'):
+    """Skeletonize labeled instances, preserving label identity.
+
+    When *dt* is provided, pixels with ``dt > dt_thresh`` are used as the
+    interior mask (fast path).  Otherwise, boundaries are removed first and
+    missing labels are re-attached after thinning.
+
+    Parameters
+    ----------
+    labels : ndarray
+        Integer label matrix.
+    dt_thresh : float
+        Distance-transform threshold for the fast path.
+    dt : ndarray, optional
+        Pre-computed distance transform.
+    method : str
+        Skeletonization method (``'zhang'`` or ``'lee'``).
+    """
+    if dt is not None:
+        inner = dt > dt_thresh
+        skel = skimage.morphology.skeletonize(inner, method=method)
+        return skel * labels
+
+    bd = find_boundaries(labels, connectivity=2)
+    inner = np.logical_xor(labels > 0, bd)
+    skel = skimage.morphology.skeletonize(inner, method=method)
+    skeleton = skel * labels
+
+    original_labels = fastremap.unique(labels)
+    original_labels = original_labels[original_labels != 0]
+
+    skeleton_labels = fastremap.unique(skeleton)
+    skeleton_labels = skeleton_labels[skeleton_labels != 0]
+
+    missing_labels = np.setdiff1d(original_labels, skeleton_labels)
+    missing_labels_mask = fastremap.mask_except(labels, list(missing_labels))
+    skeleton += missing_labels_mask
+
+    return skeleton
+
+
+def masks_to_outlines(masks, omni=False, mode="inner", connectivity=None):
+    """Return a 0/1 outline mask for labeled instances.
+
+    Parameters
+    ----------
+    masks : ndarray
+        2D or 3D label matrix. 3D inputs are processed slice-by-slice.
+    omni : bool
+        If True, use the fast native :func:`find_boundaries` (treats the
+        label matrix directly). If False, use the legacy per-label
+        erosion path (slower, iterates over each label individually).
+    mode : str
+        Forwarded to :func:`find_boundaries` (omni=True) or
+        ``skimage.segmentation.find_boundaries``-equivalent semantics
+        (omni=False, only ``"inner"`` is supported).
+    connectivity : int, optional
+        Forwarded to :func:`find_boundaries` when ``omni=True``. Defaults
+        to ``masks.ndim``.
+    """
+    if masks.ndim > 3 or masks.ndim < 2:
+        raise ValueError(
+            f"masks_to_outlines takes 2D or 3D array, not {masks.ndim}D array"
+        )
+
+    outlines = np.zeros(masks.shape, bool)
+
+    if masks.ndim == 3:
+        for i in range(masks.shape[0]):
+            outlines[i] = masks_to_outlines(
+                masks[i], omni=omni, mode=mode, connectivity=connectivity,
+            )
+        return outlines
+
+    if omni:
+        if connectivity is None:
+            connectivity = masks.ndim
+        return find_boundaries(masks, connectivity=connectivity).astype(bool)
+
+    # Legacy per-label erosion path.
+    from scipy.ndimage import binary_erosion, find_objects
+    slices = find_objects(masks.astype(int))
+    for i, si in enumerate(slices):
+        if si is not None:
+            sr, sc = si
+            mask = (masks[sr, sc] == (i + 1))
+            boundary = mask & ~binary_erosion(mask)
+            pvr, pvc = np.nonzero(boundary)
+            vr, vc = pvr + sr.start, pvc + sc.start
+            outlines[vr, vc] = 1
+    return outlines
+
+
+def hysteresis_threshold(image, low, high):
+    """PyTorch implementation of ``skimage.filters.apply_hysteresis_threshold``.
+
+    Supports 2D and 3D spatial inputs (expects batch+channel dims, i.e.
+    ``(B, C, *spatial)``).  Minor discrepancies vs. skimage occur for very
+    high thresholds on thin objects.
+    """
+    import torch
+
+    if not isinstance(image, torch.Tensor):
+        image = torch.tensor(image)
+
+    mask_low = image > low
+    mask_high = image > high
+    thresholded = mask_low.clone()
+
+    spatial_dims = len(image.shape) - 2
+    kernel_size = [3] * spatial_dims
+    hysteresis_kernel = torch.ones([1, 1] + kernel_size, device=image.device, dtype=image.dtype)
+
+    thresholded_old = torch.zeros_like(thresholded)
+    while (thresholded_old != thresholded).any():
+        if spatial_dims == 2:
+            hysteresis_magnitude = torch.nn.functional.conv2d(thresholded.float(), hysteresis_kernel, padding=1)
+        elif spatial_dims == 3:
+            hysteresis_magnitude = torch.nn.functional.conv3d(thresholded.float(), hysteresis_kernel, padding=1)
+        else:
+            raise ValueError(f'Unsupported number of spatial dimensions: {spatial_dims}')
+        thresholded_old.copy_(thresholded)
+        thresholded = ((hysteresis_magnitude > 0) & mask_low) | mask_high
+
+    return thresholded.bool()