flexray-mlx 0.2.2__tar.gz → 0.2.3__tar.gz

{flexray_mlx-0.2.2 → flexray_mlx-0.2.3}/PKG-INFO

@@ -1,14 +1,15 @@
 Metadata-Version: 2.5
 Name: flexray-mlx
-Version: 0.2.2
+Version: 0.2.3
 Summary: Native MLX implementation of FleXray
+Project-URL: Homepage, https://huggingface.co/MedMLX/flexray-mlx
 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
 Classifier: Operating System :: MacOS :: MacOS X
 Requires-Python: <3.15,>=3.13
-Requires-Dist: medmlx-core==0.2.4
+Requires-Dist: medmlx-core==0.2.5
 Requires-Dist: mlx<1,>=0.32.4
 Requires-Dist: numpy<3,>=2.0
 Requires-Dist: pillow<14,>=13
@@ -25,11 +26,7 @@ Description-Content-Type: text/markdown
 
 Segment anatomical structures in radiographs with FleXray and native MLX.
 
-[Docs](docs/technical.md) · [MedMLX](https://github.com/MedMLX)
-
-![One chest radiograph → separate lung, heart and individual rib masks.](https://raw.githubusercontent.com/MedMLX/flexray-mlx/main/docs/showcase.png)
-
-*One chest radiograph → separate lung, heart and individual rib masks.*
+[Model weights](https://huggingface.co/MedMLX/flexray-mlx)
 
 | Input | Output |
 | :--- | :--- |
@@ -42,28 +39,27 @@ Segment anatomical structures in radiographs with FleXray and native MLX.
 ```bash
 uv venv --python 3.13
 source .venv/bin/activate
-repo="https://github.com/MedMLX/flexray-mlx"
-uv pip install \
-  "flexray-mlx @ git+$repo" \
-  huggingface-hub
+uv pip install "flexray-mlx>=0.2.3"
 ```
 
-Download the released weights and run:
-
 ```python
-from huggingface_hub import snapshot_download
 from flexray_mlx.medmlx import run
 
-weights = snapshot_download("VictorButoi/flexray", local_dir="flexray", ignore_patterns=["*.onnx"])
-
-result = run(image_path="chest.png", weights_path=weights, output_dir="out")
+result = run(image_path="chest.png", output_dir="out")
 print(result["outputs"]["mask"])
 ```
 
-Use `label_name` to select a structure and `ensemble=True` to average all five members.
+Omitting `weights_path` downloads `MedMLX/flexray-mlx` once and reuses its local cache.
+An explicit `weights_path` selects your own model folder, containing `ensemble.json`
+and `members/flux0375/{config.yml,label_schema.json,model.safetensors,preprocessing.json}`.
+Use `label_name` to select a structure. Automatic setup prepares the flagship
+(`ensemble=False`); `ensemble=True` requires all five members in an explicit folder.
+
+Downloaded commits are returned in `details.weights_revision`; explicit local
+weights report `null` unless the folder carries a MedMLX download revision.
 
 ## License
 
-[Apache-2.0](LICENSE). See [third-party notices](THIRD_PARTY_NOTICES.md). Model weights retain their publisher terms. FleXray weights are noncommercial (CC-BY-NC-4.0).
+Apache-2.0; license and third-party notices are included in the distribution. Model weights retain their publisher terms. FleXray weights are noncommercial (CC-BY-NC-4.0).
 
 For research and education. Outputs require human review.

flexray_mlx-0.2.3/README.md

@@ -0,0 +1,41 @@
+# flexray-mlx
+
+Segment anatomical structures in radiographs with FleXray and native MLX.
+
+[Model weights](https://huggingface.co/MedMLX/flexray-mlx)
+
+| Input | Output |
+| :--- | :--- |
+| Raster image or single-frame DICOM | Structure masks and probabilities |
+
+**Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
+
+## Quick start
+
+```bash
+uv venv --python 3.13
+source .venv/bin/activate
+uv pip install "flexray-mlx>=0.2.3"
+```
+
+```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` once and reuses its local cache.
+An explicit `weights_path` selects your own model folder, containing `ensemble.json`
+and `members/flux0375/{config.yml,label_schema.json,model.safetensors,preprocessing.json}`.
+Use `label_name` to select a structure. Automatic setup prepares the flagship
+(`ensemble=False`); `ensemble=True` requires all five members in an explicit folder.
+
+Downloaded commits are returned in `details.weights_revision`; explicit local
+weights report `null` unless the folder carries a MedMLX download revision.
+
+## License
+
+Apache-2.0; license and third-party notices are included in the distribution. Model weights retain their publisher terms. FleXray weights are noncommercial (CC-BY-NC-4.0).
+
+For research and education. Outputs require human review.

{flexray_mlx-0.2.2 → flexray_mlx-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "flexray-mlx"
-version = "0.2.2"
+version = "0.2.3"
 description = "Native MLX implementation of FleXray"
 readme = "README.md"
 requires-python = ">=3.13,<3.15"
@@ -13,7 +13,7 @@ license = "Apache-2.0 AND MIT AND BSD-3-Clause"
 license-files = ["LICENSE", "THIRD_PARTY_NOTICES.md"]
 authors = [{ name = "Joseph Sandoval" }, { name = "Alfredo Sandoval Jr." }]
 dependencies = [
-  "medmlx-core==0.2.4",
+  "medmlx-core==0.2.5",
   "mlx>=0.32.4,<1",
   "numpy>=2.0,<3",
   "pillow>=13,<14",
@@ -22,6 +22,9 @@ dependencies = [
   "safetensors>=0.6,<1",
 ]
 
+[project.urls]
+Homepage = "https://huggingface.co/MedMLX/flexray-mlx"
+
 [project.entry-points."medmlx.models"]
 flexray = "flexray_mlx.medmlx:MODEL"
 
@@ -56,7 +59,7 @@ select = ["E", "F", "I", "UP", "B", "SIM", "RUF"]
 
 [tool.pyright]
 include = ["scripts","src","tests"]
-extraPaths = ["src", "tests"]
+extraPaths = ["src", "tests", "scripts"]
 pythonVersion = "3.13"
 typeCheckingMode = "strict"
 enableTypeIgnoreComments = false
@@ -66,11 +69,8 @@ reportUnnecessaryTypeIgnoreComment = "error"
 ban-relative-imports = "all"
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
-"typing.Any" = {msg = "Use a precise type and parse untyped values at the boundary."}
-"typing_extensions.Any" = {msg = "Use a precise type and parse untyped values at the boundary."}
-
-[tool.uv.sources]
-medmlx-core = { git = "https://github.com/MedMLX/medmlx.git", rev = "f98e0924f50228e2b9509474947b9672054b3217", subdirectory = "packages/core" }
+"typing.Any" = { msg = "Use a precise type at the boundary." }
+"typing_extensions.Any" = { msg = "Use a precise type at the boundary." }
 
 [tool.hatch.build.targets.sdist]
 only-include = [

{flexray_mlx-0.2.2 → flexray_mlx-0.2.3}/src/flexray_mlx/__init__.py

@@ -9,11 +9,12 @@ if TYPE_CHECKING:
 
 __all__ = ["FleXrayMlxModel"]
 
-__version__ = "0.2.2"
+__version__ = "0.2.3"
 
 
 def __getattr__(name: str) -> object:
     # MLX loads with the network, so importing the medmlx declaration stays light.
+    """Load the native model lazily when its public name is requested."""
     if name == "FleXrayMlxModel":
         from flexray_mlx.mlx_model import FleXrayMlxModel
 

{flexray_mlx-0.2.2 → flexray_mlx-0.2.3}/src/flexray_mlx/_kernels.py

@@ -48,6 +48,7 @@ out[i] = ((1.0f - fy) * (1.0f - fx)) * a
 
 @lru_cache(maxsize=1)
 def _resize_concat_kernel(mx: NativeRuntime) -> MetalKernel:
+    """Compile and cache the Metal kernel for corner-aligned resize and concatenation."""
     return cast(
         MetalKernel,
         mx.fast.metal_kernel(

{flexray_mlx-0.2.2 → flexray_mlx-0.2.3}/src/flexray_mlx/_random.py

@@ -31,9 +31,11 @@ class TorchRandom:
         self._next_normal: float | None = None
 
     def _words(self, count: int) -> UIntArray:
+        """Draw raw 32-bit words from the isolated MT19937 stream."""
         return self._engine.randint(0, 2**32, size=count, dtype=np.uint32)
 
     def _units(self, count: int) -> FloatArray:
+        """Convert raw words to Torch-compatible float32 unit uniforms."""
         return (self._words(count) & np.uint32(0xFFFFFF)).astype(np.float32) * np.float32(2**-24)
 
     def uniform(self, low: float = 0.0, high: float = 1.0) -> float:
@@ -49,11 +51,13 @@ class TorchRandom:
         return self.uniform() < np.float32(probability)
 
     def _double_unit(self) -> float:
+        """Combine two words into a 53-bit unit uniform."""
         words = self._words(2).astype(np.uint64)
         value = (int(words[0]) << 32) | int(words[1])
         return (value & (2**53 - 1)) * 2**-53
 
     def _short_normal(self) -> float:
+        """One double-precision Box-Muller normal; the pair's sine half is cached."""
         if self._next_normal is not None:
             value, self._next_normal = self._next_normal, None
             return value
@@ -64,6 +68,7 @@ class TorchRandom:
 
     @staticmethod
     def _normal_blocks(uniforms: FloatArray) -> FloatArray:
+        """Convert groups of 16 uniforms to Torch-compatible normal draws."""
         blocks = uniforms.reshape(-1, 16)
         radius = np.sqrt(np.float32(-2.0) * np.log(np.float32(1.0) - blocks[:, :8]))
         # The scalar CPU kernel multiplies by double pi before rounding theta.

{flexray_mlx-0.2.2 → flexray_mlx-0.2.3}/src/flexray_mlx/_types.py

@@ -2,16 +2,15 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-from os import PathLike
-from typing import TYPE_CHECKING, Literal, Protocol, TypedDict
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal, NotRequired, Protocol, TypedDict
 
 import numpy as np
-from medmlx_core.typing import FastRuntime
+from medmlx_core.typing import MlxRuntime
 from numpy.typing import NDArray
 
 if TYPE_CHECKING:
-    from mlx.core import Dtype, array
+    from mlx.core import StreamOrDevice, array
 
 type HostArray = NDArray[np.generic]
 type RealArray = NDArray[np.bool_ | np.integer | np.floating]
@@ -24,52 +23,69 @@ type JsonValue = bool | int | float | str | list[JsonValue] | dict[str, JsonValu
 type JsonObject = dict[str, JsonValue]
 
 
-class NativeRuntime(Protocol):
-    fast: FastRuntime
-    float32: Dtype
-
-    # Constructor inputs here are the NumPy arrays prepared by this network.
-    def array(self, value: FloatArray) -> array: ...
-    def eval(self, *args: array) -> None: ...
-    def conv2d(self, input: array, weight: array, /, *, padding: int) -> array: ...
-    def mean(self, a: array, /, *, axis: Sequence[int], keepdims: bool) -> array: ...
-    def square(self, a: array, /) -> array: ...
-    def rsqrt(self, a: array, /) -> array: ...
-    def maximum(self, a: array, b: array, /) -> array: ...
-    def concatenate(self, arrays: list[array], *, axis: int) -> array: ...
+class NativeRuntime(MlxRuntime, Protocol):
+    def square(self, x: array, /, *, stream: StreamOrDevice = None) -> array:
+        """Square an MLX array elementwise."""
+        ...
 
 
 class NativeOperation(Protocol):
-    def __call__(self, x: array) -> array: ...
+    def __call__(self, x: array) -> array:
+        """Apply a native operation to an MLX array."""
+        ...
 
 
 class NeuralRuntime(Protocol):
-    def MaxPool2d(self, *, kernel_size: int, stride: int) -> NativeOperation: ...
-    def Upsample(self, *, scale_factor: int, mode: str, align_corners: bool) -> NativeOperation: ...
+    def MaxPool2d(
+        self, kernel_size: int | tuple[int, int], stride: int | tuple[int, int] | None = None,
+        padding: int | tuple[int, int] | None = 0,
+    ) -> NativeOperation:
+        """Construct a two-dimensional MLX max-pooling layer."""
+        ...
 
 
 class TensorReader(Protocol):
-    def __enter__(self) -> TensorReader: ...
-    def __exit__(self, *args: object) -> None: ...
-    def keys(self) -> list[str]: ...
-    def get_tensor(self, name: str) -> HostArray: ...
+    def __enter__(self) -> TensorReader:
+        """Open the tensor-reader context."""
+        ...
+    def __exit__(self, *args: object) -> None:
+        """Close the tensor-reader context."""
+        ...
+    def keys(self) -> list[str]:
+        """List the named tensors in the checkpoint."""
+        ...
+    def get_tensor(self, name: str) -> HostArray:
+        """Read one named tensor as a host array."""
+        ...
 
 
 class SafetensorsRuntime(Protocol):
-    def safe_open(self, filename: PathLike[str], *, framework: Literal["np"]) -> TensorReader: ...
+    def safe_open(self, filename: str | Path, framework: Literal["np"]) -> TensorReader:
+        """Open named safetensors as NumPy arrays."""
+        ...
 
 
 class DicomDataset(Protocol):
     pixel_array: RealArray
 
-    def __contains__(self, name: str) -> bool: ...
+    def __contains__(self, name: str) -> bool:
+        """Return whether the dataset contains a named DICOM element."""
+        ...
 
 
 class DicomRuntime(Protocol):
-    def dcmread(self, fp: str, *, force: bool) -> DicomDataset: ...
+    def dcmread(self, fp: str, *, force: bool = False) -> DicomDataset:
+        """Read one DICOM dataset from a file path."""
+        ...
 
 
 class Normalization(TypedDict):
+    """Normalization settings parsed from ``preprocessing.json``.
+
+    ``percentiles`` applies only to ``percentile_minmax``; ``mean``/``std`` are the
+    legacy scalar form and exclude the per-sample schemes.
+    """
+
     scheme: str
     percentiles: tuple[float, float] | None
     eps: float
@@ -103,6 +119,7 @@ class RunDetails(TypedDict):
     in model-grid pixels with exclusive bottom/right.
     """
 
+    weights_revision: NotRequired[str | None]
     selected_label: str
     selected_channel: int
     threshold: float

{flexray_mlx-0.2.2 → flexray_mlx-0.2.3}/src/flexray_mlx/medmlx.py

@@ -14,6 +14,7 @@ import math
 import warnings
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, replace
+from itertools import combinations
 from pathlib import Path
 from typing import TYPE_CHECKING, SupportsFloat, SupportsIndex, SupportsInt, cast
 
@@ -24,6 +25,7 @@ from medmlx_core.errors import (
     MissingDependencyError,
     ModelExecutionError,
 )
+from medmlx_core.weights import weights_dir, weights_revision
 from PIL import Image, ImageOps, UnidentifiedImageError
 
 from flexray_mlx._types import FloatArray, JsonValue, RealArray, RunResult
@@ -91,12 +93,14 @@ class AdmittedBundle:
 
 
 def _missing(path: Path) -> AssetNotReadyError:
+    """Construct the standard missing-model-asset error for a path."""
     return AssetNotReadyError(
         f"Missing FleXray asset: {path}", reason=f"missing {path}", hint=_WEIGHTS_HINT
     )
 
 
 def _read_json(path: Path) -> JsonValue:
+    """Read a JSON model asset and classify missing or malformed files."""
     if not path.is_file():
         raise _missing(path)
     try:
@@ -114,7 +118,8 @@ def admit_bundle(root: Path, *, ensemble: bool) -> AdmittedBundle:
 
     Only the members that will run are inspected: the flagship alone, or all five
     when ``ensemble`` is true. Tensor contents are not read here; the network
-    checks them when it loads.
+    loads them and validates tensor use during inference. These checks do not
+    establish the identity or origin of the bundle or checkpoint.
 
     Raises:
         AssetNotReadyError: A file is missing or unreadable, the ensemble layout
@@ -181,6 +186,7 @@ def admit_bundle(root: Path, *, ensemble: bool) -> AdmittedBundle:
 
 
 def _pydicom() -> DicomRuntime:
+    """Import the DICOM reader or raise an actionable dependency error."""
     try:
         import pydicom
     except ImportError as error:
@@ -210,9 +216,11 @@ def is_dicom_path(path: str | Path) -> bool:
 def load_dicom_grayscale(path: str | Path) -> FloatArray:
     """Read one single-frame DICOM into a 2D float32 image min-max scaled to [0, 1].
 
-    Pixels are rescaled with ``RescaleSlope``/``RescaleIntercept`` when present,
-    RGB(A) is reduced to luma, and MONOCHROME1 is inverted. No VOI LUT or
-    windowing is applied, so the full stored intensity range maps to [0, 1].
+    Pixels are rescaled with ``RescaleSlope``/``RescaleIntercept`` when present
+    (a false-valued slope is treated as 1), RGB(A) is reduced to luma, and
+    MONOCHROME1 is inverted. No VOI LUT or windowing is applied, so the full
+    stored intensity range maps to [0, 1]. Declared multiple frames are rejected;
+    without a frame count a 3D array with last dimension 3 or 4 is treated as color.
 
     Raises:
         ValueError: The file has no pixel data or is not a single-frame 2D image.
@@ -221,6 +229,8 @@ def load_dicom_grayscale(path: str | Path) -> FloatArray:
     dataset = _pydicom().dcmread(str(path), force=True)
     if "PixelData" not in dataset:
         raise ValueError(f"DICOM file has no pixel data: {path}.")
+    if int(getattr(dataset, "NumberOfFrames", 1) or 1) > 1:
+        raise ValueError(f"Only single-frame 2D DICOM images are supported: {path}.")
     pixels: RealArray = np.asarray(dataset.pixel_array)
     if pixels.ndim == 3 and pixels.shape[-1] in {3, 4}:
         rgb = pixels[..., :3].astype(np.float64, copy=False)
@@ -339,6 +349,7 @@ class ImagePreprocessing:
     eps: float = DEFAULT_NORMALIZATION_EPS
 
     def __post_init__(self) -> None:
+        """Validate image geometry, scaling, and normalization settings."""
         _validate_image_size(self.image_size)
         if self.color_mode != "grayscale":
             raise ValueError(
@@ -391,10 +402,11 @@ class ImagePreprocessing:
     def content_box(self, height: int, width: int) -> tuple[int, int, int, int]:
         """``(top, left, bottom, right)`` of the input footprint on the prepared canvas.
 
-        Coordinates are model-grid pixels with exclusive ``bottom`` and ``right``,
-        so ``[top:bottom, left:right]`` crops the part of the model grid that
-        holds image content rather than zero padding. ``height`` and ``width``
-        are the source image size in pixels.
+        Coordinates are model-grid pixels with exclusive ``bottom`` and ``right``.
+        The box approximates the resized source footprint, not a guarantee that
+        outside pixels are zero: bilinear interpolation can mix the boundary and
+        normalization can change padding values. ``height`` and ``width`` are
+        the source image size in pixels.
 
         Raises:
             ValueError: ``height`` or ``width`` is not positive.
@@ -418,8 +430,10 @@ class ImagePreprocessing:
     def prepare(self, path: str | Path) -> FloatArray:
         """Load one image or DICOM file as a normalized float32 ``1x1xHxW`` array.
 
-        ``H`` and ``W`` are ``image_size``. Non-finite pixels become 0 with a
-        ``RuntimeWarning``; a constant image also warns.
+        ``H`` and ``W`` are ``image_size``. Non-finite values that reach the
+        preprocessing check become 0 with a ``RuntimeWarning``; a constant image
+        also warns. DICOM min-max scaling occurs first and may turn non-finite
+        decoded values into a constant zero image before that check.
         """
 
         image = _load_image_file(path)
@@ -428,11 +442,13 @@ class ImagePreprocessing:
         return self._prepare_array(image)
 
     def _prepare_pillow(self, image: Image.Image) -> FloatArray:
+        """Convert a raster image to grayscale and prepare it for inference."""
         array = _grayscale_pixels(image)
         return self._prepare_grayscale_array(array, source_dtype=array.dtype)
 
     def _prepare_array(self, array: FloatArray) -> FloatArray:
         # File inputs reach this only as decoded 2D DICOM pixels.
+        """Validate decoded DICOM pixels and prepare the grayscale array."""
         np_array = np.asarray(array)
         if np_array.size == 0:
             raise ValueError("Image inputs must be non-empty.")
@@ -443,6 +459,7 @@ class ImagePreprocessing:
     def _prepare_grayscale_array(
         self, array: RealArray, *, source_dtype: np.dtype[np.generic]
     ) -> FloatArray:
+        """Scale, pad, resize, and normalize one grayscale image."""
         if array.size == 0:
             raise ValueError("Image inputs must be non-empty.")
         if self.scale == "zero_one":
@@ -495,6 +512,7 @@ class ImagePreprocessing:
 
 
 def _load_image_file(path: str | Path) -> Image.Image | FloatArray:
+    """Load DICOM pixels or a detached Pillow raster image."""
     if is_dicom_path(path):
         return load_dicom_grayscale(path)
     with Image.open(path) as opened:
@@ -544,6 +562,7 @@ def _standardize_mean(batch: FloatArray) -> FloatArray:
 
 
 def _validate_image_size(size: Sequence[SupportsInt]) -> tuple[int, int]:
+    """Require positive height and width and return them as integers."""
     if len(size) != 2:
         raise ValueError(f"Image size must be an H W pair, got {tuple(size)}.")
     height, width = int(size[0]), int(size[1])
@@ -575,6 +594,12 @@ def _scale_numpy_integer_range(
 
 
 def _normalization_values(data: Mapping[str, object]) -> Normalization:
+    """Canonicalize ``preprocessing.json`` normalization fields.
+
+    Accepts the current ``normalization: {scheme, percentiles, eps}`` mapping (with
+    top-level ``mean``/``std``) and the legacy forms that carry only scalar
+    ``mean``/``std``, either top-level or inside ``normalization``.
+    """
     normalization = data.get("normalization")
     if isinstance(normalization, Mapping):
         normalization = cast(Mapping[str, object], normalization)
@@ -633,6 +658,7 @@ def _validate_normalization(
     mean: float | None,
     std: float | None,
 ) -> None:
+    """Reject unsupported, incompatible, or non-finite normalization settings."""
     if not isinstance(scheme, str) or scheme not in _NORMALIZATION_SCHEMES:
         raise ValueError(
             "preprocessing normalization scheme must be one of "
@@ -656,6 +682,7 @@ def _validate_normalization(
 
 
 def _validate_percentiles(values: object) -> tuple[float, float]:
+    """Require two finite increasing percentiles in the range zero to 100."""
     if isinstance(values, str | bytes) or not isinstance(values, Sequence):
         raise TypeError("preprocessing percentiles must be a two-value sequence.")
     values = cast(Sequence[str | SupportsFloat | SupportsIndex], values)
@@ -675,6 +702,7 @@ def _validate_percentiles(values: object) -> tuple[float, float]:
 
 
 def _validate_eps(value: object) -> float:
+    """Require and return a positive finite normalization stabilizer."""
     if isinstance(value, bool) or not isinstance(value, int | float):
         raise TypeError("preprocessing normalization eps must be a number.")
     eps = float(value)
@@ -938,6 +966,7 @@ def model_forward(model: Forward, bundle: Path) -> Forward:
     """Keep deferred tensor-key/architecture errors classified as incompatible assets."""
 
     def forward(images: FloatArray) -> FloatArray:
+        """Run a member and classify deferred checkpoint errors as asset errors."""
         try:
             return model(images)
         except (KeyError, ValueError) as error:
@@ -975,6 +1004,7 @@ def _predict(
 
 
 def _input_errors() -> tuple[type[BaseException], ...]:
+    """Exceptions meaning an unreadable input; pydicom's parse error only if installed."""
     errors: tuple[type[BaseException], ...] = (
         EOFError,
         OSError,
@@ -989,6 +1019,7 @@ def _input_errors() -> tuple[type[BaseException], ...]:
 
 
 def _require_bool(name: str, value: object) -> bool:
+    """Require a boolean runner argument or raise an input error."""
     if not isinstance(value, bool):
         raise InvalidInputError(f"FleXray {name} must be true or false, got {value!r}")
     return value
@@ -997,7 +1028,7 @@ def _require_bool(name: str, value: object) -> bool:
 def run(
     *,
     image_path: str | Path,
-    weights_path: str | Path,
+    weights_path: str | Path | None = None,
     output_dir: str | Path,
     label_name: str = "lungs",
     ensemble: bool = False,
@@ -1009,40 +1040,51 @@ def run(
 ) -> RunResult:
     """Segment one conventional 2D radiograph with the released FleXray network.
 
-    The model natively predicts on a 256x256 grid after zero-padding the image
-    to a square. All 61 independent sigmoid channels (background plus 60
-    structures) are saved on that padded grid. The selected structure's binary
+    The model natively predicts on a 256x256 grid after the bundle's preprocessing
+    (normally zero-padding the image to a square). All 61 independent sigmoid
+    channels (background plus 60 structures) are saved on that model grid. The
+    selected structure's binary
     mask is thresholded on the model grid, cropped, and resized with nearest neighbor,
     so it keeps the source image's pixel dimensions.
 
     Args:
         image_path: One radiograph: a raster image Pillow reads (PNG, TIFF, ...)
-            or a single-frame 2D DICOM file.
-        weights_path: The FleXray release snapshot directory with
-            ``ensemble.json`` and ``members/<name>/`` bundles
+            or a single-frame 2D DICOM file. DICOM detection uses ``.dcm`` or
+            ``.dicom`` (case-insensitive), or ``DICM`` at byte offset 128 for a
+            file without a suffix.
+        weights_path: Optional model folder. Omit to download ``MedMLX/flexray-mlx``.
+            The flagship snapshot at revision
+            ``722d8341570e8b9e1196cf50e57892da84974fcc``, containing
+            ``ensemble.json`` and the ``members/flux0375/`` bundle
             (``config.yml``, ``model.safetensors``, ``label_schema.json``,
             ``preprocessing.json``).
         output_dir: Directory for the outputs; created if missing.
         label_name: Structure to threshold into the mask, e.g. ``lungs``,
             ``heart``, ``rib_5`` or ``background`` (any name in the label schema).
         ensemble: Average the probabilities of all five released members
-            instead of running only the ``flux0375`` flagship.
+            instead of running only the ``flux0375`` flagship (default False).
+            Requires all five member bundles: ``flux0375``, ``flux000``,
+            ``flux025``, ``flux050``, ``flux075``. The published MedMLX snapshot
+            contains only the flagship; use the full ``VictorButoi/flexray``
+            snapshot at ``6b63ddf09cfdff99aea494e3738cf243a4e22cdb`` for an ensemble.
         threshold: Probability cutoff in [0, 1] for the selected structure; pixels
             at or above it are foreground.
-        tta_samples: Total passes per member; 0 or 1 means one plain pass.
-        tta_seed: Integer seed for the isolated stream that reproduces Torch CPU
-            augmentation draws; None draws fresh views, as in fxr.
+        tta_samples: Non-negative integer total passes per member; 0 or 1 means
+            one plain pass.
+        tta_seed: Integer in ``[-2**63, 2**64 - 1]`` for the isolated stream that
+            reproduces Torch CPU augmentation draws; None is unseeded, as in fxr.
         normalization_scheme: Override the bundle's normalization with ``none``,
             ``minmax``, ``percentile_minmax`` or ``standardize``. None keeps the
             bundle contract. Standardize uses sample std and adds bundle epsilon.
         overwrite: Replace existing outputs instead of refusing. Output paths
-            must remain separate from the input radiograph and its file aliases.
+            must remain separate from each other, the input radiograph and the
+            documented model bundle files, including their file aliases.
 
     Returns:
         ``outputs`` maps ``mask`` to ``source_image_label.png`` (uint8 0/1 on the
         source pixel grid), ``display_source`` to ``source_image.png`` (the 8-bit
         grayscale source), ``probabilities`` to ``probabilities.npy`` (float32
-        ``61x256x256`` on the padded model grid) and ``prediction`` to
+        ``61x256x256`` on the model grid) and ``prediction`` to
         ``prediction.json``. ``details`` holds the selected label and channel,
         threshold, members, TTA settings, normalization scheme, source size as
         ``[width, height]``, the content box as ``[top, left, bottom, right]`` in
@@ -1051,9 +1093,10 @@ def run(
     Raises:
         InvalidInputError: Missing or unreadable input, a non-2D DICOM, an
             unknown label, an out-of-range threshold, invalid TTA or
-            normalization settings, outputs that alias the input, or existing
-            outputs without ``overwrite``.
-        AssetNotReadyError: Missing, unpinned, corrupt or incompatible weights.
+            normalization settings, outputs that alias an input, model asset or
+            each other, or existing outputs without ``overwrite``.
+        AssetNotReadyError: Missing or unreadable assets, an incompatible bundle
+            layout or network, or inconsistent member metadata.
         MissingDependencyError: MLX or pydicom is unavailable.
         ModelExecutionError: Inference failed, or the network returned
             non-finite or wrongly shaped probabilities. Raised before any
@@ -1095,6 +1138,26 @@ def run(
         for path in destinations
     ):
         raise InvalidInputError("FleXray outputs must be separate from the input radiograph")
+    weights_root = (
+        weights_dir("flexray-mlx") if weights_path is None else Path(weights_path).expanduser()
+    )
+    model_assets = (
+        weights_root / "ensemble.json",
+        *(weights_root / "members" / member / name for member in MEMBERS for name in _BUNDLE_FILES),
+    )
+    if any(
+        path.resolve() == asset.resolve()
+        or (path.exists() and asset.exists() and path.samefile(asset))
+        for path in destinations
+        for asset in model_assets
+    ):
+        raise InvalidInputError("FleXray outputs must be separate from each model asset")
+    if any(
+        first.resolve() == second.resolve()
+        or (first.exists() and second.exists() and first.samefile(second))
+        for first, second in combinations(destinations, 2)
+    ):
+        raise InvalidInputError("FleXray outputs must not refer to the same file")
     existing = [path.as_posix() for path in destinations if path.exists()]
     if existing and not overwrite:
         raise InvalidInputError(
@@ -1102,7 +1165,7 @@ def run(
             "pass overwrite=True to replace them"
         )
 
-    admitted = admit_bundle(Path(weights_path).resolve(), ensemble=ensemble)
+    admitted = admit_bundle(weights_root, ensemble=ensemble)
     labels = admitted.labels
     if label_name not in labels:
         raise InvalidInputError(f"Unknown FleXray anatomy label: {label_name}")
@@ -1190,6 +1253,7 @@ def run(
             "prediction": prediction_path.as_posix(),
         },
         "details": {
+            "weights_revision": weights_revision(weights_root),
             "selected_label": label_name,
             "selected_channel": channel,
             "threshold": threshold,

{flexray_mlx-0.2.2 → flexray_mlx-0.2.3}/src/flexray_mlx/mlx_model.py

@@ -1,7 +1,8 @@
 """Native MLX execution of the pinned FleXray v1.0.2 U-Net checkpoint.
 
 The released network is fixed to seven encoder stages and six decoder stages.
-PyTorch is used by the source preprocessing and TTA code, never for this forward pass.
+Upstream FleXray uses PyTorch for preprocessing and test-time augmentation; this
+port implements those steps separately and uses MLX for the network forward pass.
 """
 
 from __future__ import annotations
@@ -95,15 +96,18 @@ class FleXrayMlxModel:
         self._pool = nn.MaxPool2d(kernel_size=2, stride=2)
 
     def _weight(self, key: str) -> array:
+        """Record and return a named checkpoint tensor."""
         self._consumed.add(key)
         return self.weights[key]
 
     def _conv(self, x: array, prefix: str, *, padding: int) -> array:
+        """Apply a learned convolution and bias using the named prefix."""
         weight = self._weight(f"{prefix}.weight")
         bias = self._weight(f"{prefix}.bias")
         return mx.conv2d(x, weight, padding=padding) + bias
 
     def _norm(self, x: array, prefix: str) -> array:
+        """Apply per-instance spatial normalization with learned scale and bias."""
         gamma = self._weight(f"{prefix}.weight")
         beta = self._weight(f"{prefix}.bias")
         # Instance norm on NHWC: per sample and channel over H and W. eps=1e-5 is
@@ -113,6 +117,7 @@ class FleXrayMlxModel:
         return (x - mean) * mx.rsqrt(variance + 1e-5) * gamma + beta
 
     def _block(self, x: array, prefix: str) -> array:
+        """Three conv, LeakyReLU, instance-norm steps plus a residual shortcut."""
         residual = x
         out = x
         for index in range(3):

{flexray_mlx-0.2.2 → flexray_mlx-0.2.3}/src/flexray_mlx/tta.py

@@ -29,6 +29,7 @@ class TtaSettings:
     seed: int | None = None
 
     def __post_init__(self) -> None:
+        """Reject invalid augmentation counts and Torch-compatible seed values."""
         if (
             isinstance(self.samples, bool)
             or not isinstance(cast(object, self.samples), int)
@@ -69,8 +70,10 @@ def predict_with_tta(
             float32 ``(B, C, H, W)`` images to float32 ``(B, L, H, W)`` logits.
         images: Prepared float32 ``(B, C, H, W)`` batch with 1 or 3 channels,
             already normalized.
-        tta_samples: Total passes per member; 0 or 1 means one plain pass.
-        seed: Seed for the isolated Torch-CPU-compatible stream; None is unseeded.
+        tta_samples: Non-negative integer total passes per member; 0 or 1 means
+            one plain pass.
+        seed: Integer in ``[-2**63, 2**64 - 1]`` for the isolated Torch-CPU-compatible
+            stream; None is unseeded.
         normalizer: Applied to each augmented view (not the plain input) so views
             are normalized as the plain input was.
 
@@ -124,7 +127,8 @@ def augment_sample(image: FloatArray, random: TorchRandom) -> tuple[FloatArray,
 
     ``image`` is one prepared ``(C, H, W)`` float32 sample. Random draws are
     consumed in fxr's order. The CLAHE, gamma, contrast and sharpness steps clip
-    to [0, 1]; invert and noise do not.
+    to [0, 1]; invert and noise do not. Fractional gamma is applied before its
+    clip, so negative values can produce NaN; clipping does not remove that NaN.
     Returns the view and whether it was flipped (so logits can be unflipped).
     """
 
@@ -200,11 +204,13 @@ def clahe(image: FloatArray, *, clip_limit: float, grid_size: tuple[int, int]) -
     interpolates bilinearly between the four nearest tile centres. ``grid_size``
     is ``(rows, cols)`` of tiles, capped at the image size; tiles are rounded up
     to even sizes and the image is padded to fit (reflect, or edge when the
-    padding reaches the image size in either direction). Returns float32 ``(C, H, W)`` in [0, 1].
+    padding reaches the image size in either direction). Supply two integer tile
+    counts; the function does not check their types. Returns float32 ``(C, H, W)`` in [0, 1].
 
     Raises:
         ValueError: ``image`` is not a non-empty 3D array, ``clip_limit`` is not
-            positive and finite, or ``grid_size`` is not two positive integers.
+            positive and finite, or ``grid_size`` has the wrong length, a boolean
+            value, or a nonpositive value.
     """
 
     if image.ndim != 3 or min(image.shape) < 1:
@@ -245,4 +251,7 @@ def clahe(image: FloatArray, *, clip_limit: float, grid_size: tuple[int, int]) -
     wy, wx = wy[None, :, None], wx[None, None]
     top = v00 * (np.float32(1) - wx) + v01 * wx
     bottom = v10 * (np.float32(1) - wx) + v11 * wx
-    return ((top * (np.float32(1) - wy) + bottom * wy) / np.float32(255))[:, :height, :width]
+    return cast(
+        FloatArray,
+        ((top * (np.float32(1) - wy) + bottom * wy) / np.float32(255))[:, :height, :width],
+    )

flexray_mlx-0.2.2/README.md

@@ -1,46 +0,0 @@
-# flexray-mlx
-
-Segment anatomical structures in radiographs with FleXray and native MLX.
-
-[Docs](docs/technical.md) · [MedMLX](https://github.com/MedMLX)
-
-![One chest radiograph → separate lung, heart and individual rib masks.](https://raw.githubusercontent.com/MedMLX/flexray-mlx/main/docs/showcase.png)
-
-*One chest radiograph → separate lung, heart and individual rib masks.*
-
-| Input | Output |
-| :--- | :--- |
-| Raster image or single-frame DICOM | Structure masks and probabilities |
-
-**Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
-
-## Quick start
-
-```bash
-uv venv --python 3.13
-source .venv/bin/activate
-repo="https://github.com/MedMLX/flexray-mlx"
-uv pip install \
-  "flexray-mlx @ git+$repo" \
-  huggingface-hub
-```
-
-Download the released weights and run:
-
-```python
-from huggingface_hub import snapshot_download
-from flexray_mlx.medmlx import run
-
-weights = snapshot_download("VictorButoi/flexray", local_dir="flexray", ignore_patterns=["*.onnx"])
-
-result = run(image_path="chest.png", weights_path=weights, output_dir="out")
-print(result["outputs"]["mask"])
-```
-
-Use `label_name` to select a structure and `ensemble=True` to average all five members.
-
-## License
-
-[Apache-2.0](LICENSE). See [third-party notices](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.