senoquant 1.0.0b1__py3-none-any.whl

senoquant/tabs/segmentation/stardist_onnx_utils/onnx_framework/inspect/probe.py

@@ -0,0 +1,100 @@
+"""Utilities for creating valid probe inputs for ONNX inspection."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+from .divisibility import infer_div_by
+from .valid_sizes import infer_valid_size_patterns_from_path, snap_size
+
+
+def make_probe_image(
+    image: np.ndarray,
+    *,
+    model_path: Path | None = None,
+    input_layout: str | None = None,
+    div_by_cache: dict[Path, tuple[int, ...]] | None = None,
+    valid_size_cache: dict[Path, list[object] | None] | None = None,
+) -> np.ndarray:
+    """Create a small probe image aligned with ONNX size constraints.
+
+    Parameters
+    ----------
+    image : numpy.ndarray
+        Input image array used to derive probe size.
+    model_path : pathlib.Path or None, optional
+        ONNX model path used for inspecting size constraints.
+    input_layout : str or None, optional
+        Model input layout (e.g., "NHWC", "NDHWC") used for size inspection.
+    div_by_cache : dict or None, optional
+        Cache for divisibility requirements keyed by model path.
+    valid_size_cache : dict or None, optional
+        Cache for valid size patterns keyed by model path.
+
+    Returns
+    -------
+    numpy.ndarray
+        Probe image padded/cropped to a valid spatial size.
+    """
+    target = 256 if image.ndim == 2 else 64
+    probe_shape = []
+    for dim in image.shape:
+        size = min(dim, target)
+        if size >= 16:
+            size = size - (size % 16)
+            if size == 0:
+                size = min(dim, target)
+        probe_shape.append(max(1, size))
+
+    probe = image[tuple(slice(0, s) for s in probe_shape)]
+
+    if model_path is None or input_layout is None:
+        return probe
+
+    patterns = None
+    if valid_size_cache is not None:
+        patterns = valid_size_cache.get(model_path)
+    if patterns is None:
+        try:
+            patterns = infer_valid_size_patterns_from_path(
+                model_path,
+                input_layout,
+                image.ndim,
+            )
+        except Exception:
+            patterns = None
+        if valid_size_cache is not None:
+            valid_size_cache[model_path] = patterns
+
+    div_by = None
+    if div_by_cache is not None:
+        div_by = div_by_cache.get(model_path)
+    if div_by is None:
+        try:
+            div_by = infer_div_by(model_path, ndim=image.ndim)
+        except Exception:
+            div_by = None
+        if div_by_cache is not None and div_by is not None:
+            div_by_cache[model_path] = div_by
+
+    desired = list(probe.shape)
+    if patterns:
+        desired = [
+            max(1, snap_size(int(size), patterns[axis]))
+            for axis, size in enumerate(desired)
+        ]
+    elif div_by:
+        desired = [
+            max(int(d), (int(size) // int(d)) * int(d)) if d else int(size)
+            for size, d in zip(desired, div_by)
+        ]
+
+    desired = [max(1, int(size)) for size in desired]
+    crop_slices = tuple(slice(0, min(s, d)) for s, d in zip(probe.shape, desired))
+    probe = probe[crop_slices]
+    pads = [(0, max(0, d - s)) for s, d in zip(probe.shape, desired)]
+    if any(pad_after > 0 for _, pad_after in pads):
+        probe = np.pad(probe, pads, mode="reflect")
+    return probe

senoquant/tabs/segmentation/stardist_onnx_utils/onnx_framework/inspect/receptive_field.py

@@ -0,0 +1,182 @@
+"""Empirically estimate receptive field and tiling overlap for ONNX models.
+
+This module mirrors StarDist's empirical receptive-field estimation:
+run the model once on a single-pixel impulse and once on zeros, then
+measure the spatial support of the difference in the probability output.
+The measured extents define the overlap needed to avoid tile boundary
+artifacts in tiled prediction.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+
+from .divisibility import infer_div_by
+
+
+def infer_receptive_field(
+    model_path: str | Path,
+    ndim: int | None = None,
+    input_shape: tuple[int, ...] | None = None,
+    eps: float = 0.0,
+) -> tuple[tuple[int, int], ...]:
+    """Estimate the receptive field via impulse response.
+
+    This mirrors StarDist's empirical receptive-field estimation: run the model
+    on an impulse image and on zeros, then find the spatial support of the
+    difference.
+
+    Parameters
+    ----------
+    model_path : str or pathlib.Path
+        Path to the ONNX model file.
+    ndim : int or None, optional
+        Spatial dimensionality (2 or 3). If None, inferred from input rank.
+    input_shape : tuple[int, ...] or None, optional
+        Spatial shape for the probe input. If None, a power-of-two shape is
+        chosen and adjusted to satisfy the inferred divisibility.
+    eps : float, optional
+        Threshold used to detect non-zero influence in the output. Default 0.0.
+
+    Returns
+    -------
+    tuple[tuple[int, int], ...]
+        Per-axis receptive field extents as (left, right) offsets from the
+        center pixel/voxel in input coordinates.
+
+    Notes
+    -----
+    - The probe uses an impulse (single 1.0) at the spatial center.
+    - The probability output is selected by heuristics (last dim == 1).
+    - Output is mapped back to input resolution using the inferred grid.
+    """
+    import onnxruntime as ort
+    from scipy.ndimage import zoom
+
+    model_path = Path(model_path)
+    session = ort.InferenceSession(str(model_path))
+
+    input_name = session.get_inputs()[0].name
+    output_names = [out.name for out in session.get_outputs()]
+
+    if ndim is None:
+        ndim = _infer_ndim_from_input(session)
+
+    if input_shape is None:
+        # Choose a reasonable power-of-two probe size and round up to a
+        # multiple of the inferred divisibility to avoid internal mismatches.
+        base = 256 if ndim == 2 else 64
+        div_by = infer_div_by(model_path, ndim=ndim)
+        input_shape = tuple(_round_up(base, d) for d in div_by)
+
+    if len(input_shape) != ndim:
+        raise ValueError("input_shape must match ndim.")
+
+    # Build impulse and zero inputs (NHWC/NDHWC).
+    center = tuple(s // 2 for s in input_shape)
+    x = np.zeros((1, *input_shape, 1), dtype=np.float32)
+    z = np.zeros_like(x)
+    x[(0, *center, 0)] = 1.0
+
+    # Run the model and extract the probability output.
+    y = _run_prob(session, output_names, input_name, x, ndim)
+    y0 = _run_prob(session, output_names, input_name, z, ndim)
+
+    # Infer grid from input/output shapes (input / output per axis).
+    grid = tuple(
+        max(1, int(round(si / so))) for si, so in zip(input_shape, y.shape)
+    )
+    y = zoom(y, grid, order=0)
+    y0 = zoom(y0, grid, order=0)
+
+    # Measure where the response differs from zero.
+    diff = np.abs(y - y0) > eps
+    indices = np.where(diff)
+    if any(len(i) == 0 for i in indices):
+        raise RuntimeError("Failed to detect receptive field; try a larger input_shape.")
+
+    return tuple((c - int(np.min(i)), int(np.max(i)) - c) for c, i in zip(center, indices))
+
+
+def recommend_tile_overlap(
+    model_path: str | Path,
+    ndim: int | None = None,
+    input_shape: tuple[int, ...] | None = None,
+    eps: float = 0.0,
+) -> tuple[int, ...]:
+    """Return recommended tile overlap per axis from empirical RF.
+
+    Parameters
+    ----------
+    model_path : str or pathlib.Path
+        Path to the ONNX model file.
+    ndim : int or None, optional
+        Spatial dimensionality (2 or 3). If None, inferred from input rank.
+    input_shape : tuple[int, ...] or None, optional
+        Spatial probe input shape. If None, a default shape is used.
+    eps : float, optional
+        Threshold used to detect non-zero influence in the output.
+
+    Returns
+    -------
+    tuple[int, ...]
+        Per-axis overlap in input pixels.
+    """
+    rf = infer_receptive_field(
+        model_path=model_path,
+        ndim=ndim,
+        input_shape=input_shape,
+        eps=eps,
+    )
+    return tuple(max(pair) for pair in rf)
+
+
+def _run_prob(session, output_names, input_name, input_tensor, ndim: int) -> np.ndarray:
+    """Run the ONNX model and return the probability output in spatial layout."""
+    outputs = session.run(output_names, {input_name: input_tensor})
+    prob = _select_prob_output(outputs)
+    prob = _to_spatial(prob, ndim)
+    return prob
+
+
+def _select_prob_output(outputs: list[np.ndarray]) -> np.ndarray:
+    """Pick the probability output from ONNX outputs."""
+    for arr in outputs:
+        if arr.ndim >= 4 and arr.shape[-1] == 1:
+            return arr
+    return outputs[0]
+
+
+def _to_spatial(prob: np.ndarray, ndim: int) -> np.ndarray:
+    """Convert a batched prob tensor into spatial layout (YX/ZYX)."""
+    if ndim == 2:
+        if prob.ndim == 4 and prob.shape[-1] == 1:
+            return prob[0, ..., 0]
+        if prob.ndim == 4 and prob.shape[1] == 1:
+            return prob[0, 0, ...]
+    if ndim == 3:
+        if prob.ndim == 5 and prob.shape[-1] == 1:
+            return prob[0, ..., 0]
+        if prob.ndim == 5 and prob.shape[1] == 1:
+            return prob[0, 0, ...]
+    raise ValueError("Unsupported prob output layout.")
+
+
+def _infer_ndim_from_input(session) -> int:
+    """Infer spatial dimensionality from ONNX session input rank."""
+    shape = session.get_inputs()[0].shape
+    if len(shape) == 4:
+        return 2
+    if len(shape) == 5:
+        return 3
+    raise ValueError(f"Unsupported input rank {len(shape)}.")
+
+
+def _round_up(value: int, multiple: int) -> int:
+    """Round up ``value`` to the next multiple."""
+    if multiple <= 0:
+        return value
+    return int(np.ceil(value / multiple) * multiple)

senoquant/tabs/segmentation/stardist_onnx_utils/onnx_framework/inspect/rf_cli.py

@@ -0,0 +1,48 @@
+"""CLI for empirical receptive-field estimation."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from .receptive_field import infer_receptive_field, recommend_tile_overlap
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Estimate ONNX receptive field.")
+    parser.add_argument("model", type=Path, help="Path to the ONNX model.")
+    parser.add_argument("--ndim", type=int, choices=(2, 3), default=None)
+    parser.add_argument(
+        "--shape",
+        type=int,
+        nargs="+",
+        default=None,
+        help="Spatial input shape (e.g. --shape 256 256 or --shape 64 64 64).",
+    )
+    parser.add_argument("--eps", type=float, default=0.0)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    input_shape = tuple(args.shape) if args.shape else None
+    rf = infer_receptive_field(
+        model_path=args.model,
+        ndim=args.ndim,
+        input_shape=input_shape,
+        eps=args.eps,
+    )
+    overlap = recommend_tile_overlap(
+        model_path=args.model,
+        ndim=args.ndim,
+        input_shape=input_shape,
+        eps=args.eps,
+    )
+
+    print(f"Model: {args.model}")
+    print(f"Receptive field: {rf}")
+    print(f"Recommended overlap: {overlap}")
+
+
+if __name__ == "__main__":
+    main()

senoquant/tabs/segmentation/stardist_onnx_utils/onnx_framework/inspect/valid_sizes.py

@@ -0,0 +1,278 @@
+"""Empirically infer valid spatial sizes for ONNX model inputs.
+
+This module probes the ONNX runtime by running the model on small inputs and
+recording which spatial sizes succeed. It then summarizes valid sizes as
+periodic residues (e.g., sizes of the form ``16k`` or ``16k+1``).
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Iterable, Sequence
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class ValidSizePattern:
+    """Periodic validity pattern for a single spatial axis.
+
+    Attributes
+    ----------
+    period : int
+        Periodicity for valid sizes.
+    residues : tuple[int, ...]
+        Allowed ``size % period`` residues.
+    min_valid : int
+        Smallest observed valid size.
+    """
+
+    period: int
+    residues: tuple[int, ...]
+    min_valid: int
+
+
+def infer_valid_size_patterns(
+    session,
+    input_name: str,
+    output_names: Iterable[str],
+    input_layout: str,
+    ndim: int,
+    max_probe: int = 64,
+) -> list[ValidSizePattern]:
+    """Probe ONNX execution to infer valid size residues per axis.
+
+    Parameters
+    ----------
+    session : onnxruntime.InferenceSession
+        ONNX Runtime session used to execute the model.
+    input_name : str
+        Name of the ONNX input tensor.
+    output_names : Iterable[str]
+        Output tensor names to request during inference.
+    input_layout : str
+        Input layout string (e.g., ``"NHWC"`` or ``"NDHWC"``).
+    ndim : int
+        Spatial dimensionality (2 or 3).
+    max_probe : int, optional
+        Maximum size to probe per axis. Default is 64.
+
+    Returns
+    -------
+    list[ValidSizePattern]
+        One entry per axis describing periodic valid size residues.
+
+    Raises
+    ------
+    RuntimeError
+        If no valid sizes can be found within the probe range.
+    """
+    if ndim not in (2, 3):
+        raise ValueError("ndim must be 2 or 3.")
+
+    output_names = list(output_names)
+
+    base = _find_valid_base(
+        session, input_name, output_names, input_layout, ndim, max_probe
+    )
+
+    patterns: list[ValidSizePattern] = []
+    any_valid = False
+    for axis in range(ndim):
+        valid = []
+        for size in range(1, max_probe + 1):
+            shape = [base] * ndim
+            shape[axis] = size
+            if _try_run(session, input_name, output_names, input_layout, shape):
+                valid.append(size)
+        if not valid:
+            patterns.append(ValidSizePattern(period=1, residues=(0,), min_valid=1))
+            continue
+        any_valid = True
+        period, residues = _infer_period_and_residues(valid, max_probe)
+        patterns.append(
+            ValidSizePattern(
+                period=int(period),
+                residues=tuple(int(r) for r in residues),
+                min_valid=int(min(valid)),
+            )
+        )
+
+    if not any_valid:
+        raise RuntimeError(
+            f"No valid sizes found within 1..{max_probe} for any axis."
+        )
+    return patterns
+
+
+def infer_valid_size_patterns_from_path(
+    model_path,
+    input_layout: str,
+    ndim: int,
+    max_probe: int = 64,
+) -> list[ValidSizePattern]:
+    """Probe valid sizes using a temporary, quiet ONNX session.
+
+    Parameters
+    ----------
+    model_path : str or pathlib.Path
+        Path to the ONNX model file.
+    input_layout : str
+        Input layout string (e.g., ``"NHWC"`` or ``"NDHWC"``).
+    ndim : int
+        Spatial dimensionality (2 or 3).
+    max_probe : int, optional
+        Maximum size to probe per axis. Default is 64.
+
+    Returns
+    -------
+    list[ValidSizePattern]
+        One entry per axis describing periodic valid size residues.
+    """
+    import onnxruntime as ort
+
+    sess_options = ort.SessionOptions()
+    # Suppress ORT error logs during probe failures.
+    sess_options.log_severity_level = 4
+    session = ort.InferenceSession(str(model_path), sess_options=sess_options)
+    input_name = session.get_inputs()[0].name
+    output_names = [o.name for o in session.get_outputs()]
+    return infer_valid_size_patterns(
+        session,
+        input_name,
+        output_names,
+        input_layout,
+        ndim,
+        max_probe=max_probe,
+    )
+
+
+def _find_valid_base(
+    session,
+    input_name: str,
+    output_names: list[str],
+    input_layout: str,
+    ndim: int,
+    max_probe: int,
+) -> int:
+    """Return the smallest symmetric size that executes successfully."""
+    for size in range(1, max_probe + 1):
+        shape = [size] * ndim
+        if _try_run(session, input_name, output_names, input_layout, shape):
+            return size
+    raise RuntimeError("Failed to find any valid base size for probing.")
+
+
+def _try_run(
+    session,
+    input_name: str,
+    output_names: list[str],
+    input_layout: str,
+    spatial_shape: list[int],
+) -> bool:
+    """Return True if the model runs on the given spatial shape."""
+    if input_layout in ("NHWC", "NDHWC"):
+        input_tensor = np.zeros((1, *spatial_shape, 1), dtype=np.float32)
+    elif input_layout in ("NCHW", "NCDHW"):
+        input_tensor = np.zeros((1, 1, *spatial_shape), dtype=np.float32)
+    else:
+        raise ValueError(f"Unsupported input layout {input_layout}.")
+    try:
+        session.run(list(output_names), {input_name: input_tensor})
+    except Exception:
+        return False
+    return True
+
+
+def _infer_period_and_residues(
+    valid_sizes: list[int], max_probe: int
+) -> tuple[int, list[int]]:
+    """Infer periodicity and residue set for valid sizes."""
+    valid_set = set(valid_sizes)
+    if not valid_set:
+        return 1, [0]
+
+    min_valid = min(valid_set)
+    for period in range(1, max_probe + 1):
+        residues = {v % period for v in valid_set}
+        ok = True
+        for size in range(min_valid, max_probe + 1):
+            if (size % period in residues) != (size in valid_set):
+                ok = False
+                break
+        if ok:
+            return period, sorted(residues)
+
+    if len(valid_sizes) < 2:
+        return max(1, valid_sizes[0]), [valid_sizes[0] % max(1, valid_sizes[0])]
+
+    diffs = [b - a for a, b in zip(valid_sizes, valid_sizes[1:]) if b > a]
+    period = diffs[0]
+    for d in diffs[1:]:
+        period = math.gcd(period, d)
+    residues = sorted({v % period for v in valid_set})
+    return max(1, period), residues
+
+
+def snap_size(size: int, pattern: ValidSizePattern) -> int:
+    """Adjust a size to the nearest valid residue at or below ``size``.
+
+    Parameters
+    ----------
+    size : int
+        Proposed size.
+    pattern : ValidSizePattern
+        Valid size pattern for the axis.
+
+    Returns
+    -------
+    int
+        Snapped valid size.
+    """
+    period = max(1, int(pattern.period))
+    residues = set(int(r) for r in pattern.residues)
+    min_valid = int(pattern.min_valid)
+    if size <= min_valid:
+        return min_valid
+    for delta in range(period + 1):
+        candidate = size - delta
+        if candidate < min_valid:
+            break
+        if candidate % period in residues:
+            return candidate
+    candidate = size
+    while candidate % period not in residues:
+        candidate += 1
+    return candidate
+
+
+def snap_shape(
+    shape: Sequence[int],
+    patterns: Sequence[ValidSizePattern],
+    *,
+    skip_axes: Sequence[int] = (),
+) -> tuple[int, ...]:
+    """Snap each axis of a shape to the nearest valid size.
+
+    Parameters
+    ----------
+    shape : Sequence[int]
+        Proposed spatial shape.
+    patterns : Sequence[ValidSizePattern]
+        Per-axis valid size patterns.
+    skip_axes : Sequence[int], optional
+        Axes to leave unchanged (e.g., skip Z for 3D models).
+
+    Returns
+    -------
+    tuple[int, ...]
+        Snapped spatial shape.
+    """
+    snapped = []
+    for axis, size in enumerate(shape):
+        if axis in skip_axes:
+            snapped.append(int(size))
+            continue
+        snapped.append(snap_size(int(size), patterns[axis]))
+    return tuple(snapped)

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

@@ -0,0 +1,8 @@
+"""Post-processing utilities for ONNX StarDist inference."""
+
+from .core import instances_from_prediction_2d, instances_from_prediction_3d
+
+__all__ = [
+    "instances_from_prediction_2d",
+    "instances_from_prediction_3d",
+]