mlx-segment-ct

Segment CT anatomy automatically or with anatomy names and clicks, using NVIDIA VISTA-3D in native MLX.

Docs · MedMLX

Automatic CT segmentation with selected organs and pelvic bones in 3D.

Apple Silicon · Metal · Python 3.13 / 3.14 · uv

Quick start

uv venv --python 3.13
source .venv/bin/activate
uv pip install "mlx-segment-ct[conversion]" huggingface-hub

Download and convert the checkpoint once, then run:

from pathlib import Path

from huggingface_hub import hf_hub_download
from medmlx_core.runtime import import_mlx
from mlx_segment_ct import mlx_correction_weights, mlx_weights

mx = import_mlx()
checkpoint = Path(hf_hub_download(
    "nvidia/NV-Segment-CT", "vista3d_pretrained_model/model.pt",
    revision="afb51518689f71e6abb367ee6301b2cd0225c66a",
))
out = Path("weights")
out.mkdir(exist_ok=True)
mx.save_safetensors(
    str(out / "automatic.safetensors"), mlx_weights.convert_weights(checkpoint, mx=mx)
)
mx.save_safetensors(
    str(out / "correction.safetensors"),
    {k: mx.array(v) for k, v in mlx_correction_weights.convert_weights(checkpoint).items()},
)

from mlx_segment_ct.medmlx import run

common = dict(image_path="ct.nii.gz", weights_path="weights")

run(**common, output_dir="all")  # all 117 classes
run(**common, output_dir="organs", prompt="liver, spleen")
run(
    **common,
    output_dir="points",  # clicks: (x, y, z) voxels
    points=[[40, 50, 30], [44, 51, 30]],
    point_labels=[1, 0],
)

Docs cover click prompts, interactive correction and the NV-Segment-CTMR checkpoint for CT and MRI.

License

Apache-2.0. See third-party notices. Model weights retain their publisher terms.

For research and education. Outputs require human review.