Metadata-Version: 2.5
Name: totalsegmentator-mlx
Version: 0.1.6
Summary: Segment 117 CT structures with TotalSegmentator v2 in native MLX
Project-URL: Homepage, https://medmlx.org
Project-URL: Documentation, https://github.com/MedMLX/totalsegmentator-mlx/blob/v0.1.6/README.md
Project-URL: Repository, https://github.com/MedMLX/totalsegmentator-mlx
Project-URL: Issues, https://github.com/MedMLX/totalsegmentator-mlx/issues
Project-URL: Weights, https://huggingface.co/MedMLX/totalsegmentator-mlx
Project-URL: MedMLX, https://github.com/MedMLX/medmlx
Author: Joseph Sandoval, Alfredo Sandoval Jr.
License-Expression: Apache-2.0
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES.md
Keywords: apple-silicon,ct,medical-imaging,metal,mlx,segmentation,totalsegmentator
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: 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: nnunet-mlx==0.6.4
Requires-Dist: numpy<3,>=2.0
Requires-Dist: scipy==1.18.1
Provides-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'
Description-Content-Type: text/markdown

<div align="center">

# totalsegmentator-mlx

Segment 117 CT structures with TotalSegmentator v2 in native MLX. Includes standard, fast and coarse modes.

[![PyPI](https://img.shields.io/pypi/v/totalsegmentator-mlx)](https://pypi.org/project/totalsegmentator-mlx/)
[![Python](https://img.shields.io/pypi/pyversions/totalsegmentator-mlx)](https://pypi.org/project/totalsegmentator-mlx/)
[![License](https://img.shields.io/badge/License-Apache_2.0-blue)](https://github.com/MedMLX/totalsegmentator-mlx/blob/v0.1.6/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/totalsegmentator-mlx)
[![MedMLX](https://img.shields.io/badge/MedMLX-runtime-A51C30)](https://github.com/MedMLX/medmlx)

[Model weights](https://huggingface.co/MedMLX/totalsegmentator-mlx) · [MedMLX](https://github.com/MedMLX/medmlx)

<img src="docs/example.png" alt="TotalSegmentator standard-mode (1.5 mm) 3D surfaces and axial CT segmentation" width="100%">

Standard mode (1.5 mm): 3D surfaces from [autoPET-III](https://huggingface.co/datasets/YongchengYAO/autoPET-III-Lite) and axial CT from [AMOS22](https://zenodo.org/records/7262581). Surfaces and contours are smoothed for display.

</div>

| Input | Output |
| :--- | :--- |
| 3D CT · NIfTI | Multilabel mask on the input grid |

## 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 model. It installs this runtime, downloads its weights on first use, and lets an AI agent describe, check and run `totalsegmentator` 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 totalsegmentator`.

> **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 "totalsegmentator-mlx>=0.1.6"
```

## Run

```python
from totalsegmentator_mlx.medmlx import run

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

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

### Inputs and outputs

| Argument | Meaning |
| --- | --- |
| `image_path` | One finite 3D CT in NIfTI-1 or NIfTI-2, in Hounsfield units |
| `output_dir` | Directory for `mask.nii.gz` |
| `weights_path` | Optional root of the model folders below |
| `mode` | `standard` (default, five 1.5 mm parts), `fast` (3 mm model, upstream `--fast`) or `coarse` (6 mm model, upstream `--fastest`) |
| `prompt` | Structure names to keep after inference; label IDs stay global |
| `roi_subset` | Locate these structures, crop to them plus 20 mm and drop other labels |
| `robust_crop` | Locate `roi_subset` with the 3 mm model instead of the 6 mm one |
| `body_seg` | Crop to the Dataset300 body model plus 3 mm |
| `overwrite` | `False` by default; an existing mask is never replaced |

With `roi_subset`, standard mode runs only the parts that contain those structures.
`roi_subset` cannot be combined with `prompt` or `body_seg`, and `robust_crop`
requires `roi_subset`. Cropped regions are restored by zero filling; an empty
localization returns an empty mask. Spatial units must be millimetres or unspecified.

The mask is a uint8 NIfTI-1 with global labels 1–117 (0 background) on the input
grid, with the upstream label table as a header extension. `run` returns
`{"outputs": {"mask": path}, "details": {...}}`, with the weights revision, mode,
arithmetic, folds and label table. Invalid inputs raise `InvalidInputError`; missing
or mismatched weights raise `AssetNotReadyError`.

### Model folders

A local `weights_path` holds one folder per dataset, each with `plans.json`,
`dataset.json` and `fold_0/checkpoint_final/` containing `state.json` and
`weights.safetensors`. Only the folders the chosen options use are required, and a
single model folder can be passed directly when only that dataset is needed.

| Folder | Model | Used by |
| --- | --- | --- |
| `291/`–`295/` | five 1.5 mm parts | `standard` |
| `297/` | 3 mm total model | `fast`; `roi_subset` with `robust_crop` |
| `298/` | 6 mm total model | `coarse`; `roi_subset` |
| `300/` | 6 mm body model | `body_seg` |

### Command line

Download the [safetensors model bundle](https://huggingface.co/MedMLX/totalsegmentator-mlx),
then segment:

```bash
python - <<'PY'
from huggingface_hub import snapshot_download

snapshot_download(
    "MedMLX/totalsegmentator-mlx",
    revision="6c03d6f90ccf53f224c3985cbdecf1eaa96555af",
    local_dir="models",
)
PY
totalsegmentator-mlx predict -i ct.nii.gz -o mask.nii.gz --model-dir models
```

`predict` takes the runner options as flags (`--mode`, `--organ` for `prompt`,
`--roi-subset`, `--robust-crop`, `--body-seg`, `--overwrite`) and adds three more:
`--model-part` runs one standard part (`organs`, `vertebrae`, `cardiac`, `muscles`
or `ribs`), `--force-split` uses upstream's three overlapping slabs, and
`--probabilities-path probabilities.npz` also exports local-label probabilities for
one unsplit, uncropped model. `totalsegmentator-mlx labels` prints the label IDs.

### Python API

```python
from totalsegmentator_mlx import segment

labels = segment("ct.nii.gz", model_dir="models", mode="fast")  # nibabel image, uint8
```

`segment` takes the command-line options as keyword arguments, plus an optional
`output_path`. It never downloads weights.

### Limits

Only the upstream `total` CT task is included: Dataset291–295, 297 and 298 with 117
labels, plus the Dataset300 body model for cropping, all `3d_fullres`, fold 0,
without mirroring. DICOM input, DICOM-SEG and RT-STRUCT output, other upstream tasks
(including `total_mr` and licensed tasks), per-class output files, statistics,
radiomics, previews, `--remove_small_blobs`, `--higher_order_resampling`,
`--v1_order`, test-time augmentation and training are rejected or absent. Inference
needs Apple Silicon with Metal and has no CPU or PyTorch fallback.

## License

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

For research and education. Outputs require human review.
