diff --git a/docs/concepts/transforms.md b/docs/concepts/transforms.md index 7bdd66ee3..7dd618c0b 100644 --- a/docs/concepts/transforms.md +++ b/docs/concepts/transforms.md @@ -10,42 +10,123 @@ and always return the same type. Internally, **all inputs are converted to a `SubjectsBatch`** before the transform runs. A single `Image` becomes a batch of size 1; a `SubjectsBatch` from a `DataLoader` passes through -directly. This means transform authors write **one method** that -works identically for single samples and batches: +directly. This means transform authors write **one batch-oriented +application method** that works identically for single samples and batches +(`apply_transform`), plus `make_params` when parameter construction is +needed: - ```python -class Flip(SpatialTransform): - def apply_transform(self, batch, params): - dims = [a - 3 for a in params["axes"]] - for name, img_batch in self._get_images(batch).items(): - img_batch._data = torch.flip(img_batch.data, dims) +from typing import Any + +import torch +import torchio as tio + + +class AddValue(tio.Transform): + """Add a fixed value to every image in a batch.""" + + def __init__(self, value: float) -> None: + super().__init__() + self.value = value + + def make_params(self, batch: tio.SubjectsBatch) -> dict[str, Any]: + """Return the value to add.""" + return {"value": self.value} + + def apply_transform( + self, + batch: tio.SubjectsBatch, + params: dict[str, Any], + ) -> tio.SubjectsBatch: + """Add the value to each 5D image tensor.""" + for image_batch in batch.images.values(): + image_batch.data = image_batch.data + params["value"] return batch + + +subject = tio.Subject( + image=tio.ScalarImage(torch.zeros(1, 2, 3, 4)), + site="A", +) +batch = tio.SubjectsBatch.from_subjects([subject]) +assert subject.image.data.shape == (1, 2, 3, 4) +assert batch.image.data.shape == (1, 1, 2, 3, 4) + +transformed = AddValue(2)(subject) +assert isinstance(transformed, tio.Subject) +assert transformed.image.data.shape == (1, 2, 3, 4) +assert torch.all(transformed.image.data == 2) +``` + +The public call performs the complete round trip: + +```text +Subject -> SubjectsBatch -> apply_transform -> Subject ``` -The negative dim indexing (`-3`, `-2`, `-1` for spatial axes) -works for both 5D `(B, C, I, J, K)` batch tensors and 5D -`(1, C, I, J, K)` single-sample tensors. +An image tensor shaped `(C, I, J, K)` therefore reaches +`apply_transform` as `(B, C, I, J, K)`. For a single `Subject`, +`B` is 1. Negative dimension indices (`-3`, `-2`, `-1`) identify +the spatial axes for both single-element and multi-element batches. When a `SubjectsBatch` is passed (e.g., from `SubjectsLoader`), transforms that support it sample **independent parameters per batch element** by default, so a single call produces diverse augmentations (see [Per-instance augmentation](per-instance-augmentation.md)). Pass `per_instance=False` to share one sampled parameter set across all -elements. Single inputs are unaffected. +elements. Fixed parameters are not sampled and therefore remain shared. +Single inputs are unaffected. -## The `make_params` / `apply` split +## The `make_params` / `apply_transform` split Every transform has two methods: -- **`make_params(subject)`**: sample random parameters (called once - per transform invocation). -- **`apply(subject, params)`**: apply the transform using those - parameters. +- **`make_params(batch)`**: create or sample parameters for the + `SubjectsBatch`. +- **`apply_transform(batch, params)`**: apply those parameters to the + `SubjectsBatch`. This separation (inspired by Torchvision V2) means the same random -parameters are applied consistently to all images, points, and bounding -boxes in a Subject. Params are saved in history for replay. +parameters are applied consistently to all images in a `Subject`. +Parameters are saved in history for inspection and inversion. + +!!! warning "`apply_transform` is a low-level kernel" + Application code should call the transform itself, for example + `result = transform(subject)`. Calling `apply_transform` directly + bypasses input wrapping, copying, probability handling, history + recording, and output-type restoration. It requires a + `SubjectsBatch`, not a `Subject`. + +### Metadata in a batch + +`Subject.metadata` is a `dict[str, Any]`. After batching, +`batch.metadata` is a `dict[str, list[Any]]`, with one value per batch +element: + +```python +import torch +import torchio as tio + +subjects = [ + tio.Subject( + image=tio.ScalarImage(torch.zeros(1, 2, 3, 4)), + site="A", + age=30, + ), + tio.Subject( + image=tio.ScalarImage(torch.ones(1, 2, 3, 4)), + site="B", + age=40, + ), +] +batch = tio.SubjectsBatch.from_subjects(subjects) +assert batch.metadata == {"site": ["A", "B"], "age": [30, 40]} +``` + +The first subject defines the image-name and metadata-key order of the +batch. All subjects must have the same schema, although their local +key order may differ. A custom transform should preserve that shared +schema and keep every metadata list aligned with the batch dimension. ## Scalar, range, or distribution: one class for both @@ -162,24 +243,22 @@ augment = tio.SomeOf( Every transform records an `AppliedTransform` in the Subject's `applied_transforms` list: - ```python -result = pipeline(subject) -for trace in result.applied_transforms: - print(trace.name, trace.params) +import torch +import torchio as tio + +subject = tio.Subject(image=tio.ScalarImage(torch.zeros(1, 2, 3, 4))) +result = tio.Noise(std=0.1)(subject) +trace = result.applied_transforms[-1] +assert trace.name == "Noise" +assert trace.params["std"] == 0.1 ``` -**Replay** applies the exact same augmentation to different data: - - -```python -# Get the params from history -params = result.applied_transforms[0].params - -# Replay on a new subject -noise = tio.Noise(std=0.1) -replayed = noise.apply_transform(new_subject, params) -``` +History parameters support inspection and inversion. TorchIO does not +currently expose a public API for applying an arbitrary saved parameter +dictionary to another input. In particular, do not use +`apply_transform(new_subject, params)` for replay: the method requires +an already wrapped `SubjectsBatch` and omits the public-call lifecycle. ## Hydra configuration diff --git a/docs/get-started/migration.md b/docs/get-started/migration.md index f693f1c7e..99c5a9652 100644 --- a/docs/get-started/migration.md +++ b/docs/get-started/migration.md @@ -18,6 +18,8 @@ This guide covers every breaking change between TorchIO v1 and v2. - Replace `RescaleIntensity(out_min_max=...)` with `Normalize(out_min=..., out_max=...)` - Replace `SubjectsDataset` with any `Dataset` passed to `SubjectsLoader` - Replace `GridAggregator` with `PatchAggregator` +- Rewrite custom transforms to accept a `SubjectsBatch` in + `make_params(batch)` and `apply_transform(batch, params)` ## Image construction @@ -207,6 +209,201 @@ tio.Compose([ ]) ``` +## Custom transforms + +### Rewrite the transform hooks + +In v1, custom transforms implemented a subject-level hook: + + +```python +# v1 +def apply_transform(self, subject: tio.Subject) -> tio.Subject: + ... +``` + +In v2, parameter creation and application are separate batch-level +hooks: + + +```python +# v2 +def make_params(self, batch: tio.SubjectsBatch) -> dict[str, Any]: + ... + +def apply_transform( + self, + batch: tio.SubjectsBatch, + params: dict[str, Any], +) -> tio.SubjectsBatch: + ... +``` + +Call the transform normally rather than calling either hook yourself. +For a single subject, TorchIO performs this conversion automatically: + +```text +Subject -> SubjectsBatch -> apply_transform -> Subject +``` + +The following complete transform works for both a single `Subject` and +a `SubjectsBatch`: + +```python +from typing import Any + +import torch +import torchio as tio + + +class AddValue(tio.Transform): + """Add a fixed value to every batched image.""" + + def __init__(self, value: float) -> None: + super().__init__() + self.value = value + self.received_shape: tuple[int, ...] | None = None + + def make_params(self, batch: tio.SubjectsBatch) -> dict[str, Any]: + """Return the value to add.""" + return {"value": self.value} + + def apply_transform( + self, + batch: tio.SubjectsBatch, + params: dict[str, Any], + ) -> tio.SubjectsBatch: + """Add the value to every image tensor.""" + for image_batch in batch.images.values(): + self.received_shape = tuple(image_batch.data.shape) + image_batch.data = image_batch.data + params["value"] + return batch + + +subject = tio.Subject(image=tio.ScalarImage(torch.zeros(1, 2, 3, 4))) +transform = AddValue(2) +result = transform(subject) +assert isinstance(result, tio.Subject) +assert transform.received_shape == (1, 1, 2, 3, 4) +assert result.image.data.shape == (1, 2, 3, 4) +assert torch.all(result.image.data == 2) +``` + +The image is 4D `(C, I, J, K)` before and after the public call, but it +is 5D `(B, C, I, J, K)` inside `apply_transform`. For a single subject, +`B` is 1. + +!!! warning "`apply_transform` is not a public replay method" + It is the low-level batch kernel. Calling it directly bypasses + wrapping, copying, probability handling, history recording, and + output-type restoration. Pass a supported input to the transform + itself instead. + +### Migrate metadata access + +In v1, subject metadata values were scalars or arbitrary objects. In a +v2 batch, each metadata key maps to a list containing one value per +element: + +```python +import torch +import torchio as tio + +subjects = [ + tio.Subject( + image=tio.ScalarImage(torch.zeros(1, 2, 3, 4)), + site="A", + age=30, + ), + tio.Subject( + image=tio.ScalarImage(torch.ones(1, 2, 3, 4)), + site="B", + age=40, + ), +] +batch = tio.SubjectsBatch.from_subjects(subjects) +assert batch.metadata == {"site": ["A", "B"], "age": [30, 40]} +``` + +Treat `batch.metadata` as `dict[str, list[Any]]`. Metadata transforms +must keep each list aligned with the batch dimension. Subjects in one +batch must have equivalent image names and metadata keys. The first +subject determines the shared key order; later subjects may use a +different local order, but custom transforms should preserve the batch +schema rather than adding, removing, or renaming keys for only some +elements. + +### Choose deterministic or per-instance behavior + +A fixed scalar is not sampled: transforms such as `Gamma` use that +value for every batch element. For built-in stochastic transforms that +support per-instance sampling, ranges and distributions produce +independent parameters for each batch element by default. Set +`per_instance=False` to share one sampled parameter set: + +```python +import torchio as tio + +independent = tio.Gamma(log_gamma=(-0.3, 0.3)) +shared = tio.Gamma(log_gamma=(-0.3, 0.3), per_instance=False) +deterministic = tio.Gamma(log_gamma=0.2) +``` + +Custom transforms do not gain per-instance sampling automatically. +Unless a transform explicitly implements and advertises that +capability, its parameters remain batch-shared. See +[Per-instance augmentation](../concepts/per-instance-augmentation.md) +for the capability contract and stochastic-realisation caveats. + +### Migrate inherently per-subject logic + +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: + +```python +from typing import Any + +import torch +import torchio as tio + + +class StripIdentifier(tio.Transform): + """Strip whitespace from subject identifiers.""" + + def make_params(self, batch: tio.SubjectsBatch) -> dict[str, Any]: + """Return no parameters.""" + return {} + + def apply_transform( + self, + batch: tio.SubjectsBatch, + params: dict[str, Any], + ) -> tio.SubjectsBatch: + """Process metadata one subject at a time.""" + subjects = batch.unbatch() + 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 + + +subject = tio.Subject( + image=tio.ScalarImage(torch.zeros(1, 2, 3, 4)), + identifier=" sub-01 ", +) +result = StripIdentifier()(subject) +assert result.identifier == "sub-01" +``` + +This pattern is more expensive than vectorized code and requires every +resulting subject to retain a compatible image and metadata schema. A +supported mapping utility is planned, but it is not part of the current +API. + ## New features ### Choice