Data Augmentation for Computer Vision

In vision, data augmentation is often the single most effective regulariser. By transforming training images in ways that preserve their labels, we teach the model the invariances we know the task has — a leaf is still diseased when rotated, a car is still a car in dim light. Good augmentation can be worth as much as doubling the dataset.

Categories of augmentation#

Geometric

• Random resized crops (scale and aspect-ratio jitter) — the most important single augmentation for classification.
• Horizontal (sometimes vertical) flips, small rotations, translations, shear.
• Perspective warps and elastic deformations (useful for handwriting and medical images).

Photometric

• Brightness, contrast, saturation and hue jitter.
• Gaussian blur, noise, JPEG compression artefacts.
• Grayscale conversion, gamma changes, simulated shadows, fog or rain.

Occlusion

• Random erasing / Cutout: blank out rectangles, forcing the model to use multiple cues rather than one discriminative patch.

Mixing

• Mixup: blend two images and their labels.
• CutMix: paste a patch from one image into another; mix labels by area.
• Mosaic (popular in YOLO): tile four images into one, exposing objects at varied scales and contexts.

Choosing augmentations: preserve the label#

Automated augmentation policies#

Choosing magnitudes by hand is tedious. AutoAugment (Cubuk et al., 2019) searched for augmentation policies with reinforcement learning — effective but expensive. RandAugment (2020) simplified this to two hyperparameters: apply $N$ random operations from a fixed list, each with global magnitude $M$. TrivialAugment (2021) goes further: one random operation with a random magnitude per image — and performs comparably. These are standard in modern training recipes.

from torchvision import transforms

train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.5, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.TrivialAugmentWide(),              # or transforms.RandAugment(num_ops=2, magnitude=9)
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    transforms.RandomErasing(p=0.25),
])

Mixup and CutMix in practice#

import numpy as np
import torch

def cutmix(x, y, num_classes, alpha=1.0):
    lam = np.random.beta(alpha, alpha)
    idx = torch.randperm(x.size(0))
    H, W = x.shape[2:]
    rh, rw = int(H * np.sqrt(1 - lam)), int(W * np.sqrt(1 - lam))
    cy, cx = np.random.randint(H), np.random.randint(W)
    y1, y2 = np.clip(cy - rh // 2, 0, H), np.clip(cy + rh // 2, 0, H)
    x1, x2 = np.clip(cx - rw // 2, 0, W), np.clip(cx + rw // 2, 0, W)
    x = x.clone(); x[:, :, y1:y2, x1:x2] = x[idx, :, y1:y2, x1:x2]
    lam = 1 - (y2 - y1) * (x2 - x1) / (H * W)            # actual area ratio
    y1h = torch.nn.functional.one_hot(y, num_classes).float()
    return x, lam * y1h + (1 - lam) * y1h[idx]

Train with soft-target cross-entropy: loss = -(targets * logits.log_softmax(1)).sum(1).mean().

Augmentation for detection and segmentation#

Geometric transforms must be applied consistently to boxes, masks and keypoints. Libraries such as Albumentations and torchvision's transforms.v2 handle this:

import albumentations as A

aug = A.Compose([
    A.RandomResizedCrop(size=(512, 512), scale=(0.5, 1.0)),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.5),
    A.GaussNoise(p=0.2),
], bbox_params=A.BboxParams(format="pascal_voc", label_fields=["labels"], min_visibility=0.3))
# out = aug(image=img, bboxes=boxes, labels=labels)

min_visibility drops boxes that become mostly cropped out — otherwise the model is trained to find objects that are no longer visible.

Test-time augmentation (TTA)#

At inference, average predictions over several augmented versions (flips, a few crops). It often adds a small accuracy gain and improves calibration, at the cost of extra inference time.

Synthetic data#

When real data is scarce — rare classes, dangerous scenarios, privacy constraints — synthetic images from simulation, 3-D rendering or generative models can help. Watch for the domain gap: models may learn artefacts of synthetic images. Mixing synthetic with real data and applying domain randomisation reduces this.

Designing an augmentation strategy#

1. Start with random resized crop + flip (if valid).
2. Add photometric jitter matching real-world variation in your deployment environment.
3. Add RandAugment/TrivialAugment and random erasing for larger training runs.
4. Add Mixup/CutMix for longer training from scratch or when overfitting persists.
5. Validate each addition; stronger augmentation often needs longer training.