Redesign batch construction and schema

src/torchio/data/batch_schema.py

@@ -0,0 +1,217 @@
+"""Schemas used to validate image and subject batches."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from .bboxes import BoundingBoxes
+from .image import Image
+from .points import Points
+from .subject import Subject
+
+
+@dataclass(frozen=True)
+class _AnnotationSchema:
+    """Describe one named annotation field."""
+
+    value_type: type[Points] | type[BoundingBoxes]
+    metadata_keys: tuple[str, ...]
+
+    @classmethod
+    def from_value(cls, value: Points | BoundingBoxes) -> _AnnotationSchema:
+        """Build a schema from one annotation value."""
+        return cls(type(value), tuple(value.metadata))
+
+    def validate(
+        self,
+        value: Points | BoundingBoxes,
+        *,
+        index: int,
+        context: str,
+    ) -> None:
+        """Validate one annotation against this schema."""
+        if type(value) is not self.value_type:
+            msg = (
+                f"{context} at index {index} has type {type(value).__name__},"
+                f" expected {self.value_type.__name__}"
+            )
+            raise ValueError(msg)
+        _validate_keys(
+            self.metadata_keys,
+            value.metadata,
+            index=index,
+            context=f"{context} metadata",
+        )
+
+
+@dataclass(frozen=True)
+class _ImageSchema:
+    """Describe one named image field."""
+
+    value_type: type[Image]
+    shape: tuple[int, ...]
+    dtype: str
+    device: torch.device
+    metadata_keys: tuple[str, ...]
+    points: dict[str, _AnnotationSchema]
+    bounding_boxes: dict[str, _AnnotationSchema]
+
+    @classmethod
+    def from_image(cls, image: Image) -> _ImageSchema:
+        """Build a schema from one image."""
+        return cls(
+            value_type=type(image),
+            shape=tuple(image.shape),
+            dtype=_normalize_dtype(image.dtype),
+            device=image.device,
+            metadata_keys=tuple(image.metadata),
+            points={
+                name: _AnnotationSchema.from_value(value)
+                for name, value in image.points.items()
+            },
+            bounding_boxes={
+                name: _AnnotationSchema.from_value(value)
+                for name, value in image.bounding_boxes.items()
+            },
+        )
+
+    def validate(self, image: Image, *, index: int, name: str) -> None:
+        """Validate one image against this schema."""
+        context = f"Image {name!r}"
+        if type(image) is not self.value_type:
+            msg = (
+                f"{context} at index {index} has type {type(image).__name__},"
+                f" expected {self.value_type.__name__}"
+            )
+            raise ValueError(msg)
+        for attribute in ("shape", "device"):
+            expected = getattr(self, attribute)
+            actual = getattr(image, attribute)
+            if actual != expected:
+                msg = (
+                    f"{context} at index {index} has {attribute} {actual},"
+                    f" expected {expected}"
+                )
+                raise ValueError(msg)
+        actual_dtype = _normalize_dtype(image.dtype)
+        if actual_dtype != self.dtype:
+            msg = (
+                f"{context} at index {index} has dtype {actual_dtype},"
+                f" expected {self.dtype}"
+            )
+            raise ValueError(msg)
+        _validate_keys(
+            self.metadata_keys,
+            image.metadata,
+            index=index,
+            context=f"{context} metadata",
+        )
+        _validate_annotations(
+            self.points,
+            image.points,
+            index=index,
+            context=f"{context} points",
+        )
+        _validate_annotations(
+            self.bounding_boxes,
+            image.bounding_boxes,
+            index=index,
+            context=f"{context} bounding boxes",
+        )
+
+
+@dataclass(frozen=True)
+class _SubjectSchema:
+    """Describe the fields shared by all subjects in a batch."""
+
+    images: dict[str, _ImageSchema]
+    metadata_keys: tuple[str, ...]
+    points: dict[str, _AnnotationSchema]
+    bounding_boxes: dict[str, _AnnotationSchema]
+
+    @classmethod
+    def from_subject(cls, subject: Subject) -> _SubjectSchema:
+        """Build a schema from the first subject in a batch."""
+        return cls(
+            images={
+                name: _ImageSchema.from_image(image)
+                for name, image in subject.images.items()
+            },
+            metadata_keys=tuple(subject.metadata),
+            points={
+                name: _AnnotationSchema.from_value(value)
+                for name, value in subject.points.items()
+            },
+            bounding_boxes={
+                name: _AnnotationSchema.from_value(value)
+                for name, value in subject.bounding_boxes.items()
+            },
+        )
+
+    def validate(self, subject: Subject, *, index: int) -> None:
+        """Validate one subject against this schema."""
+        _validate_keys(
+            self.images, subject.images, index=index, context="Subject images"
+        )
+        _validate_keys(
+            self.metadata_keys,
+            subject.metadata,
+            index=index,
+            context="Subject metadata",
+        )
+        _validate_annotations(
+            self.points,
+            subject.points,
+            index=index,
+            context="Subject points",
+        )
+        _validate_annotations(
+            self.bounding_boxes,
+            subject.bounding_boxes,
+            index=index,
+            context="Subject bounding boxes",
+        )
+        for name, schema in self.images.items():
+            schema.validate(subject.images[name], index=index, name=name)
+
+
+def _validate_annotations(
+    reference: dict[str, _AnnotationSchema],
+    current: dict[str, Points] | dict[str, BoundingBoxes],
+    *,
+    index: int,
+    context: str,
+) -> None:
+    """Validate a named annotation store."""
+    _validate_keys(reference, current, index=index, context=context)
+    for name, schema in reference.items():
+        schema.validate(current[name], index=index, context=f"{context} {name!r}")
+
+
+def _validate_keys(
+    reference: Any,
+    current: Any,
+    *,
+    index: int,
+    context: str,
+) -> None:
+    """Validate equivalent key sets while allowing reordered keys."""
+    reference_set = set(reference)
+    current_set = set(current)
+    if reference_set == current_set:
+        return
+    missing = sorted(reference_set - current_set)
+    unexpected = sorted(current_set - reference_set)
+    msg = (
+        f"{context} at index {index} has incompatible keys:"
+        f" missing {missing}, unexpected {unexpected}"
+    )
+    raise ValueError(msg)
+
+
+def _normalize_dtype(dtype: Any) -> str:
+    """Return one comparable dtype name for Torch and NumPy dtypes."""
+    return str(dtype).removeprefix("torch.")

src/torchio/data/bboxes.py

@@ -310,14 +310,19 @@ def device(self) -> torch.device:
         return self._data.device
 
     def to(self, *args: Any, **kwargs: Any) -> Self:
-        """Move bounding box data to a device and/or cast to a dtype.
+        """Move bounding box data and affine to a device or dtype.
+
+        Args:
+            *args: Positional arguments forwarded to `torch.Tensor.to`.
+            **kwargs: Keyword arguments forwarded to `torch.Tensor.to`.
 
         Returns:
             `self` (modified in-place).
         """
         self._data = self._data.to(*args, **kwargs)
         if self._labels is not None:
-            self._labels = self._labels.to(*args, **kwargs)
+            self._labels = self._labels.to(device=self._data.device)
+        self._affine.to(*args, **kwargs)
         return self
 
     # --- Methods ---

src/torchio/data/image.py

@@ -642,16 +642,24 @@ def device(self) -> torch.device:
         return self.data.device
 
     def to(self, *args: Any, **kwargs: Any) -> Self:
-        """Move image data and affine to a device and/or cast to a dtype.
+        """Move image data, affine, and annotations to a device or dtype.
 
         Accepts the same arguments as `torch.Tensor.to()`.
 
+        Args:
+            *args: Positional arguments forwarded to `torch.Tensor.to`.
+            **kwargs: Keyword arguments forwarded to `torch.Tensor.to`.
+
         Returns:
             `self` (modified in-place).
         """
         self._data = self.data.to(*args, **kwargs)
         if self._affine is not None:
             self._affine.to(*args, **kwargs)
+        for points in self._points.values():
+            points.to(*args, **kwargs)
+        for boxes in self._bounding_boxes.values():
+            boxes.to(*args, **kwargs)
         self._refresh_backend_from_data()
         return self
 

src/torchio/data/points.py

@@ -108,12 +108,17 @@ def device(self) -> torch.device:
         return self._data.device
 
     def to(self, *args: Any, **kwargs: Any) -> Self:
-        """Move point data to a device and/or cast to a dtype.
+        """Move point data and affine to a device or dtype.
+
+        Args:
+            *args: Positional arguments forwarded to `torch.Tensor.to`.
+            **kwargs: Keyword arguments forwarded to `torch.Tensor.to`.
 
         Returns:
             `self` (modified in-place).
         """
         self._data = self._data.to(*args, **kwargs)
+        self._affine.to(*args, **kwargs)
         return self
 
     # --- Methods ---

src/torchio/transforms/intensity/labels_to_image.py

@@ -10,7 +10,6 @@
 from torch import Tensor
 
 from ...data.batch import SubjectsBatch
-from ...data.image import LabelMap
 from ...data.image import ScalarImage
 from ..parameter_range import to_range
 from ..transform import Transform
@@ -173,7 +172,7 @@ def _find_label_batch(self, batch: SubjectsBatch) -> Any:
             return batch.images[self.label_key]
         # Auto-detect first LabelMap.
         for _name, img_batch in batch.images.items():
-            if issubclass(img_batch._image_class, LabelMap):
+            if img_batch.is_label:
                 return img_batch
         msg = "No LabelMap found in the subject"
         raise KeyError(msg)

src/torchio/transforms/intensity/mask.py

@@ -9,7 +9,6 @@
 from torch import Tensor
 
 from ...data.batch import SubjectsBatch
-from ...data.image import LabelMap
 from ..transform import IntensityTransform
 
 
@@ -85,7 +84,7 @@ def _resolve_mask(self, batch: SubjectsBatch) -> Tensor:
                 )
                 raise KeyError(msg)
             mask_batch = batch.images[key]
-            if not issubclass(mask_batch._image_class, LabelMap):
+            if not mask_batch.is_label:
                 msg = f'Masking method "{key}" must refer to a LabelMap.'
                 raise TypeError(msg)
             mask_data = mask_batch.data[0]

src/torchio/transforms/intensity/normalize.py

@@ -12,7 +12,6 @@
 
 from ...data.batch import ImagesBatch
 from ...data.batch import SubjectsBatch
-from ...data.image import LabelMap
 from .._statistics import compute_quantile
 from ..parameter_range import Choice
 from ..parameter_range import _ParameterRange
@@ -221,7 +220,7 @@ def _get_mask(
                 )
                 raise KeyError(msg)
             mask_batch = batch.images[key]
-            if not issubclass(mask_batch._image_class, LabelMap):
+            if not mask_batch.is_label:
                 msg = f'Masking method "{key}" must refer to a LabelMap.'
                 raise TypeError(msg)
             return mask_batch.data[0].bool()

src/torchio/transforms/intensity/standardize.py

@@ -10,7 +10,6 @@
 
 from ...data.batch import ImagesBatch
 from ...data.batch import SubjectsBatch
-from ...data.image import LabelMap
 from ..transform import IntensityTransform
 
 
@@ -166,7 +165,7 @@ def _get_mask(
             )
             raise KeyError(msg)
         mask_batch = batch.images[masking_method]
-        if not issubclass(mask_batch._image_class, LabelMap):
+        if not mask_batch.is_label:
             msg = f'Masking method "{masking_method}" must refer to a LabelMap.'
             raise TypeError(msg)
         return mask_batch.data[0].bool()

src/torchio/transforms/intensity/swap.py

@@ -10,7 +10,6 @@
 from torch import Tensor
 
 from ...data.batch import SubjectsBatch
-from ...data.image import LabelMap
 from ..parameter_range import to_nonneg_range
 from ..transform import IntensityTransform
 
@@ -63,7 +62,7 @@ def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
         """Sample swap locations (per element when batched)."""
         # Warn if label maps are present.
         for _name, img_batch in batch.images.items():
-            if issubclass(img_batch._image_class, LabelMap):
+            if img_batch.is_label:
                 warnings.warn(
                     "Swap is applied to a subject containing LabelMap "
                     "images. The spatial rearrangement will make labels "

src/torchio/transforms/label/contour.py

@@ -8,7 +8,6 @@
 import torch.nn.functional as functional
 
 from ...data.batch import SubjectsBatch
-from ...data.image import LabelMap
 from ..transform import Transform
 
 
@@ -43,7 +42,7 @@ def apply_transform(
     ) -> SubjectsBatch:
         """Replace each label map with its boundary voxels."""
         for _name, img_batch in batch.images.items():
-            if not issubclass(img_batch._image_class, LabelMap):
+            if not img_batch.is_label:
                 continue
             img_batch.data = _extract_contour(img_batch.data)
         return batch