segresnet-mlx 0.3.0__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.0 → segresnet_mlx-0.3.1}/PKG-INFO +10 -5
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/README.md +9 -4
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/pyproject.toml +1 -1
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/__init__.py +1 -1
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_nifti.py +12 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_records.py +1 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_types.py +7 -1
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/medmlx.py +123 -17
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/predictor.py +17 -4
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/telemetry.py +2 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/transforms.py +3 -2
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/.gitignore +0 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/LICENSE +0 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/THIRD_PARTY_NOTICES.md +0 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_arrays.py +0 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_runtime.py +0 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/inference.py +0 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/network.py +0 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/py.typed +0 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/topology.py +0 -0
- {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/weights.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: segresnet-mlx
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.1
|
|
4
4
|
Summary: Native MLX SegResNet with NumPy sliding-window logit inference
|
|
5
5
|
Project-URL: Homepage, https://huggingface.co/MedMLX/segresnet-renal-cect-mlx
|
|
6
6
|
Author: Joseph Sandoval, Alfredo Sandoval Jr.
|
|
@@ -26,7 +26,7 @@ Description-Content-Type: text/markdown
|
|
|
26
26
|
|
|
27
27
|
Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegResNet models in native MLX.
|
|
28
28
|
|
|
29
|
-
[BraTS weights](https://huggingface.co/MedMLX/segresnet-brats-mlx) · [Renal weights](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx)
|
|
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
30
|
|
|
31
31
|
| Input | Output |
|
|
32
32
|
| :--- | :--- |
|
|
@@ -34,14 +34,19 @@ Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegR
|
|
|
34
34
|
|
|
35
35
|
**Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
|
|
36
36
|
|
|
37
|
-
##
|
|
37
|
+
## Install
|
|
38
38
|
|
|
39
39
|
```bash
|
|
40
40
|
uv venv --python 3.13
|
|
41
41
|
source .venv/bin/activate
|
|
42
|
-
uv pip install "segresnet-mlx
|
|
42
|
+
uv pip install "segresnet-mlx @ git+https://github.com/MedMLX/segresnet-mlx.git@f60afe77518f5e617145cc4ccfa8eaa59e9eeb3e"
|
|
43
43
|
```
|
|
44
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
|
+
|
|
45
50
|
```python
|
|
46
51
|
from segresnet_mlx.medmlx import brats
|
|
47
52
|
|
|
@@ -70,6 +75,6 @@ weights report `null` unless the folder carries a MedMLX download revision.
|
|
|
70
75
|
|
|
71
76
|
## License
|
|
72
77
|
|
|
73
|
-
Apache-2.0
|
|
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.
|
|
74
79
|
|
|
75
80
|
For research and education. Outputs require human review.
|
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
|
|
3
3
|
Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegResNet models in native MLX.
|
|
4
4
|
|
|
5
|
-
[BraTS weights](https://huggingface.co/MedMLX/segresnet-brats-mlx) · [Renal weights](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx)
|
|
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
6
|
|
|
7
7
|
| Input | Output |
|
|
8
8
|
| :--- | :--- |
|
|
@@ -10,14 +10,19 @@ Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegR
|
|
|
10
10
|
|
|
11
11
|
**Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
|
|
12
12
|
|
|
13
|
-
##
|
|
13
|
+
## Install
|
|
14
14
|
|
|
15
15
|
```bash
|
|
16
16
|
uv venv --python 3.13
|
|
17
17
|
source .venv/bin/activate
|
|
18
|
-
uv pip install "segresnet-mlx
|
|
18
|
+
uv pip install "segresnet-mlx @ git+https://github.com/MedMLX/segresnet-mlx.git@f60afe77518f5e617145cc4ccfa8eaa59e9eeb3e"
|
|
19
19
|
```
|
|
20
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
|
+
|
|
21
26
|
```python
|
|
22
27
|
from segresnet_mlx.medmlx import brats
|
|
23
28
|
|
|
@@ -46,6 +51,6 @@ weights report `null` unless the folder carries a MedMLX download revision.
|
|
|
46
51
|
|
|
47
52
|
## License
|
|
48
53
|
|
|
49
|
-
Apache-2.0
|
|
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.
|
|
50
55
|
|
|
51
56
|
For research and education. Outputs require human review.
|
|
@@ -19,7 +19,7 @@ if TYPE_CHECKING:
|
|
|
19
19
|
from segresnet_mlx.weights import load_model as load_model
|
|
20
20
|
from segresnet_mlx.weights import required_shapes as required_shapes
|
|
21
21
|
|
|
22
|
-
__version__ = "0.3.
|
|
22
|
+
__version__ = "0.3.1"
|
|
23
23
|
|
|
24
24
|
_EXPORTS = {
|
|
25
25
|
"BUNDLE_CONFIGS": "segresnet_mlx.predictor",
|
|
@@ -32,10 +32,22 @@ class NiftiHeader(Protocol):
|
|
|
32
32
|
"""Set voxel spacing in the header."""
|
|
33
33
|
...
|
|
34
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
|
+
|
|
35
43
|
def __getitem__(self, key: str) -> HostArray:
|
|
36
44
|
"""Return a named NIfTI header field."""
|
|
37
45
|
...
|
|
38
46
|
|
|
47
|
+
def __setitem__(self, key: str, value: int) -> None:
|
|
48
|
+
"""Set an integer-valued NIfTI header field."""
|
|
49
|
+
...
|
|
50
|
+
|
|
39
51
|
|
|
40
52
|
class NiftiImage(Protocol):
|
|
41
53
|
shape: tuple[int, ...]
|
|
@@ -6,14 +6,19 @@ 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):
|
|
@@ -51,6 +56,7 @@ if TYPE_CHECKING:
|
|
|
51
56
|
with_coord: bool = False,
|
|
52
57
|
predictor_input: Literal["numpy"],
|
|
53
58
|
preserve_window_strides: bool = False,
|
|
59
|
+
workspace: VolumeWorkspace | None = None,
|
|
54
60
|
) -> Float32Array:
|
|
55
61
|
"""Blend predictor outputs from overlapping windows into host logits."""
|
|
56
62
|
...
|
|
@@ -6,6 +6,7 @@ import os
|
|
|
6
6
|
import tempfile
|
|
7
7
|
import time
|
|
8
8
|
from collections.abc import Mapping
|
|
9
|
+
from contextlib import nullcontext
|
|
9
10
|
from dataclasses import dataclass
|
|
10
11
|
from pathlib import Path
|
|
11
12
|
from typing import TYPE_CHECKING
|
|
@@ -14,6 +15,8 @@ from segresnet_mlx._arrays import Float32Array, Float64Array, MaskArray
|
|
|
14
15
|
from segresnet_mlx._records import RunResult, Runtime
|
|
15
16
|
|
|
16
17
|
if TYPE_CHECKING:
|
|
18
|
+
from medmlx_core.volume_workspace import VolumeWorkspace
|
|
19
|
+
|
|
17
20
|
from segresnet_mlx.network import SegResNet
|
|
18
21
|
|
|
19
22
|
import numpy as np
|
|
@@ -300,6 +303,8 @@ def run(
|
|
|
300
303
|
image_path: str | None = None,
|
|
301
304
|
overwrite: bool = False,
|
|
302
305
|
arithmetic_profile: ArithmeticProfile = "torch-order-fp32",
|
|
306
|
+
streaming: bool = False,
|
|
307
|
+
scratch_dir: str | None = None,
|
|
303
308
|
) -> RunResult:
|
|
304
309
|
"""Segment one case with a MONAI SegResNet bundle and write the label map.
|
|
305
310
|
|
|
@@ -353,6 +358,10 @@ def run(
|
|
|
353
358
|
overwrite: Replace an existing output mask when true; defaults to false.
|
|
354
359
|
arithmetic_profile: ``torch-order-fp32`` (default), ``compensated-fp32``,
|
|
355
360
|
or ``compensated-fp48``; the last supports only the renal bundle.
|
|
361
|
+
streaming: Use bounded blending intermediates and temporary mapped logits.
|
|
362
|
+
The final label volume remains resident; defaults to False.
|
|
363
|
+
scratch_dir: Parent directory for temporary maps when streaming is enabled;
|
|
364
|
+
None uses the system temporary directory. Files are cleaned up on exit.
|
|
356
365
|
|
|
357
366
|
Returns:
|
|
358
367
|
``{"outputs": {"mask": path}, "details": {...}}``; ``details`` is described
|
|
@@ -368,6 +377,8 @@ def run(
|
|
|
368
377
|
ModelExecutionError: Inference failed or produced a mask off the source grid
|
|
369
378
|
or with labels outside the bundle's set.
|
|
370
379
|
"""
|
|
380
|
+
if type(streaming) is not bool:
|
|
381
|
+
raise InvalidInputError("streaming must be boolean")
|
|
371
382
|
values = {
|
|
372
383
|
"arterial_path": arterial_path,
|
|
373
384
|
"venous_path": venous_path,
|
|
@@ -419,16 +430,29 @@ def run(
|
|
|
419
430
|
else:
|
|
420
431
|
image, inverse = preprocess_brats(volumes), None
|
|
421
432
|
prepared = time.perf_counter()
|
|
422
|
-
|
|
423
|
-
|
|
424
|
-
|
|
425
|
-
|
|
426
|
-
mask = postprocess_renal(logits[0], inverse)
|
|
427
|
-
elif bundle in _WHOLE_BODY_PIXDIM:
|
|
428
|
-
assert inverse is not None
|
|
429
|
-
mask = postprocess_wholebody(logits[0], inverse)
|
|
433
|
+
if streaming:
|
|
434
|
+
from medmlx_core.volume_workspace import VolumeWorkspace
|
|
435
|
+
|
|
436
|
+
context = VolumeWorkspace(scratch_dir)
|
|
430
437
|
else:
|
|
431
|
-
|
|
438
|
+
context = nullcontext(None)
|
|
439
|
+
with context as workspace:
|
|
440
|
+
logits, runtime = (
|
|
441
|
+
predict(network, image, config)
|
|
442
|
+
if workspace is None
|
|
443
|
+
else predict(network, image, config, workspace=workspace)
|
|
444
|
+
)
|
|
445
|
+
predicted = time.perf_counter()
|
|
446
|
+
if workspace is not None:
|
|
447
|
+
mask = postprocess_streaming(logits[0], inverse, bundle, workspace)
|
|
448
|
+
elif bundle == "renalStructures_CECT_segmentation":
|
|
449
|
+
assert inverse is not None
|
|
450
|
+
mask = postprocess_renal(logits[0], inverse)
|
|
451
|
+
elif bundle in _WHOLE_BODY_PIXDIM:
|
|
452
|
+
assert inverse is not None
|
|
453
|
+
mask = postprocess_wholebody(logits[0], inverse)
|
|
454
|
+
else:
|
|
455
|
+
mask = postprocess_brats(logits[0])
|
|
432
456
|
_validate_mask(mask, volumes[0], config)
|
|
433
457
|
output_root.mkdir(parents=True, exist_ok=True)
|
|
434
458
|
write_atomic(mask_path, mask, volumes[0].affine, config.output_dtype)
|
|
@@ -440,6 +464,7 @@ def run(
|
|
|
440
464
|
"weights_revision": weights_revision(weights.parent),
|
|
441
465
|
"bundle": bundle,
|
|
442
466
|
"arithmetic_profile": arithmetic_profile,
|
|
467
|
+
"streaming": streaming,
|
|
443
468
|
"labels": {str(value): name for value, name in config.labels.items()},
|
|
444
469
|
"labels_present": present,
|
|
445
470
|
"input_shape": list(volumes[0].shape),
|
|
@@ -482,9 +507,21 @@ def brats(
|
|
|
482
507
|
``segresnet-brats.safetensors``; an explicit safetensors file is also accepted.
|
|
483
508
|
Automatic downloads use ``MedMLX/segresnet-brats-mlx``. All arithmetic and
|
|
484
509
|
preprocessing are delegated unchanged to ``run``.
|
|
510
|
+
|
|
511
|
+
Args:
|
|
512
|
+
t1c_path: Contrast-enhanced T1 MRI NIfTI on the shared input grid.
|
|
513
|
+
t1_path: T1 MRI NIfTI on the T1c grid.
|
|
514
|
+
t2_path: T2 MRI NIfTI on the T1c grid.
|
|
515
|
+
flair_path: FLAIR MRI NIfTI on the T1c grid.
|
|
516
|
+
output_dir: Directory for the brain-tumor label map.
|
|
517
|
+
weights_path: Optional folder containing ``segresnet-brats.safetensors``;
|
|
518
|
+
omit to download and reuse the BraTS bundle. Explicit files remain accepted.
|
|
519
|
+
overwrite: Replace an existing output mask when true.
|
|
520
|
+
arithmetic_profile: ``torch-order-fp32`` (default) or ``compensated-fp32``;
|
|
521
|
+
``compensated-fp48`` requires the renal bundle.
|
|
485
522
|
"""
|
|
486
523
|
folder = (
|
|
487
|
-
weights_dir("segresnet-brats
|
|
524
|
+
weights_dir("segresnet-brats")
|
|
488
525
|
if weights_path is None else Path(weights_path).expanduser()
|
|
489
526
|
)
|
|
490
527
|
checkpoint = folder / "segresnet-brats.safetensors" if folder.is_dir() else folder
|
|
@@ -512,9 +549,20 @@ def renal_cect(
|
|
|
512
549
|
this bundle; an explicit safetensors file is also accepted. Automatic downloads
|
|
513
550
|
use ``MedMLX/segresnet-renal-cect-mlx``. Arithmetic and preprocessing are
|
|
514
551
|
delegated unchanged to ``run``.
|
|
552
|
+
|
|
553
|
+
Args:
|
|
554
|
+
arterial_path: Arterial-phase CT NIfTI on the shared input grid.
|
|
555
|
+
venous_path: Venous-phase CT NIfTI on the arterial grid.
|
|
556
|
+
excretory_path: Excretory-phase CT NIfTI on the arterial grid.
|
|
557
|
+
output_dir: Directory for the renal label map.
|
|
558
|
+
weights_path: Optional folder containing ``segresnet-renal-cect.safetensors``;
|
|
559
|
+
omit to download and reuse the renal bundle. Explicit files remain accepted.
|
|
560
|
+
overwrite: Replace an existing output mask when true.
|
|
561
|
+
arithmetic_profile: ``torch-order-fp32`` (default), ``compensated-fp32``
|
|
562
|
+
or ``compensated-fp48``; see ``run`` for their arithmetic contracts.
|
|
515
563
|
"""
|
|
516
564
|
folder = (
|
|
517
|
-
weights_dir("segresnet-renal-cect
|
|
565
|
+
weights_dir("segresnet-renal-cect")
|
|
518
566
|
if weights_path is None else Path(weights_path).expanduser()
|
|
519
567
|
)
|
|
520
568
|
checkpoint = folder / "segresnet-renal-cect.safetensors" if folder.is_dir() else folder
|
|
@@ -636,6 +684,48 @@ def postprocess_wholebody(logits: Float32Array, inverse: SpatialInverse) -> Floa
|
|
|
636
684
|
return values[0]
|
|
637
685
|
|
|
638
686
|
|
|
687
|
+
def postprocess_streaming(
|
|
688
|
+
logits: Float32Array,
|
|
689
|
+
inverse: SpatialInverse | None,
|
|
690
|
+
bundle: str,
|
|
691
|
+
workspace: VolumeWorkspace,
|
|
692
|
+
) -> MaskArray:
|
|
693
|
+
"""Retain each bundle's activation/restoration/label order with bounded class storage."""
|
|
694
|
+
from medmlx_core.volume_workspace import restored_argmax
|
|
695
|
+
|
|
696
|
+
if bundle in _WHOLE_BODY_PIXDIM:
|
|
697
|
+
assert inverse is not None
|
|
698
|
+
probabilities = workspace.map_channels(logits, transforms.softmax_channel)
|
|
699
|
+
labels = restored_argmax(probabilities, lambda channel: channel)[None]
|
|
700
|
+
values, affine = transforms.invert_spacing(
|
|
701
|
+
labels, inverse.affine, inverse.resampling, "nearest"
|
|
702
|
+
)
|
|
703
|
+
return transforms.invert_orientation(values, affine, inverse.orientation)[0][0]
|
|
704
|
+
if bundle == "renalStructures_CECT_segmentation":
|
|
705
|
+
assert inverse is not None
|
|
706
|
+
labels: Float32Array | None = None
|
|
707
|
+
for channel in range(1, logits.shape[0]):
|
|
708
|
+
values, affine = transforms.invert_spacing(
|
|
709
|
+
logits[channel : channel + 1], inverse.affine, inverse.resampling
|
|
710
|
+
)
|
|
711
|
+
values, _ = transforms.invert_orientation(values, affine, inverse.orientation)
|
|
712
|
+
active = transforms.threshold(transforms.sigmoid(values), 0.5)[0] != 0
|
|
713
|
+
if labels is None:
|
|
714
|
+
labels = np.zeros(active.shape, dtype=np.float32)
|
|
715
|
+
labels[(labels == 0) & active] = channel
|
|
716
|
+
assert labels is not None
|
|
717
|
+
return labels
|
|
718
|
+
if bundle != "brats_mri_segmentation":
|
|
719
|
+
raise InvalidInputError(f"Unsupported streaming SegResNet bundle: {bundle}")
|
|
720
|
+
flat = logits.reshape(logits.shape[0], -1)
|
|
721
|
+
output = np.empty(flat.shape[1], dtype=np.int64)
|
|
722
|
+
for start in range(0, flat.shape[1], 65536):
|
|
723
|
+
stop = min(start + 65536, flat.shape[1])
|
|
724
|
+
binary = transforms.threshold(transforms.sigmoid(flat[:, start:stop]), 0.5)
|
|
725
|
+
output[start:stop] = transforms.brats_labels(binary)[0]
|
|
726
|
+
return output.reshape(logits.shape[1:])
|
|
727
|
+
|
|
728
|
+
|
|
639
729
|
def _weights_file(weights_path: str) -> Path:
|
|
640
730
|
"""Validate a safetensors file path while retaining any Hub-cache symlink."""
|
|
641
731
|
# MLX identifies the file format from the supplied name. Hub cache paths are
|
|
@@ -678,9 +768,8 @@ def load_inputs(roles: list[tuple[str, str | None]]) -> list[Volume]:
|
|
|
678
768
|
def load_volume(path: Path, role: str) -> Volume:
|
|
679
769
|
"""Read 3D voxels using the affine selected by the upstream nibabel reader.
|
|
680
770
|
|
|
681
|
-
|
|
682
|
-
|
|
683
|
-
and affine invertibility are not checked here.
|
|
771
|
+
Voxels must be real, finite and representable as float32. Spatial units must be
|
|
772
|
+
mm or unspecified (assumed mm), and the corrected affine must be invertible.
|
|
684
773
|
"""
|
|
685
774
|
from segresnet_mlx._nifti import ImageFileError, nib
|
|
686
775
|
|
|
@@ -690,15 +779,28 @@ def load_volume(path: Path, role: str) -> Volume:
|
|
|
690
779
|
image = nib.load(path.as_posix())
|
|
691
780
|
except (OSError, ImageFileError) as error:
|
|
692
781
|
raise InvalidInputError(f"Invalid {role} NIfTI: {path}") from error
|
|
782
|
+
try:
|
|
783
|
+
spatial_units, _ = image.header.get_xyzt_units()
|
|
784
|
+
except (AttributeError, KeyError, ValueError) as error:
|
|
785
|
+
raise InvalidInputError(f"{role} has invalid NIfTI spatial units") from error
|
|
786
|
+
if spatial_units not in ("mm", "unknown"):
|
|
787
|
+
raise InvalidInputError(f"{role} spatial units must be mm or unspecified")
|
|
693
788
|
raw = np.asanyarray(image.dataobj)
|
|
694
789
|
if raw.ndim != 3:
|
|
695
790
|
raise InvalidInputError(f"{role} must be a 3D NIfTI; got shape {tuple(raw.shape)}")
|
|
791
|
+
if np.iscomplexobj(raw):
|
|
792
|
+
raise InvalidInputError(f"{role} must contain real-valued voxels")
|
|
696
793
|
if not np.all(np.isfinite(raw)):
|
|
697
794
|
raise InvalidInputError(f"{role} contains non-finite voxels")
|
|
795
|
+
limit = np.finfo(np.float32).max
|
|
796
|
+
if np.any(raw > limit) or np.any(raw < -limit):
|
|
797
|
+
raise InvalidInputError(f"{role} contains voxels outside the finite float32 range")
|
|
698
798
|
affine = np.asarray(image.affine, dtype=np.float64)
|
|
699
799
|
if affine.shape != (4, 4) or not np.all(np.isfinite(affine)):
|
|
700
800
|
raise InvalidInputError(f"{role} affine must be a finite 4x4 matrix")
|
|
701
801
|
data, reader_affine = transforms.image_arrays(image)
|
|
802
|
+
if not np.all(np.isfinite(reader_affine)) or np.linalg.matrix_rank(reader_affine[:3, :3]) < 3:
|
|
803
|
+
raise InvalidInputError(f"{role} affine must define an invertible 3D grid")
|
|
702
804
|
shape = (int(raw.shape[0]), int(raw.shape[1]), int(raw.shape[2]))
|
|
703
805
|
return Volume(path=path, shape=shape, affine=reader_affine, data=data)
|
|
704
806
|
|
|
@@ -723,7 +825,11 @@ def _load_network(
|
|
|
723
825
|
|
|
724
826
|
|
|
725
827
|
def predict(
|
|
726
|
-
network: SegResNet,
|
|
828
|
+
network: SegResNet,
|
|
829
|
+
image: Float32Array,
|
|
830
|
+
config: Bundle,
|
|
831
|
+
*,
|
|
832
|
+
workspace: VolumeWorkspace | None = None,
|
|
727
833
|
) -> tuple[Float32Array, Runtime]:
|
|
728
834
|
"""Run the bundle's ``SlidingWindowInferer`` settings around the MLX network.
|
|
729
835
|
|
|
@@ -741,9 +847,9 @@ def predict(
|
|
|
741
847
|
try:
|
|
742
848
|
host = np.ascontiguousarray(image, dtype=np.float32)
|
|
743
849
|
logits = (
|
|
744
|
-
predict_wide_logits(network, host, window, mx=mx)
|
|
850
|
+
predict_wide_logits(network, host, window, mx=mx, workspace=workspace)
|
|
745
851
|
if network.arithmetic_profile == "compensated-fp48"
|
|
746
|
-
else BundlePredictor(network, mx=mx).predict_logits(host, window)
|
|
852
|
+
else BundlePredictor(network, mx=mx).predict_logits(host, window, workspace=workspace)
|
|
747
853
|
)
|
|
748
854
|
except (RuntimeError, ValueError) as error:
|
|
749
855
|
raise ModelExecutionError(f"SegResNet MLX inference failed: {error}") from error
|
|
@@ -13,8 +13,9 @@ from segresnet_mlx._runtime import import_mlx
|
|
|
13
13
|
|
|
14
14
|
if TYPE_CHECKING:
|
|
15
15
|
from ignite.engine import Engine, Events
|
|
16
|
+
from medmlx_core.volume_workspace import VolumeWorkspace
|
|
16
17
|
|
|
17
|
-
from segresnet_mlx._types import MLX, Array, CoreSlidingWindow
|
|
18
|
+
from segresnet_mlx._types import MLX, Array, CoreSlidingWindow, WorkspaceOptions
|
|
18
19
|
from segresnet_mlx.network import SegResNet
|
|
19
20
|
|
|
20
21
|
import numpy as np
|
|
@@ -169,7 +170,11 @@ class BundlePredictor:
|
|
|
169
170
|
return cast(Float32Array, host)
|
|
170
171
|
|
|
171
172
|
def predict_logits(
|
|
172
|
-
self,
|
|
173
|
+
self,
|
|
174
|
+
values: HostArray | Array,
|
|
175
|
+
configuration: SlidingWindowConfig,
|
|
176
|
+
*,
|
|
177
|
+
workspace: VolumeWorkspace | None = None,
|
|
173
178
|
) -> Float32Array:
|
|
174
179
|
"""Stitch preprocessed ``(N, C, D, H, W)`` float32 data into full-volume logits.
|
|
175
180
|
|
|
@@ -205,6 +210,7 @@ class BundlePredictor:
|
|
|
205
210
|
# MLX reduction kernels depend on layout. Match MONAI's sliced
|
|
206
211
|
# batch-one input without a second padded volume or patch copy.
|
|
207
212
|
preserve_window_strides=True,
|
|
213
|
+
workspace=workspace,
|
|
208
214
|
)
|
|
209
215
|
except BaseException as error:
|
|
210
216
|
if self.telemetry is not None:
|
|
@@ -233,7 +239,12 @@ def predict_logits(
|
|
|
233
239
|
|
|
234
240
|
|
|
235
241
|
def predict_wide_logits(
|
|
236
|
-
network: SegResNet,
|
|
242
|
+
network: SegResNet,
|
|
243
|
+
values: Float32Array,
|
|
244
|
+
configuration: SlidingWindowConfig,
|
|
245
|
+
*,
|
|
246
|
+
mx: MLX,
|
|
247
|
+
workspace: VolumeWorkspace | None = None,
|
|
237
248
|
) -> Float32Array:
|
|
238
249
|
"""Carry paired Metal logits into FP64 host blending, rounding only at export.
|
|
239
250
|
|
|
@@ -258,6 +269,7 @@ def predict_wide_logits(
|
|
|
258
269
|
del result
|
|
259
270
|
return output
|
|
260
271
|
|
|
272
|
+
workspace_options: WorkspaceOptions = {} if workspace is None else {"workspace": workspace}
|
|
261
273
|
try:
|
|
262
274
|
logits = sliding_window_inference(
|
|
263
275
|
values.astype(np.float64),
|
|
@@ -268,8 +280,9 @@ def predict_wide_logits(
|
|
|
268
280
|
mode=configuration.mode,
|
|
269
281
|
padding_mode=configuration.padding_mode,
|
|
270
282
|
predictor_input="numpy",
|
|
283
|
+
**workspace_options,
|
|
271
284
|
)
|
|
272
|
-
return logits.astype(np.float32)
|
|
285
|
+
return logits.astype(np.float32) if workspace is None else workspace.float32(logits)
|
|
273
286
|
finally:
|
|
274
287
|
# Reuse same-shaped Metal allocations across windows, including when
|
|
275
288
|
# paired arithmetic requires a wider host accumulator.
|
|
@@ -20,6 +20,8 @@ import numpy as np
|
|
|
20
20
|
|
|
21
21
|
|
|
22
22
|
class EventEngine(Protocol):
|
|
23
|
+
"""Structural type for the Ignite engine that receives telemetry handlers."""
|
|
24
|
+
|
|
23
25
|
def add_event_handler(self, event_name: Events, handler: Callable[..., None]) -> object:
|
|
24
26
|
"""Register a handler for an Ignite lifecycle event."""
|
|
25
27
|
...
|
|
@@ -11,7 +11,7 @@ as ``Spacing`` does and returned as float32.
|
|
|
11
11
|
|
|
12
12
|
Image arrays are channel-first float32 ``(C, *spatial)`` and affines are 4x4 float64
|
|
13
13
|
voxel-to-world matrices as stored in the NIfTI header. Bundle ``pixdim`` targets
|
|
14
|
-
are in mm;
|
|
14
|
+
are in mm; the runner validates spatial units, and these functions assume mm.
|
|
15
15
|
"""
|
|
16
16
|
|
|
17
17
|
from __future__ import annotations
|
|
@@ -333,7 +333,7 @@ def brats_labels(binary: Float32Array) -> npt.NDArray[np.int64]:
|
|
|
333
333
|
|
|
334
334
|
|
|
335
335
|
def write_nifti(path: Path, data: HostArray, affine: Float64Array, dtype: npt.DTypeLike) -> None:
|
|
336
|
-
"""Cast and write a NIfTI with sform and qform code 1.
|
|
336
|
+
"""Cast and write a millimetre NIfTI with sform and qform code 1.
|
|
337
337
|
|
|
338
338
|
The sform stores the affine within NIfTI precision. The qform cannot
|
|
339
339
|
represent shear, which NiBabel removes when encoding that form.
|
|
@@ -341,6 +341,7 @@ def write_nifti(path: Path, data: HostArray, affine: Float64Array, dtype: npt.DT
|
|
|
341
341
|
array = np.asarray(data).astype(dtype, copy=False)
|
|
342
342
|
matrix = to_affine_nd(3, affine)
|
|
343
343
|
image = nib.Nifti1Image(array, matrix)
|
|
344
|
+
image.header.set_xyzt_units("mm")
|
|
344
345
|
image.set_sform(matrix, code=1)
|
|
345
346
|
image.set_qform(matrix, code=1)
|
|
346
347
|
nib.save(image, str(path))
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|