Data Augmentation — Training on Limited Image Data

Teaching a child to recognise dogs. You show them 100 dog photos — all golden retrievers, all photographed outdoors in sunlight. The child learns "dog = golden retriever outdoors." Now show them a black poodle indoors and they fail. If you had shown them photos from different angles, lighting, and backgrounds — they would generalise. Augmentation is artificially creating that variety.

The model does not know you flipped the image. It just sees a slightly different training example each epoch. Over 50 epochs with random augmentation, the model effectively trains on 50× more data than you actually collected.

import torch
import torchvision.transforms as T
import torchvision.transforms.functional as TF
import numpy as np
from PIL import Image

# ── Create a test image ───────────────────────────────────────────────
img = Image.fromarray(
    np.random.randint(50, 200, (300, 200, 3), dtype=np.uint8)
)
print(f"Original size: {img.size} (W×H)")

# ── Geometric augmentations ────────────────────────────────────────────
geometric_transforms = {
    'HorizontalFlip':   T.RandomHorizontalFlip(p=1.0),
    'VerticalFlip':     T.RandomVerticalFlip(p=1.0),
    'Rotation 15°':     T.RandomRotation(degrees=15),
    'Rotation 90°':     T.RandomRotation(degrees=90),
    'RandomResizedCrop':T.RandomResizedCrop(224, scale=(0.5, 1.0)),
    'Perspective':      T.RandomPerspective(distortion_scale=0.3, p=1.0),
    'Affine':           T.RandomAffine(degrees=10, translate=(0.1, 0.1),
                                        scale=(0.9, 1.1), shear=5),
}

print("
Geometric transform output sizes:")
for name, transform in geometric_transforms.items():
    out  = transform(img)
    size = out.size if isinstance(out, Image.Image) else tuple(out.shape)
    print(f"  {name:<20}: {size}")

# ── RandomResizedCrop — most important geometric augmentation ─────────
# Shows different crops of the same image — simulates different scales
rrc = T.RandomResizedCrop(224, scale=(0.08, 1.0), ratio=(0.75, 1.33))
to_tensor = T.ToTensor()

crops = [to_tensor(rrc(img)) for _ in range(4)]
for i, crop in enumerate(crops):
    print(f"  Crop {i+1}: shape={tuple(crop.shape)}"
          f"  mean={crop.mean():.3f}  std={crop.std():.3f}")

# ── Visualise augmentation effect numerically ─────────────────────────
original_tensor = to_tensor(img)
print(f"
Original tensor:")
print(f"  Mean: {original_tensor.mean():.4f}  Std: {original_tensor.std():.4f}")

flip_tensor = to_tensor(TF.hflip(img))
print(f"After horizontal flip:")
print(f"  Mean: {flip_tensor.mean():.4f}  ← identical mean (pixel values same)")
print(f"  Pixel at [0,0,0]: {original_tensor[0,0,0]:.3f} vs {flip_tensor[0,0,-1]:.3f}  ← mirrored")

import torch
import torchvision.transforms as T
import numpy as np
from PIL import Image

img = Image.fromarray(np.random.randint(50, 200, (224, 224, 3), dtype=np.uint8))
to_tensor = T.ToTensor()

# ── ColorJitter — the most important colour augmentation ──────────────
# brightness: multiply pixel values by random factor in [1-b, 1+b]
# contrast:   stretch/compress pixel value range
# saturation: increase/decrease colour intensity
# hue:        shift all hues by ±h (fraction of full hue circle)

colour_transforms = {
    'Brightness only':   T.ColorJitter(brightness=0.5),
    'Contrast only':     T.ColorJitter(contrast=0.5),
    'Saturation only':   T.ColorJitter(saturation=0.5),
    'Hue only':          T.ColorJitter(hue=0.2),
    'All moderate':      T.ColorJitter(brightness=0.3, contrast=0.3,
                                        saturation=0.2, hue=0.1),
    'All aggressive':    T.ColorJitter(brightness=0.8, contrast=0.8,
                                        saturation=0.5, hue=0.3),
}

original = to_tensor(img)
print("Colour augmentation effect on pixel statistics:")
print(f"  {'Transform':<20} {'Mean':>8} {'Std':>8} {'Min':>8} {'Max':>8}")
print("  " + "─" * 50)
print(f"  {'Original':<20} {original.mean():>8.4f} {original.std():>8.4f} "
      f"{original.min():>8.4f} {original.max():>8.4f}")

for name, transform in colour_transforms.items():
    t = to_tensor(transform(img))
    print(f"  {name:<20} {t.mean():>8.4f} {t.std():>8.4f} "
          f"{t.min():>8.4f} {t.max():>8.4f}")

# ── Greyscale augmentation ────────────────────────────────────────────
# Randomly convert to greyscale — forces model to use texture not just colour
grey_aug = T.RandomGrayscale(p=0.1)   # 10% chance each image
grey_out = grey_aug(img)
grey_t   = to_tensor(grey_out)
print(f"
RandomGrayscale (p=0.1):")
print(f"  Channels: {grey_t.shape[0]}  (still 3 channels — R=G=B when greyscale)")

# ── Gaussian blur — simulate out-of-focus photography ─────────────────
blur = T.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0))
blur_t = to_tensor(blur(img))
print(f"
Gaussian blur: mean_diff={abs(original.mean()-blur_t.mean()):.4f}  "
      f"std_diff={(original.std()-blur_t.std()):.4f}")

# ── Production colour augmentation pipeline ───────────────────────────
# Strong augmentation for Shopify product images
product_colour_aug = T.Compose([
    T.ColorJitter(brightness=0.4, contrast=0.3, saturation=0.3, hue=0.1),
    T.RandomGrayscale(p=0.05),
    T.GaussianBlur(kernel_size=3, sigma=(0.1, 1.5)),
])

augmented = to_tensor(product_colour_aug(img))
print(f"
Product colour pipeline:")
print(f"  Original:  mean={original.mean():.4f}  std={original.std():.4f}")
print(f"  Augmented: mean={augmented.mean():.4f}  std={augmented.std():.4f}")

import torch
import torch.nn as nn
import torchvision.transforms as T
import numpy as np
from PIL import Image

# ── Cutout / RandomErasing ────────────────────────────────────────────
# Built into torchvision as RandomErasing — applied AFTER ToTensor
random_erasing = T.RandomErasing(
    p=0.5,             # 50% chance of erasing
    scale=(0.02, 0.2), # erase 2–20% of image area
    ratio=(0.3, 3.3),  # aspect ratio of erased region
    value=0,           # fill with zeros (or 'random' for noise)
)

img = Image.fromarray(np.random.randint(50, 200, (224, 224, 3), dtype=np.uint8))
tensor = T.ToTensor()(img)
erased = random_erasing(tensor)

n_zeros = (erased == 0).float().mean().item()
print(f"RandomErasing: {n_zeros*100:.1f}% of pixels set to zero")

# ── MixUp — implemented manually ─────────────────────────────────────
def mixup_batch(images: torch.Tensor, labels: torch.Tensor,
                alpha: float = 0.2, num_classes: int = 6):
    """
    Apply MixUp to a batch.
    images: (B, C, H, W) float tensor
    labels: (B,) long tensor of class indices
    Returns: mixed images, soft label matrix (B, num_classes)
    """
    B = images.size(0)

    # Sample mixing coefficient from Beta distribution
    lam = np.random.beta(alpha, alpha) if alpha > 0 else 1.0

    # Random permutation for pairing
    idx = torch.randperm(B)

    # Mix images
    mixed_images = lam * images + (1 - lam) * images[idx]

    # Convert to one-hot and mix labels
    one_hot    = torch.zeros(B, num_classes)
    one_hot.scatter_(1, labels.unsqueeze(1), 1)
    mixed_labels = lam * one_hot + (1 - lam) * one_hot[idx]

    return mixed_images, mixed_labels

# Test MixUp
batch_images = torch.randn(8, 3, 224, 224)
batch_labels = torch.randint(0, 6, (8,))
mixed_imgs, soft_labels = mixup_batch(batch_images, batch_labels, alpha=0.2)

print(f"
MixUp batch:")
print(f"  Original labels:  {batch_labels.tolist()}")
print(f"  Soft label sample: {soft_labels[0].tolist()}")
print(f"  Sum (should be 1): {soft_labels[0].sum().item():.4f}")

# ── MixUp training step ───────────────────────────────────────────────
def mixup_criterion(criterion, pred, mixed_labels):
    """Cross-entropy with soft labels from MixUp."""
    log_probs = torch.log_softmax(pred, dim=1)
    loss      = -(mixed_labels * log_probs).sum(dim=1).mean()
    return loss

# ── CutMix ────────────────────────────────────────────────────────────
def cutmix_batch(images: torch.Tensor, labels: torch.Tensor,
                  alpha: float = 1.0, num_classes: int = 6):
    B, C, H, W = images.shape
    lam = np.random.beta(alpha, alpha)

    # Random bounding box
    cut_ratio = np.sqrt(1 - lam)
    cut_h     = int(H * cut_ratio)
    cut_w     = int(W * cut_ratio)
    cx        = np.random.randint(W)
    cy        = np.random.randint(H)
    x1 = max(cx - cut_w // 2, 0)
    x2 = min(cx + cut_w // 2, W)
    y1 = max(cy - cut_h // 2, 0)
    y2 = min(cy + cut_h // 2, H)

    idx       = torch.randperm(B)
    mixed     = images.clone()
    mixed[:, :, y1:y2, x1:x2] = images[idx, :, y1:y2, x1:x2]

    # Adjust lambda based on actual box area
    lam_actual = 1 - (x2 - x1) * (y2 - y1) / (W * H)

    one_hot    = torch.zeros(B, num_classes)
    one_hot.scatter_(1, labels.unsqueeze(1), 1)
    mixed_labels = lam_actual * one_hot + (1 - lam_actual) * one_hot[idx]

    return mixed, mixed_labels

mixed_cm, labels_cm = cutmix_batch(batch_images, batch_labels)
print(f"
CutMix batch:")
print(f"  Mixed image shape: {tuple(mixed_cm.shape)}")
print(f"  Soft label sample: {labels_cm[0].tolist()}")

import torch
import torchvision.transforms as T
from torch.utils.data import Dataset, DataLoader
import numpy as np
from PIL import Image

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

# ── Complete augmentation pipelines for different dataset sizes ────────

def get_transforms(dataset_size: int, image_size: int = 224):
    """Return appropriate train/val transforms based on dataset size."""

    val_transform = T.Compose([
        T.Resize(int(image_size * 256 / 224)),
        T.CenterCrop(image_size),
        T.ToTensor(),
        T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    ])

    if dataset_size < 1000:
        # Maximum augmentation
        train_transform = T.Compose([
            T.RandomResizedCrop(image_size, scale=(0.5, 1.0)),
            T.RandomHorizontalFlip(p=0.5),
            T.RandomVerticalFlip(p=0.1),
            T.RandomRotation(degrees=30),
            T.ColorJitter(brightness=0.8, contrast=0.8,
                           saturation=0.5, hue=0.2),
            T.RandomGrayscale(p=0.1),
            T.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
            T.ToTensor(),
            T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
            T.RandomErasing(p=0.5, scale=(0.02, 0.2)),
        ])

    elif dataset_size < 10_000:
        # Strong augmentation
        train_transform = T.Compose([
            T.RandomResizedCrop(image_size, scale=(0.6, 1.0)),
            T.RandomHorizontalFlip(p=0.5),
            T.ColorJitter(brightness=0.4, contrast=0.4,
                           saturation=0.3, hue=0.1),
            T.RandomGrayscale(p=0.05),
            T.ToTensor(),
            T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
            T.RandomErasing(p=0.3, scale=(0.02, 0.15)),
        ])

    elif dataset_size < 100_000:
        # Moderate augmentation
        train_transform = T.Compose([
            T.RandomResizedCrop(image_size, scale=(0.7, 1.0)),
            T.RandomHorizontalFlip(p=0.5),
            T.ColorJitter(brightness=0.2, contrast=0.2,
                           saturation=0.2, hue=0.1),
            T.ToTensor(),
            T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
        ])

    else:
        # Light augmentation
        train_transform = T.Compose([
            T.RandomResizedCrop(image_size, scale=(0.8, 1.0)),
            T.RandomHorizontalFlip(p=0.5),
            T.ColorJitter(brightness=0.1, contrast=0.1),
            T.ToTensor(),
            T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
        ])

    return train_transform, val_transform

# ── Demonstrate transform count ───────────────────────────────────────
for size in [500, 5000, 50000, 500000]:
    train_t, _ = get_transforms(size)
    n_steps = len(train_t.transforms)
    print(f"Dataset size {size:>7,}: {n_steps} transform steps")

# ── Verify: same image looks different each epoch ─────────────────────
train_t, _ = get_transforms(5000)
img = Image.fromarray(np.random.randint(50, 200, (300, 300, 3), dtype=np.uint8))

tensors = [train_t(img) for _ in range(4)]
print(f"
Same image, 4 different augmentations:")
for i, t in enumerate(tensors):
    print(f"  Epoch {i+1}: mean={t.mean():.4f}  std={t.std():.4f}"
          f"  corner_pixel={t[0,0,0]:.4f}")

import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np

# ── Albumentations pipeline — image AND bounding boxes transformed together ──
# bbox_params tells Albumentations the box format and which field holds
# the class label, so every geometric transform below also updates the boxes.
train_transform = A.Compose([
    A.RandomResizedCrop(height=640, width=640, scale=(0.5, 1.0)),
    A.HorizontalFlip(p=0.5),
    A.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1, p=0.7),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2(),
], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))

# ── Simulated defect-detection image with two bounding boxes ─────────────
image = np.random.randint(0, 256, (720, 1280, 3), dtype=np.uint8)
boxes = [[100, 150, 300, 400], [600, 200, 900, 500]]   # [x_min, y_min, x_max, y_max]
labels = ['scratch', 'dent']

result = train_transform(image=image, bboxes=boxes, class_labels=labels)
print(f"Transformed image shape: {tuple(result['image'].shape)}")
print(f"Transformed boxes:       {result['bboxes']}")
print(f"Labels (unchanged):      {result['class_labels']}")
print("The crop above moved and resized the boxes to match — this is the")
print("entire reason to use bbox_params instead of a plain torchvision Compose.")

# ── What silently breaks without bbox_params ──────────────────────────
# If you instead crop the image with a plain transform and leave the
# original box coordinates untouched, the pipeline still runs without
# error. Training loss still decreases, because the model is learning
# to fit boxes to whatever region of the image now happens to be there —
# it just is no longer the scratch or the dent. This kind of bug does
# not crash anything. It quietly caps the model's accuracy for the
# entire length of a training run, and the only way to catch it before
# burning GPU hours is to look at a handful of augmented examples with
# their boxes drawn on top, before training starts.

Before a training run starts on a changed augmentation pipeline, run a small fixed set of representative images through it — a dozen or so, chosen to cover edge cases like small objects, objects near the image border, and dense clusters — and render the results with boxes or masks drawn on top. A person looks at those dozen images for under a minute. That single manual check catches misaligned boxes, masks that got flipped without their image, or crops that cut an object in half, all of which run perfectly cleanly through training and only show up later as an unexplained accuracy ceiling. The augmentation config itself is saved alongside the model checkpoint — treated as a versioned part of the experiment, not a detail that lives only in whichever training script happened to be checked out that day.

Augmentation recombines and perturbs the information already present in your dataset — it cannot introduce information your images never captured in the first place. Flipping, cropping, and colour-jittering 1,000 kurta photos teaches the model that orientation, position, and lighting do not change the category, but if none of those 1,000 photos were taken at night or against a cluttered background, no amount of augmentation manufactures that missing condition. Augmentation is the right tool for teaching known invariances cheaply; genuinely new data is the only tool for covering conditions your existing images never represented at all.

As this module's geometric augmentations section notes, a horizontal flip is only label-preserving when left-right orientation genuinely does not matter for the task. It silently corrupts labels for text-containing images, since flipped text is unreadable and a model trained on it learns nothing useful about real text. It is worse than useless for medical images where laterality carries clinical meaning — a chest X-ray flipped left-to-right can turn a real finding on the correct side into training data that teaches the model the wrong side. None of these failures throw an error; the pipeline runs, the loss goes down, and the model quietly learns an invariance that is false for the actual task.

Past a certain point, aggressive augmentation stops teaching invariance and starts destroying the training signal — this module's errors section shows RandomResizedCrop with scale=(0.05, 1.0) sometimes handing the model 5% of a product photo as its entire input, which is not enough to learn from. There is a second, subtler failure: a model can effectively overfit to the augmentation pipeline itself, picking up on artifacts a specific resize interpolation or a particular colour-jitter implementation introduces, rather than becoming genuinely more robust. Augmentation strength is a hyperparameter to tune against validation performance, not a dial that only helps as you turn it up.

This module's callout that augmentation is applied only during training refers to avoiding random transforms on the validation set so metrics stay comparable between runs — it is not the whole story about augmentation and inference. Test-time augmentation (TTA) applies several augmented views of the same input image at prediction time — the original plus a horizontal flip plus a few crops, say — runs all of them through the trained model, and averages the resulting predictions. This is reproducible run to run, because the exact set of views is fixed rather than randomly resampled, and it typically buys a small but real accuracy improvement at the cost of several times more inference compute per prediction.

The augmentation strategy section of this module is explicit that the right policy scales inversely with how much data you have. Under 1,000 images, heavy augmentation is close to mandatory — it is the main defence against overfitting when there simply are not enough unique examples. Past 100,000 images, that same heavy policy actively slows convergence without improving generalisation, because the real data already provides the variation augmentation would otherwise manufacture. A policy tuned for a 1,000-image proof of concept is very likely wrong once the dataset has grown two orders of magnitude — it needs to be revisited, not carried forward unchanged.

Augmentation is the cheap first line of defence against overfitting when you already have a reasonably representative dataset but not a large one — it multiplies the effective training signal by teaching invariances (position, scale, lighting) the model would otherwise have to infer from too few examples. It stops being sufficient when the gap is not "not enough examples of what I have" but "missing conditions entirely" — a product category never photographed, a lighting condition never captured, a demographic never represented. Augmentation reshuffles existing information; it cannot manufacture information the dataset never contained, so a validation failure mode tied to a genuinely unseen condition needs new data, not a stronger transform pipeline.

Horizontal flip is the clearest case: applied to a dataset of receipts or screenshots for OCR, it turns every piece of text unreadable while the image still looks like a perfectly valid training example to any pipeline that just checks shape and dtype. Applied to chest X-rays where laterality is clinically meaningful, it can flip a genuine finding onto the wrong side without any code raising an error. The way to catch this before it costs a training run is to manually inspect a batch of actually augmented images, not just the raw source images, before training starts, and to ask explicitly for every transform in the pipeline: does this operation still produce a label the image genuinely shows, for this specific task?

Training augmentation randomly perturbs each image differently every epoch so the model never sees the exact same input twice, which is what teaches invariance. Test-time augmentation (TTA) instead generates a small, fixed set of augmented views of one input image at prediction time — say the original, a horizontal flip, and a couple of crops — runs all of them through the already-trained model, and averages the predictions or logits. The two solve different problems: training augmentation shapes what the model learns, TTA squeezes a bit more accuracy out of a model that is already trained, at the cost of running inference several times per prediction instead of once. TTA does not retrain anything; it is purely an inference-time technique.

Start from dataset size as the first-order signal — heavy augmentation for datasets under a few thousand images, progressively lighter as the dataset grows past tens and then hundreds of thousands — and then tune against validation performance rather than trusting the starting point blindly. The clearest failure signal is training loss that oscillates and never smoothly decreases, which usually means a transform like RandomResizedCrop's scale range or ColorJitter's magnitude is destroying too much of the actual training signal per image. A second, subtler signal is a large gap where training loss looks fine but validation accuracy stalls or regresses — worth checking whether the model is picking up on artifacts specific to the augmentation pipeline rather than becoming more robust.

Both techniques produce a training image that is genuinely a blend of two classes — a weighted pixel combination for MixUp, a pasted region for CutMix — so the ground truth for that image is honestly a mixture, not a single hard label. Standard nn.CrossEntropyLoss expects one correct class index and pushes the model to be confident about it; applying it to a 70/30 blended image as if it were 100% one class sends a misleading gradient that fights against the very regularisation MixUp and CutMix are meant to provide. The fix is cross-entropy computed against the mixed soft-label vector directly — the negative sum of the mixed label times the log softmax — so the loss reflects that the model's output should also be an honest blend proportional to the mixing ratio.