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

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: segresnet-mlx
-Version: 0.2.2
+Version: 0.2.3
 Summary: Native MLX SegResNet with NumPy sliding-window logit inference
 Project-URL: Homepage, https://github.com/MedMLX/segresnet-mlx
 Project-URL: Documentation, https://github.com/MedMLX/segresnet-mlx/blob/main/docs/technical.md
@@ -59,7 +59,10 @@ from huggingface_hub import hf_hub_download
 from segresnet_mlx.medmlx import run
 
 bundle = "brats_mri_segmentation"
-weights = hf_hub_download("MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors")
+weights = hf_hub_download(
+    "MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors",
+    revision="14f9cda3649965729663c7365c49fd839c81a01b",
+)
 
 result = run(
     bundle=bundle,

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/README.md

@@ -30,7 +30,10 @@ from huggingface_hub import hf_hub_download
 from segresnet_mlx.medmlx import run
 
 bundle = "brats_mri_segmentation"
-weights = hf_hub_download("MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors")
+weights = hf_hub_download(
+    "MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors",
+    revision="14f9cda3649965729663c7365c49fd839c81a01b",
+)
 
 result = run(
     bundle=bundle,

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "segresnet-mlx"
-version = "0.2.2"
+version = "0.2.3"
 description = "Native MLX SegResNet with NumPy sliding-window logit inference"
 readme = "README.md"
 requires-python = ">=3.13,<3.15"

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/src/segresnet_mlx/__init__.py

@@ -21,7 +21,7 @@ if TYPE_CHECKING:
     from segresnet_mlx.weights import load_model as load_model
     from segresnet_mlx.weights import required_shapes as required_shapes
 
-__version__ = "0.2.2"
+__version__ = "0.2.3"
 
 _EXPORTS = {
     "BUNDLE_CONFIGS": "segresnet_mlx.predictor",

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/src/segresnet_mlx/medmlx.py

@@ -292,16 +292,23 @@ def run(
     ``brats_mri_segmentation`` (``t1c_path``, ``t1_path``, ``t2_path``,
     ``flair_path``: co-registered MRI on one grid), or ``wholeBody_ct_segmentation``
     / ``wholeBody_ct_segmentation_lowres`` (``image_path``: one CT volume, segmented
-    at 1.5 mm or 3 mm respectively). ``weights_path`` is the bundle's
-    ``models/model.pt`` (or ``models/model_lowres.pt`` for the lowres whole-body
-    variant) converted with :func:`segresnet_mlx.convert_checkpoint`.
+    at 1.5 mm or 3 mm respectively). ``weights_path`` is a native ``.safetensors``
+    file: ``segresnet-brats.safetensors`` from ``MedMLX/segresnet-brats-mlx`` for
+    BraTS, or ``segresnet-renal-cect.safetensors`` from
+    ``MedMLX/segresnet-renal-cect-mlx`` for renal CECT. These published files need
+    no conversion. For other checkpoints, convert the original ``models/model.pt``
+    (``models/model_lowres.pt`` for lowres whole-body) with
+    :func:`segresnet_mlx.convert_checkpoint` and pass the resulting file.
 
     Inputs must be finite 3D NIfTI volumes; multi-input bundles need identical shape
     and affine (atol 1e-5). Writes ``<bundle>_mask.nii.gz`` in ``output_dir`` on the
     first input's shape and reader-selected affine (see ``Volume``): renal labels 0-5
     and whole-body labels 0-104 as float32, BraTS labels 0, 1, 2, 4 as uint8, as the
     bundles' ``SaveImaged`` wrote them. An existing mask is replaced only when
-    ``overwrite`` is true; output aliases of input images are always refused.
+    ``overwrite`` is true; output aliases of input images are refused by a
+    preflight check. Spatial affine values are assumed to be in millimeters; the
+    NIfTI spatial-unit field is not checked or converted. Paths are not locked
+    against concurrent replacement after the alias check.
 
     ``arithmetic_profile="compensated-fp32"`` uses Metal float pairs for
     convolution and GroupNorm accumulation, with float32 layer boundaries.
@@ -312,6 +319,27 @@ def run(
     requires transpose-convolution upsampling and compares with the unchanged
     original graph evaluated in float64 throughout its network and inferer.
 
+    Args:
+        bundle: ``brats_mri_segmentation``, ``renalStructures_CECT_segmentation``,
+            ``wholeBody_ct_segmentation`` or ``wholeBody_ct_segmentation_lowres``.
+        weights_path: Path to the matching native ``.safetensors`` file, such as
+            ``segresnet-brats.safetensors`` or ``segresnet-renal-cect.safetensors``.
+            Hugging Face cache file paths are accepted; directories and Torch ``.pt``
+            checkpoints are not.
+        output_dir: Directory for ``<bundle>_mask.nii.gz``; created if missing.
+        arterial_path: Renal arterial-phase 3D NIfTI; required for the renal bundle.
+        venous_path: Renal venous-phase 3D NIfTI on the arterial grid; required for renal.
+        excretory_path: Renal excretory-phase 3D NIfTI on the arterial grid; required
+            for renal.
+        t1c_path: BraTS contrast-enhanced T1 3D NIfTI; required for BraTS.
+        t1_path: BraTS T1 3D NIfTI on the T1c grid; required for BraTS.
+        t2_path: BraTS T2 3D NIfTI on the T1c grid; required for BraTS.
+        flair_path: BraTS FLAIR 3D NIfTI on the T1c grid; required for BraTS.
+        image_path: Single 3D CT NIfTI; required for either whole-body bundle.
+        overwrite: Replace an existing output mask when true; defaults to false.
+        arithmetic_profile: ``torch-order-fp32`` (default), ``compensated-fp32``,
+            or ``compensated-fp48``; the last supports only the renal bundle.
+
     Returns:
         ``{"outputs": {"mask": path}, "details": {...}}``; ``details`` is described
         by ``segresnet_mlx._records.Details`` and always has
@@ -534,7 +562,9 @@ def postprocess_wholebody(logits: Float32Array, inverse: SpatialInverse) -> Floa
 
 
 def _weights_file(weights_path: str) -> Path:
-    path = Path(weights_path).expanduser().resolve()
+    # MLX identifies the file format from the supplied name. Hub cache paths are
+    # named .safetensors symlinks whose underlying content-addressed blobs have no suffix.
+    path = Path(weights_path).expanduser().absolute()
     hint = (
         "Convert the bundle's models/model.pt once with "
         "segresnet_mlx.convert_checkpoint(model_pt, destination, spec=...) "
@@ -556,8 +586,9 @@ def _weights_file(weights_path: str) -> Path:
 def load_inputs(roles: list[tuple[str, str | None]]) -> list[Volume]:
     """Load each ``(role, path)`` and require one shared voxel grid.
 
-    Every volume must match the first one's shape and affine (``atol=1e-5``); the
-    runner does not resample, so a mismatch raises ``InvalidInputError``.
+    Every volume must match the first one's shape and affine (``atol=1e-5``).
+    This check does not align mismatched inputs; bundle preprocessing may still
+    resample the admitted grid afterward.
     """
     volumes = [load_volume(Path(str(path)).expanduser().resolve(), role) for role, path in roles]
     reference = volumes[0]
@@ -573,7 +604,12 @@ def load_inputs(roles: list[tuple[str, str | None]]) -> list[Volume]:
 
 
 def load_volume(path: Path, role: str) -> Volume:
-    """Read finite 3D voxels using the affine selected by the upstream nibabel reader."""
+    """Read 3D voxels using the affine selected by the upstream nibabel reader.
+
+    Source voxels are checked for finiteness before conversion to float32; values
+    outside the float32 range can overflow during that conversion. Spatial units
+    and affine invertibility are not checked here.
+    """
     from segresnet_mlx._nifti import ImageFileError, nib
 
     if not path.is_file():

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/src/segresnet_mlx/network.py

@@ -21,7 +21,8 @@ def group_mean(values: Array, mx: MLX) -> Array:
     """Mean over the last axis of ``(N, G, S)`` values, summed in 256-element tiles.
 
     Tiling bounds the float32 accumulation depth of a group that spans a whole
-    volume. Returns ``(N, G, 1)``. ``SegResNet`` does not call this (its GroupNorm
+    volume. The last dimension must be nonempty. Returns ``(N, G, 1)``. ``SegResNet``
+    does not call this (its GroupNorm
     comes from the ``medmlx_core`` operators); it is exercised by
     ``tests/test_equivalence.py``.
     """
@@ -34,16 +35,20 @@ def group_mean(values: Array, mx: MLX) -> Array:
 class SegResNet:
     """Raw-logit forward pass for MONAI's 3-D, group-normalized SegResNet.
 
-    Tensors are channels-first ``(N, C, D, H, W)`` float32 at every layer boundary.
+    Tensors are channels-first ``(N, C, D, H, W)``; the supported inference
+    contract uses float32 inputs and returned logits. This class checks rank,
+    channel count and spatial divisibility but does not itself check input dtype
+    or finiteness; the public patch and sliding-window helpers do.
     ``weights`` uses MONAI's state-dict keys and Torch parameter layouts
     (convolution weights are ``(out, in, kD, kH, kW)``); ``required_shapes`` in
     ``segresnet_mlx.weights`` lists them. Inference only: dropout is omitted and
     no activation is applied to the output.
 
-    ``arithmetic_profile`` selects the accumulation scheme (see the README):
+    ``arithmetic_profile`` selects the accumulation scheme (see ``docs/technical.md``):
     ``"torch-order-fp32"`` (default), ``"compensated-fp32"``, or
     ``"compensated-fp48"``, which needs ``upsample_mode="deconv"`` and raises
-    ``ValueError`` otherwise.
+    ``ValueError`` otherwise. The first two profiles use float32 layer boundaries;
+    ``compensated-fp48`` carries paired float32 values through the hidden layers.
     """
 
     def __init__(

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/src/segresnet_mlx/predictor.py

@@ -50,8 +50,9 @@ class SlidingWindowConfig:
     def from_bundle_config(cls, config: InfererConfig | BundleConfig) -> SlidingWindowConfig:
         """Read the staged JSON's inferer section without evaluating bundle expressions.
 
-        Missing keys fall back to MONAI's defaults. Values are taken literally, so a
-        reference such as ``@displayable_configs#sw_overlap`` is not resolved.
+        ``roi_size`` is required; the other listed keys default to MONAI's
+        values. Values are taken literally, so a reference such as
+        ``@displayable_configs#sw_overlap`` is not resolved.
         """
         inferer = cast(InfererConfig, config.get("inferer", config))
         return cls(

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/src/segresnet_mlx/topology.py

@@ -6,7 +6,7 @@ import math
 from dataclasses import dataclass
 from typing import Literal
 
-# Accumulation scheme for convolution and GroupNorm; see the README for the three
+# Accumulation scheme for convolution and GroupNorm; see docs/technical.md for the three
 # profiles. ``compensated-fp48`` additionally requires transpose-convolution upsampling.
 type ArithmeticProfile = Literal["torch-order-fp32", "compensated-fp32", "compensated-fp48"]
 

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/src/segresnet_mlx/transforms.py

@@ -10,7 +10,8 @@ bundle's ``MergeClassesd`` and the NIfTI writer. Spatial resampling follows MONA
 as ``Spacing`` does and returned as float32.
 
 Image arrays are channel-first float32 ``(C, *spatial)`` and affines are 4x4 float64
-voxel-to-world matrices as stored in the NIfTI header; ``pixdim`` values are in mm.
+voxel-to-world matrices as stored in the NIfTI header. Bundle ``pixdim`` targets
+are in mm; these functions do not inspect or convert the input's spatial units.
 """
 
 from __future__ import annotations
@@ -373,7 +374,11 @@ def brats_labels(binary: Float32Array) -> npt.NDArray[np.int64]:
 
 
 def write_nifti(path: Path, data: HostArray, affine: Float64Array, dtype: npt.DTypeLike) -> None:
-    """MONAI ``NibabelWriter``: cast, then set sform and qform (code 1) from the affine."""
+    """Cast and write a NIfTI with sform and qform code 1.
+
+    The sform stores the affine within NIfTI precision. The qform cannot
+    represent shear, which NiBabel removes when encoding that form.
+    """
     array = np.asarray(data).astype(dtype, copy=False)
     matrix = to_affine_nd(3, affine)
     image = nib.Nifti1Image(array, matrix)

{segresnet_mlx-0.2.2 → segresnet_mlx-0.2.3}/src/segresnet_mlx/weights.py

@@ -78,6 +78,8 @@ def convert_state_dict(
 
     ``state`` is either a raw state dict or a checkpoint holding it under the
     ``"model"`` key. Arrays keep their Torch layout and are returned C-contiguous.
+    Validation checks the parameter inventory and values, not the checkpoint's
+    identity or origin; other checkpoints with the same topology are admissible.
 
     Raises:
         InvalidInputError: If a required parameter is missing, an unsupported one is
@@ -117,11 +119,26 @@ def convert_checkpoint(checkpoint_path: Path, destination: Path, *, spec: SegRes
 
     The checkpoint is read with the shared loader (``torch.load`` with its default
     ``weights_only=True``, so Torch from the ``conversion`` extra is required) and
-    checked by ``convert_state_dict``. Parent directories of ``destination`` are
-    created. MLX is also needed to write the file, so this runs only on a supported
-    Apple Silicon host.
+    checked by ``convert_state_dict``. ``destination`` must end in
+    ``.safetensors``. Parent directories of ``destination`` are
+    created. An existing destination is replaced, except when it aliases the
+    source checkpoint. Alias checks occur before conversion and do not lock out
+    subsequent filesystem changes. MLX is also needed to write the file, so this
+    runs only on a supported Apple Silicon host.
+
+    Raises:
+        InvalidInputError: The destination has a different suffix or aliases the
+            source checkpoint, or checkpoint contents fail ``convert_state_dict``.
     """
 
+    if destination.suffix != ".safetensors":
+        raise InvalidInputError("Converted weights destination must end in .safetensors")
+    if checkpoint_path.resolve() == destination.resolve() or (
+        checkpoint_path.is_file()
+        and destination.is_file()
+        and checkpoint_path.samefile(destination)
+    ):
+        raise InvalidInputError("Converted weights must not overwrite their source checkpoint")
     mx = import_mlx()
     converted = convert_state_dict(load_torch_checkpoint(checkpoint_path), spec)
     destination.parent.mkdir(parents=True, exist_ok=True)
@@ -141,8 +158,9 @@ def load_model(
     """Load converted float32 safetensors into a ``SegResNet`` for ``spec``.
 
     Raises:
-        InvalidInputError: If the file is not a mapping, its keys or shapes differ
+        InvalidInputError: If the loaded value is not a mapping, its keys or shapes differ
             from ``required_shapes(spec)``, or a parameter is not finite float32.
+        RuntimeError: If MLX cannot open or decode the weights file.
     """
     mx = import_mlx()
     raw: object = mx.load(str(weights_path))