segresnet-mlx 0.2.3__tar.gz → 0.3.1__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.1/PKG-INFO +80 -0
- segresnet_mlx-0.3.1/README.md +56 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/pyproject.toml +6 -12
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/__init__.py +3 -7
- segresnet_mlx-0.3.1/src/segresnet_mlx/_nifti.py +121 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_records.py +2 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_types.py +16 -5
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/medmlx.py +204 -26
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/network.py +9 -3
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/predictor.py +22 -4
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/telemetry.py +11 -1
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/topology.py +1 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/transforms.py +32 -118
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/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.1}/.gitignore +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/LICENSE +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/THIRD_PARTY_NOTICES.md +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_arrays.py +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_runtime.py +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/inference.py +0 -0
- {segresnet_mlx-0.2.3 → segresnet_mlx-0.3.1}/src/segresnet_mlx/py.typed +0 -0
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: segresnet-mlx
|
|
3
|
+
Version: 0.3.1
|
|
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
|
+
[Docs](https://github.com/MedMLX/segresnet-mlx/blob/main/docs/technical.md) · [BraTS weights](https://huggingface.co/MedMLX/segresnet-brats-mlx) · [Renal weights](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx) · [MedMLX](https://github.com/MedMLX)
|
|
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
|
+
## Install
|
|
38
|
+
|
|
39
|
+
```bash
|
|
40
|
+
uv venv --python 3.13
|
|
41
|
+
source .venv/bin/activate
|
|
42
|
+
uv pip install "segresnet-mlx @ git+https://github.com/MedMLX/segresnet-mlx.git@f60afe77518f5e617145cc4ccfa8eaa59e9eeb3e"
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
This source revision validates NIfTI units, voxel values and geometry before
|
|
46
|
+
inference; PyPI 0.3.0 predates those checks.
|
|
47
|
+
|
|
48
|
+
## Run
|
|
49
|
+
|
|
50
|
+
```python
|
|
51
|
+
from segresnet_mlx.medmlx import brats
|
|
52
|
+
|
|
53
|
+
result = brats(
|
|
54
|
+
t1c_path="t1c.nii.gz", t1_path="t1.nii.gz", t2_path="t2.nii.gz",
|
|
55
|
+
flair_path="flair.nii.gz", output_dir="out",
|
|
56
|
+
)
|
|
57
|
+
print(result["outputs"]["mask"])
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
The registered models are `segresnet-brats` and `segresnet-renal-cect`. Their `brats`
|
|
61
|
+
and `renal_cect` runners select the matching bundle internally and download
|
|
62
|
+
`MedMLX/segresnet-brats-mlx` or `MedMLX/segresnet-renal-cect-mlx` when `weights_path`
|
|
63
|
+
is omitted. An explicit folder contains `segresnet-brats.safetensors` for BraTS or
|
|
64
|
+
`segresnet-renal-cect.safetensors` for renal CECT. The renal runner takes
|
|
65
|
+
`arterial_path`, `venous_path` and `excretory_path`. Keep each model's input volumes
|
|
66
|
+
on one voxel grid.
|
|
67
|
+
|
|
68
|
+
The standalone `run(bundle=..., weights_path=...)` interface remains available for
|
|
69
|
+
explicit local checkpoints, including the whole-body CT variants. Published native
|
|
70
|
+
weights are also available for [renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx)
|
|
71
|
+
and [whole-body CT](https://huggingface.co/MedMLX/segresnet-wholebody-mlx).
|
|
72
|
+
|
|
73
|
+
Downloaded commits are returned in `details.weights_revision`; explicit local
|
|
74
|
+
weights report `null` unless the folder carries a MedMLX download revision.
|
|
75
|
+
|
|
76
|
+
## License
|
|
77
|
+
|
|
78
|
+
[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.
|
|
79
|
+
|
|
80
|
+
For research and education. Outputs require human review.
|
|
@@ -0,0 +1,56 @@
|
|
|
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) · [BraTS weights](https://huggingface.co/MedMLX/segresnet-brats-mlx) · [Renal weights](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx) · [MedMLX](https://github.com/MedMLX)
|
|
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
|
+
## Install
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
uv venv --python 3.13
|
|
17
|
+
source .venv/bin/activate
|
|
18
|
+
uv pip install "segresnet-mlx @ git+https://github.com/MedMLX/segresnet-mlx.git@f60afe77518f5e617145cc4ccfa8eaa59e9eeb3e"
|
|
19
|
+
```
|
|
20
|
+
|
|
21
|
+
This source revision validates NIfTI units, voxel values and geometry before
|
|
22
|
+
inference; PyPI 0.3.0 predates those checks.
|
|
23
|
+
|
|
24
|
+
## Run
|
|
25
|
+
|
|
26
|
+
```python
|
|
27
|
+
from segresnet_mlx.medmlx import brats
|
|
28
|
+
|
|
29
|
+
result = brats(
|
|
30
|
+
t1c_path="t1c.nii.gz", t1_path="t1.nii.gz", t2_path="t2.nii.gz",
|
|
31
|
+
flair_path="flair.nii.gz", output_dir="out",
|
|
32
|
+
)
|
|
33
|
+
print(result["outputs"]["mask"])
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
The registered models are `segresnet-brats` and `segresnet-renal-cect`. Their `brats`
|
|
37
|
+
and `renal_cect` runners select the matching bundle internally and download
|
|
38
|
+
`MedMLX/segresnet-brats-mlx` or `MedMLX/segresnet-renal-cect-mlx` when `weights_path`
|
|
39
|
+
is omitted. An explicit folder contains `segresnet-brats.safetensors` for BraTS or
|
|
40
|
+
`segresnet-renal-cect.safetensors` for renal CECT. The renal runner takes
|
|
41
|
+
`arterial_path`, `venous_path` and `excretory_path`. Keep each model's input volumes
|
|
42
|
+
on one voxel grid.
|
|
43
|
+
|
|
44
|
+
The standalone `run(bundle=..., weights_path=...)` interface remains available for
|
|
45
|
+
explicit local checkpoints, including the whole-body CT variants. Published native
|
|
46
|
+
weights are also available for [renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx)
|
|
47
|
+
and [whole-body CT](https://huggingface.co/MedMLX/segresnet-wholebody-mlx).
|
|
48
|
+
|
|
49
|
+
Downloaded commits are returned in `details.weights_revision`; explicit local
|
|
50
|
+
weights report `null` unless the folder carries a MedMLX download revision.
|
|
51
|
+
|
|
52
|
+
## License
|
|
53
|
+
|
|
54
|
+
[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.
|
|
55
|
+
|
|
56
|
+
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.1"
|
|
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.1"
|
|
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,121 @@
|
|
|
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 get_xyzt_units(self) -> tuple[str, str]:
|
|
36
|
+
"""Return the NIfTI spatial and temporal unit names."""
|
|
37
|
+
...
|
|
38
|
+
|
|
39
|
+
def set_xyzt_units(self, xyz: str) -> None:
|
|
40
|
+
"""Set the NIfTI spatial unit name."""
|
|
41
|
+
...
|
|
42
|
+
|
|
43
|
+
def __getitem__(self, key: str) -> HostArray:
|
|
44
|
+
"""Return a named NIfTI header field."""
|
|
45
|
+
...
|
|
46
|
+
|
|
47
|
+
def __setitem__(self, key: str, value: int) -> None:
|
|
48
|
+
"""Set an integer-valued NIfTI header field."""
|
|
49
|
+
...
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
class NiftiImage(Protocol):
|
|
53
|
+
shape: tuple[int, ...]
|
|
54
|
+
affine: Float64Array
|
|
55
|
+
dataobj: ArrayBuffer | HostArray
|
|
56
|
+
header: NiftiHeader
|
|
57
|
+
|
|
58
|
+
def __init__(self, dataobj: HostArray, affine: HostArray) -> None:
|
|
59
|
+
"""Create an image from voxel data and its affine."""
|
|
60
|
+
...
|
|
61
|
+
|
|
62
|
+
def get_data_dtype(self) -> np.dtype[np.generic]:
|
|
63
|
+
"""Return the stored voxel-data dtype."""
|
|
64
|
+
...
|
|
65
|
+
|
|
66
|
+
def get_sform(self) -> Float64Array:
|
|
67
|
+
"""Return the header's sform affine."""
|
|
68
|
+
...
|
|
69
|
+
|
|
70
|
+
def get_qform(self) -> Float64Array:
|
|
71
|
+
"""Return the header's qform affine."""
|
|
72
|
+
...
|
|
73
|
+
|
|
74
|
+
def set_sform(self, affine: Float64Array, code: int | None = None) -> None:
|
|
75
|
+
"""Set the header's sform affine and optional code."""
|
|
76
|
+
...
|
|
77
|
+
|
|
78
|
+
def set_qform(self, affine: Float64Array, code: int | None = None) -> None:
|
|
79
|
+
"""Set the header's qform affine and optional code."""
|
|
80
|
+
...
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
class Orientations(Protocol):
|
|
84
|
+
def axcodes2ornt(self, axcodes: str, labels: tuple[tuple[str, str], ...]) -> Float64Array:
|
|
85
|
+
"""Convert axis codes to an orientation array."""
|
|
86
|
+
...
|
|
87
|
+
|
|
88
|
+
def ornt_transform(self, start_ornt: Float64Array, end_ornt: Float64Array) -> Float64Array:
|
|
89
|
+
"""Return the transform between two orientation arrays."""
|
|
90
|
+
...
|
|
91
|
+
|
|
92
|
+
def inv_ornt_aff(self, ornt: Float64Array, shape: Sequence[int]) -> Float64Array:
|
|
93
|
+
"""Return the inverse voxel-coordinate affine for an orientation change."""
|
|
94
|
+
...
|
|
95
|
+
|
|
96
|
+
def aff2axcodes(
|
|
97
|
+
self, aff: Float64Array, labels: tuple[tuple[str, str], ...]
|
|
98
|
+
) -> tuple[str | None, ...]:
|
|
99
|
+
"""Return axis codes implied by an affine."""
|
|
100
|
+
...
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
class Nibabel(Protocol):
|
|
104
|
+
Nifti1Image: type[NiftiImage]
|
|
105
|
+
orientations: Orientations
|
|
106
|
+
|
|
107
|
+
def load(self, filename: PathArg) -> NiftiImage:
|
|
108
|
+
"""Load an image from a NIfTI path."""
|
|
109
|
+
...
|
|
110
|
+
|
|
111
|
+
def save(self, img: NiftiImage, filename: PathArg) -> None:
|
|
112
|
+
"""Write an image to a NIfTI path."""
|
|
113
|
+
...
|
|
114
|
+
|
|
115
|
+
def io_orientation(self, affine: Float64Array) -> Float64Array:
|
|
116
|
+
"""Return the input-to-output orientation implied by an affine."""
|
|
117
|
+
...
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
# Bound to the real module; this runner admits only NIfTI images.
|
|
121
|
+
nib = cast(Nibabel, nibabel)
|
|
@@ -57,8 +57,10 @@ 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
|
|
63
|
+
streaming: NotRequired[bool]
|
|
62
64
|
labels: dict[str, str]
|
|
63
65
|
labels_present: list[int]
|
|
64
66
|
input_shape: list[int]
|
|
@@ -6,20 +6,27 @@ from typing import TYPE_CHECKING
|
|
|
6
6
|
|
|
7
7
|
if TYPE_CHECKING:
|
|
8
8
|
from collections.abc import Callable, Mapping, Sequence
|
|
9
|
-
from typing import Literal, Protocol
|
|
9
|
+
from typing import Literal, Protocol, TypedDict
|
|
10
10
|
|
|
11
11
|
import mlx.core as mx
|
|
12
12
|
from medmlx_core.typing import MlxRuntime
|
|
13
|
+
from medmlx_core.volume_workspace import VolumeWorkspace
|
|
13
14
|
|
|
14
15
|
from segresnet_mlx._arrays import Float32Array, HostArray
|
|
15
16
|
|
|
16
17
|
type Array = mx.array
|
|
18
|
+
|
|
19
|
+
class WorkspaceOptions(TypedDict, total=False):
|
|
20
|
+
workspace: VolumeWorkspace
|
|
21
|
+
|
|
17
22
|
type ArrayValue = bool | int | float | complex | list[ArrayValue] | tuple[ArrayValue, ...]
|
|
18
23
|
|
|
19
24
|
class ArrayConstructor(Protocol):
|
|
20
25
|
def __call__(
|
|
21
26
|
self, value: HostArray | Array | ArrayValue, dtype: mx.Dtype | None = None
|
|
22
|
-
) -> Array:
|
|
27
|
+
) -> Array:
|
|
28
|
+
"""Construct an MLX array with an optional dtype."""
|
|
29
|
+
...
|
|
23
30
|
|
|
24
31
|
class MLX(MlxRuntime, Protocol):
|
|
25
32
|
all = staticmethod(mx.all)
|
|
@@ -32,7 +39,6 @@ if TYPE_CHECKING:
|
|
|
32
39
|
load = staticmethod(mx.load)
|
|
33
40
|
|
|
34
41
|
def eval(self, *args: Array | Mapping[str, Array] | Sequence[Array]) -> None: ...
|
|
35
|
-
|
|
36
42
|
def save_safetensors(self, file: str, arrays: dict[str, Array]) -> None: ...
|
|
37
43
|
|
|
38
44
|
class SlidingWindowInference(Protocol):
|
|
@@ -50,10 +56,15 @@ if TYPE_CHECKING:
|
|
|
50
56
|
with_coord: bool = False,
|
|
51
57
|
predictor_input: Literal["numpy"],
|
|
52
58
|
preserve_window_strides: bool = False,
|
|
53
|
-
|
|
59
|
+
workspace: VolumeWorkspace | None = None,
|
|
60
|
+
) -> Float32Array:
|
|
61
|
+
"""Blend predictor outputs from overlapping windows into host logits."""
|
|
62
|
+
...
|
|
54
63
|
|
|
55
64
|
class CoreSlidingWindow(Protocol):
|
|
56
65
|
sliding_window_inference: SlidingWindowInference
|
|
57
66
|
|
|
58
67
|
class Metal(Protocol):
|
|
59
|
-
def is_available(self) -> bool:
|
|
68
|
+
def is_available(self) -> bool:
|
|
69
|
+
"""Report whether the Metal backend is available."""
|
|
70
|
+
...
|