From af78c59b66b806aff436ee9ed5b450276505689c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 22:20:09 +0100 Subject: [PATCH 1/6] Integrate mapping replay and adapters Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/data/batch.py | 42 +++ src/torchio/transforms/compose.py | 6 + src/torchio/transforms/cornucopia_adapter.py | 50 +++- src/torchio/transforms/monai_adapter.py | 79 +++++- src/torchio/transforms/spatial/crop_or_pad.py | 2 + .../spatial/ensure_shape_multiple.py | 2 + src/torchio/transforms/transform.py | 253 ++++++++++++++++-- tests/test_batch.py | 100 +++++++ tests/test_cornucopia_adapter.py | 70 +++++ tests/test_crop_or_pad.py | 3 +- tests/test_monai_adapter.py | 93 +++++++ tests/test_transforms_base.py | 218 +++++++++++++++ 12 files changed, 880 insertions(+), 38 deletions(-) diff --git a/src/torchio/data/batch.py b/src/torchio/data/batch.py index 15035df1c..cec2053dd 100644 --- a/src/torchio/data/batch.py +++ b/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) diff --git a/src/torchio/transforms/compose.py b/src/torchio/transforms/compose.py index 244dca1b0..0c2c3a79d 100644 --- a/src/torchio/transforms/compose.py +++ b/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, diff --git a/src/torchio/transforms/cornucopia_adapter.py b/src/torchio/transforms/cornucopia_adapter.py index 1edef558f..ef88dd40d 100644 --- a/src/torchio/transforms/cornucopia_adapter.py +++ b/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( diff --git a/src/torchio/transforms/monai_adapter.py b/src/torchio/transforms/monai_adapter.py index 21491c775..992589408 100644 --- a/src/torchio/transforms/monai_adapter.py +++ b/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,78 @@ def _apply_dict_transform( monai_transform: Callable, monai: ModuleType, ) -> None: + monai_dict = _build_monai_dict(subject, monai) + result = _validate_dict_result(monai_transform(monai_dict), monai_dict) + _update_subject_images(subject, result, monai) + _update_subject_metadata(subject, result, monai_dict) + + +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, + monai_dict: dict[str, Any], +) -> 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(monai_dict) - 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, + monai_dict: dict[str, Any], +) -> 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 monai_dict: + 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 diff --git a/src/torchio/transforms/spatial/crop_or_pad.py b/src/torchio/transforms/spatial/crop_or_pad.py index 0ea453891..5073daeae 100644 --- a/src/torchio/transforms/spatial/crop_or_pad.py +++ b/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, diff --git a/src/torchio/transforms/spatial/ensure_shape_multiple.py b/src/torchio/transforms/spatial/ensure_shape_multiple.py index d2d6e61e4..19db70f61 100644 --- a/src/torchio/transforms/spatial/ensure_shape_multiple.py +++ b/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, diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py index 40cc7a1eb..559425fe1 100644 --- a/src/torchio/transforms/transform.py +++ b/src/torchio/transforms/transform.py @@ -88,6 +88,84 @@ def _data_has_annotations(data: Any) -> bool: return False +def _has_batched_param_metadata(params: dict[str, Any]) -> bool: + return any(key in params for key in ("_batch_size", "_batched_keys", "_keep")) + + +def _get_expected_batch_size( + params: dict[str, Any], + batch: SubjectsBatch, +) -> int: + if "_batch_size" not in params or "_batched_keys" not in params: + msg = "Batched params must define both _batch_size and _batched_keys." + raise ValueError(msg) + expected_size = params["_batch_size"] + if ( + isinstance(expected_size, bool) + or not isinstance(expected_size, int) + or expected_size < 1 + ): + msg = f"_batch_size must be a positive integer, got {expected_size!r}" + raise ValueError(msg) + if expected_size != batch.batch_size: + msg = ( + f"Parameter batch size {expected_size} does not match" + f" input batch size {batch.batch_size}" + ) + raise ValueError(msg) + return expected_size + + +def _get_batched_keys(params: dict[str, Any]) -> list[str]: + batched_keys = params["_batched_keys"] + if not isinstance(batched_keys, list) or not all( + isinstance(key, str) for key in batched_keys + ): + msg = "_batched_keys must be a list of parameter names" + raise ValueError(msg) + if len(set(batched_keys)) != len(batched_keys): + msg = "_batched_keys contains duplicate parameter names" + raise ValueError(msg) + return batched_keys + + +def _validate_batched_value( + params: dict[str, Any], + key: str, + expected_size: int, +) -> None: + if key not in params: + msg = f"Batched parameter {key!r} is missing" + raise ValueError(msg) + value = params[key] + if not isinstance(value, list): + msg = f"Batched parameter {key!r} must be a list" + raise ValueError(msg) + if len(value) != expected_size: + msg = ( + f"Batched parameter {key!r} must contain" + f" {expected_size} values, got {len(value)}" + ) + raise ValueError(msg) + + +def _validate_keep(params: dict[str, Any], expected_size: int) -> None: + if "_keep" not in params: + return + keep = params["_keep"] + if not isinstance(keep, list) or len(keep) != expected_size: + msg = f"_keep must contain {expected_size} boolean values, got {keep!r}" + raise ValueError(msg) + if not all(type(value) is bool for value in keep): + msg = "_keep values must be booleans" + raise ValueError(msg) + + +def _params_apply_to_any_element(params: dict[str, Any]) -> bool: + keep = params.get("_keep") + return keep is None or any(keep) + + class Transform(nn.Module): """Abstract class for all TorchIO transforms. @@ -132,6 +210,8 @@ class Transform(nn.Module): which the transform will *not* be applied. """ + _supports_apply_with_params = True + def __init__( self, *, @@ -239,43 +319,168 @@ def forward(self, data: Any) -> Any: Args: data: Input data to transform. """ + return self._execute(data, params=None, sample_params=True) + + @overload + def apply_with_params( + self, + data: Subject, + params: dict[str, Any], + ) -> Subject: ... + @overload + def apply_with_params( + self, + data: Image, + params: dict[str, Any], + ) -> Image: ... + @overload + def apply_with_params( + self, + data: Tensor, + params: dict[str, Any], + ) -> Tensor: ... + @overload + def apply_with_params( + self, + data: np.ndarray, + params: dict[str, Any], + ) -> np.ndarray: ... + @overload + def apply_with_params( + self, + data: sitk.Image, + params: dict[str, Any], + ) -> sitk.Image: ... + @overload + def apply_with_params( + self, + data: nib.Nifti1Image, + params: dict[str, Any], + ) -> nib.Nifti1Image: ... + @overload + def apply_with_params( + self, + data: dict, + params: dict[str, Any], + ) -> dict: ... + @overload + def apply_with_params( + self, + data: ImagesBatch, + params: dict[str, Any], + ) -> ImagesBatch: ... + @overload + def apply_with_params( + self, + data: SubjectsBatch, + params: dict[str, Any], + ) -> SubjectsBatch: ... + + def apply_with_params( + self, + data: Any, + params: dict[str, Any], + ) -> Any: + """Apply an exact parameter set without sampling. + + Args: + data: Input data to transform. + params: Exact parameters accepted by `apply_transform`. + + Returns: + Transformed data with the same type as the input. + + Raises: + TypeError: If `params` is not a dictionary. + NotImplementedError: If the transform does not expose a + compatible exact-parameter kernel. + ValueError: If per-instance parameter dimensions do not + match the input batch. + """ + if not self._supports_apply_with_params: + msg = ( + f"{type(self).__name__}.apply_with_params() is not supported" + " because this transform does not expose a compatible exact" + " parameter kernel." + ) + raise NotImplementedError(msg) + if not isinstance(params, dict): + msg = f"Expected params to be a dict, got {type(params).__name__}" + raise TypeError(msg) + return self._execute( + data, + params=_copy.deepcopy(params), + sample_params=False, + ) + + def _execute( + self, + data: Any, + *, + params: dict[str, Any] | None, + sample_params: bool, + ) -> Any: + """Run the shared transform lifecycle.""" if self.copy: data = _copy.deepcopy(data) batch, unwrap = self._wrap(data) - # When per-element gating is active, the transform handles the - # probability itself (masked-out elements get identity params), - # so skip the batch-wide coin flip here. Apply iff rand < p, so - # p=0 is always a no-op and p=1 always applies. - if not self._per_instance_p_active(batch) and torch.rand(1).item() >= self.p: + if sample_params and self._should_skip(batch): return unwrap(batch) - params = self.make_params(batch) - traces = self._build_history_traces(params, batch.batch_size) + resolved_params = self._resolve_execution_params( + batch, + params, + sample_params, + ) + traces = self._build_history_traces( + resolved_params, + batch.batch_size, + ) if any(trace is not None for trace in traces): self._check_spatial_annotations(batch) - batch = self.apply_transform(batch, params) + batch = self.apply_transform(batch, resolved_params) batch._append_history(traces) result = unwrap(batch) - # Propagate history to outputs that can carry it - if not isinstance( - result, - (ImagesBatch, SubjectsBatch, Tensor, np.ndarray), - ) and not isinstance(result, dict): - with contextlib.suppress(AttributeError): - result.applied_transforms = list(batch.applied_transforms) + self._propagate_history(batch, result) return result + def _should_skip(self, batch: SubjectsBatch) -> bool: + """Return whether batch-wide probability skips this application.""" + return not self._per_instance_p_active(batch) and torch.rand(1).item() >= self.p + + def _resolve_execution_params( + self, + batch: SubjectsBatch, + params: dict[str, Any] | None, + sample_params: bool, + ) -> dict[str, Any]: + """Sample parameters or validate an exact supplied set.""" + if sample_params: + return self.make_params(batch) + assert params is not None + self._validate_batched_params(params, batch) + return params + + @staticmethod + def _propagate_history(batch: SubjectsBatch, result: Any) -> None: + """Copy history to non-batch output types that can carry it.""" + excluded_types = (ImagesBatch, SubjectsBatch, Tensor, np.ndarray, dict) + if isinstance(result, excluded_types): + return + with contextlib.suppress(AttributeError): + result.applied_transforms = list(batch.applied_transforms) + def _build_history_traces( self, params: dict[str, Any], batch_size: int, ) -> list[AppliedTransform | None]: """Build one clean optional history trace per element.""" - batched_keys = params.get("_batched_keys") - if batched_keys is None: + if "_batched_keys" not in params: return [ self._make_applied_transform(_copy.deepcopy(params)) for _ in range(batch_size) ] + batched_keys = _get_batched_keys(params) expected_size = params.get("_batch_size") if expected_size != batch_size: msg = ( @@ -293,6 +498,20 @@ def _build_history_traces( traces.append(self._make_applied_transform(element_params)) return traces + @staticmethod + def _validate_batched_params( + params: dict[str, Any], + batch: SubjectsBatch, + ) -> None: + """Validate transient per-instance bookkeeping.""" + if not _has_batched_param_metadata(params): + return + expected_size = _get_expected_batch_size(params, batch) + batched_keys = _get_batched_keys(params) + for key in batched_keys: + _validate_batched_value(params, key, expected_size) + _validate_keep(params, expected_size) + def _make_applied_transform( self, params: dict[str, Any], diff --git a/tests/test_batch.py b/tests/test_batch.py index 7a045ecd1..de40f4153 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -604,3 +604,103 @@ def test_uniform_applied_transforms_view_is_immutable(self) -> None: assert isinstance(result.applied_transforms, tuple) with pytest.raises(AttributeError): result.applied_transforms.append("invalid") # type: ignore[attr-defined] + +class TestMapSubjects: + def _batch(self) -> SubjectsBatch: + return SubjectsBatch.from_subjects( + [ + tio.Subject( + t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)), + identifier=f" subject-{index} ", + index=index, + ) + for index in range(2) + ] + ) + + def test_maps_text_metadata(self) -> None: + batch = self._batch() + + def strip_identifier(subject: tio.Subject) -> tio.Subject: + subject.metadata["identifier"] = subject.identifier.strip() + return subject + + result = batch.map_subjects(strip_identifier) + + assert result.metadata["identifier"] == ["subject-0", "subject-1"] + assert batch.metadata["identifier"] == [" subject-0 ", " subject-1 "] + + def test_preserves_annotations(self) -> None: + batch = SubjectsBatch.from_subjects( + [ + tio.Subject( + landmarks=tio.Points(torch.rand(index + 1, 3)), + identifier=f"subject-{index}", + ) + for index in range(2) + ] + ) + + result = batch.map_subjects(lambda subject: subject) + + assert result.points["landmarks"][0].num_points == 1 + assert result.points["landmarks"][1].num_points == 2 + + def test_allows_uniform_schema_change(self) -> None: + def add_site(subject: tio.Subject) -> tio.Subject: + subject.metadata["site"] = "A" + return subject + + result = self._batch().map_subjects(add_site) + + assert result.metadata["site"] == ["A", "A"] + + def test_rejects_divergent_schema_change(self) -> None: + def add_site_to_first(subject: tio.Subject) -> tio.Subject: + if subject.index == 0: + subject.metadata["site"] = "A" + return subject + + with pytest.raises(ValueError, match=r"metadata.*index 1.*missing.*site"): + self._batch().map_subjects(add_site_to_first) + + def test_rejects_non_subject_result(self) -> None: + def return_dict(subject: tio.Subject) -> dict: + return {"subject": subject} + + with pytest.raises(TypeError, match=r"index 0.*Subject.*dict"): + self._batch().map_subjects(return_dict) # type: ignore[arg-type] + + def test_retains_callback_history(self) -> None: + def flip_by_index(subject: tio.Subject) -> tio.Subject: + return tio.Flip(axes=(subject.index,))(subject) + + result = self._batch().map_subjects(flip_by_index) + + assert result.has_divergent_history + assert result.history(0)[0].params["axes"] == (0,) + assert result.history(1)[0].params["axes"] == (1,) + + def test_in_place_callback_does_not_mutate_input(self) -> None: + batch = self._batch() + + def add_in_place(subject: tio.Subject) -> tio.Subject: + subject.t1.data.add_(1) + return subject + + result = batch.map_subjects(add_in_place) + + assert torch.count_nonzero(batch.t1.data) == 0 + assert torch.all(result.t1.data == 1) + + def test_copy_false_allows_in_place_mutation(self) -> None: + batch = self._batch() + + def add_in_place(subject: tio.Subject) -> tio.Subject: + subject.t1.data.add_(1) + return subject + + result = batch.map_subjects(add_in_place, copy=False) + + assert torch.all(batch.t1.data == 1) + assert torch.all(result.t1.data == 1) diff --git a/tests/test_cornucopia_adapter.py b/tests/test_cornucopia_adapter.py index 6d54f581f..ca2902be2 100644 --- a/tests/test_cornucopia_adapter.py +++ b/tests/test_cornucopia_adapter.py @@ -81,6 +81,76 @@ def test_in_compose(self) -> None: result = pipeline(subject) assert result.t1.data.shape == subject.t1.data.shape + def test_preserves_prior_history_and_annotations(self) -> None: + subject = tio.Gamma(log_gamma=0.2)( + tio.Subject( + t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 1), + landmarks=tio.Points(torch.rand(2, 3)), + ) + ) + + result = tio.CornucopiaAdapter(lambda tensor: tensor)(subject) + + assert [trace.name for trace in result.applied_transforms] == ["Gamma"] + assert set(result.points) == {"landmarks"} + + def test_probability_zero_when_random_draw_is_zero( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + subject = _make_subject() + original = subject.t1.data.clone() + monkeypatch.setattr(torch, "rand", lambda *args, **kwargs: torch.zeros(1)) + + result = tio.CornucopiaAdapter(lambda tensor: tensor + 1, p=0)(subject) + + torch.testing.assert_close(result.t1.data, original) + + def test_rejects_non_tensor_result(self) -> None: + subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 4, 4, 4))) + + with pytest.raises(TypeError, match=r"torch.Tensor"): + tio.CornucopiaAdapter(lambda tensor: "not a tensor")(subject) + + def test_copy_false_allows_in_place_transform(self) -> None: + batch = tio.SubjectsBatch.from_subjects( + [tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)))] + ) + + result = tio.CornucopiaAdapter( + lambda tensor: tensor.add_(1), + copy=False, + )(batch) + + assert torch.all(batch.t1.data == 1) + assert torch.all(result.t1.data == 1) + + @pytest.mark.parametrize( + "transform", + [ + lambda tensor: (tensor,), + lambda tensor: [tensor], + ], + ) + def test_single_image_accepts_sequence_result(self, transform) -> None: + subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 4, 4, 4))) + + result = tio.CornucopiaAdapter(transform)(subject) + + torch.testing.assert_close(result.t1.data, subject.t1.data) + + def test_multiple_images_reject_single_tensor_result(self) -> None: + subject = _make_subject() + + with pytest.raises(TypeError, match="tuple or list with 2 image results"): + tio.CornucopiaAdapter(lambda *tensors: tensors[0])(subject) + + def test_single_image_rejects_wrong_result_count(self) -> None: + subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 4, 4, 4))) + + with pytest.raises(ValueError, match="Expected 1 image result, got 2"): + tio.CornucopiaAdapter(lambda tensor: (tensor, tensor))(subject) + # ── Real Cornucopia transforms ─────────────────────────────────────── diff --git a/tests/test_crop_or_pad.py b/tests/test_crop_or_pad.py index 5382f10f1..f653f8fef 100644 --- a/tests/test_crop_or_pad.py +++ b/tests/test_crop_or_pad.py @@ -392,8 +392,9 @@ def test_batch_pad(self) -> None: class TestProbability: - def test_p_zero_is_no_op(self) -> None: + def test_p_zero_is_no_op(self, monkeypatch: pytest.MonkeyPatch) -> None: subject = _make_subject((20, 20, 20)) + monkeypatch.setattr(torch, "rand", lambda *args, **kwargs: torch.zeros(1)) result = tio.CropOrPad(target_shape=10, p=0)(subject) assert result.t1.shape == (1, 20, 20, 20) diff --git a/tests/test_monai_adapter.py b/tests/test_monai_adapter.py index f2be5ac79..4c004c462 100644 --- a/tests/test_monai_adapter.py +++ b/tests/test_monai_adapter.py @@ -82,6 +82,56 @@ def test_dict_only_modifies_specified_keys(self) -> None: result = adapter(subject) torch.testing.assert_close(result.t2.data, original_t2) + def test_dict_updates_metadata(self) -> None: + from monai.transforms import MapTransform + + class UpdateMetadata(MapTransform): + def __init__(self) -> None: + super().__init__(keys=["t1"]) + + def __call__(self, data): + return {**data, "site": data["site"].lower()} + + subject = tio.Subject( + t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), + site="ABC", + ) + + result = tio.MonaiAdapter(UpdateMetadata())(subject) + + assert result.site == "abc" + + def test_dict_rejects_added_tensor_key(self) -> None: + from monai.transforms import MapTransform + + class AddTensor(MapTransform): + def __init__(self) -> None: + super().__init__(keys=["t1"]) + + def __call__(self, data): + return {**data, "new_image": torch.zeros(1, 8, 8, 8)} + + subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8))) + + with pytest.raises(ValueError, match=r"new tensor field.*new_image"): + tio.MonaiAdapter(AddTensor())(subject) + + def test_dict_preserves_added_metadata_order(self) -> None: + from monai.transforms import MapTransform + + class AddMetadata(MapTransform): + def __init__(self) -> None: + super().__init__(keys=["t1"]) + + def __call__(self, data): + return {**data, "zeta": 1, "alpha": 2} + + subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8))) + + result = tio.MonaiAdapter(AddMetadata())(subject) + + assert list(result.metadata) == ["zeta", "alpha"] + @pytest.mark.skipif(not HAS_MONAI, reason="MONAI not installed") class TestMonaiAdapterGeneral: @@ -116,3 +166,46 @@ def test_in_compose(self) -> None: pipeline = tio.Compose([tio.MonaiAdapter(NormalizeIntensity())]) result = pipeline(subject) assert isinstance(result, tio.Subject) + + def test_preserves_prior_history_and_annotations(self) -> None: + from monai.transforms import NormalizeIntensity + + subject = tio.Gamma(log_gamma=0.2)( + tio.Subject( + t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 1), + landmarks=tio.Points(torch.rand(2, 3)), + ) + ) + + result = tio.MonaiAdapter(NormalizeIntensity())(subject) + + assert [trace.name for trace in result.applied_transforms] == ["Gamma"] + assert set(result.points) == {"landmarks"} + + def test_probability_zero_when_random_draw_is_zero( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from monai.transforms import NormalizeIntensity + + subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 1)) + original = subject.t1.data.clone() + monkeypatch.setattr(torch, "rand", lambda *args, **kwargs: torch.zeros(1)) + + result = tio.MonaiAdapter(NormalizeIntensity(), p=0)(subject) + + torch.testing.assert_close(result.t1.data, original) + + def test_copy_false_allows_in_place_transform(self) -> None: + class AddInPlace: + def __call__(self, tensor): + return tensor.add_(1) + + batch = tio.SubjectsBatch.from_subjects( + [tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)))] + ) + + result = tio.MonaiAdapter(AddInPlace(), copy=False)(batch) + + assert torch.all(batch.t1.data == 1) + assert torch.all(result.t1.data == 1) diff --git a/tests/test_transforms_base.py b/tests/test_transforms_base.py index 7965224d5..8f73c83bb 100644 --- a/tests/test_transforms_base.py +++ b/tests/test_transforms_base.py @@ -65,6 +65,31 @@ def apply_transform(self, batch: Any, params: dict) -> Any: return batch +class _AddFromParams(tio.Transform): + """Add an explicitly supplied value.""" + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.make_params_calls = 0 + + def make_params(self, batch: Any) -> dict[str, Any]: + self.make_params_calls += 1 + return {"value": 1.0} + + def apply_transform(self, batch: Any, params: dict[str, Any]) -> Any: + for image_batch in batch.images.values(): + image_batch.data = image_batch.data + params["value"] + return batch + + +class _MutateParams(tio.Transform): + """Mutate params to test defensive copying.""" + + def apply_transform(self, batch: Any, params: dict[str, Any]) -> Any: + params["value"] = "mutated" + return batch + + # ── Transform base ─────────────────────────────────────────────────── @@ -195,6 +220,199 @@ def test_invalid_input_type(self) -> None: _IdentityTransform()("not a valid input") +class TestApplyWithParams: + def test_bypasses_probability_and_sampling(self) -> None: + subject = tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))) + transform = _AddFromParams(p=0) + + result = transform.apply_with_params(subject, {"value": 2.0}) + + assert transform.make_params_calls == 0 + torch.testing.assert_close(result.t1.data, torch.full_like(result.t1.data, 2)) + + def test_records_supplied_params(self) -> None: + subject = tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))) + + result = _AddFromParams().apply_with_params(subject, {"value": 2.0}) + + trace = result.applied_transforms[-1] + assert trace.name == "_AddFromParams" + assert trace.params == {"value": 2.0} + + def test_preserves_prior_history(self) -> None: + subject = tio.Flip(axes=(0,))( + tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))) + ) + + result = _AddFromParams().apply_with_params(subject, {"value": 2.0}) + + assert [trace.name for trace in result.applied_transforms] == [ + "Flip", + "_AddFromParams", + ] + + def test_does_not_mutate_caller_params(self) -> None: + subject = tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))) + params = {"value": "original"} + + _MutateParams().apply_with_params(subject, params) + + assert params == {"value": "original"} + + def test_copy_true_preserves_batch(self) -> None: + batch = tio.SubjectsBatch.from_subjects( + [tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)))] + ) + + result = _AddFromParams(copy=True).apply_with_params( + batch, + {"value": 2.0}, + ) + + assert result is not batch + assert torch.count_nonzero(batch.t1.data) == 0 + + def test_copy_false_mutates_batch(self) -> None: + batch = tio.SubjectsBatch.from_subjects( + [tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)))] + ) + + result = _AddFromParams(copy=False).apply_with_params( + batch, + {"value": 2.0}, + ) + + assert result is batch + assert torch.all(batch.t1.data == 2) + + def test_replays_stochastic_transform_exactly(self) -> None: + subject = tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))) + transform = tio.Noise(mean=(-1, 1), std=(0.1, 0.5)) + transformed = transform(subject) + params = transformed.applied_transforms[-1].params + + replayed = transform.apply_with_params(subject, params) + + torch.testing.assert_close(replayed.t1.data, transformed.t1.data) + + @pytest.mark.parametrize( + "data", + [ + tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))), + tio.ScalarImage(torch.zeros(1, 4, 4, 4)), + torch.zeros(1, 4, 4, 4), + np.zeros((1, 4, 4, 4), dtype=np.float32), + sitk.Image(4, 4, 4, sitk.sitkFloat32), + nib.Nifti1Image(np.zeros((4, 4, 4)), np.eye(4)), + {"t1": torch.zeros(1, 4, 4, 4), "age": 42}, + tio.ImagesBatch.from_images([tio.ScalarImage(torch.zeros(1, 4, 4, 4))]), + tio.SubjectsBatch.from_subjects( + [tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)))] + ), + ], + ) + def test_preserves_input_type(self, data: Any) -> None: + result = _IdentityTransform().apply_with_params(data, {}) + + assert type(result) is type(data) + + @pytest.mark.parametrize( + ("params", "message"), + [ + ( + { + "value": [1, 2], + "_batch_size": 3, + "_batched_keys": ["value"], + }, + "batch size", + ), + ( + { + "_batch_size": 2, + "_batched_keys": ["missing"], + }, + "missing", + ), + ( + { + "value": [1], + "_batch_size": 2, + "_batched_keys": ["value"], + }, + "2 values", + ), + ( + { + "value": (1, 2), + "_batch_size": 2, + "_batched_keys": ["value"], + }, + "must be a list", + ), + ( + { + "value": [1, 2], + "_batch_size": 2, + "_batched_keys": ["value"], + "_keep": [True], + }, + "_keep", + ), + ( + { + "value": [1, 2], + "_batch_size": 2, + "_batched_keys": ["value"], + "_keep": [True, 1], + }, + "booleans", + ), + ], + ) + def test_validates_batched_params( + self, + params: dict[str, Any], + message: str, + ) -> None: + batch = tio.SubjectsBatch.from_subjects( + [tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))) for _ in range(2)] + ) + + with pytest.raises(ValueError, match=message): + _IdentityTransform().apply_with_params(batch, params) + + def test_rejects_reserved_fields_without_batched_keys(self) -> None: + subject = tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))) + + with pytest.raises(ValueError, match="_batched_keys"): + _IdentityTransform().apply_with_params( + subject, + {"_batch_size": 1}, + ) + + @pytest.mark.parametrize( + "transform", + [ + tio.Compose([]), + tio.OneOf([_IdentityTransform()]), + tio.SomeOf([_IdentityTransform()]), + tio.MonaiAdapter(lambda value: value), + tio.CornucopiaAdapter(lambda value: value), + tio.CropOrPad(4), + tio.EnsureShapeMultiple(2), + ], + ) + def test_rejects_transforms_without_param_kernel( + self, + transform: tio.Transform, + ) -> None: + subject = tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))) + + with pytest.raises(NotImplementedError, match="apply_with_params"): + transform.apply_with_params(subject, {}) + + # ── include/exclude ────────────────────────────────────────────────── From e3ce01fe90f4bfee25bd115f3fd0198cc37680e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 22:56:33 +0100 Subject: [PATCH 2/6] Finalize replay annotation gating Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/transforms/transform.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py index 559425fe1..871fbe158 100644 --- a/src/torchio/transforms/transform.py +++ b/src/torchio/transforms/transform.py @@ -161,11 +161,6 @@ def _validate_keep(params: dict[str, Any], expected_size: int) -> None: raise ValueError(msg) -def _params_apply_to_any_element(params: dict[str, Any]) -> bool: - keep = params.get("_keep") - return keep is None or any(keep) - - class Transform(nn.Module): """Abstract class for all TorchIO transforms. From b8b3327e47bfe3c7f9b8f04c95976e1f82a1fcf5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 23:16:40 +0100 Subject: [PATCH 3/6] Format subject mapping tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- tests/test_batch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_batch.py b/tests/test_batch.py index de40f4153..c61e6f1b6 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -605,6 +605,7 @@ def test_uniform_applied_transforms_view_is_immutable(self) -> None: with pytest.raises(AttributeError): result.applied_transforms.append("invalid") # type: ignore[attr-defined] + class TestMapSubjects: def _batch(self) -> SubjectsBatch: return SubjectsBatch.from_subjects( From 5d3208cef3a2259834a5eb27e86b0226737f55ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 20 Jul 2026 23:39:55 +0100 Subject: [PATCH 4/6] Reuse canonical batch parameter keys Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/transforms/transform.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py index 871fbe158..697ba2b04 100644 --- a/src/torchio/transforms/transform.py +++ b/src/torchio/transforms/transform.py @@ -19,6 +19,7 @@ from torch import Tensor from torch import nn +from ..data.batch import _BATCH_META_KEYS from ..data.batch import ImagesBatch from ..data.batch import SubjectsBatch from ..data.batch import _slice_params @@ -89,7 +90,7 @@ def _data_has_annotations(data: Any) -> bool: def _has_batched_param_metadata(params: dict[str, Any]) -> bool: - return any(key in params for key in ("_batch_size", "_batched_keys", "_keep")) + return any(key in params for key in _BATCH_META_KEYS) def _get_expected_batch_size( From f915c86932fe985e6affb4dffe35b072648e1ea3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Tue, 21 Jul 2026 00:20:48 +0100 Subject: [PATCH 5/6] Snapshot MONAI input keys before transforms Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/transforms/monai_adapter.py | 13 +++++++------ tests/test_monai_adapter.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/src/torchio/transforms/monai_adapter.py b/src/torchio/transforms/monai_adapter.py index 992589408..0376c88d5 100644 --- a/src/torchio/transforms/monai_adapter.py +++ b/src/torchio/transforms/monai_adapter.py @@ -170,9 +170,10 @@ def _apply_dict_transform( monai: ModuleType, ) -> None: monai_dict = _build_monai_dict(subject, monai) - result = _validate_dict_result(monai_transform(monai_dict), monai_dict) + 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, monai_dict) + _update_subject_metadata(subject, result, original_keys) def _build_monai_dict(subject: Subject, monai: ModuleType) -> dict[str, Any]: @@ -187,13 +188,13 @@ def _build_monai_dict(subject: Subject, monai: ModuleType) -> dict[str, Any]: def _validate_dict_result( result: Any, - monai_dict: dict[str, 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(monai_dict) - set(result) + missing = set(original_keys) - set(result) if missing: msg = f"MONAI dictionary transform removed fields: {sorted(missing)}" raise ValueError(msg) @@ -220,13 +221,13 @@ def _update_subject_images( def _update_subject_metadata( subject: Subject, result: Mapping, - monai_dict: dict[str, Any], + 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 monai_dict: + if key in original_keys: continue _add_subject_metadata(subject, key, result[key]) diff --git a/tests/test_monai_adapter.py b/tests/test_monai_adapter.py index 4c004c462..fdba2aa1f 100644 --- a/tests/test_monai_adapter.py +++ b/tests/test_monai_adapter.py @@ -116,6 +116,25 @@ def __call__(self, data): with pytest.raises(ValueError, match=r"new tensor field.*new_image"): tio.MonaiAdapter(AddTensor())(subject) + def test_dict_rejects_in_place_field_removal(self) -> None: + from monai.transforms import MapTransform + + class RemoveMetadata(MapTransform): + def __init__(self) -> None: + super().__init__(keys=["t1"]) + + def __call__(self, data): + del data["site"] + return data + + subject = tio.Subject( + t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), + site="A", + ) + + with pytest.raises(ValueError, match=r"removed fields.*site"): + tio.MonaiAdapter(RemoveMetadata())(subject) + def test_dict_preserves_added_metadata_order(self) -> None: from monai.transforms import MapTransform From 1c8b78503d30bbc84466251a9f5767581dec611c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Tue, 21 Jul 2026 00:30:50 +0100 Subject: [PATCH 6/6] Isolate per-element history parameters Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: baf5f11a-f67c-44dc-804e-6be849fa9160 --- src/torchio/transforms/transform.py | 2 +- tests/test_transforms_base.py | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py index 697ba2b04..8364d7e7b 100644 --- a/src/torchio/transforms/transform.py +++ b/src/torchio/transforms/transform.py @@ -491,7 +491,7 @@ def _build_history_traces( traces.append(None) continue element_params = _slice_params(params, index, batched_keys) - traces.append(self._make_applied_transform(element_params)) + traces.append(self._make_applied_transform(_copy.deepcopy(element_params))) return traces @staticmethod diff --git a/tests/test_transforms_base.py b/tests/test_transforms_base.py index 8f73c83bb..307cf4767 100644 --- a/tests/test_transforms_base.py +++ b/tests/test_transforms_base.py @@ -90,6 +90,14 @@ def apply_transform(self, batch: Any, params: dict[str, Any]) -> Any: return batch +class _MutateBatchedParams(tio.Transform): + """Mutate nested batched params to test history isolation.""" + + def apply_transform(self, batch: Any, params: dict[str, Any]) -> Any: + params["value"][0]["nested"] = "mutated" + return batch + + # ── Transform base ─────────────────────────────────────────────────── @@ -259,6 +267,20 @@ def test_does_not_mutate_caller_params(self) -> None: assert params == {"value": "original"} + def test_batched_history_does_not_alias_kernel_params(self) -> None: + batch = tio.SubjectsBatch.from_subjects( + [tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4))) for _ in range(2)] + ) + params = { + "value": [{"nested": "first"}, {"nested": "second"}], + "_batch_size": 2, + "_batched_keys": ["value"], + } + + result = _MutateBatchedParams().apply_with_params(batch, params) + + assert result.history(0)[-1].params["value"] == {"nested": "first"} + def test_copy_true_preserves_batch(self) -> None: batch = tio.SubjectsBatch.from_subjects( [tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)))]