totalsegmentator-mlx

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

Model weights · MedMLX

Standard mode (1.5 mm): 3D surfaces from autoPET-III and axial CT from AMOS22. Surfaces and contours are smoothed for display.

Use with MedMLX

The MedMLX MCP server 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.

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

uv venv --python 3.13
source .venv/bin/activate
uv pip install "totalsegmentator-mlx>=0.1.6"

Run

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

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.

Command line

Download the safetensors model bundle, then segment:

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

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. See third-party notices. Model weights retain their publisher terms.

For research and education. Outputs require human review.