unest-mlx 0.2.3__py3-none-any.whl

unest_mlx/__init__.py

@@ -0,0 +1,43 @@
+"""Native MLX implementation of UNesT.
+
+Public names load on first access, so importing the package (or the medmlx
+declaration in ``unest_mlx.medmlx``) does not import MLX.
+"""
+
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from unest_mlx.mlx_model import UNesT as UNesT
+    from unest_mlx.mlx_predictor import UNesTPredictor as UNesTPredictor
+    from unest_mlx.mlx_weights import convert_state_dict as convert_state_dict
+    from unest_mlx.mlx_weights import required_shapes as required_shapes
+
+__version__ = "0.2.2"
+
+_EXPORTS = {
+    "UNesT": "unest_mlx.mlx_model",
+    "UNesTPredictor": "unest_mlx.mlx_predictor",
+    "convert_state_dict": "unest_mlx.mlx_weights",
+    "required_shapes": "unest_mlx.mlx_weights",
+}
+
+__all__ = [
+    "UNesT",
+    "UNesTPredictor",
+    "convert_state_dict",
+    "required_shapes",
+]
+
+
+def __getattr__(name: str) -> object:
+    module = _EXPORTS.get(name)
+    if module is None:
+        raise AttributeError(f"module 'unest_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})

unest_mlx/_typing.py

@@ -0,0 +1,156 @@
+"""Native tensor, image geometry and result contracts for the renal bundle."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, Sequence
+from os import PathLike
+from typing import TYPE_CHECKING, Literal, Protocol, TypedDict, overload
+
+import numpy as np
+from medmlx_core.typing import MlxRuntime
+from numpy.typing import ArrayLike, NDArray
+
+if TYPE_CHECKING:
+    from mlx.core import array
+
+
+type HostArray = NDArray[np.generic]
+type RealArray = NDArray[np.bool_ | np.integer | np.floating]
+type FloatArray = NDArray[np.float32]
+type DoubleArray = NDArray[np.float64]
+type IntArray = NDArray[np.int_]
+type NativeNetwork = Callable[[array], array]
+
+
+class NativeRuntime(MlxRuntime, Protocol):
+    """The subset of ``mlx.core`` this package calls, typed structurally."""
+
+    def eval(self, *args: array | Mapping[str, array] | Sequence[array]) -> None: ...
+    def synchronize(self) -> None: ...
+    def load(self, file: str) -> array | dict[str, array]: ...
+    def save_safetensors(self, file: str, arrays: dict[str, array]) -> None: ...
+    def softmax(self, x: array, *, axis: int) -> array: ...
+
+
+class NativeOperation(Protocol):
+    def __call__(self, x: array) -> array: ...
+
+
+class NeuralRuntime(Protocol):
+    def MaxPool3d(self, *, kernel_size: int, stride: int, padding: int) -> NativeOperation: ...
+    def leaky_relu(self, x: array, *, negative_slope: float) -> array: ...
+    def gelu(self, x: array) -> array: ...
+
+
+class SlidingWindow(Protocol):
+    def __call__(
+        self,
+        inputs: FloatArray,
+        *,
+        roi_size: tuple[int, int, int],
+        sw_batch_size: int,
+        predictor: NativeNetwork,
+        overlap: float,
+    ) -> FloatArray: ...
+
+
+class SlidingWindowNamespace(Protocol):
+    sliding_window_inference: SlidingWindow
+
+
+class NiftiHeader(Protocol):
+    def __getitem__(self, key: str) -> HostArray: ...
+    def __setitem__(self, key: str, value: int | HostArray) -> None: ...
+    def get(self, key: str) -> HostArray | None: ...
+    def get_zooms(self) -> tuple[float, ...]: ...
+    def set_zooms(self, zooms: ArrayLike) -> None: ...
+    def set_slope_inter(self, slope: float, inter: float) -> None: ...
+
+
+class NiftiImage(Protocol):
+    affine: DoubleArray
+    shape: tuple[int, ...]
+    dataobj: ArrayLike
+    header: NiftiHeader
+
+    @overload
+    def get_sform(self, coded: Literal[True]) -> tuple[DoubleArray | None, int]: ...
+    @overload
+    def get_sform(self, coded: Literal[False] = False) -> DoubleArray: ...
+    @overload
+    def get_qform(self, coded: Literal[True]) -> tuple[DoubleArray | None, int]: ...
+    @overload
+    def get_qform(self, coded: Literal[False] = False) -> DoubleArray: ...
+    def set_sform(self, affine: DoubleArray | None, code: int | None = None) -> None: ...
+    def set_qform(self, affine: DoubleArray | None, code: int | None = None) -> None: ...
+    def get_data_dtype(self) -> np.dtype[np.generic]: ...
+
+
+class Orientations(Protocol):
+    def aff2axcodes(self, affine: DoubleArray) -> tuple[str, ...]: ...
+    def axcodes2ornt(self, axcodes: str | tuple[str, ...]) -> DoubleArray: ...
+    def ornt_transform(self, start: DoubleArray, end: DoubleArray) -> DoubleArray: ...
+    def inv_ornt_aff(self, ornt: DoubleArray, shape: tuple[int, ...]) -> DoubleArray: ...
+
+
+class NiftiIo(Protocol):
+    orientations: Orientations
+
+    def load(self, filename: str | PathLike[str]) -> NiftiImage: ...
+    def save(self, image: NiftiImage, filename: str | PathLike[str]) -> None: ...
+    def Nifti1Image(self, data: ArrayLike, affine: ArrayLike | None) -> NiftiImage: ...
+    def aff2axcodes(self, affine: DoubleArray) -> tuple[str, ...]: ...
+    def io_orientation(self, affine: DoubleArray) -> DoubleArray: ...
+
+
+class NumpyFlip(Protocol):
+    def __call__[T: np.generic](self, values: NDArray[T], *, axis: list[int]) -> NDArray[T]: ...
+
+
+class Label(TypedDict):
+    """One entry of the mask's label table: integer voxel value and structure name."""
+
+    value: int
+    name: str
+
+
+class SlidingWindowDetails(TypedDict):
+    roi_size: list[int]
+    sw_batch_size: int
+    overlap: float
+
+
+class RunDetails(TypedDict):
+    """JSON-compatible run metadata returned under ``details`` by ``medmlx.run``.
+
+    ``voxel_counts`` is keyed by label name. ``shape`` and ``dtype`` describe the
+    written mask (the input grid, float32). ``input_orientation`` is the axis-code string
+    of the reader affine before reorientation to RAS (for example ``"LAS"``).
+    ``network_shape`` is the voxel grid the network ran on after RAS reorientation and
+    1 mm resampling. ``mlx_device`` is ``str(mx.default_device())`` and
+    ``duration_seconds`` is wall time from the start of the call until the mask has been
+    written, so it includes image reading and weight loading.
+    """
+
+    labels: list[Label]
+    voxel_counts: dict[str, int]
+    shape: list[int]
+    dtype: str
+    input_orientation: str
+    network_shape: list[int]
+    sliding_window: SlidingWindowDetails
+    mlx_device: str
+    duration_seconds: float
+
+
+class RunOutputs(TypedDict):
+    """Files written by a run; ``mask`` is the absolute path of ``mask.nii.gz``."""
+
+    mask: str
+
+
+class RunResult(TypedDict):
+    """Return value of ``unest_mlx.medmlx.run``."""
+
+    outputs: RunOutputs
+    details: RunDetails

unest_mlx/conversion.py

@@ -0,0 +1,90 @@
+"""Convert the bundle's Torch checkpoint to the safetensors file ``run`` loads.
+
+Requires the ``conversion`` extra (Torch). Run once per checkpoint:
+
+    python -m unest_mlx.conversion \\
+        renalStructures_UNEST_segmentation/models/model.pt unest.safetensors
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+from typing import cast
+
+from medmlx_core.conversion import load_torch_checkpoint
+from medmlx_core.errors import InvalidInputError
+from medmlx_core.runtime import import_mlx
+
+from unest_mlx import mlx_weights
+from unest_mlx._typing import NativeRuntime
+
+
+def convert_checkpoint(
+    source: str | os.PathLike[str],
+    destination: str | os.PathLike[str],
+    *,
+    overwrite: bool = False,
+) -> Path:
+    """Write graph-compatible float32 tensors, with their original keys, as safetensors.
+
+    The checkpoint is read with ``weights_only=True`` and admitted against the
+    upstream network's parameter names and shapes by ``convert_state_dict``. Tensors
+    keep the Torch layout and values; nothing is transposed, cast or quantized. The
+    inventory check does not establish checkpoint identity or origin.
+    Alias checks occur before writing and do not prevent concurrent path changes;
+    the destination write itself is not atomic.
+
+    Args:
+        source: The bundle's ``models/model.pt``.
+        destination: Output path; must end in ``.safetensors``. Parent directories are
+            created.
+        overwrite: Replace ``destination`` if it already exists; aliases of the
+            source checkpoint are rejected even when overwrite is enabled.
+
+    Returns:
+        The destination path.
+
+    Raises:
+        InvalidInputError: ``source`` is not a file, ``destination`` has the wrong suffix,
+            aliases the source, or exists without ``overwrite``, or the checkpoint does
+            not match the released parameter inventory.
+    """
+
+    source_path = Path(source).expanduser()
+    target = Path(destination).expanduser()
+    if not source_path.is_file():
+        raise InvalidInputError(f"UNesT checkpoint not found: {source_path}")
+    if target.suffix != ".safetensors":
+        raise InvalidInputError(f"Destination must end in .safetensors: {target}")
+    if target.resolve() == source_path.resolve() or (
+        target.exists() and target.samefile(source_path)
+    ):
+        raise InvalidInputError("UNesT conversion output must differ from the source checkpoint")
+    if target.exists() and not overwrite:
+        raise InvalidInputError(f"Refusing to replace {target}; pass overwrite=True")
+    weights = mlx_weights.convert_state_dict(load_torch_checkpoint(source_path, weights_only=True))
+    import_mlx()
+    import mlx.core as mx
+
+    target.parent.mkdir(parents=True, exist_ok=True)
+    cast(NativeRuntime, mx).save_safetensors(
+        str(target), {key: mx.array(value) for key, value in weights.items()}
+    )
+    return target
+
+
+def main(argv: list[str] | None = None) -> None:
+    """Command-line entry point: ``SOURCE DESTINATION [--overwrite]``; prints the output path."""
+
+    parser = argparse.ArgumentParser(description=cast(str, __doc__).splitlines()[0])
+    parser.add_argument("source", help="Bundle checkpoint, models/model.pt")
+    parser.add_argument("destination", help="Output .safetensors path")
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args(argv)
+    print(convert_checkpoint(args.source, args.destination, overwrite=args.overwrite))
+
+
+if __name__ == "__main__":
+    main()

unest_mlx/medmlx.py

@@ -0,0 +1,279 @@
+"""medmlx declaration and end-to-end runner for renal UNesT segmentation.
+
+Follows ``renalStructures_UNEST_segmentation`` 0.2.7's inference configuration:
+the NIfTI reader, RAS/1 mm preprocessing, batches of four 96-cubed windows,
+softmax, inverse transforms, argmax, and float32 NIfTI writing.
+
+The model is intended for contrast-enhanced CT in the arterial or portal venous phase,
+as in the upstream bundle's training data. It has failed on the non-contrast scans tried
+so far (the kidneys were missed), so check enhancement before running it. This is research
+software, not a diagnostic tool.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from pathlib import Path
+from typing import cast
+
+import nibabel as nib_runtime
+import numpy as np
+from medmlx_core.errors import AssetNotReadyError, InvalidInputError, ModelExecutionError
+from medmlx_core.runtime import import_mlx
+from nibabel.filebasedimages import ImageFileError
+
+from unest_mlx._typing import (
+    DoubleArray,
+    FloatArray,
+    HostArray,
+    Label,
+    NativeNetwork,
+    NativeRuntime,
+    NiftiImage,
+    NiftiIo,
+    RunResult,
+)
+
+nib = cast(NiftiIo, nib_runtime)
+
+MODEL = {
+    "id": "unest",
+    "label": "Renal structures UNesT segmentation",
+    "task": "segmentation",
+    "run": "unest_mlx.medmlx:run",
+    "summary": (
+        "Kidney cortex, medulla, and pelvicalyceal system segmentation for one CT "
+        "volume with the MONAI renalStructures_UNEST_segmentation 0.2.7 bundle."
+    ),
+}
+
+OUTPUT_NAME = "mask.nii.gz"
+# The bundle's configs/metadata.json network_data_format.outputs.pred.channel_def.
+LABELS: tuple[Label, ...] = (
+    {"value": 0, "name": "background"},
+    {"value": 1, "name": "kidney cortex"},
+    {"value": 2, "name": "medulla"},
+    {"value": 3, "name": "pelvicalyceal system"},
+)
+# configs/inference.json: SlidingWindowInferer(roi_size=[96, 96, 96], sw_batch_size=4,
+# overlap=0.5). The roi size and overlap are fixed in UNesTPredictor.
+SW_BATCH_SIZE = 4
+_CONVERT_HINT = (
+    "Convert the bundle checkpoint once with the conversion extra: python -m "
+    "unest_mlx.conversion <bundle>/models/model.pt unest.safetensors"
+)
+
+
+def run(
+    *,
+    image_path: str | os.PathLike[str],
+    weights_path: str | os.PathLike[str],
+    output_dir: str | os.PathLike[str],
+    overwrite: bool = False,
+) -> RunResult:
+    """Segment kidney cortex, medulla, and pelvicalyceal system in one CT volume.
+
+    Preprocessing follows the bundle: reorient to RAS, resample to 1 mm spacing
+    (bilinear), clip -175 to 250 HU and scale to 0-1. Windows of 96x96x96 voxels with 0.5
+    overlap are run through the network, softmaxed, resampled back to the input grid and
+    reduced with argmax. The model targets contrast-enhanced CT (arterial or portal venous
+    phase) and has failed on non-contrast scans.
+
+    Args:
+        image_path: 3D CT NIfTI (``.nii`` or ``.nii.gz``) in Hounsfield units, with
+            contrast enhancement. Header slope and intercept are applied on read.
+            Affine coordinates are interpreted as millimeters without converting
+            the NIfTI spatial-unit field; supply a millimeter affine.
+        weights_path: ``.safetensors`` converted from the bundle's ``models/model.pt``
+            by ``unest_mlx.conversion.convert_checkpoint``.
+        output_dir: Directory for the mask; created if missing.
+        overwrite: Replace an existing mask in ``output_dir``; input aliases are refused.
+
+    Returns:
+        ``{"outputs": {"mask": path}, "details": {...}}``. The mask keeps the input's
+        shape and reader affine; upstream header correction may replace a conflicting
+        affine with the form matching the header zooms. Like ``SaveImaged`` it stores
+        float32 label values 0-3 with sform and qform code 1. The new output header
+        leaves spatial units unknown.
+
+    Raises:
+        InvalidInputError: The image is missing, not a finite 3D NIfTI with a valid
+            affine, ``overwrite`` is not a bool, ``output_dir`` is not a directory, or the
+            mask exists and ``overwrite`` is false.
+        AssetNotReadyError: The weights are missing or do not match the release.
+        ModelExecutionError: The network produced non-finite logits or invalid labels.
+        MissingDependencyError: MLX or its Metal GPU is unavailable (no CPU fallback).
+    """
+
+    started = time.perf_counter()
+    source = _require_nifti(image_path)
+    weights = _require_weights(weights_path)
+    if not isinstance(cast(object, overwrite), bool):
+        raise InvalidInputError("overwrite must be true or false")
+    directory = Path(output_dir).expanduser().absolute()
+    if directory.exists() and not directory.is_dir():
+        raise InvalidInputError(f"output_dir is not a directory: {directory}")
+    mask_path = directory / OUTPUT_NAME
+    if mask_path.resolve() == source.resolve() or (
+        mask_path.exists() and mask_path.samefile(source)
+    ):
+        raise InvalidInputError("UNesT output must differ from the input image")
+    if mask_path.exists() and not overwrite:
+        raise InvalidInputError(f"{mask_path} already exists; pass overwrite=true to replace it")
+    from unest_mlx import transforms
+
+    image, data = read_volume(source)
+    reader_affine = transforms.correct_nifti_header_if_necessary(image).affine
+    network_input, geometry = transforms.preprocess(data, reader_affine)
+    mx = cast(NativeRuntime, import_mlx())
+    from unest_mlx.mlx_predictor import UNesTPredictor
+
+    network = _load_network(weights, mx=mx)
+    predictor = UNesTPredictor(network, mx=mx, sw_batch_size=SW_BATCH_SIZE)
+    try:
+        logits = predictor(network_input[None])
+    except RuntimeError as error:
+        raise ModelExecutionError(str(error)) from error
+    labels = transforms.postprocess(logits[0], geometry)
+    if labels.shape != geometry.original_shape:
+        raise ModelExecutionError("UNesT mask does not match the input grid")
+    if not set(np.unique(labels).tolist()) <= {label["value"] for label in LABELS}:
+        raise ModelExecutionError("UNesT exported labels outside the bundle contract")
+    write_mask(mask_path, labels, reader_affine)
+    return {
+        "outputs": {"mask": mask_path.resolve().as_posix()},
+        "details": {
+            "labels": [cast(Label, dict(label)) for label in LABELS],
+            "voxel_counts": {
+                str(label["name"]): int(np.count_nonzero(labels == label["value"]))
+                for label in LABELS
+            },
+            "shape": [int(value) for value in labels.shape],
+            "dtype": "float32",
+            "input_orientation": "".join(nib.aff2axcodes(reader_affine)),
+            "network_shape": [int(value) for value in geometry.network_shape],
+            "sliding_window": {
+                "roi_size": [96, 96, 96],
+                "sw_batch_size": SW_BATCH_SIZE,
+                "overlap": 0.5,
+            },
+            "mlx_device": str(mx.default_device()),
+            "duration_seconds": round(time.perf_counter() - started, 3),
+        },
+    }
+
+
+def _require_nifti(path: str | os.PathLike[str]) -> Path:
+    source = Path(path).expanduser().absolute()
+    if not source.name.endswith((".nii", ".nii.gz")):
+        raise InvalidInputError(f"image_path must be a .nii or .nii.gz file: {source}")
+    if not source.is_file():
+        raise InvalidInputError(f"image_path does not exist: {source}")
+    return source
+
+
+def _require_weights(path: str | os.PathLike[str]) -> Path:
+    weights = Path(path).expanduser().absolute()
+    if not weights.is_file():
+        raise AssetNotReadyError(
+            f"UNesT weights not found: {weights}; expected the .safetensors converted "
+            "from the bundle's models/model.pt",
+            reason=f"missing {weights}",
+            hint=_CONVERT_HINT,
+        )
+    if weights.suffix != ".safetensors":
+        raise AssetNotReadyError(
+            f"UNesT weights must be a converted .safetensors file, got {weights.name}",
+            reason=f"unconverted {weights.name}",
+            hint=_CONVERT_HINT,
+        )
+    return weights
+
+
+def read_volume(path: Path) -> tuple[NiftiImage, HostArray]:
+    """Read a finite 3D NIfTI volume and return the nibabel image and its voxel array.
+
+    The array is ``np.asanyarray(image.dataobj)`` in file axis order, with the header's
+    slope and intercept applied. Which affine (sform or qform) defines the grid is decided
+    later by ``transforms.correct_nifti_header_if_necessary``, as in the upstream reader.
+
+    Raises:
+        InvalidInputError: The file cannot be read as NIfTI, is not 3D, has non-finite
+            voxels, has a non-finite or singular affine, or has non-positive header
+            spacing.
+    """
+
+    try:
+        image = nib.load(path.as_posix())
+        data = np.asanyarray(image.dataobj)
+    except (OSError, ValueError, ImageFileError) as error:
+        raise InvalidInputError(f"Invalid UNesT input NIfTI: {path}") from error
+    if not isinstance(image, nib_runtime.Nifti1Image | nib_runtime.Nifti2Image):
+        raise InvalidInputError(f"UNesT input must be NIfTI: {path}")
+    if data.ndim != 3:
+        raise InvalidInputError(f"UNesT input must be a 3D NIfTI; got shape {data.shape}")
+    if not np.issubdtype(data.dtype, np.number) or not np.all(np.isfinite(data)):
+        raise InvalidInputError("UNesT input contains non-finite voxels")
+    affine = np.asarray(image.affine, dtype=np.float64)
+    if affine.shape != (4, 4) or not np.all(np.isfinite(affine)):
+        raise InvalidInputError("UNesT input affine must be a finite 4x4 matrix")
+    if np.linalg.det(affine[:3, :3]) == 0:
+        raise InvalidInputError("UNesT input affine must be invertible")
+    if any(float(value) <= 0.0 for value in image.header.get_zooms()[:3]):
+        raise InvalidInputError("UNesT input spacing_mm values must be positive")
+    return cast(NiftiImage, image), data
+
+
+def _load_network(path: Path, *, mx: NativeRuntime) -> NativeNetwork:
+    from unest_mlx import mlx_weights
+    from unest_mlx.mlx_model import UNesT
+
+    try:
+        loaded = mx.load(str(path))
+    except (OSError, RuntimeError, ValueError) as error:
+        raise AssetNotReadyError(
+            f"Cannot read UNesT weights {path}: {error}",
+            reason=f"unreadable {path.name}",
+            hint=_CONVERT_HINT,
+        ) from error
+    if not isinstance(loaded, dict):
+        raise AssetNotReadyError(
+            f"UNesT weights are not a tensor mapping: {path}",
+            reason=f"unreadable {path.name}",
+            hint=_CONVERT_HINT,
+        )
+    try:
+        state = {str(key): cast(HostArray, np.array(value)) for key, value in loaded.items()}
+        weights = mlx_weights.convert_state_dict(state)
+    except (InvalidInputError, RuntimeError, TypeError, ValueError) as error:
+        raise AssetNotReadyError(
+            f"UNesT weights do not match the released 0.2.7 checkpoint: {error}",
+            reason=f"mismatched {path.name}",
+            hint=_CONVERT_HINT,
+        ) from error
+    arrays = {key: mx.array(value) for key, value in weights.items()}
+    mx.eval(arrays)
+    return UNesT(arrays, mx=mx)
+
+
+def write_mask(path: Path, labels: FloatArray, affine: DoubleArray) -> None:
+    """Write a label map like MONAI's ``NibabelWriter``, replacing ``path`` atomically.
+
+    Args:
+        path: Destination; must end in ``.nii.gz``. Parent directories are created.
+        labels: Label values on the reader grid, stored as float32.
+        affine: 4x4 voxel-to-world affine, written as both sform and qform with code 1.
+    """
+
+    image = nib.Nifti1Image(np.asarray(labels, dtype=np.float32), affine)
+    image.set_sform(affine, code=1)
+    image.set_qform(affine, code=1)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    # Write beside the target and rename, so a failed run never leaves a truncated mask.
+    partial = path.with_name(f".{path.name[: -len('.nii.gz')]}.{os.getpid()}.partial.nii.gz")
+    try:
+        nib.save(image, partial)
+        os.replace(partial, path)
+    finally:
+        partial.unlink(missing_ok=True)