Data Augmentation (Computer Vision Tutorial)

By the end of this lesson you will be able to multiply a small image dataset by creating realistic variations of each photo, choose augmentations that keep the label correct, move bounding boxes along with the image, and write a torchvision augmentation pipeline. Augmentation is often the cheapest way to make a vision model more accurate.

Why augment?

A model learns only from what it sees. If every training photo of a ₹100 note was taken in daylight, upright and centred, the model may fail on a tilted note under a yellow bulb. Collecting thousands more photos is slow. Data augmentation creates new training examples by changing existing ones in ways that do not change the label: a flipped photo of a cat is still a cat.

You saw its power in the classical ML lesson: a one-pixel shift dropped accuracy from 97% to about 40%, and training on shifted copies fixed most of that.

Think of a cricket coach in Rohtak preparing a batter. She does not only bowl perfect deliveries in perfect light. She mixes pace, spin, bounce and evening practice, so the batter is ready for anything on match day.

The main kinds of augmentation

Kind	Examples	Teaches the model
Geometric	Flip, crop, shift, small rotation, zoom	Objects can be anywhere, at any size, slightly tilted
Photometric	Brightness, contrast, colour jitter	Lighting and cameras vary
Noise and blur	Gaussian noise, slight blur	Cheap cameras and motion are normal
Occlusion	Random erasing (cut out a box)	Objects can be partly hidden

The golden rule: an augmentation must keep the label true. A left-right flip is fine for a photo of a dog, but wrong for text, number plates, or the digits 2 and 5. A vertical flip makes no sense for street scenes, where the sky is always on top. Rotating a 6 by 180 degrees turns it into a 9.

Augmentations in NumPy

Each augmentation is a small function using tools from earlier lessons. A random number generator decides whether and how much to apply each one, so every epoch sees fresh variations.

Step 1: make a test image. Step 2: write the augmentation functions. Step 3: apply a random combination eight times and look.

PythonRun it yourself

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)

# Step 1: a 64x64 RGB "photo": sky, ground and a red kite
img = np.zeros((64, 64, 3))
img[:40] = [0.55, 0.75, 0.95]
img[40:] = [0.45, 0.35, 0.2]
img[12:24, 36:48] = [0.9, 0.15, 0.15]

# Step 2: augmentation functions (all work on float images in 0..1)
def random_flip(x):
    return x[:, ::-1] if rng.random() < 0.5 else x

def random_crop_resize(x, min_scale=0.7):
    h, w = x.shape[:2]
    s = rng.uniform(min_scale, 1.0)
    ch, cw = int(h * s), int(w * s)
    top, left = rng.integers(0, h - ch + 1), rng.integers(0, w - cw + 1)
    crop = x[top:top + ch, left:left + cw]
    rows = (np.arange(h) * ch / h).astype(int)       # nearest-neighbour back to h x w
    cols = (np.arange(w) * cw / w).astype(int)
    return crop[rows[:, None], cols[None, :]]

def colour_jitter(x, strength=0.25):
    bright = rng.uniform(1 - strength, 1 + strength)
    contrast = rng.uniform(1 - strength, 1 + strength)
    mean = x.mean()
    return np.clip((x - mean) * contrast + mean * bright, 0, 1)

def add_noise(x, sigma=0.04):
    return np.clip(x + rng.normal(0, sigma, x.shape), 0, 1)

def random_erase(x, p=0.5):
    if rng.random() > p:
        return x
    x = x.copy()
    h, w = x.shape[:2]
    eh, ew = rng.integers(8, 20), rng.integers(8, 20)
    top, left = rng.integers(0, h - eh), rng.integers(0, w - ew)
    x[top:top + eh, left:left + ew] = rng.random(3)      # a random-coloured patch
    return x

def augment(x):
    for f in (random_crop_resize, random_flip, colour_jitter, add_noise, random_erase):
        x = f(x)
    return x

# Step 3
fig, axes = plt.subplots(1, 9, figsize=(14, 2))
axes[0].imshow(img); axes[0].set_title("original")
for ax in axes[1:]:
    ax.imshow(augment(img))
for ax in axes:
    ax.axis("off")
plt.show()
print("each augmented image keeps the shape:", augment(img).shape)

Output

each augmented image keeps the shape: (64, 64, 3)

Every picture shows the same scene and still deserves the label "kite", but no two look identical. Run it again and you get eight new variations. During training this happens on the fly: each time the model sees a photo, it sees a new random version.

Boxes must move with the image

For object detection and segmentation, the labels are positions, not just names. If you flip the image, the boxes must flip too. A box is often stored as (x_min, y_min, x_max, y_max) in pixels. After a left-right flip of an image of width W, a pixel column x moves to W - 1 - x, so the box's left and right edges swap.

PythonRun it yourself

import numpy as np

def hflip_with_boxes(img, boxes):
    """img: (H, W) or (H, W, C); boxes: array of (x_min, y_min, x_max, y_max) in pixel indices."""
    w = img.shape[1]
    flipped = img[:, ::-1]
    b = boxes.copy()
    b[:, 0] = w - 1 - boxes[:, 2]
    b[:, 2] = w - 1 - boxes[:, 0]
    return flipped, b

img = np.zeros((10, 20), dtype=int)
img[2:5, 3:8] = 1                                   # an object at columns 3..7, rows 2..4
boxes = np.array([[3, 2, 7, 4]])

f_img, f_boxes = hflip_with_boxes(img, boxes)
print("new box:", f_boxes[0])
ys, xs = np.nonzero(f_img)
print("object pixels now span columns", xs.min(), "to", xs.max(), "and rows", ys.min(), "to", ys.max())

Output

new box: [12  2 16  4]
object pixels now span columns 12 to 16 and rows 2 to 4

The new box matches where the object's pixels really are. Getting this wrong is a classic silent bug: the model trains happily on boxes pointing at empty background. Libraries such as torchvision transforms v2 and Albumentations transform boxes and masks for you.

Does it work? A quick experiment

Here we train the same simple model twice on the digits dataset: once on the original images, once with randomly shifted copies added. Then we test both on unshifted and shifted test images.

PythonRun it yourself

import numpy as np
from sklearn.datasets import load_digits
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

def shift(imgs, dx):
    out = np.zeros_like(imgs)
    if dx > 0:
        out[:, :, dx:] = imgs[:, :, :-dx]
    else:
        out[:, :, :dx] = imgs[:, :, -dx:]
    return out

digits = load_digits()
X_tr, X_te, y_tr, y_te = train_test_split(digits.images, digits.target,
                                          test_size=0.25, random_state=0, stratify=digits.target)
flat = lambda a: a.reshape(len(a), -1) / 16

plain = LogisticRegression(max_iter=2000).fit(flat(X_tr), y_tr)
X_aug = np.concatenate([X_tr] + [shift(X_tr, d) for d in (1, -1, 2, -2)])
augmented = LogisticRegression(max_iter=3000).fit(flat(X_aug), np.tile(y_tr, 5))

print("test set        plain   augmented")
for name, test in [("unshifted", X_te), ("shift 1 px", shift(X_te, 1)), ("shift 2 px", shift(X_te, 2))]:
    print(f"{name:12}  {plain.score(flat(test), y_te):7.3f}  {augmented.score(flat(test), y_te):9.3f}")

Output

test set        plain   augmented
unshifted       0.969      0.876
shift 1 px      0.396      0.896
shift 2 px      0.056      0.838

The augmented model loses a little accuracy on perfectly centred digits but becomes far more reliable on shifted ones, which is what real-world photos look like. With a CNN, which has more capacity than this linear model, augmentation usually improves accuracy on clean images too.

Augmentation in PyTorch

torchvision's transforms.v2 module provides all of these, applied to each image as it is loaded. Run this on your computer.

Python

import torch
from torchvision.transforms import v2

train_tf = v2.Compose([
    v2.ToImage(),                                         # PIL image or HWC array -> tensor (C, H, W)
    v2.RandomResizedCrop(size=(224, 224), scale=(0.6, 1.0), antialias=True),
    v2.RandomHorizontalFlip(p=0.5),
    v2.RandomRotation(degrees=10),
    v2.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3),
    v2.ToDtype(torch.float32, scale=True),                # uint8 0..255 -> float 0..1
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    v2.RandomErasing(p=0.25),
])

val_tf = v2.Compose([                                     # no randomness for validation and test
    v2.ToImage(),
    v2.Resize(256, antialias=True),
    v2.CenterCrop(224),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

Every call train_tf(photo) returns a new random 3 x 224 x 224 tensor. For detection, wrap boxes as tv_tensors.BoundingBoxes(boxes, format="XYXY", canvas_size=(h, w)) and pass them with the image; v2 transforms move them consistently.

Common mistakes

Augmentations that change the label. Flipping text, rotating 6 into 9, or colour jitter when colour is the class (ripe vs unripe mango).
Forgetting to move boxes and masks. The image and its labels must get the same geometric change.
Augmenting the test set. Keep evaluation data clean and fixed.
Going too strong. Extreme crops can cut the object out entirely; check samples by eye before training.

Key points

Augmentation creates label-preserving variations so a model sees more variety than the raw dataset holds.
Geometric (flip, crop, rotate), photometric (brightness, contrast, colour), noise and occlusion are the main families.
Choose only augmentations that keep the label true for your task.
Boxes and masks must be transformed with the image.
In PyTorch use torchvision.transforms.v2; apply random augmentation only to training data.

Exercise

Write a function random_rotate90(img, boxes) for a square image that rotates the image by 90 degrees anticlockwise with np.rot90 with probability 0.5, and updates boxes in (x_min, y_min, x_max, y_max) format. Hint: for an anticlockwise turn of an N by N image, a pixel at (x, y) moves to (y, N - 1 - x). Test it on a 10 by 10 image with one box and check that the box covers the object's pixels.

Show solution

Python

import numpy as np

def rotate90_with_boxes(img, boxes):
    n = img.shape[0]
    out = np.rot90(img)                        # anticlockwise
    b = boxes.copy()
    b[:, 0] = boxes[:, 1]                      # new x_min = old y_min
    b[:, 2] = boxes[:, 3]                      # new x_max = old y_max
    b[:, 1] = n - 1 - boxes[:, 2]              # new y_min from old x_max
    b[:, 3] = n - 1 - boxes[:, 0]              # new y_max from old x_min
    return out, b

def random_rotate90(img, boxes, rng=np.random.default_rng()):
    return rotate90_with_boxes(img, boxes) if rng.random() < 0.5 else (img, boxes)

img = np.zeros((10, 10), dtype=int)
img[1:3, 5:9] = 1                              # x 5..8, y 1..2
out, b = rotate90_with_boxes(img, np.array([[5, 1, 8, 2]]))
ys, xs = np.nonzero(out)
print(b[0], "pixels x", xs.min(), "to", xs.max(), "y", ys.min(), "to", ys.max())

The printed box [1 1 2 4] matches the object's new pixel range (x 1 to 2, y 1 to 4). Testing on a deterministic version first, then wrapping it in a random choice, makes bugs easy to catch.

Quick Quiz

1. Which augmentation is a bad idea for a number plate reader?

2. After a left-right flip of an image 100 pixels wide, a box spans x from 10 to 30. Where does it span now?

3. Which data should receive random augmentation?

4. Why are augmentations applied "on the fly" during training?

Why augment?

The main kinds of augmentation

Augmentations in NumPy

Boxes must move with the image

Does it work? A quick experiment

Augmentation in PyTorch

Common mistakes

Key points

Exercise

Quick Quiz

Learn Computer Vision with a mentor