flexible-unet-mlx

Native MLX port of the MONAI endoscopic tool segmentation model (FlexibleUNet, EfficientNet-B2) for Apple Silicon. Segment surgical tools in endoscopy frames, locally on Mac.

• RGB endoscopy frame in, tool mask out
• Apple Silicon GPU acceleration
• FP32 inference using MLX; PyTorch is only needed once to convert the checkpoint

Part of MedMLX. Technical notes

Quick start · Input and output · Technical details · License

Requirements

Apple Silicon Mac with Metal, Python 3.13 or 3.14 and uv.

Quick start

Install with the conversion extra, and install the Hugging Face CLI as a separate tool for downloading the checkpoint:

uv pip install "flexible-unet-mlx[conversion] @ git+https://github.com/MedMLX/flexible-unet-mlx.git"
uv tool install huggingface-hub

Download the MONAI checkpoint:

hf download MONAI/endoscopic_tool_segmentation models/model.pt \
  --revision f0416a27d59516d9c1b6bec5181c9c33002c84fa --local-dir tool-bundle

Convert it once, then segment a frame:

from pathlib import Path

from flexible_unet_mlx import convert_checkpoint
from flexible_unet_mlx.medmlx import run

convert_checkpoint(Path("tool-bundle/models/model.pt"), Path("model.safetensors"))

result = run(image_path="frame.jpg", weights_path="model.safetensors", output_dir="out")
print(result["outputs"]["mask"])

Pass overwrite=True to replace existing output. The model is also discoverable by medmlx as flexible-unet.

Input and output

Input is one RGB image Pillow can read (the bundle used JPEG frames).

For frame.jpg, the mask is written to out/frame/frame_trans.png: 8-bit grayscale at the original frame size, 255 for tool and 0 for background. result["details"] holds the frame size, tool pixel counts and timings.

Technical details

Preprocessing, inference and postprocessing follow the pinned MONAI 1.4.0 bundle, including Torch's bilinear resize rounding. The tests compare against recordings of the upstream transforms and network made on one Apple Silicon Mac, using synthetic frames and seeded random weights. Agreement with the released checkpoint is a separate, optional check, and no segmentation accuracy evaluation is included. The lower-level network API, tolerances and fixture regeneration are in docs/technical.md.

Intended use

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

License

Code: Apache-2.0; see LICENSE. MONAI-derived code is Apache-2.0; see third-party notices. Checkpoint rights remain with their publishers.

Implementation mapping

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