Metadata-Version: 2.5
Name: nnunet-mlx
Version: 0.6.4
Summary: Run pretrained nnU-Net v2 models locally on Apple Silicon with native MLX
Project-URL: Homepage, https://medmlx.org
Project-URL: Documentation, https://github.com/MedMLX/nnunet-mlx/blob/v0.6.4/README.md
Project-URL: Repository, https://github.com/MedMLX/nnunet-mlx
Project-URL: Issues, https://github.com/MedMLX/nnunet-mlx/issues
Project-URL: Weights, https://huggingface.co/models?author=MedMLX&search=nnunet
Project-URL: Upstream, https://github.com/MIC-DKFZ/nnUNet
Project-URL: MedMLX, https://github.com/MedMLX/medmlx
Author: Joseph Sandoval, Alfredo Sandoval Jr.
License-Expression: Apache-2.0 AND MIT
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES.md
Keywords: apple-silicon,ct,image-segmentation,medical-imaging,metal,mlx,mri,nnunet
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Image Processing
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
Requires-Python: <3.14,>=3.13
Requires-Dist: blosc2<5,>=3.3.2
Requires-Dist: imagecodecs<2027,>=2025
Requires-Dist: medmlx-core==0.2.6
Requires-Dist: mlx==0.32.4; platform_system == 'Darwin' and platform_machine == 'arm64'
Requires-Dist: nibabel<6,>=5.3
Requires-Dist: numpy<3,>=2.0
Requires-Dist: pillow<14,>=13.0.0
Requires-Dist: safetensors<1,>=0.6
Requires-Dist: scikit-image==0.26.0
Requires-Dist: scipy==1.18.1
Requires-Dist: simpleitk==2.5.6
Requires-Dist: tifffile<2027,>=2025
Requires-Dist: typer==0.27.1
Provides-Extra: dev
Requires-Dist: hypothesis<7,>=6; extra == 'dev'
Requires-Dist: matplotlib<4,>=3.10; extra == 'dev'
Requires-Dist: pyright==1.1.414; extra == 'dev'
Requires-Dist: pytest<10,>=8.4; extra == 'dev'
Requires-Dist: ruff<1,>=0.13; extra == 'dev'
Provides-Extra: oracle
Requires-Dist: dynamic-network-architectures==0.4.4; extra == 'oracle'
Requires-Dist: nnunetv2==2.8.1; extra == 'oracle'
Requires-Dist: torch==2.13.0; extra == 'oracle'
Requires-Dist: urllib3<3,>=2.8; extra == 'oracle'
Description-Content-Type: text/markdown

<div align="center">

# nnunet-mlx

Run pretrained nnU-Net v2 models locally on Apple Silicon with native MLX.

[![PyPI](https://img.shields.io/pypi/v/nnunet-mlx)](https://pypi.org/project/nnunet-mlx/)
[![Python](https://img.shields.io/pypi/pyversions/nnunet-mlx)](https://pypi.org/project/nnunet-mlx/)
[![License](https://img.shields.io/badge/License-Apache_2.0-blue)](https://github.com/MedMLX/nnunet-mlx/blob/v0.6.4/LICENSE)
[![Apple Silicon](https://img.shields.io/badge/Apple_Silicon-MLX-black?logo=apple)](https://github.com/ml-explore/mlx)
[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97_Hugging_Face-weights-yellow)](https://huggingface.co/MedMLX/nnunet-liver-mlx)
[![MedMLX](https://img.shields.io/badge/MedMLX-runtime-A51C30)](https://github.com/MedMLX/medmlx)

[Liver model weights](https://huggingface.co/MedMLX/nnunet-liver-mlx) · [Upstream](https://github.com/MIC-DKFZ/nnUNet) · [MedMLX](https://github.com/MedMLX/medmlx)

<img src="https://raw.githubusercontent.com/MedMLX/nnunet-mlx/v0.6.4/docs/example.png" alt="nnU-Net segmentations of a liver CT and a brain tumor MRI" width="100%">

</div>

| Input | Output |
| :--- | :--- |
| NIfTI channels | Segmentation and optional probabilities |

## Use with MedMLX

The [MedMLX MCP server](https://github.com/MedMLX/medmlx/blob/main/docs/mcp.md) is the lowest-friction way to use this package. It installs this runtime, downloads published weights on first use (MRSegmentator needs a local `weights_path`), and lets an AI agent describe, check and run `nnunet-liver`, `nnunet-pancreas`, `nnunet-brats2021`, `nnunet-autopet` and `mrsegmentator` for you.

```sh
uv tool install --python 3.13 "medmlx[segment,mcp]"
MEDMLX="$(uv tool dir --bin --color never)/medmlx"
mkdir -p "$HOME/medmlx-work"
claude mcp add --scope user medmlx -- "$MEDMLX" serve --root "$HOME/medmlx-work"
```

For Codex, use `codex mcp add medmlx -- "$MEDMLX" serve --root "$HOME/medmlx-work"`. Other MCP clients take the same command and arguments. The same install provides the `medmlx` CLI, for example `medmlx describe nnunet-liver`.

To download the liver model ahead of time:

```sh
medmlx setup nnunet-liver
```

> **For AI agents:** prefer the MedMLX MCP tools (`list_models`, `describe_model`, `model_readiness`, `run_model`, `get_job`) over writing Python against this package.

## Install

```bash
uv venv --python 3.13
source .venv/bin/activate
uv pip install "nnunet-mlx>=0.6.4"
```

## Run

```python
from nnunet_mlx.medmlx import liver

result = liver(image_path="ct.nii.gz", output_dir="out")
print(result["outputs"]["segmentation"])
```

Omitting `weights_path` downloads `MedMLX/nnunet-liver-mlx` into the shared MedMLX cache
on first use; interrupted downloads resume and later runs work offline. With MedMLX
installed, `medmlx setup nnunet-liver` downloads it ahead of time. To use your own copy,
pass a published bundle folder as `weights_path`. `details.weights_revision` records the
downloaded commit, or `null` for a folder without a MedMLX download record.

Supports 2D and 3D PlainConvUNet and ResidualEncoderUNet. Supply input channels in the model’s expected order.

### Named models

| MedMLX model | Python runner | Inputs | Weights |
| --- | --- | --- | --- |
| `nnunet-liver` | `nnunet_mlx.medmlx.liver(image_path=...)` | CT | [MedMLX/nnunet-liver-mlx](https://huggingface.co/MedMLX/nnunet-liver-mlx) |
| `nnunet-autopet` | `nnunet_mlx.medmlx.autopet(image_paths=[ct, pet])` | Co-registered CT and SUV PET, in that order | [MedMLX/nnunet-autopet-mlx](https://huggingface.co/MedMLX/nnunet-autopet-mlx) |
| `nnunet-pancreas` | `nnunet_mlx.named_models.pancreas(image_path=...)` | CT | [MedMLX/nnunet-pancreas-mlx](https://huggingface.co/MedMLX/nnunet-pancreas-mlx) |
| `nnunet-brats2021` | `nnunet_mlx.named_models.brats(t1_path=..., t1ce_path=..., t2_path=..., flair_path=...)` | Four co-registered MRI | [MedMLX/nnunet-brats-mlx](https://huggingface.co/MedMLX/nnunet-brats-mlx), fold 0 only |
| `mrsegmentator` | `nnunet_mlx.named_models.mrsegmentator(image_path=..., weights_path=...)` | MRI | Local `weights_path` only |

All runners take `output_dir`, `weights_path` and `overwrite`. The liver and AutoPET
runners return the label map as `outputs["segmentation"]`; the `named_models` runners
return it as `outputs["mask"]`. Pancreas labels are background 0, pancreas 1 and cancer 2.
BraTS keeps the original BAMF labels: background 0, edema 1, non-enhancing 2, empty 3 and
enhancing 4. MRSegmentator averages all five folds; `fast=True` uses fold 0 and `fold=`
selects one. It runs without mirroring and with `torch-order-fp32` arithmetic.

### Model folders

`nnunet_mlx.medmlx.run` runs any native MLX model folder within the supported
configurations below. Keep the complete folder layout:

```text
mlx-model/
├── plans.json
├── dataset.json
├── conversion.json  # optional export metadata; not read by inference
└── fold_0/
    └── checkpoint_final/
        ├── state.json
        └── weights.safetensors
```

| Argument | Default | Meaning |
| --- | --- | --- |
| `weights_path` | required | Model folder root |
| `image_paths` | required | `.nii` or `.nii.gz` images, one per modality in `dataset.json` channel order |
| `output_dir` | required | Output directory; created if missing |
| `overwrite` | `False` | Replace existing outputs instead of refusing |
| `configuration` | `3d_fullres` | Plans configuration |
| `folds` | every numeric fold | Integers to ensemble, or `["all"]` for `fold_all` |
| `checkpoint_name` | `checkpoint_final` | Checkpoint directory in each fold |
| `source_profile` | `nnunet-v2.8.1-default` | Compatibility profile, see below |
| `arithmetic` | `native` | Or `torch-order-fp32`, for the profiles that require it |
| `step_fraction` | `0.5` | Sliding-window step, in (0, 1] |
| `use_mirroring` | `True` | Use the checkpoint's mirroring test-time augmentation |
| `save_probabilities` | `False` | Also write class or region probabilities |

Outputs, returned as absolute paths in `{"outputs": {...}, "details": {...}}`:

| Role | File | Content |
| --- | --- | --- |
| `segmentation` | `mask.nii.gz` (`mask.nii` when `dataset.json` says so) | Label map with the input image's shape and geometry |
| `probabilities` | `probabilities.npz` | Channel-first class or region probabilities (`save_probabilities` only) |
| `probability_properties` | `probabilities.pkl` | nnU-Net properties for the probability grid (`save_probabilities` only) |

Probabilities stay on the image reader's array grid; read them with their properties
sidecar. As in nnU-Net, requesting probabilities makes class models take the argmax after
float32 softmax rather than of the logits, so near ties can segment differently.
`details` holds the `dataset.json` label table, `output_kind` (`binary_mask` or
`label_map`), the selected options, checkpoint signatures and timings. The runner does
not expose cascade stages; use the CLI's `--previous-stage`.

Errors are [medmlx-core](https://github.com/MedMLX/medmlx/tree/main/packages/core) types,
each chaining the package's `ContractError`: `InvalidInputError` for bad images, options
or existing outputs; `AssetNotReadyError` for a missing, malformed or mismatched model
folder; `MissingDependencyError` when MLX 0.32.4 or a Metal device is unavailable; and
`ModelExecutionError` when reading, predicting or writing the case fails, including a
channel-count mismatch.

### Command line

Download the published liver model with `hf`, included by the runtime's
`medmlx-core` dependency, then predict:

```sh
hf download MedMLX/nnunet-liver-mlx \
  --revision 574a9c601b9b1a9e906ce985fbd46b106c5b7caa \
  --local-dir mlx-model
```

```bash
nnunet-mlx model predict \
  --model-folder mlx-model \
  --input case_0000.nii.gz \
  --output-prefix predictions/case
```

Published MLX bundles already contain the converted weights. The model cards
describe each bundle's files and terms.
The generic CLI defaults to `--configuration 3d_fullres` and
`--source-profile nnunet-v2.8.1-default`. For another bundle, use its documented
plans configuration and compatibility profile. `conversion.json` is optional
export metadata; inference does not require or read it.

Pass one `--input` per modality in channel order. Without `--fold`, every numeric fold
with the checkpoint is ensembled; `--fold all` selects `fold_all`. `--previous-stage`
supplies a cascade's previous-stage segmentation. Every command takes `--help`, and
`--json` prints machine-readable results.

- `nnunet-mlx model predict-folder` predicts a dataset folder. `--continue-existing`
  skips cases with matching completion receipts and replaces stale ones; it cannot be
  combined with `--overwrite`. Receipts compare file size and modification time.
- `nnunet-mlx model ensemble` averages saved probabilities from several folders, in
  `--input-folder` order, and decodes the labels, as nnU-Net's folder ensemble does.
  Save predictions with `--save-probabilities` first, and pass `--plans-json` and
  `--dataset-json` from the model folder. Every folder must hold the same `.npz` cases,
  each with its `.pkl` sidecar, on the same grid.
- `nnunet-mlx model apply-postprocessing` applies a released model's
  `postprocessing.pkl` to exported segmentations, as `nnUNetv2_apply_postprocessing`
  does. Prediction never applies it automatically.

### Supported configurations

| Area | Supported |
| --- | --- |
| Platform | macOS arm64 with MLX on Metal, float32, one GPU per run |
| Networks | 2D and 3D `PlainConvUNet`; residual `BasicBlockD` and `BottleneckD` encoders with up to 32 outputs in 3D and 26 in 2D |
| Configurations | `2d`, `3d_lowres`, `3d_fullres` and cascades |
| Inference | Gaussian sliding window, mirroring, fold and configuration ensembles, probability export |
| Plans | nnU-Net v2.8.1 `plans.json` and `dataset.json`, class and region heads |

There is no PyTorch, CPU or other fallback: unsupported architectures, plan fields,
checkpoints and inputs fail with an error. `dataset.json` must define `channel_names`.
Labels must be consecutive non-negative integers or ASCII decimal strings; boolean,
fractional, padded and Unicode values are rejected even where upstream's `int()` accepts
them.
The MedMLX runners and the CLI's `predict` and `predict-folder` compute the same thing:
uncompiled models, with sliding windows and folds accumulated in float16 as nnU-Net v2.8.1
does. `ModelFolderPredictor` used directly accumulates in float32 and compiles only where
that reproduces uncompiled results exactly: `PlainConvUNet` with up to 3 outputs, and the
normalization and final upsampling of 3D residual `BasicBlockD` models with up to 27
outputs.

`--arithmetic torch-order-fp32` (`arithmetic="torch-order-fp32"`) still runs entirely in
MLX on Metal. It selects local kernels that sum float32 values in the pinned PyTorch CPU
order, which some released models need to reproduce their reference labels exactly. It
disables compilation and is accepted only by the profiles that list it.

### Compatibility profiles

Models trained with the default `nnUNetTrainer` use the default profile. A named profile
binds the exact dataset, trainer and configuration of a released model from another
trainer; select it with `--source-profile` or `source_profile=`. Loading checks each
checkpoint's trainer, configuration and plans against the selected profile. Profiles
cover inference only, not training or segmentation quality.

| Profile | Architecture | Configuration | Notes |
| --- | --- | --- | --- |
| `nnunet-v2.8.1-default` ✓ | `PlainConvUNet`, residual `BasicBlockD` and `BottleneckD` | any | `nnUNetTrainer` only |
| `prostate158-2d-100epochs` ✓ | Residual `BasicBlockD` | `2d` | Prostate158; up to 3 outputs |
| `chronoroot2-arabidopsis-2000epochs` ✓ | Residual `BasicBlockD` | `2d` | ChronoRoot2 plant roots; 7 outputs |
| `cads-oar258-open` ✓ | Residual `BasicBlockD` | `3d_fullres` | CADS OAR258; 30 outputs; fold `all`, no mirroring |
| `smile-abdomenatlasf-2d` ✓ | Residual `BasicBlockD` | `2d` | SMILE AbdomenAtlas; 26 outputs; `torch-order-fp32`, fold `all` |
| `totalsegmentator-v2-ct` ✓ | `PlainConvUNet` | `3d_fullres` | TotalSegmentator v2 Dataset291–295, 297, 298 and the Dataset300 body model; `torch-order-fp32` |
| `mrsegmentator-1.2-base` | `PlainConvUNet` | `3d_fullres_bs8` | MRSegmentator 1.2; `torch-order-fp32` |
| `lion-fdg-1.0.4` | As released | `3d_fullres` | LION FDG tumor model (`nnUNetTrainerDA5`) |
| `lesiontracer-autopet3-2024` | Residual `BasicBlockD`, organ-head variant | `3d_fullres_bs3` | LesionTracer autoPET III; auxiliary organ heads are excluded at conversion |
| `prostate-mrl-no-mirroring` | As released | `3d_fullres` | Prostate MR-Linac (`nnUNetTrainerNoMirroring`) |

✓ marks profiles whose model inference has been compared with the upstream reference on
Metal for a limited set of cases. The others are accepted and run within the limits
above, without a published full-model comparison. The original TotalSegmentator
`v2.0.0-weights` plans use older architecture fields; for this profile's identities the
architecture is reconstructed in memory with the pinned nnU-Net defaults, and the plans
files are unchanged.

### Python API

The names exported from `nnunet_mlx` are the supported Python API:

```python
from nnunet_mlx import ModelFolderPredictor, load_model_bundle

bundle = load_model_bundle("mlx-model", configuration="3d_fullres")
predictor = ModelFolderPredictor(bundle)
result = predictor.predict_case(
    ["case_0000.nii.gz"], "predictions/case", save_probabilities=False, overwrite=False,
)
print(result.metrics)  # synchronized timings, tile counts and processed shapes
```

`read_conversion_manifest` optionally validates `conversion.json` and its payload files
for inspection; inference does not call it. Unsupported inputs raise `ContractError`.
Other modules are internal and may change between releases.

### Untrusted files

Inference loads safetensors weights, never PyTorch `.pth` checkpoints. nnU-Net's `.pkl`
sidecars (case properties and probability metadata) are read by a restricted loader that
accepts plain Python containers and non-object NumPy metadata up to 64 MiB, and rejects
executable globals, extension and persistent-ID opcodes, object arrays, recursion and
trailing data. A released `postprocessing.pkl` may hold one payload up to 1 MiB naming
the one supported upstream postprocessing function. These checks cover structure, not
provenance: use model folders and `.npz`/`.pkl` pairs only from producers you trust.
Report vulnerabilities through the
[MedMLX private reporting form](https://github.com/MedMLX/medmlx/security/advisories/new),
not a public issue.

## License

[Apache-2.0](https://github.com/MedMLX/nnunet-mlx/blob/v0.6.4/LICENSE). See [third-party notices](https://github.com/MedMLX/nnunet-mlx/blob/v0.6.4/THIRD_PARTY_NOTICES.md). Model weights retain their publisher terms.

For research and education. Outputs require human review.
