totalsegmentator-mlx 0.1.1__py3-none-any.whl

totalsegmentator_mlx/__init__.py

@@ -0,0 +1,23 @@
+"""TotalSegmentator v2 ``total`` CT inference on Apple Silicon (MLX), research software.
+
+``segment`` runs a CT through the native workflow and returns uint8 global label IDs
+(1-117) on the input grid; ``convert_model`` converts upstream fold-0 checkpoints to
+the layout ``segment`` reads. Names load on first access, so importing the package
+does not import the MLX engine.
+"""
+
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from totalsegmentator_mlx.conversion import convert_model
+    from totalsegmentator_mlx.pipeline import segment
+
+__all__ = ["convert_model", "segment"]
+
+
+def __getattr__(name: str) -> object:
+    owner = {"segment": "pipeline", "convert_model": "conversion"}.get(name)
+    if owner is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    return getattr(import_module(f"{__name__}.{owner}"), name)

totalsegmentator_mlx/__main__.py

@@ -0,0 +1,6 @@
+"""Run the installed command through python -m totalsegmentator_mlx."""
+
+from totalsegmentator_mlx.cli import main
+
+if __name__ == "__main__":
+    main()

totalsegmentator_mlx/cli.py

@@ -0,0 +1,128 @@
+"""Lazy command-line entry point for conversion, CT prediction and label discovery."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+from totalsegmentator_mlx.labels import LABELS, MODEL_PARTS, QUALITY_MODELS
+
+if TYPE_CHECKING:
+    from totalsegmentator_mlx.models import Mode
+
+
+class Arguments(argparse.Namespace):
+    """Typed namespace for the ``predict``, ``convert`` and ``labels`` subcommands.
+
+    Only the fields of the chosen subcommand are set by argparse.
+    """
+
+    command: str
+    input_path: Path
+    output_path: Path
+    model_dir: Path
+    mode: str
+    model_part: str | None
+    organ: list[str] | None
+    roi_subset: list[str] | None
+    robust_crop: bool
+    body_seg: bool
+    models: list[str] | None
+    force_split: bool
+    overwrite: bool
+    probabilities_path: Path | None
+    source_model_dir: Path
+    output_model_dir: Path
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Native TotalSegmentator v2 CT inference on Metal")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+    predict = subparsers.add_parser("predict", help="Segment a NIfTI CT on its original grid")
+    predict.add_argument("-i", "--input", dest="input_path", type=Path, required=True)
+    predict.add_argument("-o", "--output", dest="output_path", type=Path, required=True)
+    predict.add_argument("--model-dir", type=Path, required=True)
+    predict.add_argument(
+        "--mode",
+        choices=["standard", *QUALITY_MODELS],
+        default="standard",
+        help="standard: five 1.5 mm parts; fast: 3 mm model; coarse: 6 mm model (upstream fastest)",
+    )
+    predict.add_argument("--model-part", choices=list(MODEL_PARTS))
+    predict.add_argument("--organ", nargs="+", help="Filter canonical organ names after inference")
+    predict.add_argument(
+        "--roi-subset",
+        nargs="+",
+        help="Crop to these structures with a rough 6 mm segmentation and run only their parts",
+    )
+    predict.add_argument(
+        "--robust-crop", action="store_true", help="Locate --roi-subset with the 3 mm model"
+    )
+    predict.add_argument(
+        "--body-seg", action="store_true", help="Crop to the body with the Dataset300 model"
+    )
+    predict.add_argument("--force-split", action="store_true")
+    predict.add_argument("--overwrite", action="store_true")
+    predict.add_argument("--probabilities-path", type=Path)
+    convert = subparsers.add_parser("convert", help="Convert trusted local fold-0 checkpoints")
+    convert.add_argument("--source-model-dir", type=Path, required=True)
+    convert.add_argument("--output-model-dir", type=Path, required=True)
+    convert_models = convert.add_mutually_exclusive_group()
+    convert_models.add_argument("--model-part", choices=[*MODEL_PARTS, *QUALITY_MODELS, "body"])
+    convert_models.add_argument(
+        "--models",
+        nargs="+",
+        choices=[*MODEL_PARTS, *QUALITY_MODELS, "body"],
+        help="Models to convert into one new root (default: the five standard parts)",
+    )
+    subparsers.add_parser("labels", help="Print released global label IDs as JSON")
+    return parser
+
+
+def main(argv: list[str] | None = None) -> None:
+    """Run a subcommand; ``labels`` prints the global ID-to-name table as JSON.
+
+    Exceptions that are ``OSError``, ``ValueError``, ``RuntimeError`` or ``ImportError``
+    (which includes the ``medmlx_core`` errors) are reported on stderr with exit status 2.
+    ``predict`` prints the output path on success and ``convert`` the new model root.
+    """
+    parser = _parser()
+    args = parser.parse_args(argv, namespace=Arguments())
+    if args.command == "labels":
+        print(json.dumps(LABELS, indent=2))
+        return
+    try:
+        if args.command == "convert":
+            from totalsegmentator_mlx.conversion import convert_model
+
+            output = convert_model(
+                args.source_model_dir,
+                args.output_model_dir,
+                model_part=args.model_part,
+                models=args.models,
+            )
+            print(output)
+            return
+        from totalsegmentator_mlx.pipeline import segment
+
+        segment(
+            args.input_path,
+            model_dir=args.model_dir,
+            output_path=args.output_path,
+            mode=cast("Mode", args.mode),
+            model_part=args.model_part,
+            prompt=args.organ,
+            roi_subset=args.roi_subset,
+            robust_crop=args.robust_crop,
+            body_seg=args.body_seg,
+            force_split=args.force_split,
+            overwrite=args.overwrite,
+            probabilities_path=args.probabilities_path,
+            progress=lambda message: print(message, file=sys.stderr, flush=True),
+        )
+        print(args.output_path)
+    except (OSError, ValueError, RuntimeError, ImportError) as error:
+        parser.exit(2, f"{parser.prog}: {error}\n")

totalsegmentator_mlx/conversion.py

@@ -0,0 +1,88 @@
+"""Convert trusted local CT checkpoint folders through the public nnunet-mlx API."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from medmlx_core.errors import AssetNotReadyError, InvalidInputError
+
+from totalsegmentator_mlx.models import (
+    CONVERTIBLE,
+    MODEL_PARTS,
+    SOURCE_PROFILE,
+    model_folder,
+    validate_metadata,
+)
+
+
+def _source_folder(source: Path, task_id: int) -> Path:
+    """Find a trainer folder as ``<id>/`` or in upstream's ``Dataset<id>_*/`` results layout."""
+    folder = model_folder(source, task_id)
+    if folder.exists():
+        return folder
+    matches = sorted(source.glob(f"Dataset{task_id}_*/*__nnUNetPlans__3d_fullres"))
+    if len(matches) != 1:
+        raise AssetNotReadyError(f"Expected one Dataset{task_id} trainer folder in {source}")
+    return matches[0]
+
+
+def convert_model(
+    source_model_dir: str | Path,
+    output_model_dir: str | Path,
+    *,
+    model_part: str | None = None,
+    models: Sequence[str] | None = None,
+) -> Path:
+    """Convert fold 0 into a new model root, publishing it only after every model succeeds.
+
+    ``models`` names any of the five parts, ``fast`` (297), ``coarse`` (298)
+    and ``body`` (300); the default converts the five standard parts. Each
+    source is an nnU-Net trainer folder with plans.json, dataset.json and
+    fold_0/checkpoint_final.pth, found as ``<id>/`` or in upstream's
+    ``Dataset<id>_*/<trainer>__nnUNetPlans__3d_fullres`` results layout. A
+    single ``model_part`` also accepts its trainer folder directly. Conversion
+    needs the optional conversion extra; inference never imports Torch.
+
+    Returns:
+        ``output_model_dir``, containing one converted folder per model, named by dataset
+        ID (``291`` ... ``295``, ``297``, ``298``, ``300``).
+
+    Raises:
+        InvalidInputError: Both ``model_part`` and ``models`` are given, a name is unknown,
+            or the destination exists or lies inside the source folder.
+        AssetNotReadyError: A source folder, its metadata or its fold-0 checkpoint is
+            missing or does not match the expected dataset.
+    """
+    source, output = Path(source_model_dir), Path(output_model_dir)
+    if model_part is not None and models is not None:
+        raise InvalidInputError("Choose model_part or models, not both")
+    names = [model_part] if model_part is not None else list(models or MODEL_PARTS)
+    if not names or set(names) - set(CONVERTIBLE):
+        raise InvalidInputError(f"Models must be among {', '.join(CONVERTIBLE)}")
+    task_ids = tuple(dict.fromkeys(CONVERTIBLE[name] for name in names))
+    if output.exists() or output.resolve().is_relative_to(source.resolve()):
+        raise InvalidInputError("Conversion requires a new destination outside the source folder")
+    folders = {task_id: _source_folder(source, task_id) for task_id in task_ids}
+    for task_id, folder in folders.items():
+        validate_metadata(folder, task_id)
+        checkpoint = folder / "fold_0" / "checkpoint_final.pth"
+        if not checkpoint.is_file():
+            raise AssetNotReadyError(f"Source fold-0 checkpoint is missing: {checkpoint}")
+    from nnunet_mlx import convert_upstream_model_folder
+
+    output.parent.mkdir(parents=True, exist_ok=True)
+    with TemporaryDirectory(prefix=".totalsegmentator_convert_", dir=output.parent) as temporary:
+        staged = Path(temporary) / "models"
+        staged.mkdir()
+        for task_id, folder in folders.items():
+            convert_upstream_model_folder(
+                folder,
+                staged / str(task_id),
+                folds=[0],
+                configuration="3d_fullres",
+                source_profile=SOURCE_PROFILE,
+            )
+        staged.rename(output)
+    return output

totalsegmentator_mlx/engine.py

@@ -0,0 +1,140 @@
+"""Delegate each CT network to the public nnunet-mlx Metal inference boundary."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from functools import lru_cache
+from pathlib import Path
+from threading import RLock
+from typing import TYPE_CHECKING
+
+import nibabel as nib
+import numpy as np
+from medmlx_core.errors import ModelExecutionError
+
+from totalsegmentator_mlx.imageio import load_nifti, spatial_affine
+from totalsegmentator_mlx.labels import BODY_MODEL
+from totalsegmentator_mlx.models import ARITHMETIC, SOURCE_PROFILE, TASK_SPACING
+
+if TYPE_CHECKING:
+    from nnunet_mlx import ModelFolderPredictor
+
+_runtime_lock = RLock()
+
+
+def _model_identity(folder: Path) -> tuple[tuple[int, int, int, int, int] | None, ...]:
+    """Invalidate cached weights/metadata on replacement or an in-place update."""
+    paths = (
+        folder / "plans.json",
+        folder / "dataset.json",
+        folder / "fold_0/checkpoint_final/state.json",
+        folder / "fold_0/checkpoint_final/weights.safetensors",
+        folder / "conversion.json",
+    )
+    identities: list[tuple[int, int, int, int, int] | None] = []
+    for path in paths:
+        if path.name == "conversion.json" and not path.exists():
+            identities.append(None)
+            continue
+        info = path.stat()
+        identities.append(
+            (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns)
+        )
+    return tuple(identities)
+
+
+def _step_fraction(task_id: int) -> float:
+    """Return upstream's tile step: 0.8 below a 3 mm grid and 0.5 otherwise."""
+    return 0.8 if TASK_SPACING[task_id] < 3.0 else 0.5
+
+
+# maxsize 8 keeps all five parts, both single-model tiers and the body model loaded.
+@lru_cache(maxsize=8)
+def _load_predictor(
+    folder: Path,
+    task_id: int,
+    identity: tuple[tuple[int, int, int, int, int] | None, ...],
+) -> ModelFolderPredictor:
+    """Load one fold-0 model onto Metal; ``identity`` is only the cache key (file stats).
+
+    Raises:
+        ModelExecutionError: Metal is unavailable, the model does not take one CT channel, or a
+            non-body model carries test-time mirroring axes.
+    """
+    from nnunet_mlx import ModelFolderPredictor, load_model_bundle
+    from nnunet_mlx.runtime import runtime_info
+
+    if runtime_info().backend != "metal":
+        raise ModelExecutionError("TotalSegmentator requires an Apple Metal GPU")
+    bundle = load_model_bundle(
+        folder,
+        configuration="3d_fullres",
+        folds=[0],
+        source_profile=SOURCE_PROFILE,
+        arithmetic=ARITHMETIC,
+    )
+    if bundle.image_channels != 1 or bundle.input_channels != 1:
+        raise ModelExecutionError("CT models require exactly one CT input channel")
+    if task_id == BODY_MODEL:
+        # The body model was trained with mirroring; upstream predicts it without TTA.
+        bundle = replace(bundle, mirroring_axes=())
+    elif bundle.mirroring_axes:
+        raise ModelExecutionError("TotalSegmentator CT requires a no-mirroring model bundle")
+    # Upstream nnU-Net blends every model's tiles in float16; float32 blending
+    # changed 549 voxels of AMOS22 case amos_0008 against the upstream standard pipeline.
+    return ModelFolderPredictor(
+        bundle,
+        step_fraction=_step_fraction(task_id),
+        compile_models=False,
+        accumulation="nnunet-fp16",
+    )
+
+
+def predict_image(image: nib.Nifti1Image, folder: Path, task_id: int) -> nib.Nifti1Image:
+    """Return local labels on the supplied XYZ grid, retaining loaded networks.
+
+    ``image`` is float32 CT on the model grid (1.5, 3 or 6 mm); the returned labels are
+    this model's own IDs (``model_labels(task_id)``), not yet merged into global IDs.
+
+    The public engine NIfTI API transfers final labels from Metal. The lock owns
+    the cached predictors' single-GPU lifecycle across concurrent API callers.
+    """
+    with _runtime_lock:
+        folder = folder.resolve()
+        predictor = _load_predictor(folder, task_id, _model_identity(folder))
+        return predictor.predict_nifti(image)
+
+
+def predict_part(
+    inputs: Path,
+    output: Path,
+    folder: Path,
+    task_id: int,
+    *,
+    save_probabilities: bool,
+) -> None:
+    """Export file-backed predictions and optional local-label probabilities.
+
+    Reads ``*_0000.nii.gz`` from ``inputs`` and writes ``<case>.nii.gz`` (plus ``.npz`` and
+    ``.pkl`` when ``save_probabilities``) to ``output``, then checks that each result keeps
+    its input's shape and affine.
+    """
+    with _runtime_lock:
+        folder = folder.resolve()
+        predictor = _load_predictor(folder, task_id, _model_identity(folder))
+        predictor.predict_folder(
+            inputs,
+            output,
+            save_probabilities=save_probabilities,
+            overwrite=False,
+        )
+    for path in sorted(inputs.glob("*_0000.nii.gz")):
+        source = load_nifti(path)
+        result = load_nifti(output / path.name.replace("_0000", ""))
+        if result.shape != source.shape or not np.allclose(
+            spatial_affine(result),
+            spatial_affine(source),
+            rtol=0,
+            atol=1e-5,
+        ):
+            raise ModelExecutionError("Engine segmentation geometry differs from its CT input")

totalsegmentator_mlx/geometry.py

@@ -0,0 +1,214 @@
+"""Outer CT resampling and reversible NIfTI orientation in voxel (x, y, z) order.
+
+Volumes are 3D arrays in NIfTI storage order ``(x, y, z)``. ``canonical_ct`` reorients to
+the closest RAS+ axes, ``resample_ct`` moves to an isotropic model grid (mm) and
+``restore_labels`` returns labels to the source shape, orientation and affine.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from importlib import import_module
+from pathlib import Path
+from typing import Protocol, cast
+
+import nibabel as nib
+import numpy as np
+from medmlx_core.errors import InvalidInputError, ModelExecutionError
+from numpy.typing import NDArray
+
+from totalsegmentator_mlx.imageio import load_nifti, spatial_affine
+from totalsegmentator_mlx.resampling import cubic_zoom, nearest_labels
+
+type Labels = NDArray[np.uint8]
+
+
+class DtypeSetter(Protocol):
+    def set_data_dtype(self, datatype: object) -> None: ...
+
+
+def load_ct(source: str | Path | nib.Nifti1Image) -> nib.Nifti1Image:
+    """Load a finite, nonempty 3D CT with an invertible affine; voxel spacing is in mm.
+
+    ``source`` is a path or an in-memory NIfTI-1 image. Intensities are expected in
+    Hounsfield units (not checked) and must be integer or float and finite. Spatial units
+    must be millimeters or unspecified. Retain validated proxy values so canonicalization
+    does not decode them again. Array-backed inputs are returned unchanged; caller data
+    and headers are not modified.
+
+    Raises:
+        InvalidInputError: The volume is not 3D, is empty or non-numeric, has non-finite
+            values, has a non-invertible or non-finite affine, or has non-mm spatial units.
+    """
+    image = source if isinstance(source, nib.Nifti1Image) else load_nifti(source)
+    values = np.asanyarray(image.dataobj)
+    if values.ndim != 3 or not values.size or values.dtype.kind not in "iuf":
+        raise InvalidInputError("Input must be a nonempty three-dimensional CT")
+    if not np.isfinite(values).all():
+        raise InvalidInputError("CT intensities must be finite")
+    affine = image.affine
+    if (
+        affine is None
+        or not np.isfinite(affine).all()
+        or abs(float(np.linalg.det(affine[:3, :3]))) < 1e-12
+        or not np.allclose(affine[3], [0, 0, 0, 1])
+    ):
+        raise InvalidInputError("CT requires a finite, invertible spatial affine")
+    if image.header.get_xyzt_units()[0] not in {"unknown", "mm"}:
+        raise InvalidInputError("CT spatial units must be millimeters")
+    return (
+        image
+        if cast(object, image.dataobj) is values
+        else nib.Nifti1Image(values, affine, image.header)
+    )
+
+
+def canonical_ct(source: nib.Nifti1Image) -> nib.Nifti1Image:
+    """Reorient to the closest canonical (RAS+) axes, dropping the header.
+
+    Only the voxel array and affine are kept (axes are flipped or permuted, nothing is
+    interpolated). Restore the source orientation with ``restore_labels``.
+    """
+    # Match upstream's header-free staging before canonical reorientation.
+    clean = nib.Nifti1Image(np.asanyarray(source.dataobj), source.affine)
+    canonicalize = cast(
+        Callable[[nib.Nifti1Image], object],
+        import_module("nibabel").as_closest_canonical,
+    )
+    canonical = canonicalize(clean)
+    if not isinstance(canonical, nib.Nifti1Image):
+        raise InvalidInputError("Canonical CT must remain a NIfTI-1 image")
+    return canonical
+
+
+def resample_ct(image: nib.Nifti1Image, target: float) -> nib.Nifti1Image:
+    """Stage cubic CT values on an isotropic ``target`` mm grid as int32, as upstream does.
+
+    ``image`` is canonical-orientation CT whose header zooms are the voxel size in mm.
+    Resampling is a cubic spline (SciPy ``zoom`` semantics, nearest boundary) in float64;
+    results are then cast to int32, which truncates toward zero. When the spacing already
+    equals ``target`` the values are only cast. Raises ``InvalidInputError`` if the field
+    of view is smaller than one model voxel.
+    """
+    spacing = np.asarray(image.header.get_zooms(), dtype=np.float64)
+    factors = spacing / target
+    shape = np.rint(np.asarray(image.shape) * factors)
+    if np.any(shape < 1):
+        raise InvalidInputError(f"CT field of view is too small for the {target} mm model grid")
+    if np.array_equal(spacing, [target] * 3):
+        values = np.asanyarray(image.dataobj).astype(np.int32)
+    else:
+        values = cubic_zoom(np.asarray(image.get_fdata(), dtype=np.float64), factors).astype(
+            np.int32
+        )
+    affine = np.array(image.affine, copy=True)
+    affine[:3, :3] /= factors
+    return nib.Nifti1Image(values, affine)
+
+
+def split_ct(
+    image: nib.Nifti1Image, *, force_split: bool, multimodel: bool
+) -> tuple[nib.Nifti1Image, ...]:
+    """Use upstream's three overlapping z slabs for large multi-model volumes or on request.
+
+    Splitting happens for multi-model runs whose grid exceeds 512*512*900 voxels with more
+    than 200 slices, or when ``force_split``. The slabs cover the thirds of the z axis, each
+    extended by a margin of about 20 slices into its neighbors (float64 values);
+    ``join_slabs`` drops the margins again. Raises ``InvalidInputError`` if a split is
+    requested for fewer than 60 slices.
+    """
+    depth = image.shape[2]
+    split = force_split or (multimodel and np.prod(image.shape) > 512 * 512 * 900 and depth > 200)
+    if not split:
+        return (image,)
+    if depth < 60:
+        raise InvalidInputError("Splitting requires at least 60 slices on the model grid")
+    third, margin = depth // 3, 20
+    values = image.get_fdata()
+    # Each slab retains the upstream staging affine; merge uses voxel indices.
+    slices = (
+        slice(None, third + margin),
+        slice(third + 1 - margin, 2 * third + margin),
+        slice(2 * third + 1 - margin, None),
+    )
+    return tuple(nib.Nifti1Image(values[:, :, item], image.affine) for item in slices)
+
+
+def join_slabs(slabs: Sequence[Labels], shape: tuple[int, ...]) -> Labels:
+    """Join one volume, or three ``split_ct`` slab predictions, into labels of ``shape``.
+
+    The 20-slice overlaps are discarded, so each z slice comes from exactly one slab.
+    """
+    if len(slabs) == 1:
+        return slabs[0]
+    if len(slabs) != 3:
+        raise ModelExecutionError("Expected one volume or three prediction slabs")
+    third, margin = shape[2] // 3, 20
+    joined = np.empty(shape, dtype=np.uint8)
+    joined[:, :, :third] = slabs[0][:, :, :-margin]
+    joined[:, :, third : 2 * third] = slabs[1][:, :, margin - 1 : -margin]
+    joined[:, :, 2 * third :] = slabs[2][:, :, margin - 1 :]
+    return joined
+
+
+def restore_labels(
+    labels: Labels,
+    canonical: nib.Nifti1Image,
+    source: nib.Nifti1Image,
+) -> nib.Nifti1Image:
+    """Restore nearest-neighbor labels to the exact source shape, affine and orientation.
+
+    ``labels`` is on the resampled grid of ``canonical`` (the ``canonical_ct`` image of
+    ``source``). Labels are scaled up to the canonical shape with nearest-neighbor
+    sampling, then reoriented to match ``source``. Raises ``ModelExecutionError`` if the
+    result does not have the source's shape and affine.
+    """
+    factors = np.asarray(canonical.shape) / np.asarray(labels.shape)
+    restored = nearest_labels(labels, factors)
+    image = nib.Nifti1Image(restored, canonical.affine)
+    orientations = import_module("nibabel.orientations")
+    orientation = cast(Callable[[object], NDArray[np.float64]], orientations.io_orientation)
+    ras = cast(Callable[[str], NDArray[np.float64]], orientations.axcodes2ornt)
+    transform = cast(
+        Callable[[NDArray[np.float64], NDArray[np.float64]], NDArray[np.float64]],
+        orientations.ornt_transform,
+    )(ras("RAS"), orientation(source.affine))
+    image = image.as_reoriented(cast(Sequence[Sequence[int]], transform))
+    if image.shape != source.shape or not np.allclose(
+        spatial_affine(image), spatial_affine(source), rtol=0, atol=1e-5
+    ):
+        raise ModelExecutionError("Restored segmentation geometry differs from its CT input")
+    return source_labels(np.asanyarray(image.dataobj), source)
+
+
+def source_labels(labels: Labels, source: nib.Nifti1Image) -> nib.Nifti1Image:
+    """Wrap uint8 labels with the source affine and a copy of its header."""
+    header = source.header.copy()
+    cast(DtypeSetter, header).set_data_dtype(np.uint8)
+    return nib.Nifti1Image(labels, source.affine, header)
+
+
+def crop_ct(
+    source: nib.Nifti1Image, mask: NDArray[np.bool_], margin_mm: float
+) -> tuple[nib.Nifti1Image, tuple[slice, slice, slice]]:
+    """Crop CT to a nonempty source-grid mask's bounding box plus a margin, staged as int32.
+
+    The margin converts to whole voxels per source axis by truncation, using
+    affine column norms as upstream's header-free copy does. Cropped float64
+    values truncate to int32 before reorientation and resampling, and the
+    affine translation moves to the first retained voxel.
+    """
+    probe = nib.Nifti1Image(np.zeros((1, 1, 1), dtype=np.uint8), source.affine)
+    zooms = np.asarray(probe.header.get_zooms())
+    margin = (np.full(3, margin_mm) / zooms).astype(int)
+    occupied = [np.flatnonzero(mask.any(axis=other)) for other in ((1, 2), (0, 2), (0, 1))]
+    lower = [max(0, int(axis[0]) - pad) for axis, pad in zip(occupied, margin, strict=True)]
+    upper = [
+        min(size, int(axis[-1]) + 1 + pad)
+        for axis, pad, size in zip(occupied, margin, mask.shape, strict=True)
+    ]
+    box = (slice(lower[0], upper[0]), slice(lower[1], upper[1]), slice(lower[2], upper[2]))
+    values = np.asarray(np.asanyarray(source.dataobj)[box], dtype=np.float64).astype(np.int32)
+    affine = np.array(source.affine, dtype=np.float64, copy=True)
+    affine[:3, 3] = (affine @ np.array([*lower, 1], dtype=np.float64))[:3]
+    return nib.Nifti1Image(values, affine), box

totalsegmentator_mlx/imageio.py

@@ -0,0 +1,45 @@
+"""Typed NIfTI boundaries for NiBabel's partially annotated load/save API."""
+
+from collections.abc import Callable
+from importlib import import_module
+from pathlib import Path
+from typing import Protocol, cast
+
+import nibabel as nib
+import numpy as np
+from medmlx_core.errors import InvalidInputError
+from numpy.typing import NDArray
+
+
+class ImageWriter(Protocol):
+    def to_filename(self, filename: Path) -> None: ...
+
+
+def load_nifti(path: str | Path) -> nib.Nifti1Image:
+    """Load a NIfTI file lazily (voxels stay an array proxy).
+
+    Accepts both NIfTI-1 and NIfTI-2 files: NiBabel's ``Nifti2Image`` subclasses
+    ``Nifti1Image`` and exposes the same affine/header/dataobj interface, and
+    upstream TotalSegmentator's own loader (``nifti_io.require_nifti1``) applies the
+    identical ``isinstance(image, Nifti1Image)`` check, so this matches its behavior.
+
+    Raises ``InvalidInputError`` if the file does not load as a NIfTI image.
+    """
+    loader = cast(Callable[[str], object], import_module("nibabel").load)
+    image = loader(str(path))
+    if not isinstance(image, nib.Nifti1Image):
+        raise InvalidInputError("Input must be a NIfTI image")
+    return image
+
+
+def save_nifti(image: nib.Nifti1Image, path: Path) -> None:
+    """Write ``image`` to ``path``; compression follows the file extension (``.nii.gz``)."""
+    cast(ImageWriter, image).to_filename(path)
+
+
+def spatial_affine(image: nib.Nifti1Image) -> NDArray[np.float64]:
+    """Return the image's 4x4 voxel-to-world affine (mm) as float64."""
+    affine = image.affine
+    if affine is None:
+        raise InvalidInputError("NIfTI image requires a spatial affine")
+    return np.asarray(affine, dtype=np.float64)