TorchIO transforms are subclasses of torch.nn.Module. They are designed with a unified batch architecture: regardless of the input type (Subject, Image, Tensor, NumPy array, etc.), the transform internally converts the input into a SubjectsBatch before processing. This allows a single implementation to handle both individual samples and batches of data.
Key behaviors:
- Input/Output Consistency: Transforms return the same type they receive (e.g., a
Subject input returns a Subject). - Batching: A single
Image is treated as a batch of size 1. An image tensor of shape (C, I, J, K) reaches the internal application method as (B, C, I, J, K). - Per-instance Augmentation: When passing a
SubjectsBatch (e.g., from a DataLoader), transforms sample independent parameters for each element by default. To share the same sampled parameters across all elements in a batch, pass per_instance=False to the transform.
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)
