Add exact transform parameter replay

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,
@@ -317,46 +323,15 @@ def _rebatch_with_history(subjects: list[Any], transform_name: str) -> Any:
     """
     from ..data.batch import SubjectsBatch
 
-    _check_consistent_schema(subjects, transform_name)
     try:
         batch = SubjectsBatch.from_subjects(subjects)
-    except (RuntimeError, KeyError) as error:
+    except ValueError as error:
         msg = (
             f"Per-instance {transform_name} produced batch elements with"
             " different shapes or schemas, which cannot be re-stacked. Use"
             " only shape- and schema-preserving transforms with per-instance"
             f" {transform_name}, or pass per_instance=False."
         )
         raise RuntimeError(msg) from error
-    batch.set_per_element_history([s.applied_transforms for s in subjects])
+    batch.set_per_element_history([subject.applied_transforms for subject in subjects])
     return batch
-
-
-def _check_consistent_schema(subjects: list[Any], transform_name: str) -> None:
-    """Ensure all subjects share the same image names and classes.
-
-    Per-element branching may apply different transforms to different
-    elements; if those change the set of images (or their type), the
-    elements can no longer be re-stacked into one batch. This raises a
-    clear error instead of silently dropping data.
-
-    Args:
-        subjects: The subjects about to be re-stacked.
-        transform_name: Name of the branching transform for the message.
-
-    Raises:
-        RuntimeError: If image names or classes differ across subjects.
-    """
-    if not subjects:
-        return
-    reference = {name: type(image) for name, image in subjects[0].images.items()}
-    for subject in subjects[1:]:
-        current = {name: type(image) for name, image in subject.images.items()}
-        if current != reference:
-            msg = (
-                f"Per-instance {transform_name} produced batch elements with"
-                " different image names or types, which cannot be re-stacked."
-                " Use only schema-preserving transforms with per-instance"
-                f" {transform_name}, or pass per_instance=False."
-            )
-            raise RuntimeError(msg)

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,

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):

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,