Metadata-Version: 2.5
Name: flexible-unet-mlx
Version: 0.2.1
Summary: Native FP32 MLX FlexibleUNet released-checkpoint network
Author: Joseph Sandoval, Alfredo Sandoval Jr.
License-Expression: Apache-2.0
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES.md
Classifier: Operating System :: MacOS :: MacOS X
Requires-Python: <3.15,>=3.13
Requires-Dist: medmlx-core==0.2.3
Requires-Dist: numpy<3,>=2.0
Requires-Dist: pillow<14,>=13
Provides-Extra: conversion
Requires-Dist: torch<2.14,>=2.13; extra == 'conversion'
Provides-Extra: dev
Requires-Dist: hypothesis<7,>=6.140; extra == 'dev'
Requires-Dist: mlx==0.32.4; (platform_machine == 'arm64' and sys_platform == 'darwin') and 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'
Requires-Dist: safetensors<1,>=0.6; extra == 'dev'
Requires-Dist: torch<2.14,>=2.13; extra == 'dev'
Description-Content-Type: text/markdown

# flexible-unet-mlx

Native MLX port of the MONAI
[endoscopic tool segmentation](https://huggingface.co/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](https://github.com/MedMLX).
[Technical notes](https://github.com/MedMLX/flexible-unet-mlx/blob/main/docs/technical.md)

[Quick start](#quick-start) · [Input and output](#input-and-output) · [Technical details](#technical-details) · [License](#license)

## Requirements

Apple Silicon Mac with Metal, Python 3.13 or 3.14 and [uv](https://docs.astral.sh/uv/).

## Quick start

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

```bash
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:

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

Convert it once, then segment a frame:

```python
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](https://github.com/MedMLX/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](https://github.com/MedMLX/flexible-unet-mlx/blob/main/docs/technical.md).

## Intended use

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

## License

Code: Apache-2.0; see [LICENSE](https://github.com/MedMLX/flexible-unet-mlx/blob/main/LICENSE).
MONAI-derived code is Apache-2.0; see [third-party notices](https://github.com/MedMLX/flexible-unet-mlx/blob/main/THIRD_PARTY_NOTICES.md).
Checkpoint rights remain with their publishers.

## Implementation mapping

[Central implementation map](https://github.com/MedMLX/Med-MLX/blob/main/implementation-maps/flexible-unet-mlx.json) records upstream-to-native changes, source references and qualification scope.
