segresnet-mlx 0.2.3__tar.gz → 0.3.0__tar.gz

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.
Files changed (24) hide show
  1. segresnet_mlx-0.3.0/PKG-INFO +75 -0
  2. segresnet_mlx-0.3.0/README.md +51 -0
  3. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/pyproject.toml +6 -12
  4. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/__init__.py +3 -7
  5. segresnet_mlx-0.3.0/src/segresnet_mlx/_nifti.py +109 -0
  6. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/_records.py +1 -0
  7. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/_types.py +9 -4
  8. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/medmlx.py +83 -11
  9. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/network.py +9 -3
  10. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/predictor.py +5 -0
  11. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/telemetry.py +9 -1
  12. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/topology.py +1 -0
  13. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/transforms.py +29 -116
  14. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/weights.py +4 -86
  15. segresnet_mlx-0.2.3/PKG-INFO +0 -87
  16. segresnet_mlx-0.2.3/README.md +0 -58
  17. segresnet_mlx-0.2.3/src/segresnet_mlx/_nifti.py +0 -63
  18. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/.gitignore +0 -0
  19. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/LICENSE +0 -0
  20. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/THIRD_PARTY_NOTICES.md +0 -0
  21. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/_arrays.py +0 -0
  22. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/_runtime.py +0 -0
  23. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/inference.py +0 -0
  24. {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/py.typed +0 -0
@@ -0,0 +1,75 @@
1
+ Metadata-Version: 2.5
2
+ Name: segresnet-mlx
3
+ Version: 0.3.0
4
+ Summary: Native MLX SegResNet with NumPy sliding-window logit inference
5
+ Project-URL: Homepage, https://huggingface.co/MedMLX/segresnet-renal-cect-mlx
6
+ Author: Joseph Sandoval, Alfredo Sandoval Jr.
7
+ License-Expression: Apache-2.0
8
+ License-File: LICENSE
9
+ License-File: THIRD_PARTY_NOTICES.md
10
+ Classifier: Operating System :: MacOS :: MacOS X
11
+ Requires-Python: <3.15,>=3.13
12
+ Requires-Dist: medmlx-core==0.2.5
13
+ Requires-Dist: nibabel<6,>=5.3
14
+ Requires-Dist: numpy<3,>=2.0
15
+ Provides-Extra: dev
16
+ Requires-Dist: hypothesis<7,>=6; extra == 'dev'
17
+ Requires-Dist: monai==1.6.1; extra == 'dev'
18
+ Requires-Dist: pyright==1.1.414; extra == 'dev'
19
+ Requires-Dist: pytest<10,>=8.4; extra == 'dev'
20
+ Requires-Dist: pytorch-ignite<0.6,>=0.4.11; extra == 'dev'
21
+ Requires-Dist: ruff<1,>=0.13; extra == 'dev'
22
+ Requires-Dist: torch<3,>=2.6; extra == 'dev'
23
+ Description-Content-Type: text/markdown
24
+
25
+ # segresnet-mlx
26
+
27
+ Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegResNet models in native MLX.
28
+
29
+ [BraTS weights](https://huggingface.co/MedMLX/segresnet-brats-mlx) · [Renal weights](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx)
30
+
31
+ | Input | Output |
32
+ | :--- | :--- |
33
+ | MRI sequences or CT volumes | Segmentation on the input grid |
34
+
35
+ **Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
36
+
37
+ ## Quick start
38
+
39
+ ```bash
40
+ uv venv --python 3.13
41
+ source .venv/bin/activate
42
+ uv pip install "segresnet-mlx>=0.3.0"
43
+ ```
44
+
45
+ ```python
46
+ from segresnet_mlx.medmlx import brats
47
+
48
+ result = brats(
49
+ t1c_path="t1c.nii.gz", t1_path="t1.nii.gz", t2_path="t2.nii.gz",
50
+ flair_path="flair.nii.gz", output_dir="out",
51
+ )
52
+ print(result["outputs"]["mask"])
53
+ ```
54
+
55
+ The registered models are `segresnet-brats` and `segresnet-renal-cect`. Their `brats`
56
+ and `renal_cect` runners select the matching bundle internally and download
57
+ `MedMLX/segresnet-brats-mlx` or `MedMLX/segresnet-renal-cect-mlx` when `weights_path`
58
+ is omitted. An explicit folder contains `segresnet-brats.safetensors` for BraTS or
59
+ `segresnet-renal-cect.safetensors` for renal CECT. The renal runner takes
60
+ `arterial_path`, `venous_path` and `excretory_path`. Keep each model's input volumes
61
+ on one voxel grid.
62
+
63
+ The standalone `run(bundle=..., weights_path=...)` interface remains available for
64
+ explicit local checkpoints, including the whole-body CT variants. Published native
65
+ weights are also available for [renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx)
66
+ and [whole-body CT](https://huggingface.co/MedMLX/segresnet-wholebody-mlx).
67
+
68
+ Downloaded commits are returned in `details.weights_revision`; explicit local
69
+ weights report `null` unless the folder carries a MedMLX download revision.
70
+
71
+ ## License
72
+
73
+ Apache-2.0; license and third-party notices are included in the distribution. Model weights retain their publisher terms.
74
+
75
+ For research and education. Outputs require human review.
@@ -0,0 +1,51 @@
1
+ # segresnet-mlx
2
+
3
+ Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegResNet models in native MLX.
4
+
5
+ [BraTS weights](https://huggingface.co/MedMLX/segresnet-brats-mlx) · [Renal weights](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx)
6
+
7
+ | Input | Output |
8
+ | :--- | :--- |
9
+ | MRI sequences or CT volumes | Segmentation on the input grid |
10
+
11
+ **Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
12
+
13
+ ## Quick start
14
+
15
+ ```bash
16
+ uv venv --python 3.13
17
+ source .venv/bin/activate
18
+ uv pip install "segresnet-mlx>=0.3.0"
19
+ ```
20
+
21
+ ```python
22
+ from segresnet_mlx.medmlx import brats
23
+
24
+ result = brats(
25
+ t1c_path="t1c.nii.gz", t1_path="t1.nii.gz", t2_path="t2.nii.gz",
26
+ flair_path="flair.nii.gz", output_dir="out",
27
+ )
28
+ print(result["outputs"]["mask"])
29
+ ```
30
+
31
+ The registered models are `segresnet-brats` and `segresnet-renal-cect`. Their `brats`
32
+ and `renal_cect` runners select the matching bundle internally and download
33
+ `MedMLX/segresnet-brats-mlx` or `MedMLX/segresnet-renal-cect-mlx` when `weights_path`
34
+ is omitted. An explicit folder contains `segresnet-brats.safetensors` for BraTS or
35
+ `segresnet-renal-cect.safetensors` for renal CECT. The renal runner takes
36
+ `arterial_path`, `venous_path` and `excretory_path`. Keep each model's input volumes
37
+ on one voxel grid.
38
+
39
+ The standalone `run(bundle=..., weights_path=...)` interface remains available for
40
+ explicit local checkpoints, including the whole-body CT variants. Published native
41
+ weights are also available for [renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx)
42
+ and [whole-body CT](https://huggingface.co/MedMLX/segresnet-wholebody-mlx).
43
+
44
+ Downloaded commits are returned in `details.weights_revision`; explicit local
45
+ weights report `null` unless the folder carries a MedMLX download revision.
46
+
47
+ ## License
48
+
49
+ Apache-2.0; license and third-party notices are included in the distribution. Model weights retain their publisher terms.
50
+
51
+ For research and education. Outputs require human review.
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "segresnet-mlx"
7
- version = "0.2.3"
7
+ version = "0.3.0"
8
8
  description = "Native MLX SegResNet with NumPy sliding-window logit inference"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.13,<3.15"
@@ -13,22 +13,19 @@ license = "Apache-2.0"
13
13
  license-files = ["LICENSE", "THIRD_PARTY_NOTICES.md"]
14
14
  authors = [{ name = "Joseph Sandoval" }, { name = "Alfredo Sandoval Jr." }]
15
15
  dependencies = [
16
- "medmlx-core==0.2.4",
16
+ "medmlx-core==0.2.5",
17
17
  "nibabel>=5.3,<6",
18
18
  "numpy>=2.0,<3",
19
19
  ]
20
20
 
21
21
  [project.entry-points."medmlx.models"]
22
- segresnet = "segresnet_mlx.medmlx:MODEL"
22
+ segresnet-brats = "segresnet_mlx.medmlx:BRATS_MODEL"
23
+ segresnet-renal-cect = "segresnet_mlx.medmlx:RENAL_CECT_MODEL"
23
24
 
24
25
  [project.urls]
25
- Homepage = "https://github.com/MedMLX/segresnet-mlx"
26
- Documentation = "https://github.com/MedMLX/segresnet-mlx/blob/main/docs/technical.md"
27
- Repository = "https://github.com/MedMLX/segresnet-mlx"
28
- Issues = "https://github.com/MedMLX/segresnet-mlx/issues"
26
+ Homepage = "https://huggingface.co/MedMLX/segresnet-renal-cect-mlx"
29
27
 
30
28
  [project.optional-dependencies]
31
- conversion = ["torch>=2.6,<3"]
32
29
  dev = [
33
30
  "pyright==1.1.414",
34
31
  "pytorch-ignite>=0.4.11,<0.6",
@@ -64,7 +61,7 @@ select = ["E", "F", "I", "UP", "B", "SIM", "RUF"]
64
61
 
65
62
  [tool.pyright]
66
63
  include = ["scripts","src","tests"]
67
- extraPaths = ["src"]
64
+ extraPaths = ["src", "scripts", "tests"]
68
65
  pythonVersion = "3.13"
69
66
  typeCheckingMode = "strict"
70
67
  enableTypeIgnoreComments = false
@@ -77,9 +74,6 @@ ban-relative-imports = "all"
77
74
  "typing.Any" = {msg = "Use a precise type and parse untyped values at the boundary."}
78
75
  "typing_extensions.Any" = {msg = "Use a precise type and parse untyped values at the boundary."}
79
76
 
80
- [tool.uv.sources]
81
- medmlx-core = { git = "https://github.com/MedMLX/medmlx.git", rev = "f98e0924f50228e2b9509474947b9672054b3217", subdirectory = "packages/core" }
82
-
83
77
  [tool.hatch.build.targets.sdist]
84
78
  only-include = [
85
79
  "src/segresnet_mlx",
@@ -16,12 +16,10 @@ if TYPE_CHECKING:
16
16
  from segresnet_mlx.predictor import predict_logits as predict_logits
17
17
  from segresnet_mlx.telemetry import WindowTelemetry as WindowTelemetry
18
18
  from segresnet_mlx.topology import SegResNetSpec as SegResNetSpec
19
- from segresnet_mlx.weights import convert_checkpoint as convert_checkpoint
20
- from segresnet_mlx.weights import convert_state_dict as convert_state_dict
21
19
  from segresnet_mlx.weights import load_model as load_model
22
20
  from segresnet_mlx.weights import required_shapes as required_shapes
23
21
 
24
- __version__ = "0.2.3"
22
+ __version__ = "0.3.0"
25
23
 
26
24
  _EXPORTS = {
27
25
  "BUNDLE_CONFIGS": "segresnet_mlx.predictor",
@@ -30,8 +28,6 @@ _EXPORTS = {
30
28
  "SegResNetSpec": "segresnet_mlx.topology",
31
29
  "SlidingWindowConfig": "segresnet_mlx.predictor",
32
30
  "WindowTelemetry": "segresnet_mlx.telemetry",
33
- "convert_checkpoint": "segresnet_mlx.weights",
34
- "convert_state_dict": "segresnet_mlx.weights",
35
31
  "forward_patch": "segresnet_mlx.inference",
36
32
  "load_model": "segresnet_mlx.weights",
37
33
  "predict_logits": "segresnet_mlx.predictor",
@@ -45,8 +41,6 @@ __all__ = [
45
41
  "SegResNetSpec",
46
42
  "SlidingWindowConfig",
47
43
  "WindowTelemetry",
48
- "convert_checkpoint",
49
- "convert_state_dict",
50
44
  "forward_patch",
51
45
  "load_model",
52
46
  "predict_logits",
@@ -55,6 +49,7 @@ __all__ = [
55
49
 
56
50
 
57
51
  def __getattr__(name: str) -> object:
52
+ """Load a public export on first access and cache it in the package."""
58
53
  module = _EXPORTS.get(name)
59
54
  if module is None:
60
55
  raise AttributeError(f"module 'segresnet_mlx' has no attribute {name!r}")
@@ -64,4 +59,5 @@ def __getattr__(name: str) -> object:
64
59
 
65
60
 
66
61
  def __dir__() -> list[str]:
62
+ """Include lazy exports in interactive package discovery."""
67
63
  return sorted({*globals(), *_EXPORTS})
@@ -0,0 +1,109 @@
1
+ """Typed protocols for the NiBabel NIfTI, header and orientation calls used here."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from collections.abc import Sequence
6
+ from os import PathLike
7
+ from typing import Protocol, cast
8
+
9
+ import nibabel
10
+ import numpy as np
11
+ from nibabel.filebasedimages import ImageFileError as ImageFileError
12
+
13
+ from segresnet_mlx._arrays import Float64Array, HostArray
14
+
15
+ type PathArg = str | PathLike[str]
16
+
17
+
18
+ class ArrayBuffer(Protocol):
19
+ def __array__(
20
+ self, dtype: np.dtype[np.generic] | None = None, copy: bool | None = None
21
+ ) -> HostArray:
22
+ """Expose image data through NumPy's array conversion protocol."""
23
+ ...
24
+
25
+
26
+ class NiftiHeader(Protocol):
27
+ def get_zooms(self) -> tuple[float, ...]:
28
+ """Return the voxel spacing stored in the header."""
29
+ ...
30
+
31
+ def set_zooms(self, zooms: Sequence[float] | Float64Array) -> None:
32
+ """Set voxel spacing in the header."""
33
+ ...
34
+
35
+ def __getitem__(self, key: str) -> HostArray:
36
+ """Return a named NIfTI header field."""
37
+ ...
38
+
39
+
40
+ class NiftiImage(Protocol):
41
+ shape: tuple[int, ...]
42
+ affine: Float64Array
43
+ dataobj: ArrayBuffer | HostArray
44
+ header: NiftiHeader
45
+
46
+ def __init__(self, dataobj: HostArray, affine: HostArray) -> None:
47
+ """Create an image from voxel data and its affine."""
48
+ ...
49
+
50
+ def get_data_dtype(self) -> np.dtype[np.generic]:
51
+ """Return the stored voxel-data dtype."""
52
+ ...
53
+
54
+ def get_sform(self) -> Float64Array:
55
+ """Return the header's sform affine."""
56
+ ...
57
+
58
+ def get_qform(self) -> Float64Array:
59
+ """Return the header's qform affine."""
60
+ ...
61
+
62
+ def set_sform(self, affine: Float64Array, code: int | None = None) -> None:
63
+ """Set the header's sform affine and optional code."""
64
+ ...
65
+
66
+ def set_qform(self, affine: Float64Array, code: int | None = None) -> None:
67
+ """Set the header's qform affine and optional code."""
68
+ ...
69
+
70
+
71
+ class Orientations(Protocol):
72
+ def axcodes2ornt(self, axcodes: str, labels: tuple[tuple[str, str], ...]) -> Float64Array:
73
+ """Convert axis codes to an orientation array."""
74
+ ...
75
+
76
+ def ornt_transform(self, start_ornt: Float64Array, end_ornt: Float64Array) -> Float64Array:
77
+ """Return the transform between two orientation arrays."""
78
+ ...
79
+
80
+ def inv_ornt_aff(self, ornt: Float64Array, shape: Sequence[int]) -> Float64Array:
81
+ """Return the inverse voxel-coordinate affine for an orientation change."""
82
+ ...
83
+
84
+ def aff2axcodes(
85
+ self, aff: Float64Array, labels: tuple[tuple[str, str], ...]
86
+ ) -> tuple[str | None, ...]:
87
+ """Return axis codes implied by an affine."""
88
+ ...
89
+
90
+
91
+ class Nibabel(Protocol):
92
+ Nifti1Image: type[NiftiImage]
93
+ orientations: Orientations
94
+
95
+ def load(self, filename: PathArg) -> NiftiImage:
96
+ """Load an image from a NIfTI path."""
97
+ ...
98
+
99
+ def save(self, img: NiftiImage, filename: PathArg) -> None:
100
+ """Write an image to a NIfTI path."""
101
+ ...
102
+
103
+ def io_orientation(self, affine: Float64Array) -> Float64Array:
104
+ """Return the input-to-output orientation implied by an affine."""
105
+ ...
106
+
107
+
108
+ # Bound to the real module; this runner admits only NIfTI images.
109
+ nib = cast(Nibabel, nibabel)
@@ -57,6 +57,7 @@ class Details(TypedDict):
57
57
  ``requires_human_review`` is always true: outputs are research results.
58
58
  """
59
59
 
60
+ weights_revision: NotRequired[str | None]
60
61
  bundle: str
61
62
  arithmetic_profile: ArithmeticProfile
62
63
  labels: dict[str, str]
@@ -19,7 +19,9 @@ if TYPE_CHECKING:
19
19
  class ArrayConstructor(Protocol):
20
20
  def __call__(
21
21
  self, value: HostArray | Array | ArrayValue, dtype: mx.Dtype | None = None
22
- ) -> Array: ...
22
+ ) -> Array:
23
+ """Construct an MLX array with an optional dtype."""
24
+ ...
23
25
 
24
26
  class MLX(MlxRuntime, Protocol):
25
27
  all = staticmethod(mx.all)
@@ -32,7 +34,6 @@ if TYPE_CHECKING:
32
34
  load = staticmethod(mx.load)
33
35
 
34
36
  def eval(self, *args: Array | Mapping[str, Array] | Sequence[Array]) -> None: ...
35
-
36
37
  def save_safetensors(self, file: str, arrays: dict[str, Array]) -> None: ...
37
38
 
38
39
  class SlidingWindowInference(Protocol):
@@ -50,10 +51,14 @@ if TYPE_CHECKING:
50
51
  with_coord: bool = False,
51
52
  predictor_input: Literal["numpy"],
52
53
  preserve_window_strides: bool = False,
53
- ) -> Float32Array: ...
54
+ ) -> Float32Array:
55
+ """Blend predictor outputs from overlapping windows into host logits."""
56
+ ...
54
57
 
55
58
  class CoreSlidingWindow(Protocol):
56
59
  sliding_window_inference: SlidingWindowInference
57
60
 
58
61
  class Metal(Protocol):
59
- def is_available(self) -> bool: ...
62
+ def is_available(self) -> bool:
63
+ """Report whether the Metal backend is available."""
64
+ ...
@@ -22,6 +22,7 @@ from medmlx_core.errors import (
22
22
  InvalidInputError,
23
23
  ModelExecutionError,
24
24
  )
25
+ from medmlx_core.weights import weights_dir, weights_revision
25
26
 
26
27
  from segresnet_mlx import transforms
27
28
  from segresnet_mlx.topology import ArithmeticProfile, SegResNetSpec
@@ -39,6 +40,21 @@ MODEL = {
39
40
  ),
40
41
  }
41
42
 
43
+ BRATS_MODEL = {
44
+ "id": "segresnet-brats",
45
+ "label": "SegResNet BraTS MRI",
46
+ "task": "segmentation",
47
+ "run": "segresnet_mlx.medmlx:brats",
48
+ "summary": "BraTS brain tumor segmentation from co-registered MRI sequences.",
49
+ }
50
+ RENAL_CECT_MODEL = {
51
+ "id": "segresnet-renal-cect",
52
+ "label": "SegResNet renal CECT",
53
+ "task": "segmentation",
54
+ "run": "segresnet_mlx.medmlx:renal_cect",
55
+ "summary": "Renal structures from arterial, venous and excretory phase CT.",
56
+ }
57
+
42
58
  # Label values of the 104 TotalSegmentator structures plus background, as listed in the
43
59
  # whole-body bundle's metadata.json ``channel_def``.
44
60
  _WHOLE_BODY_LABELS: Mapping[int, str] = {
@@ -295,10 +311,8 @@ def run(
295
311
  at 1.5 mm or 3 mm respectively). ``weights_path`` is a native ``.safetensors``
296
312
  file: ``segresnet-brats.safetensors`` from ``MedMLX/segresnet-brats-mlx`` for
297
313
  BraTS, or ``segresnet-renal-cect.safetensors`` from
298
- ``MedMLX/segresnet-renal-cect-mlx`` for renal CECT. These published files need
299
- no conversion. For other checkpoints, convert the original ``models/model.pt``
300
- (``models/model_lowres.pt`` for lowres whole-body) with
301
- :func:`segresnet_mlx.convert_checkpoint` and pass the resulting file.
314
+ ``MedMLX/segresnet-renal-cect-mlx`` for renal CECT, or the corresponding
315
+ high- or low-resolution file from ``MedMLX/segresnet-wholebody-mlx``.
302
316
 
303
317
  Inputs must be finite 3D NIfTI volumes; multi-input bundles need identical shape
304
318
  and affine (atol 1e-5). Writes ``<bundle>_mask.nii.gz`` in ``output_dir`` on the
@@ -423,6 +437,7 @@ def run(
423
437
  return {
424
438
  "outputs": {"mask": mask_path.as_posix()},
425
439
  "details": {
440
+ "weights_revision": weights_revision(weights.parent),
426
441
  "bundle": bundle,
427
442
  "arithmetic_profile": arithmetic_profile,
428
443
  "labels": {str(value): name for value, name in config.labels.items()},
@@ -450,6 +465,66 @@ def run(
450
465
  }
451
466
 
452
467
 
468
+ def brats(
469
+ *,
470
+ t1c_path: str,
471
+ t1_path: str,
472
+ t2_path: str,
473
+ flair_path: str,
474
+ output_dir: str,
475
+ weights_path: str | None = None,
476
+ overwrite: bool = False,
477
+ arithmetic_profile: ArithmeticProfile = "torch-order-fp32",
478
+ ) -> RunResult:
479
+ """Segment BraTS MRI; omit weights_path to download its cached model folder.
480
+
481
+ The four input volumes must share a voxel grid. An explicit folder contains
482
+ ``segresnet-brats.safetensors``; an explicit safetensors file is also accepted.
483
+ Automatic downloads use ``MedMLX/segresnet-brats-mlx``. All arithmetic and
484
+ preprocessing are delegated unchanged to ``run``.
485
+ """
486
+ folder = (
487
+ weights_dir("segresnet-brats-mlx")
488
+ if weights_path is None else Path(weights_path).expanduser()
489
+ )
490
+ checkpoint = folder / "segresnet-brats.safetensors" if folder.is_dir() else folder
491
+ return run(
492
+ bundle="brats_mri_segmentation", weights_path=str(checkpoint), output_dir=output_dir,
493
+ t1c_path=t1c_path, t1_path=t1_path, t2_path=t2_path, flair_path=flair_path,
494
+ overwrite=overwrite, arithmetic_profile=arithmetic_profile,
495
+ )
496
+
497
+
498
+ def renal_cect(
499
+ *,
500
+ arterial_path: str,
501
+ venous_path: str,
502
+ excretory_path: str,
503
+ output_dir: str,
504
+ weights_path: str | None = None,
505
+ overwrite: bool = False,
506
+ arithmetic_profile: ArithmeticProfile = "torch-order-fp32",
507
+ ) -> RunResult:
508
+ """Segment renal CECT; omit weights_path to download its cached model folder.
509
+
510
+ The arterial, venous and excretory CT volumes must share a voxel grid. An
511
+ explicit folder contains ``segresnet-renal-cect.safetensors`` as published for
512
+ this bundle; an explicit safetensors file is also accepted. Automatic downloads
513
+ use ``MedMLX/segresnet-renal-cect-mlx``. Arithmetic and preprocessing are
514
+ delegated unchanged to ``run``.
515
+ """
516
+ folder = (
517
+ weights_dir("segresnet-renal-cect-mlx")
518
+ if weights_path is None else Path(weights_path).expanduser()
519
+ )
520
+ checkpoint = folder / "segresnet-renal-cect.safetensors" if folder.is_dir() else folder
521
+ return run(
522
+ bundle="renalStructures_CECT_segmentation", weights_path=str(checkpoint),
523
+ output_dir=output_dir, arterial_path=arterial_path, venous_path=venous_path,
524
+ excretory_path=excretory_path, overwrite=overwrite, arithmetic_profile=arithmetic_profile,
525
+ )
526
+
527
+
453
528
  def preprocess_renal(volumes: list[Volume]) -> tuple[Float32Array, SpatialInverse]:
454
529
  """The renal bundle's ``preprocessing`` from ``EnsureChannelFirstd`` onward.
455
530
 
@@ -562,14 +637,11 @@ def postprocess_wholebody(logits: Float32Array, inverse: SpatialInverse) -> Floa
562
637
 
563
638
 
564
639
  def _weights_file(weights_path: str) -> Path:
640
+ """Validate a safetensors file path while retaining any Hub-cache symlink."""
565
641
  # MLX identifies the file format from the supplied name. Hub cache paths are
566
642
  # named .safetensors symlinks whose underlying content-addressed blobs have no suffix.
567
643
  path = Path(weights_path).expanduser().absolute()
568
- hint = (
569
- "Convert the bundle's models/model.pt once with "
570
- "segresnet_mlx.convert_checkpoint(model_pt, destination, spec=...) "
571
- "(requires the conversion extra)."
572
- )
644
+ hint = "Download the matching safetensors file from the MedMLX weight repositories."
573
645
  if not path.is_file():
574
646
  raise AssetNotReadyError(
575
647
  f"SegResNet weights not found: {path}", reason=f"missing {path}", hint=hint
@@ -646,7 +718,7 @@ def _load_network(
646
718
  raise AssetNotReadyError(
647
719
  f"SegResNet weights {weights} do not match {spec}: {error}",
648
720
  reason=f"mismatched {weights.name}",
649
- hint="Convert the matching bundle's models/model.pt with convert_checkpoint.",
721
+ hint="Download the matching safetensors file from the MedMLX weight repositories.",
650
722
  ) from error
651
723
 
652
724
 
@@ -702,4 +774,4 @@ def write_atomic(path: Path, mask: MaskArray, affine: Float64Array, dtype: str)
702
774
  Path(name).unlink(missing_ok=True)
703
775
 
704
776
 
705
- __all__ = ["BUNDLES", "MODEL", "run"]
777
+ __all__ = ["BRATS_MODEL", "BUNDLES", "MODEL", "RENAL_CECT_MODEL", "brats", "renal_cect", "run"]
@@ -22,9 +22,8 @@ def group_mean(values: Array, mx: MLX) -> Array:
22
22
 
23
23
  Tiling bounds the float32 accumulation depth of a group that spans a whole
24
24
  volume. The last dimension must be nonempty. Returns ``(N, G, 1)``. ``SegResNet``
25
- does not call this (its GroupNorm
26
- comes from the ``medmlx_core`` operators); it is exercised by
27
- ``tests/test_equivalence.py``.
25
+ does not call this (its GroupNorm comes from the ``medmlx_core`` operators); it is
26
+ exercised by ``tests/test_equivalence.py``.
28
27
  """
29
28
  size = int(values.shape[-1])
30
29
  padded = mx.pad(values, [(0, 0), (0, 0), (0, (-size) % 256)])
@@ -59,6 +58,7 @@ class SegResNet:
59
58
  mx: MLX,
60
59
  arithmetic_profile: ArithmeticProfile = "torch-order-fp32",
61
60
  ) -> None:
61
+ """Bind the topology and weights and select their arithmetic operators."""
62
62
  self.spec = spec
63
63
  self._weights = weights
64
64
  self._mx = mx
@@ -105,6 +105,7 @@ class SegResNet:
105
105
  return self._forward(operators.pack(values))
106
106
 
107
107
  def _forward(self, values: Array) -> Array:
108
+ """Run the encoder and decoder, retaining skip connections between stages."""
108
109
  require_ncdhw(values, name="SegResNet input")
109
110
  if int(values.shape[1]) != self.spec.in_channels:
110
111
  raise ValueError(
@@ -163,6 +164,8 @@ class SegResNet:
163
164
  return self._add(hidden, values)
164
165
 
165
166
  def _relu(self, values: Array) -> Array:
167
+ """Apply ReLU using paired arithmetic when the fp48 profile is active."""
168
+ # fp48 activations are paired float32 words, so ReLU must act on the pair.
166
169
  return (
167
170
  cast("CompensatedFloat48Operators", self._operators).relu(values)
168
171
  if self.arithmetic_profile == "compensated-fp48"
@@ -170,6 +173,8 @@ class SegResNet:
170
173
  )
171
174
 
172
175
  def _add(self, values: Array, other: Array) -> Array:
176
+ """Add residual or skip values using the selected arithmetic profile."""
177
+ # Residual/skip add; for fp48 it is a paired-precision add, not a plain float32 sum.
173
178
  return (
174
179
  cast("CompensatedFloat48Operators", self._operators).add(values, other)
175
180
  if self.arithmetic_profile == "compensated-fp48"
@@ -177,6 +182,7 @@ class SegResNet:
177
182
  )
178
183
 
179
184
  def _norm(self, values: Array, prefix: str) -> Array:
185
+ """Apply group normalization using the weights at the given prefix."""
180
186
  return self._operators.norm(
181
187
  values, self._weights[f"{prefix}.weight"], self._weights[f"{prefix}.bias"]
182
188
  )
@@ -92,6 +92,7 @@ class BundlePredictor:
92
92
  telemetry: WindowTelemetry | None = None,
93
93
  mx: MLX | None = None,
94
94
  ) -> None:
95
+ """Initialize the network runtime and optional per-window telemetry."""
95
96
  self.network = network
96
97
  self.telemetry = telemetry
97
98
  try:
@@ -124,6 +125,7 @@ class BundlePredictor:
124
125
  self.telemetry.attach(engine, events=events)
125
126
 
126
127
  def __call__(self, values: HostArray | Array) -> Array:
128
+ """Run one patch through the same validated path as ``forward``."""
127
129
  return self.forward(values)
128
130
 
129
131
  def forward(self, values: HostArray | Array) -> Array:
@@ -140,6 +142,7 @@ class BundlePredictor:
140
142
  raise
141
143
 
142
144
  def _forward_window(self, values: HostArray | Array) -> Float32Array:
145
+ """Run and synchronize one window, returning finite host logits."""
143
146
  host = self._host_input(values)
144
147
  self.mx.synchronize()
145
148
  started = time.perf_counter()
@@ -159,6 +162,7 @@ class BundlePredictor:
159
162
  return output
160
163
 
161
164
  def _host_input(self, values: HostArray | Array) -> Float32Array:
165
+ """Admit finite float32 input as a NumPy array or raise an input error."""
162
166
  host = np.asarray(values)
163
167
  if host.dtype != np.float32 or not np.isfinite(host).all():
164
168
  raise InvalidInputError(f"{self.network_label} input must be finite float32")
@@ -242,6 +246,7 @@ def predict_wide_logits(
242
246
  def window(
243
247
  patch: np.ndarray[tuple[int, ...], np.dtype[np.float64]],
244
248
  ) -> np.ndarray[tuple[int, ...], np.dtype[np.float64]]:
249
+ """Combine paired Metal output words into float64 host logits."""
245
250
  result = network.paired_forward(cast(ArrayFactory, mx.array)(patch.astype(np.float32)))
246
251
  mx.eval(result)
247
252
  mx.synchronize()
@@ -20,7 +20,9 @@ import numpy as np
20
20
 
21
21
 
22
22
  class EventEngine(Protocol):
23
- def add_event_handler(self, event_name: Events, handler: Callable[..., None]) -> object: ...
23
+ def add_event_handler(self, event_name: Events, handler: Callable[..., None]) -> object:
24
+ """Register a handler for an Ignite lifecycle event."""
25
+ ...
24
26
 
25
27
 
26
28
  SCHEMA_VERSION = "medmlx.segresnet.mlx.v1"
@@ -37,6 +39,7 @@ class WindowTelemetry:
37
39
  """
38
40
 
39
41
  def __init__(self, path: Path, record: dict[str, Json]) -> None:
42
+ """Create the output directory and initialize the caller-owned progress record."""
40
43
  self.path = path
41
44
  self.record = record
42
45
  self._metadata = {"schema_version": SCHEMA_VERSION, **record}
@@ -141,20 +144,25 @@ class WindowTelemetry:
141
144
  attached.add_event_handler(event, self._interrupted)
142
145
 
143
146
  def _started(self, engine: object) -> None:
147
+ """Start a fresh record when the attached engine starts."""
144
148
  self.begin()
145
149
 
146
150
  def _completed(self, engine: object) -> None:
151
+ """Finalize a completed engine run."""
147
152
  self.finish()
148
153
 
149
154
  def _failed(self, engine: object, error: BaseException) -> None:
155
+ """Record and re-raise an engine failure."""
150
156
  self.fail(error)
151
157
  # Registering EXCEPTION_RAISED otherwise makes Ignite swallow the error.
152
158
  raise error
153
159
 
154
160
  def _interrupted(self, engine: object, **kwargs: object) -> None:
161
+ """Finalize an interrupted engine run."""
155
162
  self.finish("interrupted")
156
163
 
157
164
  def _write(self) -> None:
165
+ """Replace the progress file atomically with the current JSON record."""
158
166
  # Atomic visibility for progress readers, not a crash-durability promise.
159
167
  descriptor, name = tempfile.mkstemp(prefix=f".{self.path.name}.", dir=self.path.parent)
160
168
  try:
@@ -33,6 +33,7 @@ class SegResNetSpec:
33
33
  upsample_mode: str = "nontrainable"
34
34
 
35
35
  def __post_init__(self) -> None:
36
+ """Reject unsupported channel, stage, normalization, and upsampling settings."""
36
37
  if min(self.in_channels, self.out_channels, self.init_filters) < 1:
37
38
  raise ValueError("SegResNet channel counts must be positive")
38
39
  if len(self.blocks_down) != len(self.blocks_up) + 1:
@@ -18,12 +18,35 @@ from __future__ import annotations
18
18
 
19
19
  from collections.abc import Sequence
20
20
  from dataclasses import dataclass
21
- from fractions import Fraction
22
21
  from pathlib import Path
23
22
  from typing import cast
24
23
 
25
24
  import numpy as np
26
25
  import numpy.typing as npt
26
+ from medmlx_core.geometry import (
27
+ _bmm as _bmm,
28
+ )
29
+ from medmlx_core.geometry import (
30
+ _fma as _fma,
31
+ )
32
+ from medmlx_core.geometry import (
33
+ _linspace_from_neg_one as _linspace_from_neg_one,
34
+ )
35
+ from medmlx_core.geometry import (
36
+ _mm as _mm,
37
+ )
38
+ from medmlx_core.geometry import (
39
+ affine_to_spacing as affine_to_spacing,
40
+ )
41
+ from medmlx_core.geometry import (
42
+ compute_shape_offset as compute_shape_offset,
43
+ )
44
+ from medmlx_core.geometry import (
45
+ to_affine_nd as to_affine_nd,
46
+ )
47
+ from medmlx_core.geometry import (
48
+ zoom_affine as zoom_affine,
49
+ )
27
50
 
28
51
  from segresnet_mlx._arrays import HostArray, Int64Array
29
52
  from segresnet_mlx._nifti import NiftiImage, nib
@@ -83,76 +106,12 @@ def correct_nifti_header(image: NiftiImage) -> None:
83
106
  image.header.set_zooms(affine_to_spacing(image.affine, 3))
84
107
 
85
108
 
86
- def to_affine_nd(rank: int, affine: npt.ArrayLike) -> Float64Array:
87
- """``monai.data.utils.to_affine_nd`` for an integer rank."""
88
- source = np.array(affine, dtype=np.float64, copy=True)
89
- result = np.eye(rank + 1, dtype=np.float64)
90
- size = max(min(len(result) - 1, len(source) - 1), 1)
91
- result[:size, :size] = source[:size, :size]
92
- if size > 1:
93
- result[:size, -1] = source[:size, -1]
94
- return result
95
109
 
96
110
 
97
- def affine_to_spacing(affine: Float64Array, rank: int) -> Float64Array:
98
- """``monai.data.utils.affine_to_spacing`` with zeros suppressed to one."""
99
- part = np.asarray(affine[:rank, :rank], dtype=np.float64)
100
- spacing = np.sqrt(np.sum(part * part, axis=0))
101
- spacing[spacing == 0] = 1.0
102
- return spacing
103
-
104
-
105
- def zoom_affine(affine: Float64Array, scale: Float64Array) -> Float64Array:
106
- """``monai.data.utils.zoom_affine`` with ``diagonal=False``."""
107
- affine = np.array(affine, dtype=float, copy=True)
108
- scale_np = np.array(scale, dtype=float, copy=True)
109
- size = len(affine) - 1
110
- norm = affine_to_spacing(affine, size)
111
- if len(scale_np) < size:
112
- scale_np = np.append(scale_np, norm[len(scale_np) :])
113
- scale_np = scale_np[:size]
114
- scale_np = np.where(scale_np > 0, scale_np, norm)
115
- scale_np[scale_np == 0] = 1.0
116
- rzs = affine[:-1, :-1]
117
- zs = np.linalg.cholesky(rzs.T @ rzs).T
118
- rotation = rzs @ np.linalg.inv(zs)
119
- signed = np.sign(np.diag(zs)) * np.abs(scale_np)
120
- result = np.eye(len(affine))
121
- result[:-1, :-1] = rotation @ np.diag(signed)
122
- return result
123
111
 
124
112
 
125
- def compute_shape_offset(
126
- spatial_shape: Sequence[int],
127
- in_affine: Float64Array,
128
- out_affine: Float64Array,
129
- ) -> tuple[npt.NDArray[np.int_], Float64Array]:
130
- """``monai.data.utils.compute_shape_offset`` with ``scale_extent=False``."""
131
- shape = np.array(tuple(spatial_shape), copy=True, dtype=float)
132
- rank = len(shape)
133
- in_affine_ = to_affine_nd(rank, in_affine)
134
- out_affine_ = to_affine_nd(rank, out_affine)
135
- in_coords = [(0.0, dim - 1.0) for dim in shape]
136
- corners: Float64Array = np.asarray(np.meshgrid(*in_coords, indexing="ij")).reshape((rank, -1))
137
- corners = np.concatenate((corners, np.ones_like(corners[:1])))
138
- corners_out = np.linalg.solve(out_affine_, in_affine_) @ corners
139
- corners = in_affine_ @ corners
140
- all_dist = corners_out[:-1].copy()
141
- corners_out = corners_out[:-1] / corners_out[-1]
142
- out_shape = np.round(cast(Float64Array, np.ptp(corners_out, axis=1)) + 1.0)
143
- offset = None
144
- for index in range(corners.shape[1]):
145
- min_corner = np.min(all_dist - all_dist[:, index : index + 1], 1)
146
- if np.allclose(min_corner, 0.0, rtol=AFFINE_TOL):
147
- offset = corners[:-1, index]
148
- break
149
- if offset is None:
150
- offset = (
151
- in_affine_[:-1, :-1] @ (shape / 2.0)
152
- + in_affine_[:-1, -1]
153
- - out_affine_[:-1, :-1] @ (out_shape / 2.0)
154
- )
155
- return out_shape.astype(int, copy=False), offset
113
+
114
+
156
115
 
157
116
 
158
117
  def orient(
@@ -387,35 +346,8 @@ def write_nifti(path: Path, data: HostArray, affine: Float64Array, dtype: npt.DT
387
346
  nib.save(image, str(path))
388
347
 
389
348
 
390
- def _mm(left: Float64Array, right: Float64Array) -> Float64Array:
391
- """A MetaTensor affine update, ``torch.mm`` of two 4x4 float64 matrices.
392
-
393
- The reference Torch build's MKL computes each entry as
394
- ``fma(a2, b2, a0 * b0) + fma(a3, b3, a1 * b1)``.
395
- """
396
- if left.shape != (4, 4) or right.shape != (4, 4):
397
- return left @ right
398
- result = np.empty((4, 4), dtype=np.float64)
399
- for row in range(4):
400
- a = [float(value) for value in left[row]]
401
- for column in range(4):
402
- b = [float(value) for value in right[:, column]]
403
- result[row, column] = _fma(a[2], b[2], a[0] * b[0]) + _fma(a[3], b[3], a[1] * b[1])
404
- return result
405
349
 
406
350
 
407
- def _bmm(left: Float64Array, right: Float64Array) -> Float64Array:
408
- """Torch's small-matrix ``bmm`` kernel: unfused sequential float64 accumulation."""
409
- rows, inner = left.shape
410
- columns = right.shape[1]
411
- result = np.empty((rows, columns), dtype=np.float64)
412
- for row in range(rows):
413
- for column in range(columns):
414
- total = 0.0
415
- for index in range(inner):
416
- total += float(left[row, index]) * float(right[index, column])
417
- result[row, column] = total
418
- return result
419
351
 
420
352
 
421
353
  def _normalize_transform(shape: Sequence[int]) -> Float64Array:
@@ -429,29 +361,8 @@ def _normalize_transform(shape: Sequence[int]) -> Float64Array:
429
361
  return result
430
362
 
431
363
 
432
- def _fma(left: float, right: float, addend: float) -> float:
433
- """Correctly rounded ``left * right + addend``."""
434
- return float(Fraction(left) * Fraction(right) + Fraction(addend))
435
-
436
364
 
437
- def _linspace_from_neg_one(steps: int) -> Float64Array:
438
- """ATen ``linspace(-1, 1, steps) * (steps - 1) / steps`` in float64.
439
365
 
440
- ATen evaluates ``start + step * i`` below the midpoint and
441
- ``end - step * (steps - i - 1)`` above it, each as one fused multiply-add.
442
- """
443
- if steps <= 1:
444
- return np.zeros(1, dtype=np.float64)
445
- step = (1.0 - (-1.0)) / (steps - 1)
446
- halfway = steps // 2
447
- values = np.array(
448
- [
449
- _fma(step, index, -1.0) if index < halfway else _fma(-step, steps - index - 1, 1.0)
450
- for index in range(steps)
451
- ],
452
- dtype=np.float64,
453
- )
454
- return values * (steps - 1) / steps
455
366
 
456
367
 
457
368
  def _theta(
@@ -461,7 +372,8 @@ def _theta(
461
372
  ) -> Float64Array:
462
373
  """``AffineTransform`` normalization and reverse indexing of a voxel transform."""
463
374
  theta = _bmm(
464
- _bmm(_normalize_transform(in_size), xform), np.linalg.inv(_normalize_transform(out_size))
375
+ _bmm(_normalize_transform(in_size), xform),
376
+ np.linalg.inv(_normalize_transform(out_size)),
465
377
  )
466
378
  rank = len(in_size)
467
379
  reverse = list(range(rank - 1, -1, -1))
@@ -545,6 +457,7 @@ def _product_error(
545
457
  """Dekker's exact remainder ``left * right - product``."""
546
458
 
547
459
  def split[T: (Float64Array, np.float64)](value: T) -> tuple[T, T]:
460
+ """Split a float64 value into high and low words for Dekker multiplication."""
548
461
  scaled = 134217729.0 * value
549
462
  high = scaled - (scaled - value)
550
463
  return cast(T, high), cast(T, value - high)
@@ -1,16 +1,11 @@
1
- """Strict Torch checkpoint conversion for MONAI SegResNet MLX inference."""
1
+ """Strict safetensors loading for MONAI SegResNet MLX inference."""
2
2
 
3
3
  from __future__ import annotations
4
4
 
5
- from collections.abc import Mapping
6
5
  from pathlib import Path
7
6
  from typing import TYPE_CHECKING, cast
8
7
 
9
- import numpy as np
10
- import numpy.typing as npt
11
- from medmlx_core.checkpoints import load_torch_checkpoint, tensor_mapping_from_payload
12
8
  from medmlx_core.errors import InvalidInputError
13
- from medmlx_core.typing import ArrayFactory
14
9
 
15
10
  from segresnet_mlx._runtime import import_mlx
16
11
 
@@ -20,8 +15,6 @@ if TYPE_CHECKING:
20
15
  from segresnet_mlx.network import SegResNet
21
16
  from segresnet_mlx.topology import ArithmeticProfile, SegResNetSpec
22
17
 
23
- type HostArray = npt.NDArray[np.float32]
24
-
25
18
 
26
19
  def required_shapes(spec: SegResNetSpec) -> dict[str, tuple[int, ...]]:
27
20
  """Return the exact ``{parameter name: shape}`` set a checkpoint must provide.
@@ -71,84 +64,6 @@ def required_shapes(spec: SegResNetSpec) -> dict[str, tuple[int, ...]]:
71
64
  return shapes
72
65
 
73
66
 
74
- def convert_state_dict(
75
- state: Mapping[str, object] | Mapping[object, object], spec: SegResNetSpec
76
- ) -> dict[str, HostArray]:
77
- """Check a Torch state dict against ``spec`` and return float32 NumPy parameters.
78
-
79
- ``state`` is either a raw state dict or a checkpoint holding it under the
80
- ``"model"`` key. Arrays keep their Torch layout and are returned C-contiguous.
81
- Validation checks the parameter inventory and values, not the checkpoint's
82
- identity or origin; other checkpoints with the same topology are admissible.
83
-
84
- Raises:
85
- InvalidInputError: If a required parameter is missing, an unsupported one is
86
- present, a shape differs from ``required_shapes(spec)``, or a value is
87
- not finite.
88
- """
89
- # The bundle's CheckpointLoader selects model and ignores other saved objects.
90
- if "model" in state:
91
- state = cast("Mapping[str, object] | Mapping[object, object]", state["model"])
92
- source = tensor_mapping_from_payload(state, what="SegResNet")
93
- required = required_shapes(spec)
94
- missing = [key for key in required if key not in source]
95
- extra = [key for key in source if key not in required]
96
- if missing:
97
- raise InvalidInputError(
98
- "SegResNet checkpoint is missing parameter(s): " + ", ".join(missing)
99
- )
100
- if extra:
101
- raise InvalidInputError(
102
- "SegResNet checkpoint has unsupported parameter(s): " + ", ".join(extra)
103
- )
104
- converted: dict[str, HostArray] = {}
105
- for key, expected_shape in required.items():
106
- array = np.asarray(source[key], dtype=np.float32)
107
- if array.shape != expected_shape:
108
- raise InvalidInputError(
109
- f"SegResNet parameter {key!r} has shape {array.shape}; expected {expected_shape}"
110
- )
111
- if not np.isfinite(array).all():
112
- raise InvalidInputError(f"SegResNet parameter {key!r} contains non-finite values")
113
- converted[key] = np.ascontiguousarray(array)
114
- return converted
115
-
116
-
117
- def convert_checkpoint(checkpoint_path: Path, destination: Path, *, spec: SegResNetSpec) -> Path:
118
- """Convert a Torch checkpoint to float32 safetensors and return ``destination``.
119
-
120
- The checkpoint is read with the shared loader (``torch.load`` with its default
121
- ``weights_only=True``, so Torch from the ``conversion`` extra is required) and
122
- checked by ``convert_state_dict``. ``destination`` must end in
123
- ``.safetensors``. Parent directories of ``destination`` are
124
- created. An existing destination is replaced, except when it aliases the
125
- source checkpoint. Alias checks occur before conversion and do not lock out
126
- subsequent filesystem changes. MLX is also needed to write the file, so this
127
- runs only on a supported Apple Silicon host.
128
-
129
- Raises:
130
- InvalidInputError: The destination has a different suffix or aliases the
131
- source checkpoint, or checkpoint contents fail ``convert_state_dict``.
132
- """
133
-
134
- if destination.suffix != ".safetensors":
135
- raise InvalidInputError("Converted weights destination must end in .safetensors")
136
- if checkpoint_path.resolve() == destination.resolve() or (
137
- checkpoint_path.is_file()
138
- and destination.is_file()
139
- and checkpoint_path.samefile(destination)
140
- ):
141
- raise InvalidInputError("Converted weights must not overwrite their source checkpoint")
142
- mx = import_mlx()
143
- converted = convert_state_dict(load_torch_checkpoint(checkpoint_path), spec)
144
- destination.parent.mkdir(parents=True, exist_ok=True)
145
- mx.save_safetensors(
146
- str(destination),
147
- {key: cast(ArrayFactory, mx.array)(value) for key, value in converted.items()},
148
- )
149
- return destination
150
-
151
-
152
67
  def load_model(
153
68
  weights_path: Path,
154
69
  *,
@@ -162,6 +77,8 @@ def load_model(
162
77
  from ``required_shapes(spec)``, or a parameter is not finite float32.
163
78
  RuntimeError: If MLX cannot open or decode the weights file.
164
79
  """
80
+ if Path(weights_path).suffix != ".safetensors":
81
+ raise InvalidInputError("SegResNet weights must be a safetensors file")
165
82
  mx = import_mlx()
166
83
  raw: object = mx.load(str(weights_path))
167
84
  if not isinstance(raw, dict):
@@ -191,6 +108,7 @@ def load_model(
191
108
 
192
109
 
193
110
  def _residual_shapes(prefix: str, channels: int) -> dict[str, tuple[int, ...]]:
111
+ """Expected parameter shapes of one residual block (Torch layouts, no conv bias)."""
194
112
  return {
195
113
  f"{prefix}.norm1.weight": (channels,),
196
114
  f"{prefix}.norm1.bias": (channels,),
@@ -1,87 +0,0 @@
1
- Metadata-Version: 2.5
2
- Name: segresnet-mlx
3
- Version: 0.2.3
4
- Summary: Native MLX SegResNet with NumPy sliding-window logit inference
5
- Project-URL: Homepage, https://github.com/MedMLX/segresnet-mlx
6
- Project-URL: Documentation, https://github.com/MedMLX/segresnet-mlx/blob/main/docs/technical.md
7
- Project-URL: Repository, https://github.com/MedMLX/segresnet-mlx
8
- Project-URL: Issues, https://github.com/MedMLX/segresnet-mlx/issues
9
- Author: Joseph Sandoval, Alfredo Sandoval Jr.
10
- License-Expression: Apache-2.0
11
- License-File: LICENSE
12
- License-File: THIRD_PARTY_NOTICES.md
13
- Classifier: Operating System :: MacOS :: MacOS X
14
- Requires-Python: <3.15,>=3.13
15
- Requires-Dist: medmlx-core==0.2.4
16
- Requires-Dist: nibabel<6,>=5.3
17
- Requires-Dist: numpy<3,>=2.0
18
- Provides-Extra: conversion
19
- Requires-Dist: torch<3,>=2.6; extra == 'conversion'
20
- Provides-Extra: dev
21
- Requires-Dist: hypothesis<7,>=6; extra == 'dev'
22
- Requires-Dist: monai==1.6.1; extra == 'dev'
23
- Requires-Dist: pyright==1.1.414; extra == 'dev'
24
- Requires-Dist: pytest<10,>=8.4; extra == 'dev'
25
- Requires-Dist: pytorch-ignite<0.6,>=0.4.11; extra == 'dev'
26
- Requires-Dist: ruff<1,>=0.13; extra == 'dev'
27
- Requires-Dist: torch<3,>=2.6; extra == 'dev'
28
- Description-Content-Type: text/markdown
29
-
30
- # segresnet-mlx
31
-
32
- Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegResNet models in native MLX.
33
-
34
- [Docs](https://github.com/MedMLX/segresnet-mlx/blob/main/docs/technical.md) · [MedMLX](https://github.com/MedMLX)
35
-
36
- ![SegResNet brain tumor segmentation in axial, coronal and sagittal views](https://raw.githubusercontent.com/MedMLX/segresnet-mlx/main/docs/example.png)
37
-
38
- *BraTS bundle output on a Medical Segmentation Decathlon brain MRI: tumor core
39
- (red), edema (green) and enhancing tumor (yellow).*
40
-
41
- | Input | Output |
42
- | :--- | :--- |
43
- | MRI sequences or CT volumes | Segmentation on the input grid |
44
-
45
- **Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
46
-
47
- ## Quick start
48
-
49
- ```bash
50
- uv venv --python 3.13
51
- source .venv/bin/activate
52
- uv pip install segresnet-mlx huggingface-hub
53
- ```
54
-
55
- Download the converted [BraTS](https://huggingface.co/MedMLX/segresnet-brats-mlx) weights, then run:
56
-
57
- ```python
58
- from huggingface_hub import hf_hub_download
59
- from segresnet_mlx.medmlx import run
60
-
61
- bundle = "brats_mri_segmentation"
62
- weights = hf_hub_download(
63
- "MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors",
64
- revision="14f9cda3649965729663c7365c49fd839c81a01b",
65
- )
66
-
67
- result = run(
68
- bundle=bundle,
69
- t1c_path="t1c.nii.gz",
70
- t1_path="t1.nii.gz",
71
- t2_path="t2.nii.gz",
72
- flair_path="flair.nii.gz",
73
- weights_path=weights,
74
- output_dir="out",
75
- )
76
- print(result["outputs"]["mask"])
77
- ```
78
-
79
- Keep input volumes on the same voxel grid. Converted weights are also available for
80
- [renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx). Docs cover
81
- checkpoint conversion and whole-body CT.
82
-
83
- ## License
84
-
85
- [Apache-2.0](https://github.com/MedMLX/segresnet-mlx/blob/main/LICENSE). See [third-party notices](https://github.com/MedMLX/segresnet-mlx/blob/main/THIRD_PARTY_NOTICES.md). Model weights retain their publisher terms.
86
-
87
- For research and education. Outputs require human review.
@@ -1,58 +0,0 @@
1
- # segresnet-mlx
2
-
3
- Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegResNet models in native MLX.
4
-
5
- [Docs](https://github.com/MedMLX/segresnet-mlx/blob/main/docs/technical.md) · [MedMLX](https://github.com/MedMLX)
6
-
7
- ![SegResNet brain tumor segmentation in axial, coronal and sagittal views](https://raw.githubusercontent.com/MedMLX/segresnet-mlx/main/docs/example.png)
8
-
9
- *BraTS bundle output on a Medical Segmentation Decathlon brain MRI: tumor core
10
- (red), edema (green) and enhancing tumor (yellow).*
11
-
12
- | Input | Output |
13
- | :--- | :--- |
14
- | MRI sequences or CT volumes | Segmentation on the input grid |
15
-
16
- **Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
17
-
18
- ## Quick start
19
-
20
- ```bash
21
- uv venv --python 3.13
22
- source .venv/bin/activate
23
- uv pip install segresnet-mlx huggingface-hub
24
- ```
25
-
26
- Download the converted [BraTS](https://huggingface.co/MedMLX/segresnet-brats-mlx) weights, then run:
27
-
28
- ```python
29
- from huggingface_hub import hf_hub_download
30
- from segresnet_mlx.medmlx import run
31
-
32
- bundle = "brats_mri_segmentation"
33
- weights = hf_hub_download(
34
- "MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors",
35
- revision="14f9cda3649965729663c7365c49fd839c81a01b",
36
- )
37
-
38
- result = run(
39
- bundle=bundle,
40
- t1c_path="t1c.nii.gz",
41
- t1_path="t1.nii.gz",
42
- t2_path="t2.nii.gz",
43
- flair_path="flair.nii.gz",
44
- weights_path=weights,
45
- output_dir="out",
46
- )
47
- print(result["outputs"]["mask"])
48
- ```
49
-
50
- Keep input volumes on the same voxel grid. Converted weights are also available for
51
- [renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx). Docs cover
52
- checkpoint conversion and whole-body CT.
53
-
54
- ## License
55
-
56
- [Apache-2.0](https://github.com/MedMLX/segresnet-mlx/blob/main/LICENSE). See [third-party notices](https://github.com/MedMLX/segresnet-mlx/blob/main/THIRD_PARTY_NOTICES.md). Model weights retain their publisher terms.
57
-
58
- For research and education. Outputs require human review.
@@ -1,63 +0,0 @@
1
- """Typed protocols for the NiBabel NIfTI, header and orientation calls used here."""
2
-
3
- from __future__ import annotations
4
-
5
- from collections.abc import Sequence
6
- from os import PathLike
7
- from typing import Protocol, cast
8
-
9
- import nibabel
10
- import numpy as np
11
- from nibabel.filebasedimages import ImageFileError as ImageFileError
12
-
13
- from segresnet_mlx._arrays import Float64Array, HostArray
14
-
15
- type PathArg = str | PathLike[str]
16
-
17
-
18
- class ArrayBuffer(Protocol):
19
- def __array__(
20
- self, dtype: np.dtype[np.generic] | None = None, copy: bool | None = None
21
- ) -> HostArray: ...
22
-
23
-
24
- class NiftiHeader(Protocol):
25
- def get_zooms(self) -> tuple[float, ...]: ...
26
- def set_zooms(self, zooms: Sequence[float] | Float64Array) -> None: ...
27
- def __getitem__(self, key: str) -> HostArray: ...
28
-
29
-
30
- class NiftiImage(Protocol):
31
- shape: tuple[int, ...]
32
- affine: Float64Array
33
- dataobj: ArrayBuffer | HostArray
34
- header: NiftiHeader
35
-
36
- def __init__(self, dataobj: HostArray, affine: HostArray) -> None: ...
37
- def get_data_dtype(self) -> np.dtype[np.generic]: ...
38
- def get_sform(self) -> Float64Array: ...
39
- def get_qform(self) -> Float64Array: ...
40
- def set_sform(self, affine: Float64Array, code: int | None = None) -> None: ...
41
- def set_qform(self, affine: Float64Array, code: int | None = None) -> None: ...
42
-
43
-
44
- class Orientations(Protocol):
45
- def axcodes2ornt(self, axcodes: str, labels: tuple[tuple[str, str], ...]) -> Float64Array: ...
46
- def ornt_transform(self, start_ornt: Float64Array, end_ornt: Float64Array) -> Float64Array: ...
47
- def inv_ornt_aff(self, ornt: Float64Array, shape: Sequence[int]) -> Float64Array: ...
48
- def aff2axcodes(
49
- self, aff: Float64Array, labels: tuple[tuple[str, str], ...]
50
- ) -> tuple[str | None, ...]: ...
51
-
52
-
53
- class Nibabel(Protocol):
54
- Nifti1Image: type[NiftiImage]
55
- orientations: Orientations
56
-
57
- def load(self, filename: PathArg) -> NiftiImage: ...
58
- def save(self, img: NiftiImage, filename: PathArg) -> None: ...
59
- def io_orientation(self, affine: Float64Array) -> Float64Array: ...
60
-
61
-
62
- # Bound to the real module; this runner admits only NIfTI images.
63
- nib = cast(Nibabel, nibabel)
File without changes
File without changes