Skip to content

Augmentations

MONAI-compatible 3D augmentations for cryoET subvolumes. They can be composed into a transform pipeline or applied directly to batches/volumes.

MixupTransform

Creates virtual training examples by blending pairs of samples (and their labels).

copick_torch.augmentations.MixupTransform

Bases: RandomizableTransform

Implements Mixup augmentation for 3D volumes based on MONAI transform interface.

Mixup is a data augmentation technique that creates virtual training examples by mixing pairs of inputs and their labels with a random proportion.

Reference: Zhang et al., "mixup: Beyond Empirical Risk Minimization", ICLR 2018 https://arxiv.org/abs/1710.09412

__init__ 

__init__(alpha: float = 0.2, prob: float = 1.0)

Initialize the Mixup augmentation.

Parameters:

  • alpha (float, default: 0.2 ) –

    Parameter for Beta distribution. Higher values result in more mixing.

  • prob (float, default: 1.0 ) –

    Probability of applying the transform.

randomize 

randomize(data=None) -> None

Randomize the transform parameters.

__call__ 

__call__(img: Tensor, randomize: bool = True) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]

Apply mixup augmentation to a batch of images and labels.

Parameters:

  • img (Tensor) –

    Tensor of shape [batch_size, channels, depth, height, width]

  • randomize (bool, default: True ) –

    Whether to execute randomize function first, default to True.

Returns:

  • Tuple[Tensor, Tensor, Tensor, float] –

    Tuple of (mixed_images, label_a, label_b, lam) where: - mixed_images: The mixup result - label_a: Original labels - label_b: Mixed-in labels - lam: Mixing coefficient from Beta distribution

mixup_criterion 

mixup_criterion(criterion, pred, y_a, y_b, lam)

Apply mixup to the loss calculation.

Parameters:

  • criterion –

    Loss function

  • pred –

    Model predictions

  • y_a –

    First labels

  • y_b –

    Second (mixed-in) labels

  • lam –

    Mixing coefficient

Returns:

  • –

    Mixed loss

FourierAugment3D

Frequency-domain augmentation: random frequency dropout, phase noise, and intensity scaling.

copick_torch.augmentations.FourierAugment3D

Bases: RandomizableTransform, Fourier

Implements Fourier-based augmentation for 3D volumes based on MONAI transform interface.

This augmentation performs operations in the frequency domain, including random frequency dropout (masking), phase noise injection, and intensity scaling.

It can help the model become more robust to various frequency distortions that may occur in tomographic data.

__init__ 

__init__(freq_mask_prob: float = 0.3, phase_noise_std: float = 0.1, intensity_scaling_range: Tuple[float, float] = (0.8, 1.2), prob: float = 1.0) -> None

Initialize the Fourier domain augmentation.

Parameters:

  • freq_mask_prob (float, default: 0.3 ) –

    Probability of masking a frequency component

  • phase_noise_std (float, default: 0.1 ) –

    Standard deviation of Gaussian noise added to the phase

  • intensity_scaling_range (Tuple[float, float], default: (0.8, 1.2) ) –

    Range for random intensity scaling (min, max)

  • prob (float, default: 1.0 ) –

    Probability of applying the transform

randomize 

randomize(spatial_shape=None) -> None

Randomize the transform parameters.

__call__ 

__call__(volume: Tensor, randomize: bool = True) -> torch.Tensor

Apply Fourier domain augmentation to a volume.

Parameters:

  • volume (Tensor) –

    Tensor of shape [depth, height, width] or [channels, depth, height, width]

  • randomize (bool, default: True ) –

    Whether to execute randomize function first, default to True.

Returns:

  • Tensor –

    Augmented volume with same shape as input