Data augmentation

TorchIO and MONAI transform wrappers for 3D medical images, serializable inference transforms and GPU-batched patch augmentation.

Transforms wrapper

---

source

CustomDictTransform

def CustomDictTransform(
    aug
):

A class that serves as a wrapper to perform an identical transformation on both the image and the target (if it’s a mask).

Vanilla transforms

---

source

do_pad_or_crop

def do_pad_or_crop(
    o, target_shape, padding_mode, mask_name, dtype:_TensorMeta=Tensor
):

Call self as a function.

---

source

PadOrCrop

def PadOrCrop(
    size, padding_mode:int=0, mask_name:NoneType=None
):

Resize image using TorchIO CropOrPad.

Args: size: Target shape, an int or three ints. padding_mode: Padding value or mode passed to tio.CropOrPad. mask_name: Ignored by this transform, which crops or pads each image and mask on its own around the volume centre. Still accepted and saved in to_spec() so older saved specs load. Only the raw tio_transform honours it, when called on a Subject that has that key.

---

source

ZNormalization

def ZNormalization(
    masking_method:NoneType=None, channel_wise:bool=True
):

Apply TorchIO ZNormalization.

---

source

RescaleIntensity

def RescaleIntensity(
    out_min_max:tuple, in_min_max:tuple
):

Apply TorchIO RescaleIntensity for robust intensity scaling.

Args: out_min_max (tuple[float, float]): Output intensity range (min, max) in_min_max (tuple[float, float]): Input intensity range (min, max)

Example for CT images: # Normalize CT from air (-1000 HU) to bone (1000 HU) into range (-1, 1) transform = RescaleIntensity(out_min_max=(-1, 1), in_min_max=(-1000, 1000))

---

source

NormalizeIntensity

def NormalizeIntensity(
    nonzero:bool=True, channel_wise:bool=True, subtrahend:float=None, divisor:float=None
):

Apply MONAI NormalizeIntensity.

Args: nonzero (bool): Only normalize non-zero values (default: True) channel_wise (bool): Apply normalization per channel (default: True) subtrahend (float, optional): Value to subtract
divisor (float, optional): Value to divide by

---

source

inference_transforms_to_specs

def inference_transforms_to_specs(
    tfms
):

Serialize the ordered transforms that run for validation/inference.

fastai transforms with split_idx == 0 are training-only and are omitted. Every remaining transform must be allow-listed and round-trippable, preventing a standard inference artifact from silently dropping preprocessing.

---

source

transforms_from_specs

def transforms_from_specs(
    specs
):

Reconstruct live transforms from JSON spec dicts.

Accepts spec dicts (rebuilt via build_transform_from_spec), pass-through live transforms, and None.

---

source

transforms_to_specs

def transforms_to_specs(
    tfms
):

Convert transforms (or specs) to JSON-serializable spec dicts.

Accepts fastMONAI transform wrappers (via .to_spec()), pass-through spec dicts, and None. Raises TypeError if a transform cannot be serialized (e.g. a custom callable); pass such transforms via the pre_inference_tfms / pre_patch_tfms override instead of the config.

---

source

build_transform_from_spec

def build_transform_from_spec(
    spec
):

Reconstruct an allow-listed transform from a JSON spec dict.

spec must contain a name key naming a registered transform; remaining keys are passed as constructor kwargs.

---

source

BraTSMaskConverter

def BraTSMaskConverter(
    enc:NoneType=None, dec:NoneType=None, split_idx:NoneType=None, order:NoneType=None
):

Convert BraTS masks.

---

source

BinaryConverter

def BinaryConverter(
    enc:NoneType=None, dec:NoneType=None, split_idx:NoneType=None, order:NoneType=None
):

Convert to binary mask.

---

source

RandomGhosting

def RandomGhosting(
    intensity:tuple=(0.5, 1), p:float=0.5
):

Apply TorchIO RandomGhosting.

---

source

RandomSpike

def RandomSpike(
    num_spikes:int=1, intensity:tuple=(1, 3), p:float=0.5
):

Apply TorchIO RandomSpike.

---

source

RandomNoise

def RandomNoise(
    mean:int=0, std:tuple=(0, 0.25), p:float=0.5
):

Apply TorchIO RandomNoise.

---

source

RandomBiasField

def RandomBiasField(
    coefficients:float=0.5, order:int=3, p:float=0.5
):

Apply TorchIO RandomBiasField.

---

source

RandomBlur

def RandomBlur(
    std:tuple=(0, 2), p:float=0.5
):

Apply TorchIO RandomBlur.

---

source

RandomGamma

def RandomGamma(
    log_gamma:tuple=(-0.3, 0.3), p:float=0.5
):

Apply TorchIO RandomGamma.

---

source

RandomIntensityScale

def RandomIntensityScale(
    scale_range:tuple=(0.5, 2.0), p:float=0.5
):

Randomly scale image intensities by a multiplicative factor.

Useful for domain generalization across different acquisition protocols with varying intensity ranges.

Args: scale_range (tuple[float, float]): Range of scale factors (min, max). Values > 1 increase intensity, < 1 decrease intensity. p (float): Probability of applying the transform (default: 0.5)

Example: # Scale intensities randomly between 0.5x and 2.0x transform = RandomIntensityScale(scale_range=(0.5, 2.0), p=0.3)

---

source

RandomMotion

def RandomMotion(
    degrees:int=10, translation:int=10, num_transforms:int=2, image_interpolation:str='linear', p:float=0.5
):

Apply TorchIO RandomMotion.

---

source

RandomAnisotropy

def RandomAnisotropy(
    axes:tuple=(0, 1, 2), downsampling:tuple=(1.5, 5), image_interpolation:str='linear', scalars_only:bool=True,
    p:float=0.5
):

Apply TorchIO RandomAnisotropy.

---

source

RandomCutout

def RandomCutout(
    holes:int=1, max_holes:int=3, spatial_size:int=8, max_spatial_size:int=16, fill:str='min', mask_only:bool=True,
    p:float=0.2
):

Randomly erase spherical regions in 3D medical images with mask-aware placement.

Simulates post-operative surgical cavities by filling random ellipsoid volumes with specified values. When mask_only=True (default), cutouts only affect voxels inside the segmentation mask, ensuring no healthy tissue is modified.

Args: holes: Minimum number of cutout regions. Default: 1. max_holes: Maximum number of regions. Default: 3. spatial_size: Minimum cutout diameter in voxels. Default: 8. max_spatial_size: Maximum cutout diameter. Default: 16. fill: Fill value - ‘min’, ‘mean’, ‘random’, or float. Default: ‘min’. mask_only: If True, cutouts only affect mask-positive voxels (tumor tissue). If False, cutouts can affect any voxel (original behavior). Default: True. p: Probability of applying transform. Default: 0.2.

Example: >>> # Simulate post-op cavities only within tumor regions >>> tfm = RandomCutout(holes=1, max_holes=2, spatial_size=10, … max_spatial_size=25, fill=‘min’, mask_only=True, p=0.2)

>>> # Original behavior - cutouts anywhere in the volume
>>> tfm = RandomCutout(mask_only=False, p=0.2)

Dictionary transforms

---

source

RandomElasticDeformation

def RandomElasticDeformation(
    num_control_points:int=7, max_displacement:float=7.5, image_interpolation:str='linear', p:float=0.5
):

Apply TorchIO RandomElasticDeformation.

---

source

RandomAffine

def RandomAffine(
    scales:int=0, degrees:int=10, translation:int=0, isotropic:bool=False, image_interpolation:str='linear',
    default_pad_value:float=0.0, p:float=0.5
):

Apply TorchIO RandomAffine.

---

source

RandomFlip

def RandomFlip(
    axes:str='LRAPIS', p:float=0.5
):

Apply TorchIO RandomFlip.

---

source

OneOf

def OneOf(
    transform_dict, p:int=1
):

Apply only one of the given transforms using TorchIO OneOf.

transform_dict maps transforms to probabilities, or is a sequence of transforms picked with equal probability. Transforms can be TorchIO transforms or fastMONAI wrappers, which are unwrapped through their tio_transform property.

Augmentation suggestion

GPU patch augmentation

---

source

GpuPatchAugmentation

def GpuPatchAugmentation(
    affine:NoneType=None, anisotropy:NoneType=None, flip:NoneType=None, gamma:NoneType=None,
    intensity_scale:NoneType=None, noise:NoneType=None, blur:NoneType=None
):

GPU-batched augmentation for patch-based training.

Operates on [B, C, D, H, W] tensors already on GPU. All operations run under torch.no_grad() since augmentation does not need gradient tracking. Inputs are cloned, so the caller’s tensors are never mutated.

Transform order: spatial (affine, anisotropy, flip) then intensity (gamma, intensity_scale, noise, blur). Spatial transforms apply the same parameters to both image and mask. Intensity transforms skip the mask.

Each transform is controlled by a parameter dict with at minimum a ‘p’ key for per-sample probability. Pass None to disable a transform. As in TorchIO, index i of a per-axis value refers to array axis i of the patch (D, H, W), and lengths are in voxels.

Args: affine: dict with keys ‘scales’: (lo, hi). Each axis gets a factor from U(lo, hi); values below 1 shrink the content, as in TorchIO. A single number is not accepted. ‘degrees’: d or per-axis maxima (d0, d1, d2). The rotation about axis i is drawn from U(-d_i, d_i). ‘translation’: t or per-axis maxima (t0, t1, t2), in voxels. The shift along axis i is drawn from U(-t_i, t_i). ‘default_pad_value’: value for voxels moved in from outside the patch (default 0). ‘p’: probability. TorchIO’s range forms for degrees and translation, (a, b) or six values, are not supported. None to disable. anisotropy: dict with keys ‘axes’ (one is picked at random), ‘downsampling’ as (lo, hi), ‘p’. None to disable. flip: dict with keys ‘axes’ and ‘p’; each axis is flipped independently with probability p. None to disable. gamma: dict with keys ‘log_gamma’ as (lo, hi), ‘p’. None to disable. intensity_scale: dict with keys ‘scale_range’ as (lo, hi), ‘p’. None to disable. noise: dict with keys ‘std’, ‘p’. Zero-mean noise whose std is drawn from U(0, s) for a number s, or from U(lo, hi) for (lo, hi). None to disable. blur: dict with keys ‘std’, ‘p’. The sigma in voxels is drawn from U(0, s) for a number s, or from U(lo, hi) for (lo, hi); a sample whose sigma is 0 is left unchanged. None to disable.

Example::

>>> gpu_aug = GpuPatchAugmentation(
...     affine={'scales': (0.7, 1.4), 'degrees': (30, 30, 30),
...             'translation': (25, 25, 10), 'default_pad_value': 0., 'p': 0.2},
...     gamma={'log_gamma': (-0.3, 0.3), 'p': 0.3},
...     flip={'axes': (0, 1, 2), 'p': 0.5},
... )
>>> img_aug, mask_aug = gpu_aug(img_gpu, mask_gpu)

Design rationale (why a from-scratch engine, not MONAI/TorchIO): Genuinely batched: one affine_grid/grid_sample over the whole [B, C, D, H, W] batch. MONAI/TorchIO transforms are strictly per-sample, so a batched alternative must loop per sample. Measured on an RTX 6000 Ada (full pipeline, 128^3): this batched GPU path ~3 ms/batch (B=2) vs MONAI per-sample-on-GPU ~77 ms (~25x slower) vs TorchIO CPU ~287 ms (~100x slower); the gap is structural and grows with batch size, so the GPU path is kept.

Differences from the CPU/TorchIO path (suggest_patch_augmentations): Each parameter means the same thing in both paths. The CPU patch path hands patches to TorchIO with an identity affine, so TorchIO works in voxels there too. The GPU path is not bit-identical:

- Rotations act in affine_grid's normalised [-1, 1] patch coordinates.
  A rotation between two axes of different patch length is stretched by
  their length ratio, so it is rigid in voxel space only between axes of
  equal length (e.g. a 256 x 256 plane). It is close to rigid in
  physical space when patch_size x target_spacing is near isotropic.
  TorchIO rotates rigidly in voxel space.
- Blur uses one isotropic sigma per sample, a kernel cut at 3 sigma and
  replicate padding. TorchIO draws a sigma per axis and uses scipy's
  gaussian_filter (cut at 4 sigma, reflect padding).
- Anisotropy follows the same recipe as TorchIO (nearest down, linear
  up, image only), but resamples with F.interpolate on the voxel grid,
  so sample positions and the downsampled size can differ slightly.
- Gamma uses one exponent per sample, while TorchIO draws one per
  channel. This only matters for multi-channel input.

---

source

gpu_patch_augmentations

def gpu_patch_augmentations(
    patch_size, target_spacing, anisotropy_threshold:float=3.0, translation_fraction:float=0.15, affine_p:float=0.2,
    anisotropy_p:float=0.25, gamma_p:float=0.3, intensity_scale_p:float=0.1, noise_p:float=0.1, blur_p:float=0.2,
    flip_p:float=0.5
):

Create GpuPatchAugmentation with nnU-Net-inspired defaults.

GPU counterpart of suggest_patch_augmentations: it shares the same parameter logic (via _compute_patch_aug_params) and probabilities, and each value means the same thing in both paths. It runs batched on GPU, much faster than the per-sample CPU/TorchIO path, but is not bit-identical. See GpuPatchAugmentation for the benchmark and the remaining differences (rotation on non-cubic patches, isotropic blur sigma, anisotropy resampling grid, per-sample gamma).

Args: patch_size: List/tuple of 3 ints, patch dimensions. target_spacing: List/tuple of 3 floats, voxel spacing. anisotropy_threshold: Ratio threshold for anisotropy detection (default 3.0). translation_fraction: Fraction of patch_size for translation (default 0.15). affine_p: Probability for RandomAffine (default 0.2). anisotropy_p: Probability for RandomAnisotropy (default 0.25). gamma_p: Probability for RandomGamma (default 0.3). intensity_scale_p: Probability for RandomIntensityScale (default 0.1). noise_p: Probability for RandomNoise (default 0.1). blur_p: Probability for RandomBlur (default 0.2). flip_p: Probability for RandomFlip per axis (default 0.5).

Returns: GpuPatchAugmentation instance.

Example::

>>> gpu_aug = gpu_patch_augmentations([128, 128, 32], [0.5, 0.5, 1.5])
>>> dls = MedPatchDataLoaders.from_df(..., gpu_augmentation=gpu_aug)

---

source

suggest_patch_augmentations

def suggest_patch_augmentations(
    patch_size, target_spacing, anisotropy_threshold:float=3.0, translation_fraction:float=0.15
):

Suggest patch-based augmentations with nnU-Net-inspired defaults.

Derives rotation degrees, translation, and RandomAnisotropy axes from patch geometry and voxel spacing. Returns a list of fastMONAI transform instances ready for the patch_tfms parameter in MedPatchDataLoaders.

Anisotropy detection: if max(spacing)/min(spacing) >= threshold, rotation is restricted to 5 deg out-of-plane and 30 deg in-plane. Otherwise 30 deg symmetric. Translation is patch_size * fraction per axis.

Args: patch_size: List/tuple of 3 ints, patch dimensions. target_spacing: List/tuple of 3 floats, voxel spacing. anisotropy_threshold: Ratio threshold for anisotropy detection (default 3.0). translation_fraction: Fraction of patch_size for translation (default 0.15).

Returns: list: fastMONAI transform instances (7 normally, 6 if RandomAnisotropy omitted).

Example::

>>> patch_tfms = suggest_patch_augmentations([128, 128, 32], [0.5, 0.5, 1.5])
>>> dls = MedPatchDataLoaders.from_df(..., patch_tfms=patch_tfms)