segresnet-mlx 0.2.1__py3-none-any.whl

segresnet_mlx/__init__.py

@@ -0,0 +1,67 @@
+"""Explicit MLX port of MONAI SegResNet inference.
+
+Public names load on first access, so importing the package (or the medmlx
+declaration in ``segresnet_mlx.medmlx``) does not import MLX.
+"""
+
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from segresnet_mlx.inference import forward_patch as forward_patch
+    from segresnet_mlx.network import SegResNet as SegResNet
+    from segresnet_mlx.predictor import BUNDLE_CONFIGS as BUNDLE_CONFIGS
+    from segresnet_mlx.predictor import BundlePredictor as BundlePredictor
+    from segresnet_mlx.predictor import SlidingWindowConfig as SlidingWindowConfig
+    from segresnet_mlx.predictor import predict_logits as predict_logits
+    from segresnet_mlx.telemetry import WindowTelemetry as WindowTelemetry
+    from segresnet_mlx.topology import SegResNetSpec as SegResNetSpec
+    from segresnet_mlx.weights import convert_checkpoint as convert_checkpoint
+    from segresnet_mlx.weights import convert_state_dict as convert_state_dict
+    from segresnet_mlx.weights import load_model as load_model
+    from segresnet_mlx.weights import required_shapes as required_shapes
+
+__version__ = "0.2.1"
+
+_EXPORTS = {
+    "BUNDLE_CONFIGS": "segresnet_mlx.predictor",
+    "BundlePredictor": "segresnet_mlx.predictor",
+    "SegResNet": "segresnet_mlx.network",
+    "SegResNetSpec": "segresnet_mlx.topology",
+    "SlidingWindowConfig": "segresnet_mlx.predictor",
+    "WindowTelemetry": "segresnet_mlx.telemetry",
+    "convert_checkpoint": "segresnet_mlx.weights",
+    "convert_state_dict": "segresnet_mlx.weights",
+    "forward_patch": "segresnet_mlx.inference",
+    "load_model": "segresnet_mlx.weights",
+    "predict_logits": "segresnet_mlx.predictor",
+    "required_shapes": "segresnet_mlx.weights",
+}
+
+__all__ = [
+    "BUNDLE_CONFIGS",
+    "BundlePredictor",
+    "SegResNet",
+    "SegResNetSpec",
+    "SlidingWindowConfig",
+    "WindowTelemetry",
+    "convert_checkpoint",
+    "convert_state_dict",
+    "forward_patch",
+    "load_model",
+    "predict_logits",
+    "required_shapes",
+]
+
+
+def __getattr__(name: str) -> object:
+    module = _EXPORTS.get(name)
+    if module is None:
+        raise AttributeError(f"module 'segresnet_mlx' has no attribute {name!r}")
+    value = getattr(import_module(module), name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted({*globals(), *_EXPORTS})

segresnet_mlx/_arrays.py

@@ -0,0 +1,36 @@
+"""Host array dtypes at model, checkpoint and image boundaries."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, cast
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from segresnet_mlx._types import Array
+
+type HostArray = np.ndarray[tuple[int, ...], np.dtype[np.generic]]
+type Float32Array = np.ndarray[tuple[int, ...], np.dtype[np.float32]]
+type Float64Array = np.ndarray[tuple[int, ...], np.dtype[np.float64]]
+type Int64Array = np.ndarray[tuple[int, ...], np.dtype[np.int64]]
+type MaskArray = Float32Array | Int64Array
+
+type NumericScalar = (
+    np.bool_
+    | np.int8
+    | np.int16
+    | np.int32
+    | np.int64
+    | np.uint8
+    | np.uint16
+    | np.uint32
+    | np.uint64
+    | np.float16
+    | np.float32
+    | np.float64
+    | np.complex64
+    | np.complex128
+)
+type NumericArray = np.ndarray[tuple[int, ...], np.dtype[NumericScalar]]
+as_numpy = cast("Callable[[Array], NumericArray]", np.array)

segresnet_mlx/_nifti.py

@@ -0,0 +1,63 @@
+"""Typed protocols for the NiBabel NIfTI, header and orientation calls used here."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from os import PathLike
+from typing import Protocol, cast
+
+import nibabel
+import numpy as np
+from nibabel.filebasedimages import ImageFileError as ImageFileError
+
+from segresnet_mlx._arrays import Float64Array, HostArray
+
+type PathArg = str | PathLike[str]
+
+
+class ArrayBuffer(Protocol):
+    def __array__(
+        self, dtype: np.dtype[np.generic] | None = None, copy: bool | None = None
+    ) -> HostArray: ...
+
+
+class NiftiHeader(Protocol):
+    def get_zooms(self) -> tuple[float, ...]: ...
+    def set_zooms(self, zooms: Sequence[float] | Float64Array) -> None: ...
+    def __getitem__(self, key: str) -> HostArray: ...
+
+
+class NiftiImage(Protocol):
+    shape: tuple[int, ...]
+    affine: Float64Array
+    dataobj: ArrayBuffer | HostArray
+    header: NiftiHeader
+
+    def __init__(self, dataobj: HostArray, affine: HostArray) -> None: ...
+    def get_data_dtype(self) -> np.dtype[np.generic]: ...
+    def get_sform(self) -> Float64Array: ...
+    def get_qform(self) -> Float64Array: ...
+    def set_sform(self, affine: Float64Array, code: int | None = None) -> None: ...
+    def set_qform(self, affine: Float64Array, code: int | None = None) -> None: ...
+
+
+class Orientations(Protocol):
+    def axcodes2ornt(self, axcodes: str, labels: tuple[tuple[str, str], ...]) -> Float64Array: ...
+    def ornt_transform(self, start_ornt: Float64Array, end_ornt: Float64Array) -> Float64Array: ...
+    def inv_ornt_aff(self, ornt: Float64Array, shape: Sequence[int]) -> Float64Array: ...
+    def aff2axcodes(
+        self, aff: Float64Array, labels: tuple[tuple[str, str], ...]
+    ) -> tuple[str | None, ...]: ...
+
+
+class Nibabel(Protocol):
+    Nifti1Image: type[NiftiImage]
+    orientations: Orientations
+
+    def load(self, filename: PathArg) -> NiftiImage: ...
+    def save(self, img: NiftiImage, filename: PathArg) -> None: ...
+    def io_orientation(self, affine: Float64Array) -> Float64Array: ...
+
+
+# Bound to the real module; this runner admits only NIfTI images.
+nib = cast(Nibabel, nibabel)

segresnet_mlx/_records.py

@@ -0,0 +1,78 @@
+"""JSON schemas for the public runner and window telemetry."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import NotRequired, TypedDict
+
+from segresnet_mlx.topology import ArithmeticProfile
+
+type Json = bool | int | float | str | list[Json] | dict[str, Json] | None
+
+
+class InfererConfig(TypedDict):
+    roi_size: Sequence[int]
+    sw_batch_size: NotRequired[int]
+    overlap: NotRequired[float]
+    mode: NotRequired[str]
+    padding_mode: NotRequired[str]
+
+
+class BundleConfig(TypedDict):
+    inferer: InfererConfig
+
+
+class Runtime(TypedDict):
+    """MLX device string and peak MLX memory in bytes (from ``mx.get_peak_memory``)."""
+
+    device: str
+    peak_bytes: int | None
+
+
+class SlidingWindow(TypedDict):
+    roi_size: list[int]
+    sw_batch_size: int
+    overlap: float
+    mode: str
+    padding_mode: str
+
+
+class Seconds(TypedDict):
+    """Wall-clock seconds per runner stage, rounded to four decimals."""
+
+    load_inputs_and_weights: float
+    preprocess: float
+    inference: float
+    postprocess_and_write: float
+
+
+class Details(TypedDict):
+    """The ``details`` block of ``run``'s result.
+
+    ``labels`` maps each label value (as a string) to its name and ``labels_present``
+    lists the values found in the written mask. ``input_shape`` is the first input's
+    voxel shape (equal to the mask shape); ``network_input_shape`` is the spatial shape
+    after preprocessing, which the sliding window sees. ``output_dtype`` is the NIfTI
+    dtype of the mask. ``peak_mlx_bytes`` is the peak MLX memory in bytes.
+    ``requires_human_review`` is always true: outputs are research results.
+    """
+
+    bundle: str
+    arithmetic_profile: ArithmeticProfile
+    labels: dict[str, str]
+    labels_present: list[int]
+    input_shape: list[int]
+    network_input_shape: list[int]
+    output_dtype: str
+    sliding_window: SlidingWindow
+    mlx_device: str
+    peak_mlx_bytes: int | None
+    seconds: Seconds
+    requires_human_review: bool
+
+
+class RunResult(TypedDict):
+    """Result of ``segresnet_mlx.medmlx.run``: ``outputs["mask"]`` is the mask path."""
+
+    outputs: dict[str, str]
+    details: Details

segresnet_mlx/_runtime.py

@@ -0,0 +1,22 @@
+"""Typed native import after the shared runtime's existing admission checks."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+from medmlx_core.runtime import import_mlx as _import_mlx
+
+if TYPE_CHECKING:
+    from segresnet_mlx._types import MLX
+
+
+def import_mlx() -> MLX:
+    """Return ``mlx.core`` with the Metal GPU as the default device.
+
+    The host checks and device selection are done by ``medmlx_core``; on an
+    unsupported host it raises ``MissingDependencyError`` and there is no CPU fallback.
+    """
+    _import_mlx()
+    import mlx.core as mx
+
+    return cast("MLX", mx)

segresnet_mlx/_types.py

@@ -0,0 +1,58 @@
+"""Static native MLX contracts; runtime imports remain lazy."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Mapping, Sequence
+    from typing import Literal, Protocol
+
+    import mlx.core as mx
+    from medmlx_core.typing import MlxRuntime
+
+    from segresnet_mlx._arrays import Float32Array, HostArray
+
+    type Array = mx.array
+    type ArrayValue = bool | int | float | complex | list[ArrayValue] | tuple[ArrayValue, ...]
+
+    class ArrayConstructor(Protocol):
+        def __call__(
+            self, value: HostArray | Array | ArrayValue, dtype: mx.Dtype | None = None
+        ) -> Array: ...
+
+    class MLX(MlxRuntime, Protocol):
+        all = staticmethod(mx.all)
+        isfinite = staticmethod(mx.isfinite)
+        full = staticmethod(mx.full)
+        sum = staticmethod(mx.sum)
+        ones = staticmethod(mx.ones)
+        synchronize = staticmethod(mx.synchronize)
+        clear_cache = staticmethod(mx.clear_cache)
+        load = staticmethod(mx.load)
+
+        def eval(self, *args: Array | Mapping[str, Array] | Sequence[Array]) -> None: ...
+
+        def save_safetensors(self, file: str, arrays: dict[str, Array]) -> None: ...
+
+    class SlidingWindowInference(Protocol):
+        def __call__(
+            self,
+            inputs: Float32Array,
+            roi_size: Sequence[int],
+            sw_batch_size: int,
+            predictor: Callable[[Float32Array | Array], Float32Array]
+            | Callable[[Float32Array | Array, list[list[slice]]], Float32Array],
+            *,
+            overlap: float,
+            mode: str,
+            padding_mode: str,
+            with_coord: bool,
+            predictor_input: Literal["numpy"],
+        ) -> Float32Array: ...
+
+    class CoreSlidingWindow(Protocol):
+        sliding_window_inference: SlidingWindowInference
+
+    class Metal(Protocol):
+        def is_available(self) -> bool: ...

segresnet_mlx/inference.py

@@ -0,0 +1,37 @@
+"""Synchronized single-patch inference with NumPy inputs and raw-logit outputs."""
+
+from __future__ import annotations
+
+from typing import cast
+
+import numpy as np
+from medmlx_core.errors import InvalidInputError
+from medmlx_core.typing import ArrayFactory
+
+from segresnet_mlx._arrays import Float32Array, HostArray
+from segresnet_mlx._runtime import import_mlx
+from segresnet_mlx.network import SegResNet
+
+
+def forward_patch(network: SegResNet, values: HostArray) -> Float32Array:
+    """Run one finite float32 ``(N, C, D, H, W)`` patch without transforms or stitching.
+
+    Returns float32 NumPy logits ``(N, out_channels, D, H, W)``, synchronized and
+    checked for finiteness.
+
+    Raises:
+        InvalidInputError: If ``values`` is not finite float32.
+        RuntimeError: If the network produced non-finite logits.
+    """
+    host = np.asarray(values)
+    if host.dtype != np.float32 or not np.isfinite(host).all():
+        raise InvalidInputError("SegResNet input must be finite float32")
+    mx = import_mlx()
+    mx.synchronize()
+    result = network(cast(ArrayFactory, mx.array)(host))
+    mx.eval(result)
+    mx.synchronize()
+    output = cast(Float32Array, np.array(result))
+    if not np.isfinite(output).all():
+        raise RuntimeError("SegResNet produced non-finite logits")
+    return output