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.
Files changed (21) hide show
  1. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/PKG-INFO +10 -5
  2. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/README.md +9 -4
  3. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/pyproject.toml +1 -1
  4. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/__init__.py +1 -1
  5. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_nifti.py +12 -0
  6. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_records.py +1 -0
  7. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_types.py +7 -1
  8. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/medmlx.py +123 -17
  9. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/predictor.py +17 -4
  10. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/telemetry.py +2 -0
  11. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/transforms.py +3 -2
  12. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/.gitignore +0 -0
  13. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/LICENSE +0 -0
  14. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/THIRD_PARTY_NOTICES.md +0 -0
  15. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_arrays.py +0 -0
  16. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/_runtime.py +0 -0
  17. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/inference.py +0 -0
  18. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/network.py +0 -0
  19. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/py.typed +0 -0
  20. {segresnet_mlx-0.3.0 → segresnet_mlx-0.3.1}/src/segresnet_mlx/topology.py +0 -0
  21. {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.0
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
- ## Quick start
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>=0.3.0"
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; license and third-party notices are included in the distribution. Model weights retain their publisher terms.
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
- ## Quick start
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>=0.3.0"
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; license and third-party notices are included in the distribution. Model weights retain their publisher terms.
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.
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "segresnet-mlx"
7
- version = "0.3.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"
@@ -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.0"
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, ...]
@@ -60,6 +60,7 @@ class Details(TypedDict):
60
60
  weights_revision: NotRequired[str | None]
61
61
  bundle: str
62
62
  arithmetic_profile: ArithmeticProfile
63
+ streaming: NotRequired[bool]
63
64
  labels: dict[str, str]
64
65
  labels_present: list[int]
65
66
  input_shape: list[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
- logits, runtime = predict(network, image, config)
423
- predicted = time.perf_counter()
424
- if bundle == "renalStructures_CECT_segmentation":
425
- assert inverse is not None
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
- mask = postprocess_brats(logits[0])
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-mlx")
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-mlx")
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
- Source voxels are checked for finiteness before conversion to float32; values
682
- outside the float32 range can overflow during that conversion. Spatial units
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, image: Float32Array, config: Bundle
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, values: HostArray | Array, configuration: SlidingWindowConfig
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, values: Float32Array, configuration: SlidingWindowConfig, *, mx: MLX
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; these functions do not inspect or convert the input's spatial units.
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