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.
- segresnet_mlx-0.3.0/PKG-INFO +75 -0
- segresnet_mlx-0.3.0/README.md +51 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/pyproject.toml +6 -12
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/__init__.py +3 -7
- segresnet_mlx-0.3.0/src/segresnet_mlx/_nifti.py +109 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/_records.py +1 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/_types.py +9 -4
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/medmlx.py +83 -11
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/network.py +9 -3
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/predictor.py +5 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/telemetry.py +9 -1
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/topology.py +1 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/transforms.py +29 -116
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/weights.py +4 -86
- segresnet_mlx-0.2.3/PKG-INFO +0 -87
- segresnet_mlx-0.2.3/README.md +0 -58
- segresnet_mlx-0.2.3/src/segresnet_mlx/_nifti.py +0 -63
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/.gitignore +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/LICENSE +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/THIRD_PARTY_NOTICES.md +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/_arrays.py +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/_runtime.py +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.0}/src/segresnet_mlx/inference.py +0 -0
- {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.
|
|
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.
|
|
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:
|
|
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://
|
|
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.
|
|
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)
|
|
@@ -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
|
|
299
|
-
|
|
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="
|
|
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
|
-
|
|
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
|
-
|
|
126
|
-
|
|
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),
|
|
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
|
|
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,),
|
segresnet_mlx-0.2.3/PKG-INFO
DELETED
|
@@ -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
|
-

|
|
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.
|
segresnet_mlx-0.2.3/README.md
DELETED
|
@@ -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
|
-

|
|
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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|