Integrate subject mapping and exact replay

src/torchio/data/batch.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import copy as _copy
+from collections.abc import Callable
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 from typing import Any
@@ -433,6 +434,47 @@ def unbatch(self) -> list[Any]:
             subjects.append(sub)
         return subjects
 
+    def map_subjects(
+        self,
+        callback: Callable[[Subject], Subject],
+        *,
+        copy: bool = True,
+    ) -> Self:
+        """Apply a callback to every subject and rebuild the batch.
+
+        Each callback receives an unbatched `Subject` carrying its exact
+        history. By default, image tensors are cloned before the callback
+        so the input batch is unchanged. With `copy=False`, callbacks may
+        mutate the input batch's image tensors.
+
+        Args:
+            callback: Callable taking and returning one `Subject`.
+            copy: Clone each image tensor before invoking the callback.
+
+        Returns:
+            A new batch containing the callback results.
+
+        Raises:
+            TypeError: If the callback does not return a `Subject`.
+            ValueError: If callback results cannot be batched together.
+        """
+        from .subject import Subject
+
+        mapped = []
+        for index, subject in enumerate(self.unbatch()):
+            if copy:
+                for image in subject.images.values():
+                    image.set_data(image.data.clone())
+            result = callback(subject)
+            if not isinstance(result, Subject):
+                msg = (
+                    f"Expected callback result at index {index} to be a Subject,"
+                    f" got {type(result).__name__}"
+                )
+                raise TypeError(msg)
+            mapped.append(result)
+        return type(self).from_subjects(mapped)
+
     def _batch_items(self, items: Sequence[Any]) -> Self:
         """Rebuild a subject batch from subjects."""
         return type(self).from_subjects(items)

src/torchio/transforms/compose.py

@@ -65,6 +65,8 @@ class Compose(Transform):
         ... })
     """
 
+    _supports_apply_with_params = False
+
     def __init__(
         self,
         transforms: Sequence[Transform] | Mapping[str, Transform] | None = None,
@@ -122,6 +124,8 @@ class OneOf(Transform):
         ... })
     """
 
+    _supports_apply_with_params = False
+
     def __init__(
         self,
         transforms: Sequence[Transform] | dict[Transform, float],
@@ -208,6 +212,8 @@ class SomeOf(Transform):
         ... )
     """
 
+    _supports_apply_with_params = False
+
     def __init__(
         self,
         transforms: Sequence[Transform] | None = None,

src/torchio/transforms/cornucopia_adapter.py

@@ -49,6 +49,8 @@ class CornucopiaAdapter(Transform):
         objects are not guaranteed to be serializable.
     """
 
+    _supports_apply_with_params = False
+
     def __init__(
         self,
         cornucopia_transform: Callable,
@@ -68,14 +70,14 @@ def forward(self, data: Any) -> Any:
         batch, unwrap = self._wrap(data)
         if self.copy:
             batch = _copy.deepcopy(batch)
-        if torch.rand(1).item() > self.p:
+        if torch.rand(1).item() >= self.p:
             return unwrap(batch)
-        subjects = batch.unbatch()
-        for subject in subjects:
+
+        def apply_to_subject(subject: Subject) -> Subject:
             _apply_cornucopia(subject, self.cornucopia_transform, self)
-        from ..data.batch import SubjectsBatch
+            return subject
 
-        result = SubjectsBatch.from_subjects(subjects)
+        result = batch.map_subjects(apply_to_subject, copy=False)
         return unwrap(result)
 
     def apply_transform(
@@ -121,14 +123,40 @@ def _apply_cornucopia(
     # Cornucopia transforms accept multiple tensors as *args
     # and return the same number of tensors.
     results = cornucopia_transform(*tensors)
-
-    # If only one image, result is a single tensor (not a tuple).
-    if len(names) == 1:
-        results = (results,)
+    results = _normalize_results(results, len(names))
 
     for name, result_tensor in zip(names, results, strict=True):
-        if isinstance(result_tensor, torch.Tensor):
-            images[name].set_data(result_tensor)
+        if not isinstance(result_tensor, torch.Tensor):
+            msg = (
+                f"Expected torch.Tensor for image field {name!r},"
+                f" got {type(result_tensor).__name__}"
+            )
+            raise TypeError(msg)
+        images[name].set_data(result_tensor)
+
+
+def _normalize_results(
+    results: Any,
+    num_images: int,
+) -> tuple[Any, ...] | list[Any]:
+    """Normalize Cornucopia outputs and validate their arity."""
+    if num_images == 1:
+        if not isinstance(results, (tuple, list)):
+            return (results,)
+        if len(results) != 1:
+            msg = f"Expected 1 image result, got {len(results)}"
+            raise ValueError(msg)
+        return results
+    if not isinstance(results, (tuple, list)):
+        msg = (
+            f"Expected a tuple or list with {num_images} image results,"
+            f" got {type(results).__name__}"
+        )
+        raise TypeError(msg)
+    if len(results) != num_images:
+        msg = f"Expected {num_images} image results, got {len(results)}"
+        raise ValueError(msg)
+    return results
 
 
 def _filter_images(

src/torchio/transforms/monai_adapter.py

@@ -57,6 +57,8 @@ class MonaiAdapter(Transform):
         serializable.
     """
 
+    _supports_apply_with_params = False
+
     def __init__(self, monai_transform: Callable, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         if not callable(monai_transform):
@@ -72,12 +74,11 @@ def forward(self, data):
         batch, unwrap = self._wrap(data)
         if self.copy:
             batch = _copy.deepcopy(batch)
-        if torch.rand(1).item() > self.p:
+        if torch.rand(1).item() >= self.p:
             return unwrap(batch)
-        # MONAI transforms operate per-subject
         monai = get_monai()
-        subjects = batch.unbatch()
-        for subject in subjects:
+
+        def apply_to_subject(subject: Subject) -> Subject:
             is_dict = isinstance(
                 self.monai_transform,
                 monai.transforms.MapTransform,
@@ -87,9 +88,9 @@ def forward(self, data):
             else:
                 images = self._get_subject_images(subject)
                 _apply_array_transform(images, self.monai_transform, monai)
-        from ..data.batch import SubjectsBatch
+            return subject
 
-        result = SubjectsBatch.from_subjects(subjects)
+        result = batch.map_subjects(apply_to_subject, copy=False)
         return unwrap(result)
 
     def apply_transform(self, batch: Any, params: dict[str, Any]) -> Any:
@@ -168,18 +169,79 @@ def _apply_dict_transform(
     monai_transform: Callable,
     monai: ModuleType,
 ) -> None:
+    monai_dict = _build_monai_dict(subject, monai)
+    original_keys = tuple(monai_dict)
+    result = _validate_dict_result(monai_transform(monai_dict), original_keys)
+    _update_subject_images(subject, result, monai)
+    _update_subject_metadata(subject, result, original_keys)
+
+
+def _build_monai_dict(subject: Subject, monai: ModuleType) -> dict[str, Any]:
+    """Build a MONAI mapping from subject images and metadata."""
     monai_dict: dict[str, Any] = {}
     for name, image in subject.images.items():
         monai_dict[name] = _image_to_meta_tensor(image, monai)
     for key, value in subject.metadata.items():
         monai_dict[key] = value
+    return monai_dict
 
-    result = monai_transform(monai_dict)
 
+def _validate_dict_result(
+    result: Any,
+    original_keys: tuple[str, ...],
+) -> Mapping:
+    """Validate the mapping returned by a MONAI dictionary transform."""
     if not isinstance(result, Mapping):
         msg = f"Expected mapping from MONAI dict transform, got {type(result).__name__}"
         raise TypeError(msg)
+    missing = set(original_keys) - set(result)
+    if missing:
+        msg = f"MONAI dictionary transform removed fields: {sorted(missing)}"
+        raise ValueError(msg)
+    return result
+
 
+def _update_subject_images(
+    subject: Subject,
+    result: Mapping,
+    monai: ModuleType,
+) -> None:
+    """Update existing subject images from a MONAI result."""
     for name, image in subject.images.items():
-        if name in result and isinstance(result[name], torch.Tensor):
-            _update_image_from_result(image, result[name], monai)
+        value = result[name]
+        if not isinstance(value, torch.Tensor):
+            msg = (
+                f"Expected torch.Tensor for image field {name!r},"
+                f" got {type(value).__name__}"
+            )
+            raise TypeError(msg)
+        _update_image_from_result(image, value, monai)
+
+
+def _update_subject_metadata(
+    subject: Subject,
+    result: Mapping,
+    original_keys: tuple[str, ...],
+) -> None:
+    """Update existing metadata and append new scalar fields."""
+    for key in subject.metadata:
+        subject.metadata[key] = result[key]
+    for key in result:
+        if key in original_keys:
+            continue
+        _add_subject_metadata(subject, key, result[key])
+
+
+def _add_subject_metadata(subject: Subject, key: Any, value: Any) -> None:
+    """Add one new metadata field returned by MONAI."""
+    if not isinstance(key, str):
+        msg = f"Expected MONAI output keys to be strings, got {key!r}"
+        raise TypeError(msg)
+    if isinstance(value, torch.Tensor):
+        msg = (
+            f"MONAI dictionary transform added new tensor field {key!r}."
+            " TorchIO cannot infer its image type; add it to the Subject"
+            " before applying the adapter."
+        )
+        raise ValueError(msg)
+    subject.metadata[key] = value

src/torchio/transforms/spatial/crop_or_pad.py

@@ -431,6 +431,8 @@ class CropOrPad(SpatialTransform):
         >>> transform = tio.CropOrPad(target_shape=256, padding_mode='mean')
     """
 
+    _supports_apply_with_params = False
+
     def __init__(
         self,
         target_shape: TargetShapeParam,

src/torchio/transforms/spatial/ensure_shape_multiple.py

@@ -90,6 +90,8 @@ class EnsureShapeMultiple(SpatialTransform):
         >>> transform = tio.EnsureShapeMultiple((4, 8, 16))
     """
 
+    _supports_apply_with_params = False
+
     def __init__(
         self,
         target_multiple: TargetMultipleParam,