senoquant 1.0.0b1__py3-none-any.whl

senoquant/tabs/segmentation/stardist_onnx_utils/onnx_framework/post/core.py

@@ -0,0 +1,157 @@
+"""Post-processing helpers using StarDist geometry and NMS."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def instances_from_prediction_2d(
+    prob: np.ndarray,
+    dist: np.ndarray,
+    *,
+    grid: tuple[int, int],
+    prob_thresh: float,
+    nms_thresh: float,
+    scale: dict[str, float] | None = None,
+    img_shape: tuple[int, int] | None = None,
+) -> tuple[np.ndarray, dict]:
+    """Create 2D instance labels from StarDist outputs.
+
+    Parameters
+    ----------
+    prob : numpy.ndarray
+        Probability map with shape (Y, X).
+    dist : numpy.ndarray
+        Distance/ray map with shape (Y, X, R), where R is the number of rays.
+    grid : tuple[int, int]
+        Subsampling grid of the model (e.g., (1, 1) or (2, 2)).
+    prob_thresh : float
+        Probability threshold used to filter candidate points before NMS.
+    nms_thresh : float
+        NMS IoU/overlap threshold for suppressing nearby detections.
+    scale : dict[str, float], optional
+        Scale factors applied to the input image before inference. If
+        provided, must include ``"X"`` and ``"Y"`` so that points and
+        distances can be rescaled back to the original image space.
+    img_shape : tuple[int, int], optional
+        Original image shape (Y, X). If provided, the output label image
+        is generated in this shape instead of the scaled prediction shape.
+    Returns
+    -------
+    tuple[numpy.ndarray, dict]
+        Label image of shape (Y, X) and a metadata dict with:
+        - ``points``: center points used for each instance.
+        - ``prob``: per-instance probabilities.
+        - ``dist``: per-instance ray distances.
+
+    Notes
+    -----
+    This function performs non-maximum suppression on the probability map
+    and then rasterizes polygons using the selected points and distances.
+    """
+    from ..._stardist.nms import non_maximum_suppression
+    points, scores, distances = non_maximum_suppression(
+        dist,
+        prob,
+        grid=grid,
+        prob_thresh=prob_thresh,
+        nms_thresh=nms_thresh,
+    )
+    if scale is not None:
+        if not (isinstance(scale, dict) and "X" in scale and "Y" in scale):
+            raise ValueError("scale must be a dictionary with entries for 'X' and 'Y'")
+        rescale = (1 / scale["Y"], 1 / scale["X"])
+        points = points * np.array(rescale).reshape(1, 2)
+    else:
+        rescale = (1, 1)
+    from ..._stardist.geometry.geom2d import polygons_to_label
+    shape = img_shape if img_shape is not None else tuple(
+        s * g for s, g in zip(prob.shape, grid)
+    )
+    labels = polygons_to_label(
+        distances, points, shape=shape, prob=scores, scale_dist=rescale
+    )
+    return labels, {"points": points, "prob": scores, "dist": distances}
+
+
+def instances_from_prediction_3d(
+    prob: np.ndarray,
+    dist: np.ndarray,
+    *,
+    grid: tuple[int, int, int],
+    prob_thresh: float,
+    nms_thresh: float,
+    rays,
+    scale: dict[str, float] | None = None,
+    img_shape: tuple[int, int, int] | None = None,
+) -> tuple[np.ndarray, dict]:
+    """Create 3D instance labels from StarDist outputs.
+
+    Parameters
+    ----------
+    prob : numpy.ndarray
+        Probability map with shape (Z, Y, X).
+    dist : numpy.ndarray
+        Distance/ray map with shape (Z, Y, X, R), where R is the number of rays.
+    grid : tuple[int, int, int]
+        Subsampling grid of the model (e.g., (1, 1, 1) or (2, 2, 2)).
+    prob_thresh : float
+        Probability threshold used to filter candidate points before NMS.
+    nms_thresh : float
+        NMS IoU/overlap threshold for suppressing nearby detections.
+    rays : object
+        StarDist 3D rays object describing ray directions and sampling.
+    scale : dict[str, float], optional
+        Scale factors applied to the input image before inference. If
+        provided, must include ``"X"``, ``"Y"``, and ``"Z"`` so that points
+        and rays can be rescaled back to the original image space.
+    img_shape : tuple[int, int, int], optional
+        Original image shape (Z, Y, X). If provided, the output label
+        volume is generated in this shape instead of the scaled prediction
+        shape.
+    Returns
+    -------
+    tuple[numpy.ndarray, dict]
+        Label volume of shape (Z, Y, X) and a metadata dict with:
+        - ``points``: center points used for each instance.
+        - ``prob``: per-instance probabilities.
+        - ``dist``: per-instance ray distances.
+
+    Notes
+    -----
+    This function performs non-maximum suppression in 3D and then
+    rasterizes polyhedra using the selected points and distances. The
+    Python backend uses an axis-aligned bounding-box approximation.
+    """
+    from ..._stardist.nms import non_maximum_suppression_3d
+    points, scores, distances = non_maximum_suppression_3d(
+        dist,
+        prob,
+        rays,
+        grid=grid,
+        prob_thresh=prob_thresh,
+        nms_thresh=nms_thresh,
+    )
+    if scale is not None:
+        if not (
+            isinstance(scale, dict)
+            and "X" in scale
+            and "Y" in scale
+            and "Z" in scale
+        ):
+            raise ValueError(
+                "scale must be a dictionary with entries for 'X', 'Y', and 'Z'"
+            )
+        rescale = (1 / scale["Z"], 1 / scale["Y"], 1 / scale["X"])
+        points = points * np.array(rescale).reshape(1, 3)
+        rays = rays.copy(scale=rescale)
+    else:
+        rescale = (1, 1, 1)
+    from ..._stardist.geometry.geom3d import polyhedron_to_label
+    shape = img_shape if img_shape is not None else tuple(
+        s * g for s, g in zip(prob.shape, grid)
+    )
+    labels = polyhedron_to_label(
+        distances, points, rays=rays, shape=shape, prob=scores, verbose=False
+    )
+    return labels, {"points": points, "prob": scores, "dist": distances}

senoquant/tabs/segmentation/stardist_onnx_utils/onnx_framework/pre/__init__.py

@@ -0,0 +1,17 @@
+"""Preprocessing utilities for ONNX inference."""
+
+from .core import (
+    normalize,
+    pad_for_tiling,
+    pad_to_multiple,
+    unpad_to_shape,
+    validate_image,
+)
+
+__all__ = [
+    "normalize",
+    "pad_for_tiling",
+    "pad_to_multiple",
+    "unpad_to_shape",
+    "validate_image",
+]

senoquant/tabs/segmentation/stardist_onnx_utils/onnx_framework/pre/core.py

@@ -0,0 +1,226 @@
+"""Preprocessing helpers for ONNX StarDist inference.
+
+These helpers mirror the high-level behavior of StarDist's preprocessing:
+validate input dimensionality, optionally normalize, and pad to the grid
+required by the network. The helpers assume single-channel images with
+spatial axes ordered as YX (2D) or ZYX (3D).
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+
+def validate_image(image: np.ndarray) -> None:
+    """Validate input image dimensionality.
+
+    Parameters
+    ----------
+    image : numpy.ndarray
+        Input image array expected to be 2D (YX) or 3D (ZYX) and single-channel.
+
+    Raises
+    ------
+    ValueError
+        If the input is not 2D or 3D.
+    """
+    if image.ndim not in (2, 3):
+        raise ValueError("Input image must be 2D (YX) or 3D (ZYX).")
+
+
+def normalize(
+    image: np.ndarray,
+    pmin: float = 1.0,
+    pmax: float = 99.8,
+    eps: float = 1e-6,
+) -> np.ndarray:
+    """Percentile normalize a single-channel image.
+
+    Parameters
+    ----------
+    image : numpy.ndarray
+        Input image to normalize.
+    pmin : float, optional
+        Lower percentile for normalization. Default is 1.0.
+    pmax : float, optional
+        Upper percentile for normalization. Default is 99.8.
+    eps : float, optional
+        Small constant to avoid division by zero. Default is 1e-6.
+
+    Returns
+    -------
+    numpy.ndarray
+        Normalized float32 image with values clipped to [0, 1].
+
+    Raises
+    ------
+    ValueError
+        If ``pmax`` is not greater than ``pmin``.
+    """
+    if pmax <= pmin:
+        raise ValueError("pmax must be greater than pmin.")
+    lo = np.percentile(image, pmin, keepdims=True)
+    hi = np.percentile(image, pmax, keepdims=True)
+    scale = hi - lo
+    scale = np.where(scale < eps, 1.0, scale)
+    normalized = (image - lo) / scale
+    return np.clip(normalized, 0.0, 1.0).astype(np.float32)
+
+
+def pad_to_multiple(
+    image: np.ndarray,
+    multiples: tuple[int, ...],
+    mode: str = "reflect",
+) -> tuple[np.ndarray, tuple[tuple[int, int], ...]]:
+    """Pad each axis so its length is divisible by the corresponding multiple.
+
+    Parameters
+    ----------
+    image : numpy.ndarray
+        Input image (2D or 3D).
+    multiples : tuple[int, ...]
+        Per-axis divisibility constraints (e.g., the StarDist grid).
+    mode : str, optional
+        Padding mode passed to ``numpy.pad``. Default is "reflect".
+
+    Returns
+    -------
+    numpy.ndarray
+        Padded image as float32.
+    tuple[tuple[int, int], ...]
+        Padding applied per axis as (before, after) pairs. This implementation
+        only pads at the end of each axis (before = 0).
+
+    Raises
+    ------
+    ValueError
+        If the multiples do not match the image dimensionality or contain
+        non-positive values.
+    """
+    validate_image(image)
+    if len(multiples) != image.ndim:
+        raise ValueError("Multiples must match image dimensionality.")
+    pads = []
+    for dim, mult in zip(image.shape, multiples):
+        if mult <= 0:
+            raise ValueError("Multiples must be positive.")
+        pad_after = (mult - (dim % mult)) % mult
+        pads.append((0, pad_after))
+    padded = np.pad(image, pads, mode=mode)
+    return padded.astype(np.float32, copy=False), tuple(pads)
+
+
+def unpad_to_shape(
+    data: np.ndarray,
+    pads: tuple[tuple[int, int], ...],
+    scale: tuple[int, ...] | None = None,
+) -> np.ndarray:
+    """Crop padded output back to the unpadded shape (accounting for scale).
+
+    Parameters
+    ----------
+    data : numpy.ndarray
+        Output array to crop (e.g., probability or distance map).
+    pads : tuple[tuple[int, int], ...]
+        Padding applied to the input image as returned by ``pad_to_multiple``.
+    scale : tuple[int, ...] or None, optional
+        Per-axis scale factor between input and output spatial grids.
+        For StarDist, this is typically the grid (e.g., (1, 1) or (2, 2)).
+        If None, assumes scale 1 for all axes.
+
+    Returns
+    -------
+    numpy.ndarray
+        Cropped array with padding removed.
+
+    Raises
+    ------
+    ValueError
+        If scale dimensionality does not match pads, or if non-zero
+        pre-padding is provided (this function assumes end-padding only).
+    """
+    if scale is None:
+        scale = (1,) * len(pads)
+    if len(scale) != len(pads):
+        raise ValueError("Scale must match pad dimensionality.")
+    slices = []
+    for (before, after), mult in zip(pads, scale):
+        if before != 0:
+            raise ValueError("Only end-padding is supported.")
+        if after >= mult:
+            slices.append(slice(0, -(after // mult)))
+        else:
+            slices.append(slice(None))
+    slices.extend([slice(None)] * (data.ndim - len(pads)))
+    return data[tuple(slices)]
+
+
+def pad_for_tiling(
+    image: np.ndarray,
+    grid: tuple[int, ...],
+    tile_shape: tuple[int, ...],
+    overlap: tuple[int, ...],
+    div_by: tuple[int, ...] | None = None,
+    mode: str = "reflect",
+) -> tuple[np.ndarray, tuple[tuple[int, int], ...]]:
+    """Pad input so tiled prediction aligns with overlap and divisibility.
+
+    Parameters
+    ----------
+    image : numpy.ndarray
+        Input image (2D or 3D).
+    grid : tuple[int, ...]
+        Model output grid/stride per axis.
+    tile_shape : tuple[int, ...]
+        Spatial size of each tile in input pixels.
+    overlap : tuple[int, ...]
+        Overlap per axis in input pixels.
+    div_by : tuple[int, ...] or None, optional
+        Additional per-axis divisibility constraints. If None, uses ``grid``.
+    mode : str, optional
+        Padding mode passed to ``numpy.pad``. Default is "reflect".
+
+    Returns
+    -------
+    numpy.ndarray
+        Padded image as float32.
+    tuple[tuple[int, int], ...]
+        Padding applied per axis as (before, after) pairs. This implementation
+        only pads at the end of each axis (before = 0).
+
+    Raises
+    ------
+    ValueError
+        If shapes are inconsistent or overlap is invalid.
+    """
+    if len(grid) != image.ndim:
+        raise ValueError("Grid must match image dimensionality.")
+    if len(tile_shape) != image.ndim:
+        raise ValueError("tile_shape must match image dimensionality.")
+    if len(overlap) != image.ndim:
+        raise ValueError("overlap must match image dimensionality.")
+
+    if div_by is None:
+        div_by = grid
+    if len(div_by) != image.ndim:
+        raise ValueError("div_by must match image dimensionality.")
+
+    pads = []
+    for dim, g, ts, ov, d in zip(image.shape, grid, tile_shape, overlap, div_by):
+        step = ts - ov
+        if step <= 0:
+            raise ValueError("overlap must be smaller than tile size.")
+        if dim <= ts:
+            target = ts
+        else:
+            target = ts + math.ceil((dim - ts) / step) * step
+        div_req = math.lcm(int(g), int(d))
+        if div_req > 1:
+            while target % div_req != 0:
+                target += step
+        pads.append((0, int(max(0, target - dim))))
+
+    padded = np.pad(image, pads, mode=mode)
+    return padded.astype(np.float32, copy=False), tuple(pads)

senoquant/tabs/segmentation/stardist_onnx_utils/onnx_framework/predict/__init__.py

@@ -0,0 +1,5 @@
+"""Prediction utilities for ONNX StarDist inference."""
+
+from .core import TilingSpec, default_tiling_spec, predict_tiled
+
+__all__ = ["TilingSpec", "default_tiling_spec", "predict_tiled"]