diff --git a/docs/get-started/migration.md b/docs/get-started/migration.md index 99c5a9652..c4b61cf48 100644 --- a/docs/get-started/migration.md +++ b/docs/get-started/migration.md @@ -360,7 +360,7 @@ for the capability contract and stochastic-realisation caveats. Prefer vectorized operations on 5D tensors or metadata lists. If logic must call a subject-oriented external API, the current low-level escape hatch is to unbatch, process every subject without changing its schema, -restack, and adopt the prior history: +and restack. Exact per-element histories are preserved automatically: ```python from typing import Any @@ -386,9 +386,7 @@ class StripIdentifier(tio.Transform): for subject in subjects: identifier = subject.metadata["identifier"] subject.metadata["identifier"] = identifier.strip() - rebuilt = tio.SubjectsBatch.from_subjects(subjects) - rebuilt.adopt_history(batch, subjects) - return rebuilt + return tio.SubjectsBatch.from_subjects(subjects) subject = tio.Subject( diff --git a/src/torchio/data/batch.py b/src/torchio/data/batch.py index 6aea3579d..15035df1c 100644 --- a/src/torchio/data/batch.py +++ b/src/torchio/data/batch.py @@ -3,7 +3,6 @@ from __future__ import annotations import copy as _copy -import dataclasses from collections.abc import Sequence from typing import TYPE_CHECKING from typing import Any @@ -13,13 +12,13 @@ from typing_extensions import Self from .affine import AffineMatrix +from .batch_history import _BatchedHistoryMixin from .batch_schema import _ImageSchema from .batch_schema import _SubjectSchema from .bboxes import BoundingBoxes from .image import Image from .image import LabelMap from .image import ScalarImage -from .invertible import Invertible from .points import Points if TYPE_CHECKING: @@ -29,7 +28,7 @@ _BATCH_META_KEYS = ("_batch_size", "_batched_keys", "_keep") -class ImagesBatch(Invertible): +class ImagesBatch(_BatchedHistoryMixin): """A batch of images with per-sample affines and private prototypes. Wraps a 5D tensor `(B, C, I, J, K)` and a list of `AffineMatrix` @@ -57,6 +56,7 @@ def _initialize( data: Tensor, affines: Sequence[AffineMatrix], prototypes: Sequence[Image], + histories: Sequence[Sequence[Any]] | None = None, ) -> None: """Initialize a validated image batch.""" if data.ndim != 5: @@ -74,7 +74,7 @@ def _initialize( self._data = data self._affines = [affine.clone() for affine in affines] self._prototypes = list(prototypes) - self.applied_transforms: list[Any] = [] + self._initialize_histories(data.shape[0], histories) @classmethod def _from_parts( @@ -82,10 +82,11 @@ def _from_parts( data: Tensor, affines: Sequence[AffineMatrix], prototypes: Sequence[Image], + histories: Sequence[Sequence[Any]] | None = None, ) -> Self: """Build an image batch from validated internal parts.""" batch = cls.__new__(cls) - batch._initialize(data, affines, prototypes) + batch._initialize(data, affines, prototypes, histories) return batch @classmethod @@ -139,7 +140,8 @@ def from_images(cls, images: Sequence[Image]) -> Self: stacked = torch.stack(tensors) affines = [image.affine for image in images] prototypes = [_make_image_prototype(image) for image in images] - return cls._from_parts(stacked, affines, prototypes) + histories = [image.applied_transforms for image in images] + return cls._from_parts(stacked, affines, prototypes, histories) @property def data(self) -> Tensor: @@ -210,10 +212,7 @@ def __getitem__(self, index: int) -> Image: affine=self._affines[index].clone(), ) image._metadata = _copy.deepcopy(prototype.metadata) - image.applied_transforms = [ - *prototype.applied_transforms, - *self.applied_transforms, - ] + image.applied_transforms = list(self.history(index)) return image def __len__(self) -> int: @@ -223,6 +222,10 @@ def unbatch(self) -> list[Image]: """Split the batch into individual images.""" return [self[i] for i in range(self.batch_size)] + def _batch_items(self, items: Sequence[Any]) -> Self: + """Rebuild an image batch from images.""" + return type(self).from_images(items) + @property def has_annotations(self) -> bool: """Whether any image prototype carries annotations.""" @@ -237,7 +240,7 @@ def __repr__(self) -> str: return f"ImagesBatch({cls}, batch_size={b}, shape=({c}, {i}, {j}, {k}))" -class SubjectsBatch(Invertible): +class SubjectsBatch(_BatchedHistoryMixin): """A batch of image columns and per-element object stores. Each image field becomes an `ImagesBatch`. Metadata, points, and @@ -271,32 +274,7 @@ def __init__( self._metadata, ) self._schema: _SubjectSchema | None = None - self.applied_transforms: list[Any] = [] - # When per-element branching occurs (e.g. per-instance OneOf), - # this stores the frozen per-element history prefix. Transforms - # applied afterwards still append to `applied_transforms`, and - # `unbatch()` merges the prefix with the sliced suffix. - self._per_element_history: list[list[Any]] | None = None - - def set_per_element_history(self, histories: list[list[Any]]) -> None: - """Freeze a distinct transform history for each batch element. - - Used when different elements receive different transforms (for - example per-instance [`OneOf`][torchio.OneOf]). Resets the shared - `applied_transforms` so that subsequent transforms accumulate as - a common suffix. - - Args: - histories: One history list per batch element. - """ - if len(histories) != self.batch_size: - msg = ( - f"Expected {self.batch_size} per-element histories," - f" got {len(histories)}" - ) - raise ValueError(msg) - self._per_element_history = [list(history) for history in histories] - self.applied_transforms = [] + self._initialize_histories(self._batch_size) @classmethod def from_subjects(cls, subjects: Sequence[Any]) -> Self: @@ -316,6 +294,9 @@ def from_subjects(cls, subjects: Sequence[Any]) -> Self: metadata=_collect_subject_metadata(subjects, schema), ) batch._schema = schema + batch._set_histories( + [subject.applied_transforms for subject in subjects], + ) return batch @property @@ -448,76 +429,17 @@ def unbatch(self) -> list[Any]: for key, values in self._metadata.items(): kwargs[key] = _copy.deepcopy(values[i]) sub = Subject(**kwargs) - suffix = _slice_history(self.applied_transforms, i) - if self._per_element_history is not None: - sub.applied_transforms = list(self._per_element_history[i]) + suffix - else: - sub.applied_transforms = suffix + sub.applied_transforms = list(self.history(i)) subjects.append(sub) return subjects + def _batch_items(self, items: Sequence[Any]) -> Self: + """Rebuild a subject batch from subjects.""" + return type(self).from_subjects(items) + def __len__(self) -> int: return self.batch_size - def adopt_history(self, source: SubjectsBatch, subjects: list[Any]) -> None: - """Carry transform history from *source* after rebuilding the batch. - - Used by code that unbatches, processes, and re-stacks subjects - (for example the MONAI and Cornucopia adapters). Preserves a - per-element history if *source* had one, otherwise copies the - shared history. - - Args: - source: The batch the subjects were unbatched from. - subjects: The processed subjects, in batch order. - """ - if source._per_element_history is not None: - self.set_per_element_history([s.applied_transforms for s in subjects]) - else: - self.applied_transforms = list(source.applied_transforms) - - def clear_history(self) -> None: - """Remove all applied transform records, including per-element ones.""" - self.applied_transforms = [] - self._per_element_history = None - - def get_inverse_transform(self, **kwargs: Any) -> Any: - """Build a transform that inverts the recorded history. - - Raises: - RuntimeError: If the batch carries per-element histories (from - a per-instance `OneOf`/`SomeOf`), since a single batch - inverse is ambiguous. Call `apply_inverse_transform` - (which inverts each element) or `unbatch()` and invert - each subject. - """ - if self._per_element_history is not None: - msg = ( - "This batch has per-element transform histories from a" - " per-instance OneOf/SomeOf, so a single batch inverse is" - " ambiguous. Call apply_inverse_transform() (which inverts" - " each element) or unbatch() and invert each subject." - ) - raise RuntimeError(msg) - return super().get_inverse_transform(**kwargs) - - def apply_inverse_transform(self, **kwargs: Any) -> SubjectsBatch: - """Apply the inverse of the recorded history. - - When the batch carries per-element histories, each element is - inverted independently and the results are re-stacked. - - Args: - **kwargs: Forwarded to `get_inverse_transform`. - - Returns: - A batch with the transforms undone. - """ - if self._per_element_history is not None: - inverted = [s.apply_inverse_transform(**kwargs) for s in self.unbatch()] - return type(self).from_subjects(inverted) - return super().apply_inverse_transform(**kwargs) - def __repr__(self) -> str: fields = [] for label, store in ( @@ -704,40 +626,3 @@ def _slice_params( else: sliced[key] = value return sliced - - -def _slice_history(history: list[Any], index: int) -> list[Any]: - """Build the per-subject transform history for batch element *index*. - - Batch-shared traces are copied unchanged. Per-instance traces are - sliced to the element's own parameters, and traces whose per-element - keep mask excludes this element are dropped. - - Args: - history: The batch-level list of `AppliedTransform` records. - index: The batch element whose history to build. - - Returns: - The list of `AppliedTransform` records for the element. - """ - sliced: list[Any] = [] - for trace in history: - params = getattr(trace, "params", None) - if not isinstance(params, dict) or "_batched_keys" not in params: - sliced.append(trace) - continue - expected_size = params.get("_batch_size") - if expected_size is not None and not 0 <= index < expected_size: - msg = ( - f"Cannot extract per-instance history for element {index}:" - f" the transform was recorded for a batch of size" - f" {expected_size}" - ) - raise IndexError(msg) - keep = params.get("_keep") - if keep is not None and not keep[index]: - continue - batched_keys = params["_batched_keys"] - new_params = _slice_params(params, index, batched_keys) - sliced.append(dataclasses.replace(trace, params=new_params)) - return sliced diff --git a/src/torchio/data/batch_history.py b/src/torchio/data/batch_history.py new file mode 100644 index 000000000..f82d42c8c --- /dev/null +++ b/src/torchio/data/batch_history.py @@ -0,0 +1,155 @@ +"""Exact per-element history support for batch containers.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from typing_extensions import Self + +from .invertible import Invertible + + +class _BatchedHistoryMixin(Invertible): + """Store one exact transform history per batch element.""" + + _histories: list[list[Any]] + + @property + def batch_size(self) -> int: + """Number of batch elements.""" + raise NotImplementedError + + def unbatch(self) -> list[Any]: + """Return individual batch elements.""" + raise NotImplementedError + + def _initialize_histories( + self, + batch_size: int, + histories: Sequence[Sequence[Any]] | None = None, + ) -> None: + """Initialize exact per-element histories.""" + if histories is None: + self._histories = [[] for _ in range(batch_size)] + return + self._set_histories(histories) + + @property + def histories(self) -> tuple[tuple[Any, ...], ...]: + """Immutable view of every element's exact history.""" + return tuple(tuple(history) for history in self._histories) + + def history(self, index: int) -> tuple[Any, ...]: + """Return one element's exact history. + + Args: + index: Batch element index. + + Returns: + Immutable transform-history view for the element. + """ + if not 0 <= index < self.batch_size: + msg = ( + f"Cannot get history for element {index}:" + f" batch size is {self.batch_size}" + ) + raise IndexError(msg) + return tuple(self._histories[index]) + + @property + def has_divergent_history(self) -> bool: + """Whether element histories differ.""" + if self.batch_size <= 1: + return False + first = self._histories[0] + return any( + not _histories_equal(first, history) for history in self._histories[1:] + ) + + @property + def applied_transforms(self) -> tuple[Any, ...]: + """Immutable uniform batch history for compatibility. + + Raises: + RuntimeError: If element histories differ. + """ + if self.batch_size == 0: + return () + if self.has_divergent_history: + msg = ( + "This batch has divergent element histories. Use `histories`" + " or `history(index)` instead of `applied_transforms`." + ) + raise RuntimeError(msg) + return tuple(self._histories[0]) + + @applied_transforms.setter + def applied_transforms(self, history: Sequence[Any]) -> None: + copied = list(history) + self._histories = [list(copied) for _ in range(self.batch_size)] + + def _set_histories(self, histories: Sequence[Sequence[Any]]) -> None: + """Replace every element history.""" + if len(histories) != self.batch_size: + msg = f"Expected {self.batch_size} histories, got {len(histories)}" + raise ValueError(msg) + self._histories = [list(history) for history in histories] + + def _append_history(self, traces: Sequence[Any | None]) -> None: + """Append one optional trace per batch element.""" + if len(traces) != self.batch_size: + msg = f"Expected {self.batch_size} traces, got {len(traces)}" + raise ValueError(msg) + for history, trace in zip(self._histories, traces, strict=True): + if trace is not None: + history.append(trace) + + def clear_history(self) -> None: + """Remove every element history.""" + self._histories = [[] for _ in range(self.batch_size)] + + def get_inverse_transform(self, **kwargs: Any) -> Any: + """Build a vectorized inverse for a uniform batch history. + + Args: + **kwargs: Forwarded to `Invertible.get_inverse_transform`. + + Raises: + RuntimeError: If element histories differ. + """ + if self.has_divergent_history: + msg = ( + "This batch has divergent element histories, so one vectorized" + " inverse is ambiguous. Use `apply_inverse_transform()`." + ) + raise RuntimeError(msg) + return super().get_inverse_transform(**kwargs) + + def apply_inverse_transform(self, **kwargs: Any) -> Self: + """Apply vectorized or per-element inverse transforms. + + Args: + **kwargs: Forwarded to `get_inverse_transform`. + + Returns: + A batch with transforms undone. + """ + if not self.has_divergent_history: + return super().apply_inverse_transform(**kwargs) + inverted = [item.apply_inverse_transform(**kwargs) for item in self.unbatch()] + result = self._batch_items(inverted) + result.clear_history() + return result + + def _batch_items(self, items: Sequence[Any]) -> Self: + """Rebuild the concrete batch type from unbatched items.""" + raise NotImplementedError + + +def _histories_equal(first: Sequence[Any], second: Sequence[Any]) -> bool: + """Compare histories, treating ambiguous value equality as unequal.""" + try: + return bool(first == second) + except (RuntimeError, TypeError, ValueError): + return False diff --git a/src/torchio/data/image.py b/src/torchio/data/image.py index 000ee4366..eb6b5895f 100644 --- a/src/torchio/data/image.py +++ b/src/torchio/data/image.py @@ -157,6 +157,8 @@ class Image(Invertible): >>> image = tio.ScalarImage(nifti_image) # from nibabel (lazy) """ + applied_transforms: list[Any] + #: Source types accepted by the constructor. ImageInput = ( ImageSource # str | Path | IOBase | OpenFile diff --git a/src/torchio/data/invertible.py b/src/torchio/data/invertible.py index 65a890f6b..2267754f4 100644 --- a/src/torchio/data/invertible.py +++ b/src/torchio/data/invertible.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import Any from typing_extensions import Self @@ -13,11 +14,12 @@ class Invertible: Provides `apply_inverse_transform()` and `get_inverse_transform()` to undo recorded transforms. - Classes that inherit from this mixin must initialise - `self.applied_transforms = []` in their constructor. + Subclasses must expose an `applied_transforms` sequence. `Subject` + and `Image` use mutable lists; batch containers provide an immutable + uniform-history view backed by their per-element history store. """ - applied_transforms: list[Any] + applied_transforms: Sequence[Any] def get_inverse_transform( self, diff --git a/src/torchio/data/subject.py b/src/torchio/data/subject.py index fe0d5e543..5f5e323db 100644 --- a/src/torchio/data/subject.py +++ b/src/torchio/data/subject.py @@ -56,6 +56,8 @@ class Subject(Invertible): >>> subject.age # metadata access (returns 45) """ + applied_transforms: list[Any] + def __init__(self, **kwargs: Any) -> None: images: dict[str, Image] = {} points: dict[str, Points] = {} diff --git a/src/torchio/transforms/compose.py b/src/torchio/transforms/compose.py index 68d9a97be..244dca1b0 100644 --- a/src/torchio/transforms/compose.py +++ b/src/torchio/transforms/compose.py @@ -283,9 +283,7 @@ def to_hydra(self) -> dict[str, Any]: def _apply_to_element(subject: Any, apply_fn: Any) -> Any: """Apply a transform (or callable) to one element, preserving history. - Wrapping a `Subject` directly into a batch discards its existing - history, so the element is wrapped into a one-element batch seeded - with its prior history; the transform then appends to that history. + The one-element batch inherits the subject's exact history. Args: subject: The single subject to transform (carrying its history). @@ -298,7 +296,6 @@ def _apply_to_element(subject: Any, apply_fn: Any) -> Any: from ..data.batch import SubjectsBatch element_batch = SubjectsBatch.from_subjects([subject]) - element_batch.applied_transforms = list(subject.applied_transforms) element_batch = apply_fn(element_batch) return element_batch.unbatch()[0] @@ -317,10 +314,9 @@ 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 (TypeError, ValueError) as error: msg = ( f"Per-instance {transform_name} produced batch elements with" " different shapes or schemas, which cannot be re-stacked. Use" @@ -328,35 +324,4 @@ def _rebatch_with_history(subjects: list[Any], transform_name: str) -> Any: 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]) 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) diff --git a/src/torchio/transforms/cornucopia_adapter.py b/src/torchio/transforms/cornucopia_adapter.py index f8b0588d7..1edef558f 100644 --- a/src/torchio/transforms/cornucopia_adapter.py +++ b/src/torchio/transforms/cornucopia_adapter.py @@ -76,7 +76,6 @@ def forward(self, data: Any) -> Any: from ..data.batch import SubjectsBatch result = SubjectsBatch.from_subjects(subjects) - result.adopt_history(batch, subjects) return unwrap(result) def apply_transform( diff --git a/src/torchio/transforms/inverse.py b/src/torchio/transforms/inverse.py index 56bff7ba2..465f316f2 100644 --- a/src/torchio/transforms/inverse.py +++ b/src/torchio/transforms/inverse.py @@ -3,6 +3,7 @@ from __future__ import annotations import warnings +from collections.abc import Sequence from typing import Any from .compose import Compose @@ -13,7 +14,7 @@ def get_inverse_transform( - history: list[AppliedTransform], + history: Sequence[AppliedTransform], *, warn: bool = True, ignore_intensity: bool = False, @@ -80,15 +81,16 @@ def apply_inverse_transform( Returns: Data with transforms undone, same type as input. """ - if not hasattr(data, "applied_transforms"): - return data - # Batches with per-element histories (from per-instance OneOf/SomeOf) - # know how to invert each element; delegate to their own method. - if getattr(data, "_per_element_history", None) is not None: + from ..data.batch import ImagesBatch + from ..data.batch import SubjectsBatch + + if isinstance(data, (ImagesBatch, SubjectsBatch)): return data.apply_inverse_transform( warn=warn, ignore_intensity=ignore_intensity, ) + if not hasattr(data, "applied_transforms"): + return data inverse = get_inverse_transform( data.applied_transforms, warn=warn, diff --git a/src/torchio/transforms/monai_adapter.py b/src/torchio/transforms/monai_adapter.py index 64cb94b40..21491c775 100644 --- a/src/torchio/transforms/monai_adapter.py +++ b/src/torchio/transforms/monai_adapter.py @@ -90,7 +90,6 @@ def forward(self, data): from ..data.batch import SubjectsBatch result = SubjectsBatch.from_subjects(subjects) - result.adopt_history(batch, subjects) return unwrap(result) def apply_transform(self, batch: Any, params: dict[str, Any]) -> Any: diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py index 2fbb4882f..40cc7a1eb 100644 --- a/src/torchio/transforms/transform.py +++ b/src/torchio/transforms/transform.py @@ -21,6 +21,7 @@ from ..data.batch import ImagesBatch from ..data.batch import SubjectsBatch +from ..data.batch import _slice_params from ..data.bboxes import BoundingBoxes from ..data.image import Image from ..data.image import ScalarImage @@ -49,21 +50,6 @@ class AppliedTransform: _TRANSFORM_REGISTRY: dict[str, type[Transform]] = {} -def _all_elements_gated_out(params: dict[str, Any]) -> bool: - """Whether per-element gating masked out every batch element. - - Args: - params: The parameter dict produced by `make_params`, possibly - carrying a `_keep` mask added by `_tag_batched`. - - Returns: - `True` only when a `_keep` mask is present and none of its - elements are kept, i.e. the transform was an exact no-op. - """ - keep = params.get("_keep") - return keep is not None and not any(keep) - - def _copy_optional_list(value: list[str] | None) -> list[str] | None: return None if value is None else list(value) @@ -263,34 +249,62 @@ def forward(self, data: Any) -> Any: if not self._per_instance_p_active(batch) and torch.rand(1).item() >= self.p: return unwrap(batch) params = self.make_params(batch) - if not _all_elements_gated_out(params): + traces = self._build_history_traces(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) - # Record history on the batch, unless every element was gated out by - # per-element probability: that is an exact no-op, and recording it - # would let history replay (e.g. an invertible spatial transform) - # trigger an unnecessary identity resample. - if not _all_elements_gated_out(params): - trace = AppliedTransform( - name=type(self).__name__, - params=params, - include=_copy_optional_list(self.include), - exclude=_copy_optional_list(self.exclude), - ) - if not hasattr(batch, "applied_transforms"): - batch.applied_transforms = [] - batch.applied_transforms.append(trace) + batch._append_history(traces) result = unwrap(batch) # Propagate history to outputs that can carry it - if ( - hasattr(batch, "applied_transforms") - and not isinstance(result, (SubjectsBatch, Tensor, np.ndarray)) - and not isinstance(result, dict) - ): + 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) return result + 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: + return [ + self._make_applied_transform(_copy.deepcopy(params)) + for _ in range(batch_size) + ] + expected_size = params.get("_batch_size") + if expected_size != batch_size: + msg = ( + f"Parameter batch size {expected_size} does not match" + f" input batch size {batch_size}" + ) + raise ValueError(msg) + keep = params.get("_keep") + traces: list[AppliedTransform | None] = [] + for index in range(batch_size): + if keep is not None and not keep[index]: + traces.append(None) + continue + element_params = _slice_params(params, index, batched_keys) + traces.append(self._make_applied_transform(element_params)) + return traces + + def _make_applied_transform( + self, + params: dict[str, Any], + ) -> AppliedTransform: + """Build one history trace.""" + return AppliedTransform( + name=type(self).__name__, + params=params, + include=_copy_optional_list(self.include), + exclude=_copy_optional_list(self.exclude), + ) + def _check_spatial_annotations(self, data: Any) -> None: """Reject spatial transforms that would leave stale annotations.""" if isinstance(self, SpatialTransform) and _data_has_annotations(data): @@ -545,6 +559,7 @@ def _wrap( return data, _unwrap_subjects_batch case ImagesBatch(): sb = SubjectsBatch({"tio_default_image": data}) + sb._set_histories(data.histories) return sb, _unwrap_images_batch case Subject(): sb = SubjectsBatch.from_subjects([data]) @@ -560,6 +575,7 @@ def _wrap_single_image(img: Image, unwrap_fn: Any) -> tuple[Any, Any]: from ..data.batch import SubjectsBatch sub = Subject(tio_default_image=img) + sub.applied_transforms = list(img.applied_transforms) sb = SubjectsBatch.from_subjects([sub]) return sb, unwrap_fn @@ -655,7 +671,9 @@ def _unwrap_subjects_batch(batch: SubjectsBatch) -> SubjectsBatch: def _unwrap_images_batch(batch: SubjectsBatch) -> ImagesBatch: - return batch.images["tio_default_image"] + image_batch = batch.images["tio_default_image"] + image_batch._set_histories(batch.histories) + return image_batch def _unwrap_subject(batch: SubjectsBatch) -> Subject: @@ -664,7 +682,9 @@ def _unwrap_subject(batch: SubjectsBatch) -> Subject: def _unwrap_image(batch: SubjectsBatch) -> Image: sub = batch.unbatch()[0] - return sub.tio_default_image + image = sub.tio_default_image + image.applied_transforms = list(sub.applied_transforms) + return image def _unwrap_tensor(batch: SubjectsBatch) -> Tensor: diff --git a/tests/conftest.py b/tests/conftest.py index 95cd5456d..32177b9bb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,7 +10,6 @@ import torch from torchio.data.batch import SubjectsBatch -from torchio.data.batch import _slice_params @pytest.fixture @@ -38,10 +37,6 @@ def _assert( ) -> None: original = copy.deepcopy(batch) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params, "per-instance path was not active" - batched_keys = params["_batched_keys"] - keep = params.get("_keep") image_names = list(transform._get_images(result).keys()) result_images = transform._get_images(result) original_subjects = original.unbatch() @@ -51,10 +46,16 @@ def _assert( name: image.data.clone() for name, image in transform._get_images(single).items() } - element_params = _slice_params(params, index, batched_keys) - single = transform.apply_transform(single, element_params) + prior_history_size = len(original.history(index)) + history = result.history(index) + gated_out = len(history) == prior_history_size + if not gated_out: + assert len(history) == prior_history_size + 1 + single = transform.apply_transform( + single, + history[-1].params, + ) single_images = transform._get_images(single) - gated_out = keep is not None and not keep[index] for name in image_names: result_row = result_images[name].data[index : index + 1] torch.testing.assert_close( diff --git a/tests/test_anisotropy.py b/tests/test_anisotropy.py index 5692b28cd..5b4f17bb5 100644 --- a/tests/test_anisotropy.py +++ b/tests/test_anisotropy.py @@ -63,9 +63,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Anisotropy(downsampling=(2.0, 5.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["factor"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["factor"] for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_batch.py b/tests/test_batch.py index e9331cdc0..7a045ecd1 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -486,6 +486,14 @@ def test_from_images_empty_raises(self) -> None: with pytest.raises(ValueError, match="empty"): ImagesBatch.from_images([]) + def test_empty_history_view_is_defensive(self) -> None: + batch = object.__new__(ImagesBatch) + batch._data = torch.empty(0, 1, 1, 1, 1) + batch._histories = [] + + assert batch.has_divergent_history is False + assert batch.applied_transforms == () + def test_data_setter_non_5d_raises(self) -> None: from torchio.data.batch import ImagesBatch @@ -565,15 +573,12 @@ def _batch(self, batch_size: int = 4) -> SubjectsBatch: ] return SubjectsBatch.from_subjects(subjects) - def test_adopt_history_preserves_per_element(self) -> None: - # Simulate the adapter pattern: a per-element batch is unbatched, - # processed, and re-stacked; history must survive. + def test_restack_preserves_divergent_history(self) -> None: torch.manual_seed(0) batch = self._batch() branched = tio.OneOf([tio.Flip(axes=(0,)), tio.Flip(axes=(1,))])(batch) subjects = branched.unbatch() rebuilt = SubjectsBatch.from_subjects(subjects) - rebuilt.adopt_history(branched, subjects) for original, restored in zip( branched.unbatch(), rebuilt.unbatch(), @@ -583,13 +588,19 @@ def test_adopt_history_preserves_per_element(self) -> None: t.name for t in original.applied_transforms ] - def test_adopt_history_shared_case(self) -> None: + def test_restack_preserves_uniform_history(self) -> None: torch.manual_seed(0) batch = self._batch() transformed = tio.Gamma(log_gamma=0.3, per_instance=False)(batch) subjects = transformed.unbatch() rebuilt = SubjectsBatch.from_subjects(subjects) - rebuilt.adopt_history(transformed, subjects) - assert rebuilt._per_element_history is None + assert not rebuilt.has_divergent_history for subject in rebuilt.unbatch(): assert [t.name for t in subject.applied_transforms] == ["Gamma"] + + def test_uniform_applied_transforms_view_is_immutable(self) -> None: + result = tio.Gamma(log_gamma=0.2, per_instance=False)(self._batch()) + + assert isinstance(result.applied_transforms, tuple) + with pytest.raises(AttributeError): + result.applied_transforms.append("invalid") # type: ignore[attr-defined] diff --git a/tests/test_bias_field.py b/tests/test_bias_field.py index 9167fd7f1..f09346479 100644 --- a/tests/test_bias_field.py +++ b/tests/test_bias_field.py @@ -94,9 +94,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.BiasField(std=(0.3, 0.6))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["std"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["std"] for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_shares_std(self) -> None: diff --git a/tests/test_blur.py b/tests/test_blur.py index 101b84fbb..32374f1b1 100644 --- a/tests/test_blur.py +++ b/tests/test_blur.py @@ -50,9 +50,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Blur(std=(1.0, 4.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["std"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({tuple(item["std"]) for item in params}) > 1 data = result.t1.data assert not torch.allclose(data[0], data[1]) diff --git a/tests/test_flip.py b/tests/test_flip.py index 60b766cdf..26f7736bc 100644 --- a/tests/test_flip.py +++ b/tests/test_flip.py @@ -263,10 +263,9 @@ def test_per_instance_axes_differ_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Flip(axes=(0, 1, 2), flip_probability=0.5)(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["axes"]) == batch.batch_size - distinct = {tuple(a) for a in params["axes"]} + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + distinct = {tuple(item["axes"]) for item in params} assert len(distinct) > 1 def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_gamma.py b/tests/test_gamma.py index 8a82abb21..fb062139e 100644 --- a/tests/test_gamma.py +++ b/tests/test_gamma.py @@ -92,10 +92,9 @@ def test_per_instance_default_differs_across_batch(self) -> None: batch = self._batch() transform = tio.Gamma(log_gamma=(0.2, 0.8)) result = transform(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["log_gamma"], list) - assert len(params["log_gamma"]) == batch.batch_size - assert len(set(params["log_gamma"])) > 1 + log_gammas = [history[-1].params["log_gamma"] for history in result.histories] + assert all(isinstance(value, float) for value in log_gammas) + assert len(set(log_gammas)) > 1 def test_per_instance_false_is_shared(self) -> None: torch.manual_seed(0) @@ -117,7 +116,7 @@ def test_per_instance_values_applied(self) -> None: original = batch.t1.data.clone() transform = tio.Gamma(log_gamma=(0.2, 0.8)) result = transform(batch) - log_gammas = result.applied_transforms[-1].params["log_gamma"] + log_gammas = [history[-1].params["log_gamma"] for history in result.histories] for i, log_gamma in enumerate(log_gammas): gamma = torch.tensor(log_gamma).exp() expected = original[i].sign() * original[i].abs().pow(gamma) diff --git a/tests/test_ghosting.py b/tests/test_ghosting.py index 6c97ef233..a0d9dc2b2 100644 --- a/tests/test_ghosting.py +++ b/tests/test_ghosting.py @@ -61,9 +61,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Ghosting(intensity=(0.5, 1.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["intensity"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["intensity"] for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_labels_to_image.py b/tests/test_labels_to_image.py index 70a0942e4..bad8992a1 100644 --- a/tests/test_labels_to_image.py +++ b/tests/test_labels_to_image.py @@ -72,10 +72,9 @@ def test_per_instance_means_differ_across_batch(self) -> None: batch = self._batch() transform = tio.LabelsToImage(label_key="seg", default_mean=(0.2, 0.9)) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["means"]) == batch.batch_size - means_for_label_1 = [m[1] for m in params["means"]] + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + means_for_label_1 = [item["means"][1] for item in params] assert len(set(means_for_label_1)) > 1 assert result.image_from_labels.data.shape[0] == batch.batch_size @@ -117,18 +116,19 @@ def test_each_element_uses_its_own_label_stats(self) -> None: ) torch.manual_seed(1) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params image = result.img.data for index in range(batch.batch_size): + params = result.history(index)[-1].params region_one = image[index, 0, : size // 2] region_two = image[index, 0, size // 2 :] assert region_one.mean().item() == pytest.approx( - params["means"][index][1], abs=0.5 + params["means"][1], abs=0.5 ) assert region_two.mean().item() == pytest.approx( - params["means"][index][2], abs=0.5 + params["means"][2], abs=0.5 ) # Independent per-element sampling: means vary across the batch. - label_one_means = {round(params["means"][i][1], 3) for i in range(3)} + label_one_means = { + round(result.history(i)[-1].params["means"][1], 3) for i in range(3) + } assert len(label_one_means) > 1 diff --git a/tests/test_motion.py b/tests/test_motion.py index b3c5d1304..5979b610f 100644 --- a/tests/test_motion.py +++ b/tests/test_motion.py @@ -61,9 +61,9 @@ def test_per_instance_differs_across_batch(self) -> None: result = tio.Motion(degrees=(5, 15), translation=(5, 15), num_transforms=2)( batch ) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["transforms"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({repr(item["transforms"]) for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_noise.py b/tests/test_noise.py index e5c912603..767b0e6fb 100644 --- a/tests/test_noise.py +++ b/tests/test_noise.py @@ -244,10 +244,9 @@ def test_per_instance_default_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Noise(std=(0.5, 1.5))(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["std"], list) - assert len(params["std"]) == batch.batch_size - assert len(set(params["std"])) > 1 + stds = [history[-1].params["std"] for history in result.histories] + assert all(isinstance(value, float) for value in stds) + assert len(set(stds)) > 1 def test_per_instance_false_is_shared(self) -> None: torch.manual_seed(0) @@ -265,7 +264,7 @@ def test_per_instance_mean_applied_per_element(self) -> None: torch.manual_seed(0) batch = self._batch(batch_size=5) result = tio.Noise(mean=(5.0, 20.0), std=0.0)(batch) - means = result.applied_transforms[-1].params["mean"] + means = [history[-1].params["mean"] for history in result.histories] for i, mean in enumerate(means): torch.testing.assert_close( result.t1.data[i].mean(), diff --git a/tests/test_normalize.py b/tests/test_normalize.py index 788e605c0..f4a05b219 100644 --- a/tests/test_normalize.py +++ b/tests/test_normalize.py @@ -278,15 +278,14 @@ def test_per_instance_out_range_differs(self) -> None: batch = self._batch() transform = tio.RescaleIntensity(out_min=(-1.0, 0.0), out_max=(0.5, 1.0)) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["out_min"]) == batch.batch_size - assert len(set(params["out_min"])) > 1 + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["out_min"] for item in params}) > 1 # Each element rescaled to its own output range. for i in range(batch.batch_size): data = result.t1.data[i] - assert data.min() >= params["out_min"][i] - 1e-4 - assert data.max() <= params["out_max"][i] + 1e-4 + assert data.min() >= params[i]["out_min"] - 1e-4 + assert data.max() <= params[i]["out_max"] + 1e-4 def test_per_instance_false_shares_params(self) -> None: torch.manual_seed(0) @@ -314,7 +313,9 @@ def test_per_instance_inverse_zero_range_no_nan(self) -> None: batch = tio.SubjectsBatch.from_subjects(subjects) transform = tio.RescaleIntensity(out_min=0.0, out_max=0.0) result = transform(batch) - assert "_batched_keys" in result.applied_transforms[-1].params + assert all( + "_batched_keys" not in history[-1].params for history in result.histories + ) restored = result.apply_inverse_transform() assert not torch.isnan(restored.t1.data).any() # Identical inputs: the batch-shared input range covers every diff --git a/tests/test_one_of.py b/tests/test_one_of.py index 1d03b5fee..24db9617d 100644 --- a/tests/test_one_of.py +++ b/tests/test_one_of.py @@ -127,9 +127,10 @@ def test_functional_inverse_restores_per_element(self) -> None: def test_get_inverse_transform_raises_for_per_element(self) -> None: torch.manual_seed(0) - batch = self._batch(batch_size=4) - result = tio.OneOf([tio.Flip(axes=(0,))])(batch) - with pytest.raises(RuntimeError, match="per-element"): + batch = self._batch(batch_size=16) + result = tio.OneOf([tio.Flip(axes=(0,)), tio.Flip(axes=(1,))])(batch) + assert result.has_divergent_history + with pytest.raises(RuntimeError, match="divergent"): result.get_inverse_transform() def test_clear_history_clears_per_element(self) -> None: @@ -137,7 +138,7 @@ def test_clear_history_clears_per_element(self) -> None: batch = self._batch(batch_size=4) result = tio.OneOf([tio.Flip(axes=(0,))])(batch) result.clear_history() - assert result._per_element_history is None + assert all(not history for history in result.histories) for subject in result.unbatch(): assert subject.applied_transforms == [] @@ -149,7 +150,7 @@ def test_p_zero_is_noop_preserving_history(self) -> None: flipped = tio.Flip(axes=(0,))(batch) result = tio.OneOf([tio.Flip(axes=(1,))], p=0.0)(flipped) torch.testing.assert_close(result.t1.data, flipped.t1.data) - assert result._per_element_history is None + assert not result.has_divergent_history # The shared Flip history is intact and still invertible as a batch. restored = result.apply_inverse_transform() torch.testing.assert_close(restored.t1.data, original) @@ -180,3 +181,10 @@ def apply_transform(self, batch, params): tio.OneOf([_Spy()])(_make_subject()) assert seen == [False] + + +def test_rebatch_rejects_non_subject_elements() -> None: + from torchio.transforms.compose import _rebatch_with_history + + with pytest.raises(RuntimeError, match="cannot be re-stacked"): + _rebatch_with_history([object()], "OneOf") diff --git a/tests/test_per_instance.py b/tests/test_per_instance.py index d5e0bdef6..6e59430bf 100644 --- a/tests/test_per_instance.py +++ b/tests/test_per_instance.py @@ -82,9 +82,9 @@ def test_compose_child_is_per_instance(self) -> None: batch = _identical_batch() pipeline = tio.Compose([tio.Gamma(log_gamma=(0.2, 0.8))]) result = pipeline(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["log_gamma"], list) - assert len(set(params["log_gamma"])) > 1 + log_gammas = [history[-1].params["log_gamma"] for history in result.histories] + assert all(isinstance(value, float) for value in log_gammas) + assert len(set(log_gammas)) > 1 def test_compose_respects_per_instance_false(self) -> None: torch.manual_seed(0) @@ -100,35 +100,40 @@ def test_unbatch_slices_history(self) -> None: torch.manual_seed(0) batch = _identical_batch(batch_size=4) result = tio.Gamma(log_gamma=(0.2, 0.8))(batch) - batch_log_gammas = result.applied_transforms[-1].params["log_gamma"] + batch_log_gammas = [ + history[-1].params["log_gamma"] for history in result.histories + ] for i, subject in enumerate(result.unbatch()): trace = subject.applied_transforms[-1] assert trace.params["log_gamma"] == batch_log_gammas[i] assert "_batched_keys" not in trace.params + def test_uniform_history_uses_independent_traces(self) -> None: + batch = _identical_batch(batch_size=4) + + result = tio.Gamma(log_gamma=0.2, per_instance=False)(batch) + + traces = [history[-1] for history in result.histories] + assert all(trace.params == traces[0].params for trace in traces) + assert all(trace is not traces[0] for trace in traces[1:]) + class TestSpatialBatchSizeValidation: def test_mismatched_batch_size_raises(self) -> None: torch.manual_seed(0) batch = _identical_batch(batch_size=4) transform = tio.Affine(degrees=(20.0, 80.0), default_pad_value=0.0) - result = transform(batch) - params = result.applied_transforms[-1].params + params = transform.make_params(batch) smaller = _identical_batch(batch_size=2) with pytest.raises(RuntimeError, match="Per-instance spatial parameters"): transform.apply_transform(smaller, params) - def test_history_slice_out_of_range_raises(self) -> None: - # Slicing per-instance history for an element beyond the recorded - # batch size must fail with a clear error rather than an opaque one. - from torchio.data.batch import _slice_history - + def test_history_index_out_of_range_raises(self) -> None: torch.manual_seed(0) batch = _identical_batch(batch_size=4) result = tio.Noise(std=(0.1, 0.5))(batch) - history = result.applied_transforms - with pytest.raises(IndexError, match="batch of size 4"): - _slice_history(history, 4) + with pytest.raises(IndexError, match=r"element 4.*batch size is 4"): + result.history(4) class TestPerInstanceDtypePreservation: @@ -176,7 +181,7 @@ def test_fully_gated_records_no_history(self) -> None: torch.manual_seed(0) batch = _identical_batch(batch_size=4) result = tio.Affine(degrees=20.0, p=0.0)(batch) - assert result.applied_transforms == [] + assert not result.applied_transforms def test_fully_gated_inverse_preserves_float64(self) -> None: torch.manual_seed(0) diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 4c98800cc..eed80fa3b 100644 --- a/tests/test_spatial.py +++ b/tests/test_spatial.py @@ -410,9 +410,9 @@ def test_per_instance_rotations_differ(self) -> None: batch = self._identical_batch() transform = AffineTransform(degrees=(20.0, 80.0), default_pad_value=0.0) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["affine_matrix"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({repr(item["affine_matrix"]) for item in params}) > 1 data = result.t1.data assert not torch.allclose(data[0], data[1]) assert not torch.allclose(data[1], data[2]) @@ -502,9 +502,9 @@ def test_per_instance_elastic_differs_across_batch(self) -> None: max_displacement=(1.0, 3.0), ) result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["control_points"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({repr(item["control_points"]) for item in params}) > 1 data = result.t1.data assert not torch.allclose(data[0], data[1]) assert not torch.allclose(data[1], data[2]) diff --git a/tests/test_spike.py b/tests/test_spike.py index 3d1fc1729..aaf31b346 100644 --- a/tests/test_spike.py +++ b/tests/test_spike.py @@ -56,9 +56,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Spike(intensity=(1.0, 3.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["intensity"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({item["intensity"] for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_swap.py b/tests/test_swap.py index 45aba36d5..eb09c4905 100644 --- a/tests/test_swap.py +++ b/tests/test_swap.py @@ -63,9 +63,9 @@ def test_per_instance_differs_across_batch(self) -> None: torch.manual_seed(0) batch = self._batch() result = tio.Swap(patch_size=4, num_iterations=20)(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["locations"]) == batch.batch_size + params = [history[-1].params for history in result.histories] + assert all("_batched_keys" not in item for item in params) + assert len({repr(item["locations"]) for item in params}) > 1 assert not torch.allclose(result.t1.data[0], result.t1.data[1]) def test_per_instance_false_is_shared(self) -> None: diff --git a/tests/test_vectorization.py b/tests/test_vectorization.py index ea7364d82..860c7f892 100644 --- a/tests/test_vectorization.py +++ b/tests/test_vectorization.py @@ -79,3 +79,13 @@ def test_anisotropy_tie_rounding_matches_scalar(assert_vectorized) -> None: [tio.Subject(t1=tio.ScalarImage(data.clone() + index)) for index in range(4)] ) assert_vectorized(tio.Anisotropy(downsampling=2.0), batch) + + +def test_gating_with_prior_history(assert_vectorized) -> None: + torch.manual_seed(0) + batch = tio.Gamma(log_gamma=0.2, per_instance=False)(_batch(batch_size=6)) + + assert_vectorized( + tio.Flip(axes=(0, 1, 2), flip_probability=1.0, p=0.5), + batch, + )