Metadata-Version: 2.5
Name: flexray-mlx
Version: 0.2.6
Summary: Segment anatomical structures in radiographs with FleXray and native MLX
Project-URL: Homepage, https://medmlx.org
Project-URL: Documentation, https://github.com/MedMLX/flexray-mlx/blob/v0.2.6/README.md
Project-URL: Repository, https://github.com/MedMLX/flexray-mlx
Project-URL: Issues, https://github.com/MedMLX/flexray-mlx/issues
Project-URL: Weights, https://huggingface.co/MedMLX/flexray-mlx
Project-URL: MedMLX, https://github.com/MedMLX/medmlx
Author: Joseph Sandoval, Alfredo Sandoval Jr.
License-Expression: Apache-2.0 AND MIT AND BSD-3-Clause
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES.md
Keywords: apple-silicon,flexray,medical-imaging,metal,mlx,radiograph,segmentation,x-ray
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Image Processing
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
Requires-Python: <3.14,>=3.13
Requires-Dist: medmlx-core==0.2.6
Requires-Dist: mlx==0.32.4; platform_system == 'Darwin' and platform_machine == 'arm64'
Requires-Dist: numpy<3,>=2.0
Requires-Dist: pillow<14,>=13
Requires-Dist: pydicom<4,>=3
Requires-Dist: pyyaml<7,>=6
Requires-Dist: safetensors<1,>=0.6
Provides-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'
Description-Content-Type: text/markdown

<div align="center">

# flexray-mlx

Segment anatomical structures in radiographs with FleXray and native MLX.

[![PyPI](https://img.shields.io/pypi/v/flexray-mlx)](https://pypi.org/project/flexray-mlx/)
[![Python](https://img.shields.io/pypi/pyversions/flexray-mlx)](https://pypi.org/project/flexray-mlx/)
[![License](https://img.shields.io/badge/License-Apache_2.0-blue)](https://github.com/MedMLX/flexray-mlx/blob/v0.2.6/LICENSE)
[![Apple Silicon](https://img.shields.io/badge/Apple_Silicon-MLX-black?logo=apple)](https://github.com/ml-explore/mlx)
[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97_Hugging_Face-weights-yellow)](https://huggingface.co/MedMLX/flexray-mlx)
[![MedMLX](https://img.shields.io/badge/MedMLX-runtime-A51C30)](https://github.com/MedMLX/medmlx)

[Model weights](https://huggingface.co/MedMLX/flexray-mlx) · [MedMLX](https://github.com/MedMLX/medmlx)

<img src="https://raw.githubusercontent.com/MedMLX/flexray-mlx/v0.2.6/docs/showcase.png" alt="Chest radiograph with lung, heart and individual rib segmentations" width="100%">

</div>

| Input | Output |
| :--- | :--- |
| Raster image or single-frame DICOM | Structure masks and probabilities |

## Use with MedMLX

The [MedMLX MCP server](https://github.com/MedMLX/medmlx/blob/main/docs/mcp.md) 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 `flexray` for you.

```sh
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 flexray`.

> **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

```bash
uv venv --python 3.13
source .venv/bin/activate
uv pip install "flexray-mlx>=0.2.6"
```

## Run

```python
from flexray_mlx.medmlx import run

result = run(image_path="chest.png", output_dir="out")
print(result["outputs"]["mask"])
```

Omitting `weights_path` downloads `MedMLX/flexray-mlx` into the shared MedMLX cache on
first use; interrupted downloads resume and later runs work offline. With MedMLX
installed, `medmlx setup flexray` downloads it ahead of time. To use your own copy, pass
a folder containing `ensemble.json` and
`members/flux0375/{config.yml,label_schema.json,model.safetensors,preprocessing.json}`
as `weights_path`. Automatic setup prepares the flagship (`ensemble=False`);
`ensemble=True` requires all five members in an explicit folder.
`details.weights_revision` records the downloaded commit, or `null` for a folder without
a MedMLX download record.

Use `label_name` to select a structure.

### Inputs and outputs

| Argument | Meaning |
| --- | --- |
| `image_path` | One radiograph: an image Pillow reads (PNG, TIFF, ...) or a single-frame 2D DICOM |
| `output_dir` | Output directory, created if missing |
| `weights_path` | Optional bundle folder, as above |
| `label_name` | Structure to threshold into the mask, `lungs` by default (any label, such as `heart`, `rib_5` or `vertebra_l1`) |
| `ensemble` | `False` by default; `True` averages all five released members |
| `threshold` | Probability cutoff in [0, 1], 0.5 by default |
| `tta_samples` | Passes per member, 1 by default: one plain view plus augmented views |
| `tta_seed` | `None` by default, unseeded as in FleXray; an integer reproduces Torch CPU augmentation draws |
| `normalization_scheme` | `None` keeps the bundle's; or `none`, `minmax`, `percentile_minmax` or `standardize` |
| `overwrite` | `False` by default; existing outputs are replaced only when `True` |

DICOM files end in `.dcm` or `.dicom` (any case) or have `DICM` at byte 128. Accepting
a file does not mean the model suits its modality. Output paths must differ from each
other, the input and the bundle files, including links, even with `overwrite`.

| Output | File | Content |
| --- | --- | --- |
| `mask` | `source_image_label.png` | uint8 0/1 mask of the selected structure on the source pixel grid |
| `display_source` | `source_image.png` | 8-bit grayscale source the mask overlays |
| `probabilities` | `probabilities.npy` | float32 61×256×256 sigmoid probabilities on the padded model grid |
| `prediction` | `prediction.json` | Labels, selected channel, threshold, members and content box |

FleXray thresholds on the model grid and exports no source-grid images; the PNGs are a
MedMLX addition. The mask is thresholded first, then cropped to the content box and
restored to the source size with nearest-neighbour resampling, so no new decisions come
from interpolation. `run` returns `{"outputs": {...}, "details": {...}}`; `details`
repeats the label and channel, threshold, members, TTA passes, seed, normalization,
source size, content box and mask pixel count. Missing or incompatible assets raise
`AssetNotReadyError`; invalid inputs or existing outputs raise `InvalidInputError`;
non-finite probabilities raise `ModelExecutionError` before anything is written.

The [published bundle](https://huggingface.co/MedMLX/flexray-mlx/tree/722d8341570e8b9e1196cf50e57892da84974fcc)
holds only `flux0375`; its `ensemble.json` still lists all five members in order
(`flux0375`, `flux000`, `flux025`, `flux050`, `flux075`). For `ensemble=True`, use the
[upstream snapshot](https://huggingface.co/VictorButoi/flexray/tree/6b63ddf09cfdff99aea494e3738cf243a4e22cdb)
with all five member folders.

Augmented views follow FleXray's `tta_v3` chain: flips, inversion, CLAHE or gamma,
contrast, sharpness and Gaussian noise. Gamma on negative normalized pixels can produce
NaN, which the runner rejects. The released weights were trained with percentile
normalization; other schemes are not validated.

### Python API

```python
from pathlib import Path

import numpy as np
from flexray_mlx import FleXrayMlxModel

model = FleXrayMlxModel(Path("flexray/members/flux0375"))
logits = model(np.zeros((1, 1, 256, 256), dtype=np.float32))  # Bx1x256x256 -> Bx61x256x256
```

The member folder needs the released `config.yml` and `model.safetensors`; loading
converts convolution weights to MLX layout and every tensor to float32, keeping the
names. Inputs are prepared images; preparation, sigmoid, ensembling, TTA and
source-grid restoration are the caller's. `flexray_mlx.tta.predict_with_tta(forward_fn,
images, tta_samples=..., seed=..., normalizer=...)` takes one or more forward callables
and already normalized float32 BxCxHxW images, and returns `(logits, probabilities)`;
pass the `normalize` method of `flexray_mlx.medmlx.ImagePreprocessing` as `normalizer`
to follow the bundle. Inference needs Apple Silicon with Metal and has no CPU fallback.

## License

[Apache-2.0](https://github.com/MedMLX/flexray-mlx/blob/v0.2.6/LICENSE). See [third-party notices](https://github.com/MedMLX/flexray-mlx/blob/v0.2.6/THIRD_PARTY_NOTICES.md). Model weights retain their publisher terms. FleXray weights are noncommercial (CC-BY-NC-4.0).

For research and education. Outputs require human review.
