Image transforms and augmentation

Preprocessing versus augmentation

Image models need inputs as float32 tensors of shape (C, H, W), with values on a consistent scale. torchvision transforms do this per sample, inside the Dataset. For new code, the torchvision.transforms.v2 module is recommended; it also handles bounding boxes and masks for detection and segmentation.

A typical pipeline converts the image to a tensor (v2.ToImage() then v2.ToDtype(torch.float32, scale=True), which scales 0-255 pixels to 0-1), resizes or crops it, and normalises each channel with v2.Normalize(mean, std). For a pretrained model, use the mean and std it was trained with, commonly the ImageNet statistics.

Augmentation applies random changes such as RandomResizedCrop, RandomHorizontalFlip and ColorJitter, so each epoch sees slightly different versions of each image. That reduces overfitting. Apply augmentation only to training data: validation must be deterministic (resize and centre crop) so scores are comparable between runs. torchvision's ImageFolder builds a labelled dataset from a root/class_name/image.jpg folder layout and accepts a transform argument.

Train and validation pipelines

The two pipelines share normalisation but only training is random. ImageFolder assigns labels from sub-folder names, in sorted order.

import torch
from torchvision import datasets
from torchvision.transforms import v2

MEAN, STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]   # ImageNet statistics

train_tf = v2.Compose([
    v2.ToImage(),
    v2.RandomResizedCrop(224, antialias=True),
    v2.RandomHorizontalFlip(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=MEAN, std=STD),
])
val_tf = v2.Compose([
    v2.ToImage(),
    v2.Resize(256, antialias=True),
    v2.CenterCrop(224),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=MEAN, std=STD),
])

# Expects data/train/<class_name>/*.jpg and data/val/<class_name>/*.jpg
train_ds = datasets.ImageFolder("data/train", transform=train_tf)
val_ds = datasets.ImageFolder("data/val", transform=val_tf)
print(train_ds.classes)

Augment with care

Choose augmentations that keep the label true. A horizontal flip is fine for cats and dogs, but it turns a "b" into a "d" in handwriting and swaps left and right in medical scans. If you flip digits or text, you are teaching the model wrong labels.

त्वरित जाँच: Why does the validation pipeline leave out `RandomHorizontalFlip` and `RandomResizedCrop`?

• Validation images are already flipped
• Random transforms are not supported outside training
• Validation must be deterministic so scores are comparable and reflect real inputs
• Random transforms only work with `ImageFolder`