flexray-mlx 0.2.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of flexray-mlx might be problematic. Click here for more details.
- flexray_mlx/__init__.py +21 -0
- flexray_mlx/_kernels.py +98 -0
- flexray_mlx/_random.py +91 -0
- flexray_mlx/_types.py +126 -0
- flexray_mlx/medmlx.py +1208 -0
- flexray_mlx/mlx_model.py +168 -0
- flexray_mlx/py.typed +0 -0
- flexray_mlx/tta.py +248 -0
- flexray_mlx-0.2.1.dist-info/METADATA +69 -0
- flexray_mlx-0.2.1.dist-info/RECORD +14 -0
- flexray_mlx-0.2.1.dist-info/WHEEL +4 -0
- flexray_mlx-0.2.1.dist-info/entry_points.txt +2 -0
- flexray_mlx-0.2.1.dist-info/licenses/LICENSE +202 -0
- flexray_mlx-0.2.1.dist-info/licenses/THIRD_PARTY_NOTICES.md +173 -0
flexray_mlx/mlx_model.py
ADDED
|
@@ -0,0 +1,168 @@
|
|
|
1
|
+
"""Native MLX execution of the pinned FleXray v1.0.2 U-Net checkpoint.
|
|
2
|
+
|
|
3
|
+
The released network is fixed to seven encoder stages and six decoder stages.
|
|
4
|
+
PyTorch is used by the source preprocessing and TTA code, never for this forward pass.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
from importlib import import_module
|
|
10
|
+
from pathlib import Path
|
|
11
|
+
from typing import TYPE_CHECKING, cast
|
|
12
|
+
|
|
13
|
+
import numpy as np
|
|
14
|
+
import yaml
|
|
15
|
+
from medmlx_core import import_mlx
|
|
16
|
+
|
|
17
|
+
from flexray_mlx._kernels import upsample_concat
|
|
18
|
+
from flexray_mlx._types import FloatArray
|
|
19
|
+
|
|
20
|
+
if TYPE_CHECKING:
|
|
21
|
+
from mlx.core import array
|
|
22
|
+
|
|
23
|
+
from flexray_mlx._types import NativeRuntime, NeuralRuntime, SafetensorsRuntime
|
|
24
|
+
|
|
25
|
+
# Namespaces stay lazy until this native implementation is admitted.
|
|
26
|
+
mx = cast("NativeRuntime", import_mlx())
|
|
27
|
+
nn = cast("NeuralRuntime", import_module("mlx.nn"))
|
|
28
|
+
safetensors = cast("SafetensorsRuntime", import_module("safetensors"))
|
|
29
|
+
|
|
30
|
+
_FILTERS = (64, 128, 256, 512, 512, 720, 1024)
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class FleXrayMlxModel:
|
|
34
|
+
"""Released residual U-Net run natively in MLX.
|
|
35
|
+
|
|
36
|
+
Activations are channels-last (NHWC) internally, as MLX convolutions expect.
|
|
37
|
+
Convolution weights are stored upstream in PyTorch OIHW order and are
|
|
38
|
+
transposed to OHWI when loaded. The public input and output stay NCHW.
|
|
39
|
+
|
|
40
|
+
Attributes:
|
|
41
|
+
weights: Every tensor in ``model.safetensors`` as float32, under its
|
|
42
|
+
original key.
|
|
43
|
+
"""
|
|
44
|
+
|
|
45
|
+
def __init__(self, bundle: Path) -> None:
|
|
46
|
+
"""Load one member bundle after checking it matches the released network.
|
|
47
|
+
|
|
48
|
+
Args:
|
|
49
|
+
bundle: Directory holding the released ``config.yml`` (its ``model``
|
|
50
|
+
section) and ``model.safetensors``.
|
|
51
|
+
|
|
52
|
+
Raises:
|
|
53
|
+
ValueError: The config describes any other architecture or activation.
|
|
54
|
+
OSError: ``config.yml`` or ``model.safetensors`` cannot be read.
|
|
55
|
+
Malformed files raise the YAML or safetensors library's own errors.
|
|
56
|
+
"""
|
|
57
|
+
config = cast(
|
|
58
|
+
dict[str, dict[str, object]],
|
|
59
|
+
yaml.safe_load((bundle / "config.yml").read_text(encoding="utf-8")),
|
|
60
|
+
)["model"]
|
|
61
|
+
expected: dict[str, object] = {
|
|
62
|
+
"_class": "fxr.models.UNet",
|
|
63
|
+
"in_channels": 1,
|
|
64
|
+
"out_channels": 61,
|
|
65
|
+
"filters": list(_FILTERS),
|
|
66
|
+
"convs_per_block": 3,
|
|
67
|
+
"dropout": 0.0,
|
|
68
|
+
"norm": "instance",
|
|
69
|
+
"residual": True,
|
|
70
|
+
"norm_after_activation": True,
|
|
71
|
+
"skip_connections": True,
|
|
72
|
+
"upsample_align_corners": True,
|
|
73
|
+
"residual_projection": "auto",
|
|
74
|
+
"residual_shortcut_norm": True,
|
|
75
|
+
}
|
|
76
|
+
for key, value in expected.items():
|
|
77
|
+
if config.get(key) != value:
|
|
78
|
+
raise ValueError(f"Unsupported FleXray MLX architecture: {key}={config.get(key)!r}")
|
|
79
|
+
if str(config.get("activation", "LeakyReLU")).casefold() != "leakyrelu":
|
|
80
|
+
raise ValueError("Unsupported FleXray MLX activation")
|
|
81
|
+
if config.get("bottleneck_convs") not in (None, 3) or config.get("up_filters") is not None:
|
|
82
|
+
raise ValueError("Unsupported FleXray MLX bottleneck or decoder configuration")
|
|
83
|
+
|
|
84
|
+
self.weights: dict[str, array] = {}
|
|
85
|
+
with safetensors.safe_open(bundle / "model.safetensors", framework="np") as reader:
|
|
86
|
+
for key in reader.keys(): # noqa: SIM118 - safe_open is not iterable
|
|
87
|
+
value = np.asarray(reader.get_tensor(key), dtype=np.float32)
|
|
88
|
+
# 4D weights are convolution kernels: OIHW -> OHWI. Norm weights are 1D.
|
|
89
|
+
if key.endswith("weight") and value.ndim == 4:
|
|
90
|
+
value = value.transpose(0, 2, 3, 1)
|
|
91
|
+
self.weights[key] = mx.array(value)
|
|
92
|
+
# MLX is lazy; evaluate now so loading is complete before the first forward pass.
|
|
93
|
+
mx.eval(*self.weights.values())
|
|
94
|
+
self._consumed: set[str] = set()
|
|
95
|
+
self._pool = nn.MaxPool2d(kernel_size=2, stride=2)
|
|
96
|
+
|
|
97
|
+
def _weight(self, key: str) -> array:
|
|
98
|
+
self._consumed.add(key)
|
|
99
|
+
return self.weights[key]
|
|
100
|
+
|
|
101
|
+
def _conv(self, x: array, prefix: str, *, padding: int) -> array:
|
|
102
|
+
weight = self._weight(f"{prefix}.weight")
|
|
103
|
+
bias = self._weight(f"{prefix}.bias")
|
|
104
|
+
return mx.conv2d(x, weight, padding=padding) + bias
|
|
105
|
+
|
|
106
|
+
def _norm(self, x: array, prefix: str) -> array:
|
|
107
|
+
gamma = self._weight(f"{prefix}.weight")
|
|
108
|
+
beta = self._weight(f"{prefix}.bias")
|
|
109
|
+
# Instance norm on NHWC: per sample and channel over H and W. eps=1e-5 is
|
|
110
|
+
# PyTorch's InstanceNorm2d default; the upstream parity test covers it.
|
|
111
|
+
mean = mx.mean(x, axis=(1, 2), keepdims=True)
|
|
112
|
+
variance = mx.mean(mx.square(x - mean), axis=(1, 2), keepdims=True)
|
|
113
|
+
return (x - mean) * mx.rsqrt(variance + 1e-5) * gamma + beta
|
|
114
|
+
|
|
115
|
+
def _block(self, x: array, prefix: str) -> array:
|
|
116
|
+
residual = x
|
|
117
|
+
out = x
|
|
118
|
+
for index in range(3):
|
|
119
|
+
# Checkpoint key layout per step: conv at layers.3i, activation at
|
|
120
|
+
# 3i+1 (no tensors), norm at 3i+2 (norm_after_activation).
|
|
121
|
+
layer = f"{prefix}.layers.{index * 3}"
|
|
122
|
+
out = self._conv(out, layer, padding=1)
|
|
123
|
+
# LeakyReLU with PyTorch's default negative slope of 0.01.
|
|
124
|
+
out = mx.maximum(out, out * 0.01)
|
|
125
|
+
out = self._norm(out, f"{prefix}.layers.{index * 3 + 2}")
|
|
126
|
+
# A 1x1 projection shortcut exists only where channel counts differ.
|
|
127
|
+
if f"{prefix}.shortcut.conv.weight" in self.weights:
|
|
128
|
+
residual = self._conv(residual, f"{prefix}.shortcut.conv", padding=0)
|
|
129
|
+
residual = self._norm(residual, f"{prefix}.shortcut.instancenorm")
|
|
130
|
+
return out + residual
|
|
131
|
+
|
|
132
|
+
def __call__(self, images: FloatArray) -> FloatArray:
|
|
133
|
+
"""Run the U-Net on prepared images and return raw logits.
|
|
134
|
+
|
|
135
|
+
Args:
|
|
136
|
+
images: Float32 ``(B, 1, 256, 256)`` NCHW radiographs, already scaled
|
|
137
|
+
and normalized as the bundle's preprocessing specifies.
|
|
138
|
+
|
|
139
|
+
Returns:
|
|
140
|
+
A new float32 ``(B, 61, 256, 256)`` NCHW array of logits, one channel
|
|
141
|
+
per label in the bundle's label-schema order (channel 0 is background).
|
|
142
|
+
Apply a sigmoid for probabilities.
|
|
143
|
+
|
|
144
|
+
Raises:
|
|
145
|
+
ValueError: ``images`` has another shape, or the bundle holds tensors
|
|
146
|
+
this network does not use. A tensor the network needs but the
|
|
147
|
+
bundle lacks raises ``KeyError``.
|
|
148
|
+
"""
|
|
149
|
+
if images.ndim != 4 or images.shape[1:] != (1, 256, 256):
|
|
150
|
+
raise ValueError(f"FleXray MLX expects Bx1x256x256, got {images.shape}")
|
|
151
|
+
x = mx.array(np.asarray(images, dtype=np.float32).transpose(0, 2, 3, 1))
|
|
152
|
+
skips: list[array] = []
|
|
153
|
+
for index in range(len(_FILTERS)):
|
|
154
|
+
x = self._block(x, f"down_blocks.{index}")
|
|
155
|
+
if index < len(_FILTERS) - 1:
|
|
156
|
+
skips.append(x)
|
|
157
|
+
x = self._pool(x)
|
|
158
|
+
for index, skip in enumerate(reversed(skips)):
|
|
159
|
+
# Preserve corner alignment and resized-before-skip channel order,
|
|
160
|
+
# while avoiding the separate bilinear samples and resize tensor.
|
|
161
|
+
x = self._block(upsample_concat(x, skip, mx=mx), f"up_blocks.{index}")
|
|
162
|
+
logits = self._conv(x, "output_conv", padding=0)
|
|
163
|
+
mx.eval(logits)
|
|
164
|
+
# Fail on extra or misnamed tensors instead of silently running a different net.
|
|
165
|
+
if self._consumed != self.weights.keys():
|
|
166
|
+
missing = sorted(self.weights.keys() - self._consumed)
|
|
167
|
+
raise ValueError(f"FleXray MLX did not consume every released tensor: {missing[:3]}")
|
|
168
|
+
return cast(FloatArray, np.asarray(logits)).transpose(0, 3, 1, 2).copy()
|
flexray_mlx/py.typed
ADDED
|
File without changes
|
flexray_mlx/tta.py
ADDED
|
@@ -0,0 +1,248 @@
|
|
|
1
|
+
"""NumPy port of pinned fxr ``tta_v3`` (FleXray bf6d014, MIT).
|
|
2
|
+
|
|
3
|
+
Each view is drawn once, normalized again, shared by all members, and unflipped
|
|
4
|
+
on the logits. Accumulation is sequential in view/member order, in float32
|
|
5
|
+
probability space. The plain view is already normalized and stays untouched.
|
|
6
|
+
See root THIRD_PARTY_NOTICES.md for source notices.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import math
|
|
12
|
+
from collections.abc import Callable, Sequence
|
|
13
|
+
from dataclasses import dataclass
|
|
14
|
+
from typing import cast
|
|
15
|
+
|
|
16
|
+
import numpy as np
|
|
17
|
+
|
|
18
|
+
from flexray_mlx._random import TorchRandom
|
|
19
|
+
from flexray_mlx._types import BoolArray, FloatArray, IntArray
|
|
20
|
+
|
|
21
|
+
Forward = Callable[[FloatArray], FloatArray]
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@dataclass(frozen=True, slots=True)
|
|
25
|
+
class TtaSettings:
|
|
26
|
+
"""Validated pass count and isolated Torch CPU seed."""
|
|
27
|
+
|
|
28
|
+
samples: int = 1
|
|
29
|
+
seed: int | None = None
|
|
30
|
+
|
|
31
|
+
def __post_init__(self) -> None:
|
|
32
|
+
if (
|
|
33
|
+
isinstance(self.samples, bool)
|
|
34
|
+
or not isinstance(cast(object, self.samples), int)
|
|
35
|
+
or self.samples < 0
|
|
36
|
+
):
|
|
37
|
+
raise ValueError("tta_samples must be a non-negative integer")
|
|
38
|
+
if self.seed is not None and (
|
|
39
|
+
isinstance(self.seed, bool)
|
|
40
|
+
or not isinstance(cast(object, self.seed), int)
|
|
41
|
+
or not -(2**63) <= self.seed < 2**64
|
|
42
|
+
):
|
|
43
|
+
raise ValueError("tta_seed must be an integer in [-2**63, 2**64 - 1], or None")
|
|
44
|
+
|
|
45
|
+
@property
|
|
46
|
+
def resolved_samples(self) -> int:
|
|
47
|
+
"""Passes per member actually run; ``samples`` of 0 resolves to one plain pass."""
|
|
48
|
+
return max(self.samples, 1)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def predict_with_tta(
|
|
52
|
+
forward_fn: Forward | Sequence[Forward],
|
|
53
|
+
images: FloatArray,
|
|
54
|
+
*,
|
|
55
|
+
tta_samples: int = 1,
|
|
56
|
+
seed: int | None = None,
|
|
57
|
+
normalizer: Forward | None = None,
|
|
58
|
+
) -> tuple[FloatArray, FloatArray]:
|
|
59
|
+
"""Return multilabel ``(logits, probabilities)``, following fxr's total-pass contract.
|
|
60
|
+
|
|
61
|
+
The first pass is the plain input; each further pass is an augmented view
|
|
62
|
+
shared by all members. Probabilities are the mean of every member's sigmoid
|
|
63
|
+
over every pass. For a single plain pass of one member the logits are that
|
|
64
|
+
member's raw output; otherwise they are the logit of the mean probability
|
|
65
|
+
clipped to [1e-6, 1 - 1e-6].
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
forward_fn: One callable, or a sequence of member callables, mapping
|
|
69
|
+
float32 ``(B, C, H, W)`` images to float32 ``(B, L, H, W)`` logits.
|
|
70
|
+
images: Prepared float32 ``(B, C, H, W)`` batch with 1 or 3 channels,
|
|
71
|
+
already normalized.
|
|
72
|
+
tta_samples: Total passes per member; 0 or 1 means one plain pass.
|
|
73
|
+
seed: Seed for the isolated Torch-CPU-compatible stream; None is unseeded.
|
|
74
|
+
normalizer: Applied to each augmented view (not the plain input) so views
|
|
75
|
+
are normalized as the plain input was.
|
|
76
|
+
|
|
77
|
+
Raises:
|
|
78
|
+
ValueError: Invalid ``tta_samples``/``seed``, no forward callable, or
|
|
79
|
+
``images`` that is not float32 ``(B, C, H, W)`` with 1 or 3 channels.
|
|
80
|
+
"""
|
|
81
|
+
|
|
82
|
+
# Deferred: flexray_mlx.medmlx imports this module at load time.
|
|
83
|
+
from flexray_mlx.medmlx import multilabel_probabilities
|
|
84
|
+
|
|
85
|
+
settings = TtaSettings(tta_samples, seed)
|
|
86
|
+
forwards = (forward_fn,) if callable(forward_fn) else tuple(forward_fn)
|
|
87
|
+
if not forwards:
|
|
88
|
+
raise ValueError("predict_with_tta requires at least one forward callable")
|
|
89
|
+
if images.ndim != 4 or images.shape[1] not in {1, 3} or images.dtype != np.float32:
|
|
90
|
+
raise ValueError("FleXray TTA expects float32 BxCxHxW images with 1 or 3 channels")
|
|
91
|
+
logits = forwards[0](images)
|
|
92
|
+
probability_sum = multilabel_probabilities(logits)
|
|
93
|
+
if settings.resolved_samples == 1 and len(forwards) == 1:
|
|
94
|
+
return logits, probability_sum
|
|
95
|
+
for forward in forwards[1:]:
|
|
96
|
+
probability_sum += multilabel_probabilities(forward(images))
|
|
97
|
+
random = TorchRandom(settings.seed)
|
|
98
|
+
for _ in range(settings.resolved_samples - 1):
|
|
99
|
+
augmented, flipped = augment_batch(images, random)
|
|
100
|
+
if normalizer is not None:
|
|
101
|
+
augmented = normalizer(augmented)
|
|
102
|
+
for forward in forwards:
|
|
103
|
+
view_logits = forward(augmented)
|
|
104
|
+
unflipped = np.where(flipped[:, None, None, None], view_logits[..., ::-1], view_logits)
|
|
105
|
+
probability_sum += multilabel_probabilities(unflipped)
|
|
106
|
+
probabilities = probability_sum / np.float32(settings.resolved_samples * len(forwards))
|
|
107
|
+
clamped = np.clip(probabilities, np.float32(1e-6), np.float32(1 - 1e-6))
|
|
108
|
+
return np.log(clamped) - np.log1p(-clamped), probabilities
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
def augment_batch(images: FloatArray, random: TorchRandom) -> tuple[FloatArray, BoolArray]:
|
|
112
|
+
"""Draw the fxr chain independently per sample, in batch order.
|
|
113
|
+
|
|
114
|
+
Returns the stacked ``(B, C, H, W)`` views and a ``(B,)`` boolean mask of
|
|
115
|
+
which samples were flipped along the width axis.
|
|
116
|
+
"""
|
|
117
|
+
|
|
118
|
+
views, flips = zip(*(augment_sample(image, random) for image in images), strict=True)
|
|
119
|
+
return np.stack(views), np.asarray(flips, dtype=np.bool_)
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
def augment_sample(image: FloatArray, random: TorchRandom) -> tuple[FloatArray, np.bool_]:
|
|
123
|
+
"""Horizontal flip, invert, CLAHE/gamma, raw contrast, sharpness, then noise.
|
|
124
|
+
|
|
125
|
+
``image`` is one prepared ``(C, H, W)`` float32 sample. Random draws are
|
|
126
|
+
consumed in fxr's order. The CLAHE, gamma, contrast and sharpness steps clip
|
|
127
|
+
to [0, 1]; invert and noise do not.
|
|
128
|
+
Returns the view and whether it was flipped (so logits can be unflipped).
|
|
129
|
+
"""
|
|
130
|
+
|
|
131
|
+
flipped = random.chance(0.5)
|
|
132
|
+
if flipped:
|
|
133
|
+
image = image[..., ::-1]
|
|
134
|
+
if random.chance(0.5):
|
|
135
|
+
image = np.float32(1.0) - image
|
|
136
|
+
choice = random.uniform()
|
|
137
|
+
if choice < 0.1:
|
|
138
|
+
image = clahe(image, clip_limit=random.uniform(1.0, 2.0), grid_size=(8, 8))
|
|
139
|
+
elif choice < 0.35:
|
|
140
|
+
gamma, gain = random.uniform(0.9, 1.1), random.uniform(0.9, 1.1)
|
|
141
|
+
image = np.clip(np.float32(gain) * np.power(image, np.float32(gamma)), 0, 1)
|
|
142
|
+
if random.chance(0.25):
|
|
143
|
+
image = cast(FloatArray, np.clip(image * np.float32(random.uniform(0.7, 1.3)), 0, 1))
|
|
144
|
+
if random.chance(0.5):
|
|
145
|
+
image = adjust_sharpness(image, random.uniform(0.7, 1.3))
|
|
146
|
+
if random.chance(0.25):
|
|
147
|
+
image = cast(FloatArray, image + random.normal(image.shape) * np.float32(0.01))
|
|
148
|
+
return image, flipped
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def adjust_sharpness(image: FloatArray, factor: float) -> FloatArray:
|
|
152
|
+
"""Sharpen ``(C, H, W)`` ``image`` as ``factor * image + (1 - factor) * blurred``.
|
|
153
|
+
|
|
154
|
+
The blur is torchvision's 3x3 ``[1,1,1; 1,5,1; 1,1,1] / 13`` kernel, and the
|
|
155
|
+
one-pixel border keeps the original values. The blend is clipped to [0, 1].
|
|
156
|
+
Images with a side of 2 pixels or fewer are returned unchanged.
|
|
157
|
+
"""
|
|
158
|
+
|
|
159
|
+
if min(image.shape[-2:]) <= 2:
|
|
160
|
+
return image
|
|
161
|
+
# Deferred: flexray_mlx.medmlx imports this module at load time.
|
|
162
|
+
from flexray_mlx.medmlx import fma32
|
|
163
|
+
|
|
164
|
+
height, width = image.shape[-2:]
|
|
165
|
+
blurred = image.copy()
|
|
166
|
+
interior = np.zeros_like(image[:, 1:-1, 1:-1])
|
|
167
|
+
for y in range(3):
|
|
168
|
+
for x in range(3):
|
|
169
|
+
weight = np.float32(5.0 if (y, x) == (1, 1) else 1.0) / np.float32(13.0)
|
|
170
|
+
interior = fma32(image[:, y : y + height - 2, x : x + width - 2], weight, interior)
|
|
171
|
+
blurred[:, 1:-1, 1:-1] = interior
|
|
172
|
+
return cast(
|
|
173
|
+
FloatArray, np.clip(np.float32(factor) * image + np.float32(1.0 - factor) * blurred, 0, 1)
|
|
174
|
+
)
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
def _interpolation_axis(
|
|
178
|
+
length: int, tile_size: int, tile_count: int
|
|
179
|
+
) -> tuple[IntArray, IntArray, FloatArray]:
|
|
180
|
+
"""Per pixel: nearest lower/upper tile index and the weight of the upper tile.
|
|
181
|
+
|
|
182
|
+
Tile centres sit at ``(k + 0.5) * tile_size``; indices clamp at the edges.
|
|
183
|
+
"""
|
|
184
|
+
|
|
185
|
+
position = np.arange(length, dtype=np.float32) / np.float32(tile_size) - np.float32(0.5)
|
|
186
|
+
lower = np.floor(position).astype(np.int64)
|
|
187
|
+
return (
|
|
188
|
+
np.clip(lower, 0, tile_count - 1),
|
|
189
|
+
np.clip(lower + 1, 0, tile_count - 1),
|
|
190
|
+
(position - lower.astype(np.float32)),
|
|
191
|
+
)
|
|
192
|
+
|
|
193
|
+
|
|
194
|
+
def clahe(image: FloatArray, *, clip_limit: float, grid_size: tuple[int, int]) -> FloatArray:
|
|
195
|
+
"""Contrast-limited adaptive histogram equalization of a ``(C, H, W)`` image.
|
|
196
|
+
|
|
197
|
+
Follows fxr: values are clipped to [0, 1] and quantized to 256 bins per
|
|
198
|
+
channel, histograms are built per tile and clipped at
|
|
199
|
+
``clip_limit * pixels / 256`` with the excess redistributed, and each pixel
|
|
200
|
+
interpolates bilinearly between the four nearest tile centres. ``grid_size``
|
|
201
|
+
is ``(rows, cols)`` of tiles, capped at the image size; tiles are rounded up
|
|
202
|
+
to even sizes and the image is padded to fit (reflect, or edge when the
|
|
203
|
+
padding reaches the image size in either direction). Returns float32 ``(C, H, W)`` in [0, 1].
|
|
204
|
+
|
|
205
|
+
Raises:
|
|
206
|
+
ValueError: ``image`` is not a non-empty 3D array, ``clip_limit`` is not
|
|
207
|
+
positive and finite, or ``grid_size`` is not two positive integers.
|
|
208
|
+
"""
|
|
209
|
+
|
|
210
|
+
if image.ndim != 3 or min(image.shape) < 1:
|
|
211
|
+
raise ValueError(f"CLAHE expects a non-empty CxHxW image, got {image.shape}")
|
|
212
|
+
if not math.isfinite(clip_limit) or clip_limit <= 0:
|
|
213
|
+
raise ValueError("CLAHE clip_limit must be positive and finite")
|
|
214
|
+
if len(grid_size) != 2 or any(isinstance(n, bool) or n <= 0 for n in grid_size):
|
|
215
|
+
raise ValueError("CLAHE grid_size must contain two positive integers")
|
|
216
|
+
channels, height, width = image.shape
|
|
217
|
+
rows, cols = min(grid_size[0], height), min(grid_size[1], width)
|
|
218
|
+
th, tw = math.ceil(height / rows), math.ceil(width / cols)
|
|
219
|
+
th, tw = th + th % 2, tw + tw % 2
|
|
220
|
+
ph, pw = th * rows - height, tw * cols - width
|
|
221
|
+
work = np.clip(image, 0, 1).astype(np.float32)
|
|
222
|
+
if ph or pw:
|
|
223
|
+
mode = "reflect" if ph < height and pw < width else "edge"
|
|
224
|
+
work = np.pad(work, ((0, 0), (0, ph), (0, pw)), mode=mode)
|
|
225
|
+
tiles = work.reshape(channels, rows, th, cols, tw).transpose(1, 3, 0, 2, 4)
|
|
226
|
+
pixels = th * tw
|
|
227
|
+
bins = (tiles.reshape(-1, pixels) * np.float32(255)).astype(np.int64)
|
|
228
|
+
offsets = np.arange(len(bins), dtype=np.int64)[:, None] * 256
|
|
229
|
+
histograms = np.bincount((bins + offsets).ravel(), minlength=len(bins) * 256).reshape(-1, 256)
|
|
230
|
+
histograms = np.minimum(histograms, max(math.floor(clip_limit * pixels / 256), 1))
|
|
231
|
+
clipped = pixels - histograms.sum(axis=1)
|
|
232
|
+
residual = clipped % 256
|
|
233
|
+
histograms += ((clipped - residual) // 256)[:, None]
|
|
234
|
+
histograms += np.arange(256)[None] < residual[:, None]
|
|
235
|
+
lookup = np.floor(np.cumsum(histograms, axis=1).astype(np.float32) * np.float32(255 / pixels))
|
|
236
|
+
lookup = np.clip(lookup, 0, 255).reshape(rows, cols, channels, 256).transpose(2, 0, 1, 3)
|
|
237
|
+
pixel_bins = (work * np.float32(255)).astype(np.int64)
|
|
238
|
+
y0, y1, wy = _interpolation_axis(work.shape[1], th, rows)
|
|
239
|
+
x0, x1, wx = _interpolation_axis(work.shape[2], tw, cols)
|
|
240
|
+
channel = cast(IntArray, np.arange(channels))[:, None, None]
|
|
241
|
+
v00 = lookup[channel, y0[None, :, None], x0[None, None], pixel_bins]
|
|
242
|
+
v01 = lookup[channel, y0[None, :, None], x1[None, None], pixel_bins]
|
|
243
|
+
v10 = lookup[channel, y1[None, :, None], x0[None, None], pixel_bins]
|
|
244
|
+
v11 = lookup[channel, y1[None, :, None], x1[None, None], pixel_bins]
|
|
245
|
+
wy, wx = wy[None, :, None], wx[None, None]
|
|
246
|
+
top = v00 * (np.float32(1) - wx) + v01 * wx
|
|
247
|
+
bottom = v10 * (np.float32(1) - wx) + v11 * wx
|
|
248
|
+
return ((top * (np.float32(1) - wy) + bottom * wy) / np.float32(255))[:, :height, :width]
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: flexray-mlx
|
|
3
|
+
Version: 0.2.1
|
|
4
|
+
Summary: Native MLX implementation of FleXray
|
|
5
|
+
Author: Joseph Sandoval, Alfredo Sandoval Jr.
|
|
6
|
+
License-Expression: Apache-2.0 AND MIT AND BSD-3-Clause
|
|
7
|
+
License-File: LICENSE
|
|
8
|
+
License-File: THIRD_PARTY_NOTICES.md
|
|
9
|
+
Classifier: Operating System :: MacOS :: MacOS X
|
|
10
|
+
Requires-Python: <3.15,>=3.13
|
|
11
|
+
Requires-Dist: medmlx-core==0.2.3
|
|
12
|
+
Requires-Dist: mlx<1,>=0.32.4
|
|
13
|
+
Requires-Dist: numpy<3,>=2.0
|
|
14
|
+
Requires-Dist: pillow<14,>=13
|
|
15
|
+
Requires-Dist: pydicom<4,>=3
|
|
16
|
+
Requires-Dist: pyyaml<7,>=6
|
|
17
|
+
Requires-Dist: safetensors<1,>=0.6
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: pyright==1.1.414; extra == 'dev'
|
|
20
|
+
Requires-Dist: pytest<10,>=8.4; extra == 'dev'
|
|
21
|
+
Requires-Dist: ruff<1,>=0.13; extra == 'dev'
|
|
22
|
+
Description-Content-Type: text/markdown
|
|
23
|
+
|
|
24
|
+
# flexray-mlx
|
|
25
|
+
|
|
26
|
+
Segment anatomical structures in radiographs with FleXray and native MLX.
|
|
27
|
+
|
|
28
|
+
[Docs](docs/technical.md) · [MedMLX](https://github.com/MedMLX)
|
|
29
|
+
|
|
30
|
+

|
|
31
|
+
|
|
32
|
+
*One chest radiograph → separate lung, heart and individual rib masks.*
|
|
33
|
+
|
|
34
|
+
| Input | Output |
|
|
35
|
+
| :--- | :--- |
|
|
36
|
+
| Raster image or single-frame DICOM | Structure masks and probabilities |
|
|
37
|
+
|
|
38
|
+
**Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
|
|
39
|
+
|
|
40
|
+
## Quick start
|
|
41
|
+
|
|
42
|
+
```bash
|
|
43
|
+
uv venv --python 3.13
|
|
44
|
+
source .venv/bin/activate
|
|
45
|
+
repo="https://github.com/MedMLX/flexray-mlx"
|
|
46
|
+
uv pip install \
|
|
47
|
+
"flexray-mlx @ git+$repo" \
|
|
48
|
+
huggingface-hub
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
Download the released weights and run:
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
from huggingface_hub import snapshot_download
|
|
55
|
+
from flexray_mlx.medmlx import run
|
|
56
|
+
|
|
57
|
+
weights = snapshot_download("VictorButoi/flexray", local_dir="flexray", ignore_patterns=["*.onnx"])
|
|
58
|
+
|
|
59
|
+
result = run(image_path="chest.png", weights_path=weights, output_dir="out")
|
|
60
|
+
print(result["outputs"]["mask"])
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
Use `label_name` to select a structure and `ensemble=True` to average all five members.
|
|
64
|
+
|
|
65
|
+
## License
|
|
66
|
+
|
|
67
|
+
[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).
|
|
68
|
+
|
|
69
|
+
For research and education. Outputs require human review.
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
flexray_mlx/__init__.py,sha256=zELHZB2HOgjvicl7pfqRAdO0x8OeeG0NHG8Wnk_H25c,563
|
|
2
|
+
flexray_mlx/_kernels.py,sha256=R13EcWmb9pUolbJUICGJ6R0a3XjzKvdo7msoG6uwDxY,3604
|
|
3
|
+
flexray_mlx/_random.py,sha256=zgeYBsuLy-eKRmrWHZmvcadr_EqaTHi0RiPPjeil5v8,3921
|
|
4
|
+
flexray_mlx/_types.py,sha256=sdxMEDJJa8cP5xVqJ12_ESaO-8cGOT8fFxukxrxzCzA,3586
|
|
5
|
+
flexray_mlx/medmlx.py,sha256=uEoDBbtRB2K_oMRgv2UJw8ZxBjc3Pf4881GApnW_Kmo,49894
|
|
6
|
+
flexray_mlx/mlx_model.py,sha256=zjNuEdppSrctGLskKgaEmkjkhK_wcOM8koECVB8m85c,7678
|
|
7
|
+
flexray_mlx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
8
|
+
flexray_mlx/tta.py,sha256=yg8MDF9VvOJQ8p4rvslOWo2lAGlN9FV-9a849dRqFb8,11346
|
|
9
|
+
flexray_mlx-0.2.1.dist-info/METADATA,sha256=dBt_JlT4b918Y7bpB1V0iqP9xaNZAoU3aGtWXiNnuHw,2201
|
|
10
|
+
flexray_mlx-0.2.1.dist-info/WHEEL,sha256=2iZrJkDIkeDh0k9SikKMqFuRP-ua6mb4Ir1AsER4BEw,87
|
|
11
|
+
flexray_mlx-0.2.1.dist-info/entry_points.txt,sha256=fISmTVUJp2A8zG55t1nLC_bwQuztR_shgMVADfCLxwE,51
|
|
12
|
+
flexray_mlx-0.2.1.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
|
|
13
|
+
flexray_mlx-0.2.1.dist-info/licenses/THIRD_PARTY_NOTICES.md,sha256=HaCU3wNwRTSBx0jregZw0a1GAsi3b6-t_B59tEEABzU,7430
|
|
14
|
+
flexray_mlx-0.2.1.dist-info/RECORD,,
|