Computer Vision Pipelines

Everything starts with labelled data — images paired with the correct answer (the label). Before any training, you split that data into a train set the model learns from and a validation set you hold back to check progress honestly. A common split is 80/20.

The split below is the simplest possible version — slice a list. Run it and watch which samples land where.

# Before training you split your labelled data into two piles:
#   train -> the model learns from these
#   val   -> held back, used ONLY to check progress honestly
# A common split is 80% train, 20% validation.

samples = ["img0", "img1", "img2", "img3", "img4",
           "img5", "img6", "img7", "img8", "img9"]

def train_val_split(data, val_fraction=0.2):
    n_val = int(len(data) * val_fraction)   # 10 * 0.2 = 2 go to val
    val = data[:n_val]                       # first slice -> validati
...

Augmentation creates safe, label-preserving copies of training images so the model sees more variety: horizontal flips, random crops, and colour jitter (brightness/contrast). A flipped cat is still a cat, so the label stays the same while the pixels change.

Here's the simplest augmentation — a horizontal flip — written in plain Python so you can see exactly what changes.

# Augmentation = make safe copies of training images so the model sees
# more variety. A horizontal flip mirrors the image left<->right, which
# teaches the model that a cat facing left is still a cat.

image = [
    [1, 2, 3],
    [4, 5, 6],
]

def horizontal_flip(img):
    # Reverse the pixels WITHIN each row (columns flip), rows stay in order.
    return [list(reversed(row)) for row in img]

flipped = horizontal_flip(image)

for row in flipped:
    print(row)

# Expected output:
# [3, 2, 1]
#
...

Preprocessing makes every image look the same to the model. Two steps dominate: resize (so all images share one width and height the model expects) and normalize (rescale pixels from 0..255 down to a small range like 0..1). Normalizing keeps training stable.

Run the example below to normalize a tiny "image" by hand — every pixel divided by 255.0.

# A camera gives pixel values from 0 (black) to 255 (white).
# Models train best when inputs sit in a small, fixed range like 0..1.
# "Normalizing" = divide every pixel by 255.0 so 0->0.0 and 255->1.0.

# A tiny 2x3 grayscale "image" stored as a nested list (rows of pixels).
image = [
    [0, 128, 255],
    [64, 192, 32],
]

def normalize(img):
    # Walk each row, then each pixel, dividing by 255.0.
    return [[pixel / 255.0 for pixel in row] for row in img]

normalized = normalize(image)

for
...

In real projects you don't do this by hand — torchvision.transforms composes resize, crop, tensor-conversion and normalization into one reusable pipeline. The same recipe must run at training, validation, and inference time.

# In production you don't normalize by hand — torchvision composes the
# whole preprocessing pipeline for you. This is the standard recipe for a
# model pretrained on ImageNet.
import torch
from torchvision import transforms

# Resize -> crop to the model's input size -> to tensor -> normalize.
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),                       # pixels become 0..1 floats
    transforms.Normalize(                        # then ImageNet mean/std
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    ),
])

# Pretend we already loaded a 224x224 RGB image as a tensor of zeros.
fake_image = torch.zeros(3, 224, 224)
print("input shape :", tuple(fake_image.shape))
print("output shape:", tuple(transforms.ToTensor()(
    transforms.ToPILImage()(fake_image)).shape))

# Expected output:
# input shape : (3, 224, 224)
# output shape: (3, 224, 224)

Training a vision model from scratch needs millions of images. Transfer learning avoids that: you take a backbone (like ResNet-50) that already learned generic features — edges, textures, shapes — from ImageNet, and you only replace its final layer (the head) so it outputs your classes.

Below, Albumentations builds the augmentation pipelines (note: train augments, val does not) and a pretrained ResNet-50 has its head swapped for 5 classes.

# Albumentations is the go-to library for fast image augmentation, and
# torchvision.models gives you a pretrained backbone for transfer learning.
import albumentations as A
import torchvision.models as models
import torch.nn as nn

# Build a TRAINING augmentation pipeline (flips, crops, colour jitter).
train_aug = A.Compose([
    A.HorizontalFlip(p=0.5),                 # mirror half the time
    A.RandomResizedCrop(height=224, width=224, scale=(0.8, 1.0)),
    A.ColorJitter(brightness=0.2, contrast=0.2, p=0.5),
    A.Normalize(),                           # 0..1 then ImageNet stats
])

# The VALIDATION pipeline must NOT augment — only resize + normalize.
val_aug = A.Compose([
    A.Resize(256, 256),
    A.CenterCrop(224, 224),
    A.Normalize(),
])

# Transfer learning: take a pretrained ResNet-50 and swap its final layer
# for one that outputs YOUR number of classes (here: 5).
model = models.resnet50(weights="IMAGENET1K_V2")
model.fc = nn.Linear(model.fc.in_features, 5)   # 2048 -> 5 classes
print("ready:", type(model).__name__, "with", model.fc.out_features, "classes")

# Expected output:
# ready: ResNet with 5 classes

Training runs in epochs — one full pass over the training data. Each epoch you let the model learn from the train set (compute loss, back-propagate, update weights), then validate on the held-back set without learning from it. Watching train loss fall while val accuracy rises tells you it's working; if val accuracy stalls or drops while train keeps improving, the model is overfitting.

# The training loop: for each epoch, learn from train data, then check
# val data WITHOUT learning from it. Watch train vs val to spot overfitting.
import torch

def train_one_epoch(model, loader, optimizer, loss_fn):
    model.train()                            # enable dropout/batchnorm updates
    for images, labels in loader:
        optimizer.zero_grad()                # reset old gradients
        preds = model(images)
        loss = loss_fn(preds, labels)
        loss.backward()                      # compute gradients
        optimizer.step()                     # nudge the weights

@torch.no_grad()                             # no gradients = faster, safer
def validate(model, loader):
    model.eval()                             # freeze dropout/batchnorm
    correct = total = 0
    for images, labels in loader:
        preds = model(images).argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
    return correct / total

# A typical printout across 3 epochs:
# Epoch 1  train_loss=1.42  val_acc=0.61
# Epoch 2  train_loss=0.88  val_acc=0.74
# Epoch 3  train_loss=0.55  val_acc=0.81

# Expected output:
# Epoch 1  train_loss=1.42  val_acc=0.61
# Epoch 2  train_loss=0.88  val_acc=0.74
# Epoch 3  train_loss=0.55  val_acc=0.81

Accuracy alone lies on imbalanced data. If 95% of images are "not cancer", a model that always says "not cancer" scores 95% yet catches nothing. So you also look at precision (of the things I flagged, how many were right?), recall (of the things I should have caught, how many did I?), and F1 (their balance). A confusion matrix shows exactly which classes get mixed up.

Deployment means running the trained model on one new image at a time. The golden rule: apply the exact same preprocessing you used for validation — never the training augmentation — switch to model.eval(), and run a single forward pass. A softmax turns the raw scores into probabilities so you can report a confidence.

# Deployment = use the trained model on ONE new image. The golden rule:
# apply the EXACT same preprocessing you used for validation — never the
# training augmentation. Then run a single forward pass in eval mode.
import torch
import torch.nn.functional as F

classes = ["cat", "dog", "car", "house", "tree"]

@torch.no_grad()
def predict(model, image_tensor):
    model.eval()
    logits = model(image_tensor.unsqueeze(0))   # add a batch dimension
    probs = F.softmax(logits, dim=1)[0]         # turn scores into 0..1
    idx = int(probs.argmax())
    return classes[idx], float(probs[idx])

# label, confidence = predict(model, preprocessed_image)
# print(f"{label} ({confidence:.1%})")

# Expected output:
# dog (92.3%)