biomedparse-mlx

Native MLX port of Microsoft BioMedParse v2 for Apple Silicon. Segment medical images from text prompts, locally on Mac.

• Text-guided 3D segmentation with multiple prompts and labels
• Raw NIfTI and DICOM windowing, with geometry-preserving NIfTI mask output
• Float32 model execution on Metal; PyTorch is used for conversion and oracle tests
• A 2D image path with overlapping prompt masks; dedicated 2D weights remain unqualified

Part of MedMLX. Technical notes · Usage

Quick start · Input and output · Technical details · License

Requirements

Apple Silicon Mac with Metal, Python 3.13 or 3.14 and uv. The released checkpoint requires access to the gated Microsoft model on Hugging Face.

Quick start

Prepare the locked checkout and local checkpoint assets:

gh repo clone MedMLX/biomedparse-mlx
cd biomedparse-mlx
make env
MACOSX_DEPLOYMENT_TARGET=27.0 uv sync --locked --extra dev --extra oracle --extra conversion --extra assets

uv run --locked --no-sync hf download microsoft/BiomedParse biomedparse_v2.ckpt \
  --revision e473e5b2b1a3f44649734afd3dc7cf1770aaa9e2 --local-dir weights/source
uv run --locked --no-sync hf download openai/clip-vit-base-patch32 tokenizer.json \
  --revision 3d74acf9a28c67741b2f4f2ea7635f0aaf6f0268 --local-dir weights

uv run --locked --no-sync biomedparse-mlx convert \
  --source weights/source/biomedparse_v2.ckpt --output weights/biomedparse-v2.npz

The source checkpoint is about 4.46 GB. Conversion writes a float32 tensor bundle and a versioned receipt. Inference uses local assets without implicit downloads.

Segment a preprocessed volume:

uv run --locked --no-sync biomedparse-mlx segment \
  --image scan.npy --weights weights/biomedparse-v2.npz --tokenizer weights/tokenizer.json \
  --prompt 'Visualization of the liver in abdominal CT imaging' --label 6 \
  --output-dir outputs/liver --slice-batch-size 2 --seed 7

Repeat --prompt and --label for multiple structures. Existing outputs require --overwrite. Python usage, raw-image windowing, instance segmentation and the segment-2d command are described in docs/usage.md. The model is discoverable by medmlx as biomedparse-v2.

Input and output

Volumes are preprocessed numeric .npy arrays or .npz files containing imgs in the upstream 0..255 intensity range. Raw NIfTI volumes and DICOM slices require explicit --modality and, for CT, --ct-site windowing options. Volume slicing follows upstream's fixed shape heuristic.

The volume path writes segmentation.npz with uint16 labels on the original array grid, prompt-by-slice existence probabilities, and run.json. NIfTI inputs also produce segmentation.nii.gz on the source affine. The 2D path returns separate, potentially overlapping masks and probabilities per prompt.

Technical details

The native graph follows the pinned BioMedParse v2 architecture and checkpoint contract. Controlled oracle comparisons replay upstream attention masks; matching integer seeds alone does not align Torch and MLX sampling. Recorded qualification covers bounded synthetic and AMOS cases. Dedicated 2D weights, broader task accuracy and clinical performance remain unqualified.

Architecture, numerical tolerances, development checks and recorded measurements are in docs/technical.md.

Intended use

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

License

Original contributions: Apache-2.0; see LICENSE. Upstream MIT/BSD notices and full license texts are retained in THIRD_PARTY_NOTICES.md. Model weights have separate publisher terms: Microsoft's model card declares CC-BY-NC-SA-4.0. Weights, tokenizers, datasets and masks are excluded from Git.

Implementation mapping

Central implementation map records upstream-to-native changes, source references and qualification scope.