Metadata-Version: 2.5
Name: nnunet-mlx
Version: 0.5.1
Summary: MLX-native nnU-Net v2 inference on Apple Silicon
Project-URL: Repository, https://github.com/MedMLX/nnunet-mlx
Project-URL: Issues, https://github.com/MedMLX/nnunet-mlx/issues
Project-URL: Documentation, https://github.com/MedMLX/nnunet-mlx/tree/main/docs
Author: Joseph Sandoval
License-Expression: Apache-2.0 AND MIT
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES.md
Keywords: apple-silicon,image-segmentation,medical-imaging,mlx,nnunet
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Image Processing
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
Requires-Python: <3.15,>=3.13
Requires-Dist: blosc2<5,>=3.3.2
Requires-Dist: imagecodecs<2027,>=2025
Requires-Dist: medmlx-core==0.2.3
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: conversion
Requires-Dist: torch==2.13.0; extra == 'conversion'
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

# nnunet-mlx

[nnU-Net v2](https://github.com/MIC-DKFZ/nnUNet) inference in native MLX
for Apple Silicon. Run existing nnU-Net models locally on Mac.

- Convert and run pretrained nnU-Net v2 model folders
- 2D and 3D `PlainConvUNet` and `ResidualEncoderUNet`
- FP32 using MLX; PyTorch is only needed to convert `.pth` checkpoints

Part of [MedMLX](https://github.com/MedMLX).
[Technical notes](https://github.com/MedMLX/nnunet-mlx/blob/main/docs/technical.md) · [Compatibility](https://github.com/MedMLX/nnunet-mlx/blob/main/docs/compatibility.md)

[Quick start](#quick-start) · [Input and output](#input-and-output) · [Technical details](#technical-details) · [License](#license)

## Requirements

Apple Silicon Mac with Metal, [uv](https://docs.astral.sh/uv/), and conda or mamba
(setup uses it to build a pinned Python 3.13 environment).

## Quick start

Prepare the source checkout:

```bash
git clone https://github.com/MedMLX/nnunet-mlx.git
cd nnunet-mlx
make env                     # creates .venv with MLX on Metal
uv pip install --python .venv/bin/python -e '.[conversion]'   # to convert .pth checkpoints
```

Use `.venv/bin/nnunet-mlx`, or activate `.venv`. Every command accepts `--help`.

Convert an upstream nnU-Net model folder once:

```bash
nnunet-mlx convert model-folder \
  --source-model-folder /path/to/upstream-model \
  --output-model-folder mlx-model
```

Predict one case, passing one `--input` per modality in channel order:

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

Use `nnunet-mlx model predict-folder` for a whole folder. Add `--overwrite` to replace
existing output, or `--continue-existing` to skip finished cases.

Only convert checkpoints you trust: `.pth` files are pickles. See [SECURITY.md](https://github.com/MedMLX/nnunet-mlx/blob/main/SECURITY.md).

### Python

```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,
)
```

For a single-channel 3D class model,
`predictor.predict_nifti(image)` accepts a NiBabel NIfTI-1 image and returns local
labels on its original grid without intermediate files. Fold logits stay on Metal
except where restoration needs the host route (separate-Z or non-linear
resampling). Cascades and regions use the file APIs above.

The model is also discoverable by [medmlx](https://github.com/MedMLX/medmlx)
as `nnunet`; see [docs/technical.md](https://github.com/MedMLX/nnunet-mlx/blob/main/docs/technical.md) for the runner.

## Input and output

Supply one NIfTI file per input channel in the model's channel order. Prediction
commands write segmentation files under the selected output prefix or directory;
probability export is optional. Supported models and configurations are listed
below and in the [compatibility notes](https://github.com/MedMLX/nnunet-mlx/blob/main/docs/compatibility.md).

| Area | Supported |
| --- | --- |
| Networks | 2D and 3D `PlainConvUNet` and `ResidualEncoderUNet` |
| Configurations | `2d`, `3d_lowres`, `3d_fullres`, cascades |
| Inference | Gaussian sliding window, mirroring, fold and configuration ensembles, probabilities |
| Precision | FP32 |

Anything else fails with a clear error rather than falling back. Pretrained-model
support varies; see [Compatibility](https://github.com/MedMLX/nnunet-mlx/blob/main/docs/compatibility.md).

## Technical details

Targets nnU-Net v2.8.1 and dynamic-network-architectures v0.4.4. Networks, conversion
and the prediction path are compared against the upstream PyTorch implementation on
seeded-weight and synthetic inputs and on selected released models (see
[Compatibility](https://github.com/MedMLX/nnunet-mlx/blob/main/docs/compatibility.md)); agreement on other models or data is not implied.
Development, parity testing, the PancCTMultiTalentV2 staging recipe and the runner
contract are in [docs/technical.md](https://github.com/MedMLX/nnunet-mlx/blob/main/docs/technical.md); see also
[Architecture](https://github.com/MedMLX/nnunet-mlx/blob/main/docs/architecture.md) and [Roadmap](https://github.com/MedMLX/nnunet-mlx/blob/main/docs/roadmap.md).

## Intended use

Intended for research and education, not clinical diagnosis or treatment decisions.
Outputs require human review.

## License

Apache-2.0; see [LICENSE](https://github.com/MedMLX/nnunet-mlx/blob/main/LICENSE) and [THIRD_PARTY_NOTICES.md](https://github.com/MedMLX/nnunet-mlx/blob/main/THIRD_PARTY_NOTICES.md). Kernels derived from MLX keep
their MIT license. Pretrained nnU-Net weights have their own licenses and are not
distributed here. Please cite via [CITATION.cff](https://github.com/MedMLX/nnunet-mlx/blob/main/CITATION.cff) and cite
[nnU-Net](https://github.com/MIC-DKFZ/nnUNet).

## Implementation mapping

[Central implementation map](https://github.com/MedMLX/Med-MLX/blob/main/implementation-maps/nnunet-mlx.json) records upstream-to-native changes, source references and qualification scope.
