senoquant 1.0.0b1__py3-none-any.whl

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

@@ -0,0 +1,45 @@
+"""ONNX tiling and prediction framework for StarDist."""
+
+__all__ = [
+    "normalize",
+    "pad_for_tiling",
+    "pad_to_multiple",
+    "unpad_to_shape",
+    "TilingSpec",
+    "default_tiling_spec",
+    "predict_tiled",
+    "instances_from_prediction_2d",
+    "instances_from_prediction_3d",
+    "DEFAULT_2D_MODEL",
+    "DEFAULT_3D_MODEL",
+    "convert_model_to_onnx",
+    "convert_pretrained_2d",
+    "convert_pretrained_3d",
+    "infer_div_by",
+    "summarize_model_io",
+]
+
+
+def __getattr__(name):
+    if name in {"normalize", "pad_for_tiling", "pad_to_multiple", "unpad_to_shape"}:
+        from . import pre as _pre
+        return getattr(_pre, name)
+    if name in {"TilingSpec", "default_tiling_spec", "predict_tiled"}:
+        from . import predict as _predict
+        return getattr(_predict, name)
+    if name in {"instances_from_prediction_2d", "instances_from_prediction_3d"}:
+        from . import post as _post
+        return getattr(_post, name)
+    if name in {
+        "DEFAULT_2D_MODEL",
+        "DEFAULT_3D_MODEL",
+        "convert_model_to_onnx",
+        "convert_pretrained_2d",
+        "convert_pretrained_3d",
+    }:
+        from . import convert as _convert
+        return getattr(_convert, name)
+    if name in {"infer_div_by", "summarize_model_io"}:
+        from . import inspect as _inspect
+        return getattr(_inspect, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

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

@@ -0,0 +1,17 @@
+"""ONNX conversion helpers for StarDist models."""
+
+from .core import (
+    DEFAULT_2D_MODEL,
+    DEFAULT_3D_MODEL,
+    convert_model_to_onnx,
+    convert_pretrained_2d,
+    convert_pretrained_3d,
+)
+
+__all__ = [
+    "DEFAULT_2D_MODEL",
+    "DEFAULT_3D_MODEL",
+    "convert_model_to_onnx",
+    "convert_pretrained_2d",
+    "convert_pretrained_3d",
+]

senoquant/tabs/segmentation/stardist_onnx_utils/onnx_framework/convert/cli.py

@@ -0,0 +1,55 @@
+"""CLI for converting StarDist models to ONNX."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from .core import (
+    DEFAULT_2D_MODEL,
+    DEFAULT_3D_MODEL,
+    convert_pretrained_2d,
+    convert_pretrained_3d,
+)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Convert StarDist models to ONNX.")
+    parser.add_argument(
+        "--dim",
+        choices=("2", "3", "2d", "3d"),
+        default="2d",
+        help="Model dimensionality.",
+    )
+    parser.add_argument(
+        "--model",
+        default=None,
+        help="Pretrained model name/alias or model directory path.",
+    )
+    parser.add_argument(
+        "--output",
+        default=".",
+        help="Output directory or ONNX file path.",
+    )
+    parser.add_argument(
+        "--opset",
+        type=int,
+        default=18,
+        help="ONNX opset version to export.",
+    )
+    args = parser.parse_args()
+
+    dim = 2 if args.dim in ("2", "2d") else 3
+    model_name = args.model or (DEFAULT_2D_MODEL if dim == 2 else DEFAULT_3D_MODEL)
+    output = Path(args.output)
+
+    if dim == 2:
+        path = convert_pretrained_2d(model_name, output, opset=args.opset)
+    else:
+        path = convert_pretrained_3d(model_name, output, opset=args.opset)
+
+    print(f"Saved ONNX model to {path}")
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,285 @@
+"""Convert StarDist Keras models to ONNX."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import re
+import sys
+import types
+import importlib
+import tempfile
+
+
+DEFAULT_2D_MODEL = "2D_versatile_fluo"
+DEFAULT_3D_MODEL = "3D_demo"
+
+
+def convert_pretrained_2d(
+    model_name: str = DEFAULT_2D_MODEL,
+    output: str | Path = ".",
+    *,
+    opset: int = 18,
+) -> Path:
+    """Convert a pretrained StarDist2D model to ONNX.
+
+    Parameters
+    ----------
+    model_name : str, optional
+        Pretrained model name or alias. Defaults to ``2D_versatile_fluo``.
+    output : str or pathlib.Path, optional
+        Output directory or ONNX file path. Defaults to the current directory.
+    opset : int, optional
+        ONNX opset version to export. Defaults to 13.
+
+    Returns
+    -------
+    pathlib.Path
+        Path to the saved ONNX model.
+    """
+    model = _load_stardist_model(2, model_name)
+    output_path = _resolve_output_path(output, f"stardist2d_{_safe_name(model_name)}.onnx")
+    return convert_model_to_onnx(model, output_path, opset=opset)
+
+
+def convert_pretrained_3d(
+    model_name: str = DEFAULT_3D_MODEL,
+    output: str | Path = ".",
+    *,
+    opset: int = 18,
+) -> Path:
+    """Convert a pretrained StarDist3D model to ONNX.
+
+    Parameters
+    ----------
+    model_name : str, optional
+        Pretrained model name or alias. Defaults to ``3D_demo``.
+    output : str or pathlib.Path, optional
+        Output directory or ONNX file path. Defaults to the current directory.
+    opset : int, optional
+        ONNX opset version to export. Defaults to 13.
+
+    Returns
+    -------
+    pathlib.Path
+        Path to the saved ONNX model.
+    """
+    model = _load_stardist_model(3, model_name)
+    output_path = _resolve_output_path(output, f"stardist3d_{_safe_name(model_name)}.onnx")
+    return convert_model_to_onnx(model, output_path, opset=opset)
+
+
+def convert_model_to_onnx(model, output_path: str | Path, *, opset: int = 18) -> Path:
+    """Convert a StarDist model instance to ONNX.
+
+    Parameters
+    ----------
+    model : object
+        StarDist2D or StarDist3D instance with a ``keras_model`` attribute.
+    output_path : str or pathlib.Path
+        File path to save the ONNX model.
+    opset : int, optional
+        ONNX opset version to export. Defaults to 13.
+
+    Returns
+    -------
+    pathlib.Path
+        Path to the saved ONNX model.
+    """
+    tf = _import_tensorflow()
+    tf2onnx = _import_tf2onnx()
+
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    keras_model = model.keras_model
+    keras_model.trainable = False
+
+    input_tensor = keras_model.inputs[0]
+    input_name = input_tensor.name.split(":")[0]
+    input_shape = list(input_tensor.shape)
+    if input_shape and input_shape[0] is None:
+        input_shape[0] = 1
+    input_signature = (tf.TensorSpec(tuple(input_shape), input_tensor.dtype, name=input_name),)
+    try:
+        _convert_via_saved_model(tf2onnx, keras_model, input_signature, opset, output_path)
+    except Exception:
+        try:
+            output_names = [out.name.split(":")[0] for out in keras_model.outputs]
+            tf2onnx.convert.from_keras(
+                keras_model,
+                input_signature=input_signature,
+                opset=opset,
+                output_path=str(output_path),
+                output_names=output_names,
+            )
+        except TypeError:
+            try:
+                tf2onnx.convert.from_keras(
+                    keras_model,
+                    input_signature=input_signature,
+                    opset=opset,
+                    output_path=str(output_path),
+                )
+            except ValueError as exc:
+                if "explicit_paddings" not in str(exc):
+                    raise
+                _convert_via_frozen_graph(
+                    tf2onnx, tf, keras_model, input_signature, opset, output_path
+                )
+    return output_path
+
+
+def _load_stardist_model(ndim: int, name_or_path: str):
+    _ensure_csbdeep_on_path()
+    _ensure_stardist_stub()
+    if ndim == 2:
+        module = importlib.import_module(
+            "senoquant.tabs.segmentation.stardist_onnx_utils._stardist.models"
+        )
+        model_cls = module.StarDist2D
+    elif ndim == 3:
+        module = importlib.import_module(
+            "senoquant.tabs.segmentation.stardist_onnx_utils._stardist.models"
+        )
+        model_cls = module.StarDist3D
+    else:
+        raise ValueError("ndim must be 2 or 3.")
+
+    model_path = Path(name_or_path)
+    if model_path.is_dir():
+        return model_cls(None, name=model_path.name, basedir=str(model_path.parent))
+    model = model_cls.from_pretrained(name_or_path)
+    if model is None:
+        raise ValueError(f"Unknown pretrained model: {name_or_path}")
+    return model
+
+
+def _resolve_output_path(output: str | Path, default_name: str) -> Path:
+    output_path = Path(output)
+    if output_path.suffix.lower() != ".onnx":
+        output_path = output_path / default_name
+    return output_path
+
+
+def _safe_name(name: str) -> str:
+    return re.sub(r"[^A-Za-z0-9._-]+", "_", name).strip("_")
+
+
+def _ensure_csbdeep_on_path() -> None:
+    root = Path(__file__).resolve().parents[2]
+    csbdeep_root = root / "_csbdeep"
+    if csbdeep_root.exists():
+        csbdeep_path = str(csbdeep_root)
+        if csbdeep_path not in sys.path:
+            sys.path.insert(0, csbdeep_path)
+
+
+def _ensure_stardist_stub() -> None:
+    base_pkg = "senoquant.tabs.segmentation.stardist_onnx_utils._stardist"
+    root = Path(__file__).resolve().parents[2] / "_stardist"
+    if base_pkg not in sys.modules:
+        pkg = types.ModuleType(base_pkg)
+        pkg.__path__ = [str(root)]
+        sys.modules[base_pkg] = pkg
+    geom_name = f"{base_pkg}.geometry"
+    if geom_name not in sys.modules:
+        geom = types.ModuleType(geom_name)
+
+        def _stub(*_args, **_kwargs):
+            raise RuntimeError("StarDist geometry helpers are unavailable in converter.")
+
+        geom.star_dist = _stub
+        geom.dist_to_coord = _stub
+        geom.polygons_to_label = _stub
+        geom.star_dist3D = _stub
+        geom.polyhedron_to_label = _stub
+        sys.modules[geom_name] = geom
+
+
+def _import_tensorflow():
+    try:
+        import tensorflow as tf
+    except ImportError as exc:
+        raise RuntimeError("TensorFlow is required to export StarDist models.") from exc
+    return tf
+
+
+def _import_tf2onnx():
+    try:
+        import numpy as np
+        # tf2onnx still references deprecated numpy aliases in some versions.
+        for alias, value in {
+            "bool": np.bool_,
+            "object": np.object_,
+        }.items():
+            if not hasattr(np, alias):
+                setattr(np, alias, value)
+        import tf2onnx
+    except ImportError as exc:
+        raise RuntimeError("tf2onnx is required to export StarDist models.") from exc
+    return tf2onnx
+
+
+def _convert_via_frozen_graph(tf2onnx, tf, keras_model, input_signature, opset, output_path):
+    @tf.function
+    def _model_fn(*args):
+        return keras_model(*args, training=False)
+
+    concrete = _model_fn.get_concrete_function(*input_signature)
+
+    try:
+        from tensorflow.python.framework.convert_to_constants import (
+            convert_variables_to_constants_v2,
+        )
+    except ImportError as exc:
+        raise RuntimeError("TensorFlow constants converter is unavailable.") from exc
+
+    frozen_func = convert_variables_to_constants_v2(concrete)
+    graph_def = frozen_func.graph.as_graph_def()
+    inputs = [tensor.name for tensor in frozen_func.inputs]
+    outputs = [tensor.name for tensor in frozen_func.outputs]
+
+    _strip_empty_explicit_paddings(graph_def)
+
+    try:
+        tf2onnx.convert.from_graph_def(
+            graph_def,
+            input_names=inputs,
+            output_names=outputs,
+            opset=opset,
+            output_path=str(output_path),
+        )
+    except TypeError:
+        tf2onnx.convert.from_graph_def(
+            graph_def,
+            inputs,
+            outputs,
+            opset=opset,
+            output_path=str(output_path),
+        )
+
+
+def _convert_via_saved_model(tf2onnx, keras_model, input_signature, opset, output_path):
+    if not hasattr(keras_model, "export"):
+        raise RuntimeError("Keras model does not support export().")
+    export_dir = Path(tempfile.mkdtemp(prefix="stardist_saved_model_"))
+    keras_model.export(
+        str(export_dir),
+        format="tf_saved_model",
+        input_signature=input_signature,
+    )
+    if hasattr(tf2onnx.convert, "from_saved_model"):
+        tf2onnx.convert.from_saved_model(
+            str(export_dir),
+            output_path=str(output_path),
+            opset=opset,
+        )
+    else:
+        raise RuntimeError("tf2onnx does not support from_saved_model.")
+
+
+def _strip_empty_explicit_paddings(graph_def):
+    for node in graph_def.node:
+        attr = node.attr.get("explicit_paddings")
+        if attr is not None and len(attr.list.i) == 0:
+            del node.attr["explicit_paddings"]

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

@@ -0,0 +1,15 @@
+"""ONNX model inspection utilities."""
+
+from .divisibility import infer_div_by, summarize_model_io
+from .receptive_field import infer_receptive_field, recommend_tile_overlap
+from .valid_sizes import infer_valid_size_patterns
+from .probe import make_probe_image
+
+__all__ = [
+    "infer_div_by",
+    "summarize_model_io",
+    "infer_receptive_field",
+    "recommend_tile_overlap",
+    "infer_valid_size_patterns",
+    "make_probe_image",
+]

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

@@ -0,0 +1,36 @@
+"""CLI for inspecting StarDist ONNX models."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from .divisibility import infer_div_by, summarize_model_io
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Inspect an ONNX model.")
+    parser.add_argument("model", type=Path, help="Path to the ONNX model.")
+    parser.add_argument("--ndim", type=int, choices=(2, 3), default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    # Summarize model IO shapes to show dynamic/static dims.
+    summary = summarize_model_io(args.model)
+    # Infer the spatial divisibility required by the graph.
+    div_by = infer_div_by(args.model, ndim=args.ndim)
+
+    print(f"Model: {args.model}")
+    print("Inputs:")
+    for idx, dims in enumerate(summary["inputs"]):
+        print(f"  [{idx}] {dims}")
+    print("Outputs:")
+    for idx, dims in enumerate(summary["outputs"]):
+        print(f"  [{idx}] {dims}")
+    print(f"Inferred div_by: {div_by}")
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,193 @@
+"""Infer input divisibility constraints from an ONNX graph.
+
+This module inspects ONNX graphs to infer the minimal spatial divisibility
+required to run the model without shape mismatches through down/upsampling
+paths (e.g., U-Net skip connections).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable
+
+
+def infer_div_by(model_path: str | Path, ndim: int | None = None) -> tuple[int, ...]:
+    """Infer the spatial divisibility required by an ONNX model.
+
+    This inspects the graph to estimate the cumulative downsampling factor
+    along spatial axes. The result is the minimal per-axis multiple that the
+    model input should be divisible by to avoid internal shape mismatches
+    (e.g., concatenation of encoder/decoder feature maps).
+
+    Parameters
+    ----------
+    model_path : str or pathlib.Path
+        Path to the ONNX model file.
+    ndim : int or None, optional
+        Number of spatial dimensions (2 or 3). If ``None``, the input rank is
+        used to infer dimensionality (rank 4 -> 2D, rank 5 -> 3D).
+
+    Returns
+    -------
+    tuple[int, ...]
+        Per-axis divisibility requirement (e.g., ``(16, 16)`` or
+        ``(8, 8, 8)``).
+
+    Notes
+    -----
+    - The algorithm tracks cumulative scaling factors by propagating
+      per-axis scale values through the graph.
+    - Downsampling ops (Conv/Pool with stride > 1) increase the scale.
+    - Upsampling ops (ConvTranspose/Resize) reduce the scale.
+    - The maximum scale observed across the graph is returned.
+    """
+    # Load the ONNX graph and find the primary input tensor.
+    model = _load_onnx(model_path)
+    input_name = model.graph.input[0].name if model.graph.input else None
+    if input_name is None:
+        raise ValueError("ONNX model has no graph inputs.")
+
+    # Determine the number of spatial dimensions if not specified.
+    if ndim is None:
+        ndim = _infer_ndim(model)
+
+    # Collect initializer tensors so we can read Resize scales, etc.
+    init_map = _initializers(model)
+
+    # Map tensor name -> per-axis scale relative to the original input.
+    scale_map: dict[str, list[float]] = {input_name: [1.0] * ndim}
+    # Track the maximum cumulative downsample per axis across the graph.
+    max_scale = [1.0] * ndim
+
+    for node in model.graph.node:
+        # Resolve the input scales for this node if we have them.
+        input_scales = [scale_map[name] for name in node.input if name in scale_map]
+        # Merge multiple inputs by taking the maximum scale per axis.
+        base = (
+            [max(values) for values in zip(*input_scales)]
+            if input_scales
+            else [1.0] * ndim
+        )
+        # Default: node does not change spatial scale.
+        factor = [1.0] * ndim
+
+        # Downsampling: increase scale by stride.
+        if node.op_type in ("Conv", "MaxPool", "AveragePool"):
+            strides = _get_attr_ints(node, "strides")
+            if strides:
+                factor = [float(s) for s in strides[-ndim:]]
+        # Upsampling: reduce scale by stride.
+        elif node.op_type == "ConvTranspose":
+            strides = _get_attr_ints(node, "strides")
+            if strides:
+                factor = [1.0 / float(s) if s else 1.0 for s in strides[-ndim:]]
+        # Resize/Upsample may carry explicit scales as initializers.
+        elif node.op_type in ("Resize", "Upsample"):
+            scales = _get_resize_scales(node, init_map)
+            if scales is not None and len(scales) >= ndim:
+                spatial = scales[-ndim:]
+                factor = [
+                    1.0 / float(s) if float(s) not in (0.0, 1.0) else 1.0
+                    for s in spatial
+                ]
+
+        # Propagate the updated scale to all outputs of this node.
+        out_scale = [b * f for b, f in zip(base, factor)]
+        for output in node.output:
+            scale_map[output] = out_scale
+        # Record the maximum scale seen so far.
+        max_scale = [max(m, s) for m, s in zip(max_scale, out_scale)]
+
+    # Convert to integer divisibility requirements.
+    return tuple(int(round(s)) if s >= 1 else 1 for s in max_scale)
+
+
+def summarize_model_io(model_path: str | Path) -> dict[str, list[list[str]]]:
+    """Return a simple summary of model input/output shapes.
+
+    Parameters
+    ----------
+    model_path : str or pathlib.Path
+        Path to the ONNX model file.
+
+    Returns
+    -------
+    dict
+        Dictionary with ``inputs`` and ``outputs`` lists. Each entry is a
+        list of dimension labels (e.g., ``"1"``, ``"H (dynamic)"``).
+    """
+    # Load the graph and format the shapes for user-friendly display.
+    model = _load_onnx(model_path)
+    inputs = [_format_shape(inp.type.tensor_type.shape) for inp in model.graph.input]
+    outputs = [_format_shape(out.type.tensor_type.shape) for out in model.graph.output]
+    return {"inputs": inputs, "outputs": outputs}
+
+
+def _load_onnx(model_path: str | Path):
+    """Load an ONNX model, raising a helpful error if onnx is missing."""
+    try:
+        import onnx
+    except Exception as exc:
+        # Keep error explicit so users know to install the dependency.
+        raise RuntimeError("onnx is required for model inspection.") from exc
+    return onnx.load(str(model_path))
+
+
+def _initializers(model) -> dict[str, Iterable[float]]:
+    """Materialize ONNX initializers into a name -> numpy array map."""
+    from onnx import numpy_helper
+
+    return {
+        init.name: numpy_helper.to_array(init)
+        for init in model.graph.initializer
+    }
+
+
+def _infer_ndim(model) -> int:
+    """Infer the spatial dimensionality from the model input rank."""
+    if not model.graph.input:
+        raise ValueError("ONNX model has no graph inputs.")
+    shape = model.graph.input[0].type.tensor_type.shape
+    rank = len(shape.dim)
+    if rank == 4:
+        return 2
+    if rank == 5:
+        return 3
+    raise ValueError(f"Unsupported input rank {rank}; pass ndim explicitly.")
+
+
+def _get_attr_ints(node, name: str) -> list[int] | None:
+    """Extract INT/INTS attributes from a node."""
+    for attr in node.attribute:
+        if attr.name == name:
+            if attr.type == attr.INTS:
+                return list(attr.ints)
+            if attr.type == attr.INT:
+                return [attr.i]
+    return None
+
+
+def _get_resize_scales(node, init_map: dict[str, Iterable[float]]):
+    """Return resize scales from initializer inputs or node attributes."""
+    # Newer ONNX Resize uses a scales tensor input.
+    for input_name in reversed(node.input):
+        if input_name in init_map:
+            return init_map[input_name]
+    # Older Resize/Upsample variants may store scales as attributes.
+    for attr in node.attribute:
+        if attr.name == "scales" and attr.type == attr.FLOATS:
+            return list(attr.floats)
+    return None
+
+
+def _format_shape(shape) -> list[str]:
+    """Format an ONNX TensorShapeProto into a list of human-readable dims."""
+    dims: list[str] = []
+    for dim in shape.dim:
+        if dim.dim_param:
+            dims.append(f"{dim.dim_param} (dynamic)")
+        elif dim.dim_value:
+            dims.append(str(dim.dim_value))
+        else:
+            dims.append("? (dynamic)")
+    return dims