medaugmentx 0.2.0__py3-none-any.whl

medaugmentx/__init__.py

@@ -0,0 +1,22 @@
+"""MedAugment — clinically-aware medical image augmentation.
+
+Public surface for Phase 2.
+"""
+from medaugmentx.core import (
+    Compose,
+    MedVolume,
+    OneOf,
+    SomeOf,
+    Transform,
+)
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "__version__",
+    "MedVolume",
+    "Transform",
+    "Compose",
+    "OneOf",
+    "SomeOf",
+]

medaugmentx/core/__init__.py

@@ -0,0 +1,16 @@
+"""Core data model and pipeline primitives."""
+from medaugmentx.core.base import Transform
+from medaugmentx.core.compose import Compose, OneOf, SomeOf
+from medaugmentx.core.utils import as_float32, derive_rng, resolve_rng
+from medaugmentx.core.volume import MedVolume
+
+__all__ = [
+    "MedVolume",
+    "Transform",
+    "Compose",
+    "OneOf",
+    "SomeOf",
+    "as_float32",
+    "derive_rng",
+    "resolve_rng",
+]

medaugmentx/core/base.py

@@ -0,0 +1,81 @@
+"""Abstract base class that every MedAugment transform inherits from."""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+import numpy as np
+
+from medaugmentx.core.utils import SeedLike, resolve_rng
+from medaugmentx.core.volume import MedVolume
+
+
+class Transform(ABC):
+    """Base class for all augmentations.
+
+    Subclasses must override :meth:`apply` and accept ``p`` and ``seed``
+    through ``super().__init__``. The base class handles probabilistic
+    application, seeding, and serialisation.
+
+    Probabilistic application is gated on ``self.rng`` so that two transforms
+    in the same :class:`Compose` with the same seed do not share a random
+    stream — see :func:`medaugmentx.core.utils.derive_rng`.
+
+    Example:
+
+        class MyShift(Transform):
+            def __init__(self, max_shift=0.1, p=1.0, seed=None):
+                super().__init__(p=p, seed=seed)
+                self.max_shift = max_shift
+
+            def apply(self, volume):
+                delta = self.rng.uniform(-self.max_shift, self.max_shift)
+                return volume.replace(image=volume.image + delta)
+    """
+
+    def __init__(self, p: float = 1.0, seed: SeedLike = None) -> None:
+        if not 0.0 <= float(p) <= 1.0:
+            raise ValueError(f"p must be in [0, 1], got {p}")
+        self.p: float = float(p)
+        # Store the seed for serialisation (int or None only; Generator can't round-trip).
+        self._seed: int | None = seed if isinstance(seed, (int, type(None))) else None
+        self.rng: np.random.Generator = resolve_rng(seed)
+
+    def __call__(self, volume: MedVolume) -> MedVolume:
+        if not isinstance(volume, MedVolume):
+            raise TypeError(f"Transform expects a MedVolume, got {type(volume).__name__}")
+        if self.p < 1.0 and self.rng.random() >= self.p:
+            return volume
+        return self.apply(volume)
+
+    @abstractmethod
+    def apply(self, volume: MedVolume) -> MedVolume:
+        """Perform the transform unconditionally — already past probability gate."""
+
+    def set_rng(self, rng: np.random.Generator) -> None:
+        """Reseed this transform with a specific :class:`numpy.random.Generator`.
+
+        Used by :class:`Compose` to give each child its own deterministic stream.
+        """
+        if not isinstance(rng, np.random.Generator):
+            raise TypeError("rng must be a numpy.random.Generator")
+        self.rng = rng
+
+    def __repr__(self) -> str:
+        attrs = ", ".join(
+            f"{k}={v!r}"
+            for k, v in self.__dict__.items()
+            if k != "rng" and not k.startswith("_")
+        )
+        return f"{self.__class__.__name__}({attrs})"
+
+    def to_dict(self) -> dict[str, Any]:
+        """Best-effort dictionary form of this transform's parameters.
+
+        Phase 1 ships only this introspection helper; full YAML serialisation
+        and round-tripping arrive in Phase 2.
+        """
+        params = {
+            k: v for k, v in self.__dict__.items() if k != "rng" and not k.startswith("_")
+        }
+        return {"name": self.__class__.__name__, "params": params}

medaugmentx/core/compose.py

@@ -0,0 +1,195 @@
+"""Pipeline builders: Compose, OneOf, SomeOf."""
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+import numpy as np
+
+from medaugmentx.core.base import Transform
+from medaugmentx.core.utils import SeedLike, derive_rng
+from medaugmentx.core.volume import MedVolume
+
+
+class Compose(Transform):
+    """Apply transforms sequentially.
+
+    All children share a deterministic seeding chain derived from the
+    top-level seed, so ``Compose([...], seed=42)`` produces the same output
+    every time, on every machine, for the same NumPy version.
+    """
+
+    def __init__(
+        self,
+        transforms: Iterable[Transform],
+        p: float = 1.0,
+        seed: SeedLike = None,
+    ) -> None:
+        super().__init__(p=p, seed=seed)
+        self.transforms: list[Transform] = list(transforms)
+        for t in self.transforms:
+            if not isinstance(t, Transform):
+                raise TypeError(
+                    f"Compose expected Transform instances, got {type(t).__name__}"
+                )
+        self._reseed_children()
+
+    def _reseed_children(self) -> None:
+        if not self.transforms:
+            return
+        for t, child_rng in zip(self.transforms, derive_rng(self.rng, len(self.transforms))):
+            t.set_rng(child_rng)
+
+    def set_rng(self, rng: np.random.Generator) -> None:
+        super().set_rng(rng)
+        self._reseed_children()
+
+    def apply(self, volume: MedVolume) -> MedVolume:
+        out = volume
+        for t in self.transforms:
+            out = t(out)
+        return out
+
+    def __len__(self) -> int:
+        return len(self.transforms)
+
+    def __iter__(self):
+        return iter(self.transforms)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.__class__.__name__,
+            "params": {
+                "transforms": [t.to_dict() for t in self.transforms],
+                "p": self.p,
+                "seed": self._seed,
+            },
+        }
+
+
+class OneOf(Transform):
+    """Pick exactly one child uniformly at random and apply it.
+
+    The container's ``p`` controls whether *any* child runs at all. When
+    weights are provided they are normalised; otherwise the choice is uniform.
+    """
+
+    def __init__(
+        self,
+        transforms: Sequence[Transform],
+        weights: Sequence[float] | None = None,
+        p: float = 1.0,
+        seed: SeedLike = None,
+    ) -> None:
+        super().__init__(p=p, seed=seed)
+        self.transforms: list[Transform] = list(transforms)
+        if not self.transforms:
+            raise ValueError("OneOf requires at least one transform")
+        for t in self.transforms:
+            if not isinstance(t, Transform):
+                raise TypeError(f"OneOf expected Transform, got {type(t).__name__}")
+
+        if weights is None:
+            self.weights = np.full(len(self.transforms), 1.0 / len(self.transforms))
+        else:
+            w = np.asarray(weights, dtype=np.float64)
+            if w.shape != (len(self.transforms),):
+                raise ValueError("weights length must match number of transforms")
+            if (w < 0).any() or w.sum() <= 0:
+                raise ValueError("weights must be non-negative and sum to > 0")
+            self.weights = w / w.sum()
+
+        self._reseed_children()
+
+    def _reseed_children(self) -> None:
+        for t, child_rng in zip(self.transforms, derive_rng(self.rng, len(self.transforms))):
+            t.set_rng(child_rng)
+
+    def set_rng(self, rng: np.random.Generator) -> None:
+        super().set_rng(rng)
+        self._reseed_children()
+
+    def apply(self, volume: MedVolume) -> MedVolume:
+        idx = int(self.rng.choice(len(self.transforms), p=self.weights))
+        # Force the chosen child to run regardless of its own ``p``.
+        return self.transforms[idx].apply(volume)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.__class__.__name__,
+            "params": {
+                "transforms": [t.to_dict() for t in self.transforms],
+                "weights": self.weights.tolist(),
+                "p": self.p,
+                "seed": self._seed,
+            },
+        }
+
+
+class SomeOf(Transform):
+    """Pick ``n`` children at random (without replacement) and apply them in order.
+
+    ``n`` may be an int or a ``(low, high)`` inclusive range — when a range,
+    a value is sampled per call.
+    """
+
+    def __init__(
+        self,
+        transforms: Sequence[Transform],
+        n: int | tuple[int, int] = 1,
+        p: float = 1.0,
+        seed: SeedLike = None,
+    ) -> None:
+        super().__init__(p=p, seed=seed)
+        self.transforms: list[Transform] = list(transforms)
+        if not self.transforms:
+            raise ValueError("SomeOf requires at least one transform")
+        for t in self.transforms:
+            if not isinstance(t, Transform):
+                raise TypeError(f"SomeOf expected Transform, got {type(t).__name__}")
+
+        if isinstance(n, int):
+            lo, hi = n, n
+        else:
+            lo, hi = int(n[0]), int(n[1])
+        if not 0 <= lo <= hi <= len(self.transforms):
+            raise ValueError(f"n={n} invalid for {len(self.transforms)} transforms")
+        self.n_range: tuple[int, int] = (lo, hi)
+
+        self._reseed_children()
+
+    def _reseed_children(self) -> None:
+        for t, child_rng in zip(self.transforms, derive_rng(self.rng, len(self.transforms))):
+            t.set_rng(child_rng)
+
+    def set_rng(self, rng: np.random.Generator) -> None:
+        super().set_rng(rng)
+        self._reseed_children()
+
+    def apply(self, volume: MedVolume) -> MedVolume:
+        lo, hi = self.n_range
+        n = int(self.rng.integers(lo, hi + 1))
+        if n == 0:
+            return volume
+        idxs = self.rng.choice(len(self.transforms), size=n, replace=False)
+        idxs.sort()
+        out = volume
+        for i in idxs:
+            out = self.transforms[int(i)].apply(out)
+        return out
+
+    def to_dict(self) -> dict[str, Any]:
+        lo, hi = self.n_range
+        n: Any = lo if lo == hi else list(self.n_range)
+        return {
+            "name": self.__class__.__name__,
+            "params": {
+                "transforms": [t.to_dict() for t in self.transforms],
+                "n": n,
+                "p": self.p,
+                "seed": self._seed,
+            },
+        }
+
+
+__all__ = ["Compose", "OneOf", "SomeOf"]

medaugmentx/core/utils.py

@@ -0,0 +1,87 @@
+"""Small helpers shared across the library."""
+from __future__ import annotations
+
+from typing import Union
+
+import numpy as np
+
+SeedLike = Union[int, np.random.Generator, None]
+
+
+def resolve_rng(seed: SeedLike) -> np.random.Generator:
+    """Return a ``numpy.random.Generator`` from any accepted seed input.
+
+    Accepts ``None`` (fresh entropy), an ``int`` seed, or an existing
+    ``Generator`` (returned as-is for chaining).
+    """
+    if isinstance(seed, np.random.Generator):
+        return seed
+    return np.random.default_rng(seed)
+
+
+def derive_rng(rng: np.random.Generator, n: int) -> list[np.random.Generator]:
+    """Spawn ``n`` independent generators from ``rng`` deterministically.
+
+    Used by :class:`~medaugmentx.core.compose.Compose` to give each child
+    transform its own stream while keeping the whole pipeline reproducible
+    from a single top-level seed.
+    """
+    seeds = rng.integers(0, np.iinfo(np.uint64).max, size=n, dtype=np.uint64, endpoint=False)
+    return [np.random.default_rng(int(s)) for s in seeds]
+
+
+def as_float32(image: np.ndarray) -> np.ndarray:
+    """Cast to ``float32`` only when needed; cheap no-op otherwise."""
+    if image.dtype == np.float32:
+        return image
+    return image.astype(np.float32, copy=False)
+
+
+def normalize_axes(axes: int | tuple | list | None, ndim: int) -> tuple:
+    """Normalise an ``axes`` argument to a sorted tuple of non-negative ints.
+
+    ``None`` expands to all axes. Negative axes are wrapped relative to ndim.
+    """
+    if axes is None:
+        return tuple(range(ndim))
+    if isinstance(axes, int):
+        axes = (axes,)
+    out: list[int] = []
+    for a in axes:
+        ax = int(a)
+        if ax < 0:
+            ax += ndim
+        if not 0 <= ax < ndim:
+            raise ValueError(f"axis {a} out of range for ndim={ndim}")
+        out.append(ax)
+    return tuple(sorted(set(out)))
+
+
+def axis_label_to_index(label: str, ndim: int) -> int:
+    """Map a friendly axis label (``"x"``, ``"y"``, ``"z"``) to a NumPy axis.
+
+    Convention used throughout the library:
+
+    - 3D arrays are stored as ``(D, H, W)`` — i.e. ``(z, y, x)``.
+    - 2D arrays are stored as ``(H, W)`` — i.e. ``(y, x)``.
+
+    So for 3D ``"z"`` -> 0, ``"y"`` -> 1, ``"x"`` -> 2; for 2D ``"y"`` -> 0,
+    ``"x"`` -> 1. ``"z"`` is invalid for 2D arrays.
+    """
+    label = label.lower()
+    if ndim == 3:
+        mapping = {"z": 0, "y": 1, "x": 2}
+    elif ndim == 2:
+        mapping = {"y": 0, "x": 1}
+    else:
+        raise ValueError(f"Only 2D or 3D supported, got ndim={ndim}")
+    if label not in mapping:
+        raise ValueError(f"Unknown axis label {label!r} for ndim={ndim}")
+    return mapping[label]
+
+
+def clip_intensity(image: np.ndarray, lo: float | None = None, hi: float | None = None) -> np.ndarray:
+    """Clip in-place if writeable, otherwise return a clipped copy."""
+    if lo is None and hi is None:
+        return image
+    return np.clip(image, lo, hi)

medaugmentx/core/volume.py

@@ -0,0 +1,117 @@
+"""The MedVolume container — image + optional mask + spacing + metadata."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field, replace
+from typing import Any
+
+import numpy as np
+
+
+@dataclass
+class MedVolume:
+    """A single medical image (2D or 3D) with optional segmentation mask.
+
+    All transforms in MedAugment operate on this container so that masks and
+    metadata stay in lockstep with the image array.
+
+    Attributes:
+        image: 2D ``(H, W)`` or 3D ``(D, H, W)`` array. Recommended dtype is
+            ``float32``; integer inputs are accepted but will be cast where
+            arithmetic is required.
+        mask: Optional integer label map with the same shape as ``image``.
+        spacing: Voxel size in millimetres, one entry per spatial axis.
+            For 3D volumes this is ``(slice_thickness, row_mm, col_mm)``.
+        metadata: Free-form dictionary. Conventional keys: ``modality``
+            (``"MR" | "CT" | "DX" | "DBT"``), ``vendor``, ``patient_id``,
+            ``original_dtype``.
+    """
+
+    image: np.ndarray
+    mask: np.ndarray | None = None
+    spacing: tuple[float, ...] = ()
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.image, np.ndarray):
+            raise TypeError(f"image must be a numpy.ndarray, got {type(self.image).__name__}")
+        if self.image.ndim not in (2, 3):
+            raise ValueError(f"image must be 2D or 3D; got shape {self.image.shape}")
+
+        if self.mask is not None:
+            if not isinstance(self.mask, np.ndarray):
+                raise TypeError("mask must be a numpy.ndarray or None")
+            if self.mask.shape != self.image.shape:
+                raise ValueError(
+                    f"mask shape {self.mask.shape} does not match image shape {self.image.shape}"
+                )
+
+        if self.spacing:
+            if len(self.spacing) != self.image.ndim:
+                raise ValueError(
+                    f"spacing has {len(self.spacing)} entries but image is {self.image.ndim}D"
+                )
+            self.spacing = tuple(float(s) for s in self.spacing)
+        else:
+            self.spacing = tuple(1.0 for _ in range(self.image.ndim))
+
+        if not isinstance(self.metadata, dict):
+            raise TypeError("metadata must be a dict")
+
+    @property
+    def ndim(self) -> int:
+        return int(self.image.ndim)
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        return tuple(self.image.shape)
+
+    @property
+    def is_3d(self) -> bool:
+        return self.image.ndim == 3
+
+    @property
+    def has_mask(self) -> bool:
+        return self.mask is not None
+
+    @property
+    def modality(self) -> str | None:
+        return self.metadata.get("modality")
+
+    def replace(
+        self,
+        *,
+        image: np.ndarray | None = None,
+        mask: np.ndarray | None = None,
+        spacing: tuple[float, ...] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> MedVolume:
+        """Return a new MedVolume with selected fields swapped out.
+
+        Use ``mask=...`` only to provide a new mask; pass ``mask=None`` and
+        rely on the existing one by omitting the keyword. Metadata is shallow-
+        copied to avoid silent aliasing across volumes.
+        """
+        return replace(
+            self,
+            image=self.image if image is None else image,
+            mask=self.mask if mask is None else mask,
+            spacing=self.spacing if spacing is None else tuple(float(s) for s in spacing),
+            metadata=dict(self.metadata if metadata is None else metadata),
+        )
+
+    def copy(self) -> MedVolume:
+        """Deep copy of the underlying arrays and metadata."""
+        return MedVolume(
+            image=self.image.copy(),
+            mask=None if self.mask is None else self.mask.copy(),
+            spacing=tuple(self.spacing),
+            metadata=dict(self.metadata),
+        )
+
+    def __repr__(self) -> str:
+        mask_repr = "None" if self.mask is None else f"shape={self.mask.shape}, dtype={self.mask.dtype}"
+        return (
+            f"MedVolume(image=shape={self.image.shape}, dtype={self.image.dtype}, "
+            f"mask={mask_repr}, spacing={self.spacing}, "
+            f"modality={self.modality!r})"
+        )

medaugmentx/io/__init__.py

@@ -0,0 +1,18 @@
+"""Unified I/O for medical image formats.
+
+Each loader returns a :class:`~medaugmentx.core.volume.MedVolume` with
+``spacing`` populated in millimetres and ``metadata`` carrying the
+modality and any vendor-specific information that callers may need.
+
+Optional dependencies:
+
+- DICOM I/O requires ``pydicom`` (``pip install medaugmentx[dicom]``).
+- NIfTI I/O requires ``nibabel`` (``pip install medaugmentx[nifti]``).
+
+If a backend is missing the loader raises a clear :class:`ImportError`
+when called, not at import time.
+"""
+from medaugmentx.io.dicom import load_dicom_series
+from medaugmentx.io.nifti import load_nifti, save_nifti
+
+__all__ = ["load_dicom_series", "load_nifti", "save_nifti"]