mlx-segment-ct 0.2.1__py3-none-any.whl

mlx_segment_ct/__init__.py

@@ -0,0 +1,83 @@
+"""Native MLX implementation of NV-Segment-CT (VISTA-3D).
+
+Public names load on first access, so importing the package (or the medmlx
+declaration in ``mlx_segment_ct.medmlx``) does not import MLX.
+"""
+
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from mlx_segment_ct.correction import normalized_point_positions
+    from mlx_segment_ct.correction_inferer import (
+        point_based_window_inferer,
+    )
+    from mlx_segment_ct.interactive import (
+        interactive_labels,
+        interactive_logits,
+        transform_points,
+    )
+    from mlx_segment_ct.labels import InteractivePrompt
+    from mlx_segment_ct.mlx_network import AutomaticModel, FeatureModel, required_shapes
+    from mlx_segment_ct.mlx_point import PointModel, additional_shapes, correction_shapes
+    from mlx_segment_ct.mlx_point_window import PointWindowPredictor
+    from mlx_segment_ct.mlx_upsample import AutomaticUpsample
+    from mlx_segment_ct.mlx_worker import automatic_logits
+    from mlx_segment_ct.streaming import LabelReducer, bounded_labels, rolling_logits
+
+__version__ = "0.2.1"
+
+_EXPORTS = {
+    "AutomaticModel": "mlx_segment_ct.mlx_network",
+    "AutomaticUpsample": "mlx_segment_ct.mlx_upsample",
+    "FeatureModel": "mlx_segment_ct.mlx_network",
+    "InteractivePrompt": "mlx_segment_ct.labels",
+    "LabelReducer": "mlx_segment_ct.streaming",
+    "PointModel": "mlx_segment_ct.mlx_point",
+    "PointWindowPredictor": "mlx_segment_ct.mlx_point_window",
+    "additional_shapes": "mlx_segment_ct.mlx_point",
+    "automatic_logits": "mlx_segment_ct.mlx_worker",
+    "bounded_labels": "mlx_segment_ct.streaming",
+    "correction_shapes": "mlx_segment_ct.mlx_point",
+    "interactive_labels": "mlx_segment_ct.interactive",
+    "interactive_logits": "mlx_segment_ct.interactive",
+    "normalized_point_positions": "mlx_segment_ct.correction",
+    "point_based_window_inferer": "mlx_segment_ct.correction_inferer",
+    "required_shapes": "mlx_segment_ct.mlx_network",
+    "rolling_logits": "mlx_segment_ct.streaming",
+    "transform_points": "mlx_segment_ct.interactive",
+}
+
+__all__ = [
+    "AutomaticModel",
+    "AutomaticUpsample",
+    "FeatureModel",
+    "InteractivePrompt",
+    "LabelReducer",
+    "PointModel",
+    "PointWindowPredictor",
+    "additional_shapes",
+    "automatic_logits",
+    "bounded_labels",
+    "correction_shapes",
+    "interactive_labels",
+    "interactive_logits",
+    "normalized_point_positions",
+    "point_based_window_inferer",
+    "required_shapes",
+    "rolling_logits",
+    "transform_points",
+]
+
+
+def __getattr__(name: str) -> object:
+    module = _EXPORTS.get(name)
+    if module is None:
+        raise AttributeError(f"module 'mlx_segment_ct' has no attribute {name!r}")
+    value = getattr(import_module(module), name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted({*globals(), *_EXPORTS})

mlx_segment_ct/_types.py

@@ -0,0 +1,378 @@
+"""Consumed native MLX and host-orchestration contracts.
+
+The native graph uses real MLX array/Dtype identities. The orchestration boundary
+is generic because the existing CPU oracle tests inject NumPy arrays there.
+These partial interfaces describe the MLX and medmlx-core versions pinned in uv.lock;
+they do not replace either runtime or claim complete upstream API coverage.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, Sequence
+from os import PathLike
+from typing import TYPE_CHECKING, Literal, NotRequired, Protocol, Required, TypedDict
+
+import numpy as np
+from medmlx_core.typing import MlxRuntime
+from numpy.typing import DTypeLike, NDArray
+
+if TYPE_CHECKING:
+    from mlx.core import Dtype, array
+
+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
+)
+HostArray = NDArray[NumericScalar]
+FloatArray = NDArray[np.float32]
+DoubleArray = NDArray[np.float64]
+IntArray = NDArray[np.int64]
+BoolArray = NDArray[np.bool_]
+type ArrayValue = array | HostArray
+type NumericInput = bool | int | float | ArrayValue | Sequence[NumericInput]
+type ArrayTree = array | Sequence[ArrayTree] | Mapping[str, ArrayTree]
+type PatchCoordinates = Sequence[Sequence[slice]]
+
+
+class ArrayBoundary[ArrayT](Protocol):
+    @property
+    def array(self) -> Callable[..., ArrayT]: ...
+    def eval(self, *args: ArrayT | Sequence[ArrayT] | Mapping[str, ArrayT]) -> None: ...
+    def synchronize(self) -> None: ...
+
+
+class CoreRuntime(MlxRuntime, Protocol):
+    def eval(self, *args: array | Mapping[str, array] | Sequence[array]) -> None: ...
+    def synchronize(self) -> None: ...
+    def async_eval(self, *args: ArrayTree) -> None: ...
+    def ones(self, shape: Sequence[int], dtype: Dtype | None = None) -> array: ...
+    def sin(self, a: array) -> array: ...
+    def cos(self, a: array) -> array: ...
+    def abs(self, a: array) -> array: ...
+    def sum(
+        self, a: array, axis: int | Sequence[int] | None = None, keepdims: bool = False
+    ) -> array: ...
+    def softmax(
+        self, a: array, axis: int | Sequence[int] | None = None, *, precise: bool = False
+    ) -> array: ...
+    def addmm(
+        self, c: array, a: array, b: array, alpha: float = 1.0, beta: float = 1.0
+    ) -> array: ...
+    def meshgrid(
+        self, *arrays: array, indexing: Literal["xy", "ij"] = "xy", sparse: bool = False
+    ) -> tuple[array, ...]: ...
+    def contiguous(self, a: array, /, *, allow_col_major: bool = False) -> array: ...
+    def load(self, file: str, /, format: str | None = None) -> array | dict[str, array]: ...
+    def save_safetensors(
+        self, file: str, arrays: Mapping[str, array], metadata: Mapping[str, str] | None = None
+    ) -> None: ...
+
+
+class PointPredictor[ArrayT](Protocol):
+    def __call__(self, image: ArrayT, positions: ArrayT, codes: ArrayT, target: int) -> ArrayT: ...
+
+
+class CoordinatePredictor[ArrayT](Protocol):
+    def __call__(self, image: ArrayT, coordinates: PatchCoordinates, /) -> ArrayT: ...
+
+
+class PointOptions(TypedDict, total=False):
+    transpose: Required[bool]
+    sw_batch_size: int
+    with_coord: bool
+    mode: str
+    sw_device: object
+    device: object
+    overlap: float
+
+
+class PointWindowCall[ArrayT](Protocol):
+    def __call__(
+        self,
+        image: ArrayT,
+        *,
+        point_coords: HostArray,
+        point_labels: HostArray,
+        class_vector: HostArray | None,
+        prompt_class: HostArray | None,
+        patch_coords: PatchCoordinates,
+        prev_mask: HostArray | None,
+        transpose: bool,
+        **kwargs: object,
+    ) -> HostArray: ...
+
+
+class Metrics(TypedDict):
+    patch_seconds: list[float]
+    padded_shape: NotRequired[list[int]]
+    padding: NotRequired[list[int]]
+    streaming_axis: NotRequired[int]
+    ring_depth: NotRequired[int]
+    predicted_accumulator_bytes: NotRequired[int]
+    count_buffer_bytes: NotRequired[int]
+    label_buffer_bytes: NotRequired[int]
+    accumulator_dtype: NotRequired[str]
+
+
+class OutputLabel(TypedDict):
+    value: int
+    name: str | None
+
+
+class InteractiveParameters(TypedDict):
+    roi_size: list[int]
+    center_only: bool
+    padding_mode: str
+    overlap_reduction: str
+    threshold: int
+    uncovered_label: int
+    positive_components: bool
+
+
+class Details(TypedDict, total=False):
+    mode: Required[str]
+    label_prompt: list[int]
+    label_names: list[str]
+    automatic_branch: bool
+    point_start: int
+    output_labels: Required[list[OutputLabel]]
+    points: list[list[float]]
+    point_labels: list[int]
+    effective_point_labels: list[int]
+    preprocessed_points: list[list[float]]
+    preprocessed_affine: list[list[float]]
+    preprocessed_shape: Required[list[int]]
+    output_shape: Required[list[int]]
+    output_dtype: Required[str]
+    algorithm: Required[str]
+    parameters: InteractiveParameters
+    patch_count: Required[int]
+    load_preprocess_seconds: Required[float]
+    inference_stitch_seconds: Required[float]
+    restore_save_seconds: Required[float]
+    class_ids: list[int]
+    roi_size: list[int]
+    overlap: float
+    blend_mode: str
+    padding_mode: str
+    padded_shape: list[int]
+    modality: str
+    checkpoint: str
+    restoration: str
+    backend: str
+    mlx_version: str
+    mlx_default_device: str
+    precision: str
+    source_revision: str
+    mlx_peak_memory_bytes: int | None
+    elapsed_seconds: float
+    warning: str
+
+
+class PublishedDetails(TypedDict, total=False):
+    mode: Required[str]
+    label_prompt: list[int]
+    label_names: list[str]
+    automatic_branch: bool
+    point_start: int
+    output_labels: Required[list[OutputLabel]]
+    points: list[list[float]]
+    point_labels: list[int]
+    effective_point_labels: list[int]
+    preprocessed_points: list[list[float]]
+    preprocessed_affine: list[list[float]]
+    preprocessed_shape: Required[list[int]]
+    output_shape: Required[list[int]]
+    output_dtype: Required[str]
+    algorithm: Required[str]
+    parameters: InteractiveParameters
+    patch_count: Required[int]
+    load_preprocess_seconds: Required[float]
+    inference_stitch_seconds: Required[float]
+    restore_save_seconds: Required[float]
+    class_ids: list[int]
+    roi_size: list[int]
+    overlap: float
+    blend_mode: str
+    padding_mode: str
+    padded_shape: list[int]
+    modality: str
+    checkpoint: Required[str]
+    restoration: Required[str]
+    backend: Required[str]
+    mlx_version: Required[str]
+    mlx_default_device: Required[str]
+    precision: Required[str]
+    source_revision: Required[str]
+    mlx_peak_memory_bytes: Required[int | None]
+    elapsed_seconds: Required[float]
+    warning: Required[str]
+
+
+class Outputs(TypedDict):
+    mask: str
+
+
+class RunResult(TypedDict):
+    outputs: Outputs
+    details: PublishedDetails
+
+
+class RunArguments(TypedDict, total=False):
+    image_path: str | PathLike[str]
+    weights_path: str | PathLike[str]
+    output_dir: str | PathLike[str]
+    segment_all: bool
+    prompt: str | Sequence[str] | None
+    label_prompt: Sequence[int] | None
+    previous_logits_path: str | PathLike[str] | None
+    points: Sequence[Sequence[float]] | None
+    point_labels: Sequence[int] | None
+    point_start: int
+    interactive: bool
+    overwrite: bool
+    checkpoint: str
+    modality: str | None
+    restoration: str
+
+
+class NiftiHeader(Protocol):
+    @property
+    def binaryblock(self) -> bytes: ...
+    def __getitem__(self, key: str) -> HostArray: ...
+    def get(self, key: str) -> HostArray | None: ...
+    def get_zooms(self) -> tuple[float, ...]: ...
+    def set_zooms(self, zooms: Sequence[float] | HostArray) -> None: ...
+    def as_byteswapped(self, endianness: str | None = None) -> NiftiHeader: ...
+
+
+class NiftiArrayProxy(Protocol):
+    """File-backed NIfTI data supports conversion without ndarray methods."""
+
+    def __array__(self, dtype: DTypeLike | None = None) -> NDArray[np.generic]: ...
+
+
+class NiftiImage(Protocol):
+    @property
+    def shape(self) -> tuple[int, ...]: ...
+    def get_data_dtype(self) -> np.dtype[NumericScalar]: ...
+    @property
+    def header(self) -> NiftiHeader: ...
+    @property
+    def affine(self) -> DoubleArray: ...
+    @property
+    def dataobj(self) -> HostArray | NiftiArrayProxy: ...
+    def get_sform(self) -> DoubleArray: ...
+    def get_qform(self) -> DoubleArray: ...
+    def set_sform(self, affine: HostArray, code: int | None = None) -> None: ...
+    def set_qform(self, affine: HostArray, code: int | None = None) -> None: ...
+
+
+class NiftiConstructor(Protocol):
+    def __call__(self, dataobj: HostArray, affine: HostArray) -> NiftiImage: ...
+
+
+class NiftiRuntime(Protocol):
+    @property
+    def Nifti1Image(self) -> NiftiConstructor: ...
+    def load(self, filename: str) -> object: ...
+    def save(self, img: NiftiImage, filename: str) -> None: ...
+    def aff2axcodes(self, aff: HostArray) -> tuple[str | None, ...]: ...
+    def io_orientation(self, affine: HostArray) -> DoubleArray: ...
+
+
+class OrientationRuntime(Protocol):
+    def axcodes2ornt(self, axcodes: tuple[str | None, ...], labels: None = None) -> DoubleArray: ...
+    def ornt_transform(self, start_ornt: DoubleArray, end_ornt: DoubleArray) -> DoubleArray: ...
+    def inv_ornt_aff(self, ornt: DoubleArray, shape: Sequence[int]) -> DoubleArray: ...
+
+
+class RollingMetrics(TypedDict):
+    patch_seconds: list[float]
+    padded_shape: list[int]
+    padding: list[int]
+    streaming_axis: int
+    ring_depth: int
+    predicted_accumulator_bytes: int
+    count_buffer_bytes: int
+    label_buffer_bytes: int
+    accumulator_dtype: str
+
+
+class SlidingWindowCall(Protocol):
+    def __call__(
+        self,
+        inputs: ArrayValue,
+        roi_size: Sequence[int],
+        sw_batch_size: int,
+        predictor: Callable[..., array],
+        *,
+        overlap: float,
+        mode: str,
+        padding_mode: str,
+        roi_weight_map: ArrayValue | None,
+        with_coord: bool,
+        **kwargs: object,
+    ) -> HostArray: ...
+
+
+class SlidingWindowNamespace(Protocol):
+    @property
+    def sliding_window_inference(self) -> SlidingWindowCall: ...
+
+
+class StandInRuntime[ArrayT, DtypeT](ArrayBoundary[ArrayT], Protocol):
+    @property
+    def float32(self) -> DtypeT: ...
+    def sum(self, a: ArrayT) -> ArrayT | NumericScalar: ...
+
+
+type Float32Dtype = Dtype | type[np.float32]
+
+
+class Float32Cast[ArrayT, DtypeT](Protocol):
+    def astype(self, dtype: DtypeT) -> ArrayT: ...
+
+
+class RollingMetricsBuilder(TypedDict):
+    patch_seconds: NotRequired[list[float]]
+    padded_shape: NotRequired[list[int]]
+    padding: NotRequired[list[int]]
+    streaming_axis: NotRequired[int]
+    ring_depth: NotRequired[int]
+    predicted_accumulator_bytes: NotRequired[int]
+    count_buffer_bytes: NotRequired[int]
+    label_buffer_bytes: NotRequired[int]
+    accumulator_dtype: NotRequired[str]
+
+
+class AutomaticDetails(TypedDict):
+    class_ids: list[int]
+    padded_shape: list[int]
+    label_names: list[str] | None
+
+
+class InteractiveDetails(TypedDict):
+    automatic_branch: bool
+    effective_point_labels: list[int]
+    point_start: int
+
+
+class NumericArrayArithmetic[ArrayT](Protocol):
+    """Native and NumPy arrays accept corresponding arrays and NumPy scalars.
+
+    MLX 0.32.3 accepts the NumPy reduction scalars used by host-injected point
+    fixtures; its bundled array operator declaration omits those scalar types.
+    """
+
+    def __add__(self, other: ArrayValue | NumericScalar | float, /) -> ArrayT: ...

mlx_segment_ct/activations.py

@@ -0,0 +1,74 @@
+"""GELU with PyTorch 2.14.1's ARM float32 erf and exponential arithmetic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+from medmlx_core.typing import MetalKernel
+
+from mlx_segment_ct._types import CoreRuntime
+
+if TYPE_CHECKING:
+    from mlx.core import array
+
+_GELU = """uint i = thread_position_in_grid.x;
+if (i >= SIZE) return;
+float x = X[i], z = x * 0.70710678118654752440f;
+float t = 1.0f / fma(0.3275911f, abs(z), 1.0f);
+float r = fma(1.061405429f, t, -1.453152027f);
+r = fma(r, t, 1.421413741f);
+r = fma(r, t, -0.284496736f);
+r = fma(r, t, 0.254829592f);
+float exponent = -(z * z), exponential;
+if (exponent < -0x1.5d5e2ap+6f) {
+    exponential = exp(exponent);
+} else {
+    // PyTorch's NEON exp_u20 range reduction and polynomial.
+    float n = round(exponent * 0x1.715476p+0f);
+    float reduced = fma(-n, 0x1.62e4p-1f, exponent);
+    reduced = fma(-n, 0x1.7f7d1cp-20f, reduced);
+    float scale = as_type<float>((uint(int(n)) << 23) + 0x3f800000u);
+    float square = reduced * reduced;
+    float p = fma(reduced, 0x1.0e4020p-7f, 0x1.573e2ep-5f);
+    float q = fma(reduced, 0x1.555e66p-3f, 0x1.fffdb6p-2f);
+    q = fma(p, square, q);
+    float polynomial = fma(q, square, 0x1.ffffecp-1f * reduced);
+    exponential = fma(polynomial, scale, scale);
+}
+float erf_value = fma(-t * exponential, r, 1.0f);
+erf_value = as_type<float>(as_type<uint>(erf_value) ^ (as_type<uint>(z) & 0x80000000u));
+Y[i] = (x * 0.5f) * (1.0f + erf_value);
+"""
+
+
+class Gelu:
+    """Source-ordered ARM GELU on Metal; no reference-runtime dependency.
+
+    The kernel reproduces PyTorch's vectorized erf (a polynomial fit) and exponential
+    in the reference's operation order, with floating-point contraction disabled
+    except for the explicit ``fma`` calls.
+    """
+
+    def __init__(self, mx: CoreRuntime) -> None:
+        self.mx = mx
+        self.kernel = cast(
+            MetalKernel,
+            mx.fast.metal_kernel(
+                name="vista_arm_gelu",
+                input_names=["X"],
+                output_names=["Y"],
+                source=_GELU,
+                header="#pragma clang fp contract(off)\n",
+            ),
+        )
+
+    def __call__(self, values: array) -> array:
+        """Exact (erf-based) GELU of a float32 array; the result has the same shape."""
+        return self.kernel(
+            inputs=[values],
+            template=[("SIZE", values.size)],
+            grid=(values.size, 1, 1),
+            threadgroup=(256, 1, 1),
+            output_shapes=[tuple(values.shape)],
+            output_dtypes=[self.mx.float32],
+        )[0]