06 - Image Augmentation: Techniques to Prevent Overfitting

# Example: cat/dog classification

# CORRECT transformations (preserve semantics):
# - Horizontal flip: a horizontally mirrored cat is still a cat ✓
# - Brightness variation: a cat in dim light is still a cat ✓
# - Random crop: a detail of a cat is still recognizable ✓
# - Slight rotation: a cat rotated 15 degrees is still a cat ✓

# DANGEROUS transformations (may change semantics):
# - Vertical flip for traffic signs: "STOP" upside down loses meaning ✗
# - Rotation > 45 degrees for text/numbers: "6" rotated becomes "9" ✗
# - Extreme scale: 5% crop may lose all context ✗
# - Extreme color jitter for medical diagnostics: color is semantically relevant ✗

# Golden rule:
# "An augmentation is valid if a human, seeing the augmented image,
#  would still assign the same label"

# Practical impact on benchmarks (same ResNet-18, same hyperparameters):
# CIFAR-10 without augmentation:      ~84.3% accuracy
# + Flip + Crop:                      ~91.8% accuracy  (+7.5%)
# + Color Jitter:                     ~93.2% accuracy  (+1.4%)
# + Cutout/CoarseDropout:             ~94.1% accuracy  (+0.9%)
# + MixUp (alpha=0.2):                ~95.3% accuracy  (+1.2%)
# + CutMix (alpha=1.0):               ~95.8% accuracy  (+0.5%)
# + AutoAugment (CIFAR-10 policy):    ~97.1% accuracy  (+1.3%)
# TrivialAugment + MixUp:             ~97.4% accuracy  (best combo)

# Note: each increment is obtained WITHOUT adding any real data.
# Data augmentation = virtually infinite dataset from finite data.

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

# ---- Standard pipeline for classification ----
def get_classification_transforms(img_size: int = 224, is_train: bool = True):
    if is_train:
        return A.Compose([
            # Geometric
            A.RandomResizedCrop(img_size, img_size, scale=(0.7, 1.0),
                               ratio=(0.75, 1.33), p=1.0),
            A.HorizontalFlip(p=0.5),
            A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.15,
                               rotate_limit=15, border_mode=cv2.BORDER_REFLECT, p=0.7),
            A.OneOf([
                A.Perspective(scale=(0.05, 0.1)),
                A.GridDistortion(num_steps=5, distort_limit=0.3),
                A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50)
            ], p=0.3),

            # Photometric
            A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.7),
            A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30,
                                val_shift_limit=20, p=0.5),
            A.OneOf([
                A.GaussNoise(var_limit=(10, 50)),
                A.GaussianBlur(blur_limit=(3, 7)),
                A.MotionBlur(blur_limit=7),
                A.MedianBlur(blur_limit=5)
            ], p=0.4),
            A.ImageCompression(quality_lower=70, quality_upper=100, p=0.2),

            # Dropout / occlusion simulation
            A.CoarseDropout(max_holes=8, max_height=32, max_width=32,
                           min_holes=1, p=0.3),  # similar to Cutout
            A.RandomGridShuffle(grid=(3, 3), p=0.1),

            # ImageNet normalization
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])
    else:
        # Validation: deterministic operations only
        return A.Compose([
            A.Resize(int(img_size * 1.14), int(img_size * 1.14)),
            A.CenterCrop(img_size, img_size),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])

# ---- Detection pipeline (automatically updates bounding boxes!) ----
def get_detection_transforms(img_size: int = 640, is_train: bool = True):
    if is_train:
        return A.Compose([
            A.RandomResizedCrop(img_size, img_size, scale=(0.5, 1.0)),
            A.HorizontalFlip(p=0.5),
            A.RandomBrightnessContrast(p=0.7),
            A.HueSaturationValue(p=0.5),
            A.OneOf([
                A.GaussNoise(),
                A.GaussianBlur(blur_limit=3)
            ], p=0.3),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ],
        # CRITICAL: specify bbox format for automatic coordinate update
        bbox_params=A.BboxParams(
            format='yolo',           # or 'pascal_voc', 'coco', 'albumentations'
            label_fields=['class_labels'],
            min_visibility=0.3,      # remove bbox if visibility < 30%
            min_area=100             # remove bbox if area < 100 pixels
        ))
    else:
        return A.Compose([
            A.Resize(img_size, img_size),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ],
        bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))

# ---- Segmentation pipeline (automatically updates masks!) ----
def get_segmentation_transforms(img_size: int = 512, is_train: bool = True):
    if is_train:
        return A.Compose([
            A.RandomResizedCrop(img_size, img_size, scale=(0.7, 1.0)),
            A.HorizontalFlip(p=0.5),
            A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1,
                               rotate_limit=10, p=0.5),
            A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=0.3),
            A.RandomBrightnessContrast(p=0.5),
            A.GaussNoise(var_limit=(10, 30), p=0.3),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])
        # NOTE: mask is passed as 'mask' argument - automatically updated!
    else:
        return A.Compose([
            A.Resize(img_size, img_size),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])

# Usage with detection
transform = get_detection_transforms(is_train=True)
image = cv2.imread('image.jpg')
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
bboxes = [(0.5, 0.5, 0.3, 0.4)]  # YOLO format [x_c, y_c, w, h]
labels = [0]

result = transform(image=image_rgb, bboxes=bboxes, class_labels=labels)
transformed_image = result['image']     # tensor [3, H, W]
transformed_boxes = result['bboxes']    # automatically updated!
print(f"Original: {bboxes} -> Transformed: {transformed_boxes}")

import torch
import numpy as np

def mixup_batch(images: torch.Tensor, labels: torch.Tensor,
                alpha: float = 0.2) -> tuple:
    """
    MixUp: linearly interpolates two images and their labels.
    Output: mixed image, soft labels.

    lambda ~ Beta(alpha, alpha)
    image_mixed = lambda * image_a + (1 - lambda) * image_b
    label_mixed = lambda * label_a + (1 - lambda) * label_b

    With alpha=0.2, lambda is typically close to 0 or 1 (near-pure),
    with alpha=1.0 (uniform Beta), images are equally blended.
    """
    batch_size = images.size(0)
    lam = np.random.beta(alpha, alpha)

    # Random permutation for the second image in the batch
    perm = torch.randperm(batch_size)

    mixed_images = lam * images + (1 - lam) * images[perm]
    labels_a = labels
    labels_b = labels[perm]

    # Loss: loss = lam * CE(pred, a) + (1-lam) * CE(pred, b)
    return mixed_images, labels_a, labels_b, lam

def cutmix_batch(images: torch.Tensor, labels: torch.Tensor,
                 alpha: float = 1.0) -> tuple:
    """
    CutMix: replaces a rectangular region of one image with another.
    More effective than MixUp for detection (preserves intact regions).

    lambda ~ Beta(alpha, alpha)  # determines the cut area proportion
    """
    batch_size, C, H, W = images.size()
    lam = np.random.beta(alpha, alpha)
    perm = torch.randperm(batch_size)

    # Calculate cut box dimensions
    cut_ratio = np.sqrt(1 - lam)
    cut_w = int(W * cut_ratio)
    cut_h = int(H * cut_ratio)

    # Random box center
    cx = np.random.randint(W)
    cy = np.random.randint(H)

    # Box coordinates (clipped to image bounds)
    x1 = np.clip(cx - cut_w // 2, 0, W)
    x2 = np.clip(cx + cut_w // 2, 0, W)
    y1 = np.clip(cy - cut_h // 2, 0, H)
    y2 = np.clip(cy + cut_h // 2, 0, H)

    # Apply CutMix (immutable: creates a copy)
    mixed_images = images.clone()
    mixed_images[:, :, y1:y2, x1:x2] = images[perm, :, y1:y2, x1:x2]

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

    return mixed_images, labels, labels[perm], lam

def mosaic_augmentation(images: list, labels: list) -> tuple:
    """
    Mosaic augmentation (introduced in YOLOv5):
    Combines 4 images into a 2x2 mosaic with random central crop.
    Particularly effective for detection on small objects:
    each image in the mosaic is at 25% of original size,
    simulating objects at greater distances.
    """
    assert len(images) == 4, "Mosaic requires exactly 4 images"

    _, H, W = images[0].shape
    mosaic = torch.zeros(3, H * 2, W * 2)

    # Place the 4 images in the mosaic
    mosaic[:, 0:H, 0:W]     = images[0]  # top-left
    mosaic[:, 0:H, W:2*W]   = images[1]  # top-right
    mosaic[:, H:2*H, 0:W]   = images[2]  # bottom-left
    mosaic[:, H:2*H, W:2*W] = images[3]  # bottom-right

    # Random central crop (simulates different viewpoints)
    crop_y = np.random.randint(H // 2, H)
    crop_x = np.random.randint(W // 2, W)
    mosaic_cropped = mosaic[:, crop_y-H//2:crop_y+H//2,
                            crop_x-W//2:crop_x+W//2]

    combined_labels = []
    for lbl in labels:
        combined_labels.extend(lbl)

    return mosaic_cropped, combined_labels

# ---- Training loop with MixUp/CutMix ----
def train_with_advanced_augmentation(
    model, train_loader, optimizer, criterion, device,
    mixup_alpha: float = 0.2, cutmix_alpha: float = 1.0,
    mixup_prob: float = 0.5, cutmix_prob: float = 0.5
) -> float:
    """Training step that applies MixUp or CutMix with given probability."""
    model.train()
    total_loss = 0.0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        r = np.random.rand()
        if r < mixup_prob:
            mixed_images, labels_a, labels_b, lam = mixup_batch(images, labels, mixup_alpha)
            outputs = model(mixed_images)
            loss = lam * criterion(outputs, labels_a) + (1 - lam) * criterion(outputs, labels_b)
        elif r < mixup_prob + cutmix_prob:
            mixed_images, labels_a, labels_b, lam = cutmix_batch(images, labels, cutmix_alpha)
            outputs = model(mixed_images)
            loss = lam * criterion(outputs, labels_a) + (1 - lam) * criterion(outputs, labels_b)
        else:
            outputs = model(images)
            loss = criterion(outputs, labels)

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        total_loss += loss.item()

    return total_loss / len(train_loader)

from torchvision import transforms

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

# ---- AutoAugment (policy learned on ImageNet) ----
train_auto = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.AutoAugment(
        policy=transforms.AutoAugmentPolicy.IMAGENET,  # or CIFAR10, SVHN
        interpolation=transforms.InterpolationMode.BILINEAR
    ),
    transforms.ToTensor(),
    transforms.Normalize(mean=MEAN, std=STD)
])

# ---- RandAugment (N=2, M=9 are typical optimal values) ----
train_rand = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.RandAugment(num_ops=2, magnitude=9),
    transforms.ToTensor(),
    transforms.Normalize(mean=MEAN, std=STD)
])

# ---- TrivialAugment: even simpler, often best ----
# Selects 1 random operation with uniform random magnitude
train_trivial = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.TrivialAugmentWide(),  # PyTorch 1.13+
    transforms.ToTensor(),
    transforms.Normalize(mean=MEAN, std=STD),
    transforms.RandomErasing(p=0.1)  # adds Cutout
])

# Recommendation:
# - AutoAugment: use ONLY if you can afford the search or use pre-found policies
# - RandAugment: good default for new datasets
# - TrivialAugment + MixUp/CutMix: best practical combination for classification

import torch
import torch.nn.functional as F
import torchvision.transforms.functional as TF

@torch.no_grad()
def tta_predict(model: torch.nn.Module,
                image: torch.Tensor,  # [1, C, H, W]
                n_augmentations: int = 5) -> torch.Tensor:
    """
    Test-Time Augmentation: average predictions over augmented images.

    Practical tips:
    - Horizontal flip TTA is the most reliable
    - Avoid strong rotation TTA (can degrade performance)
    - n=5 is a good speed/performance trade-off
    """
    device = image.device
    model.eval()
    predictions = []

    # 1. Original image
    pred = F.softmax(model(image), dim=1)
    predictions.append(pred)

    # 2. Horizontal flip
    flipped = TF.hflip(image)
    pred = F.softmax(model(flipped), dim=1)
    predictions.append(pred)

    # 3-N. Slightly scaled crops
    _, C, H, W = image.shape
    scale = 0.9
    for _ in range(n_augmentations - 2):
        new_h, new_w = int(H * scale), int(W * scale)
        resized = TF.resize(image, (new_h, new_w))

        top  = torch.randint(0, H - new_h + 1, (1,)).item()
        left = torch.randint(0, W - new_w + 1, (1,)).item()
        cropped = TF.crop(resized, top, left, new_h, new_w)
        cropped = TF.resize(cropped, (H, W))

        pred = F.softmax(model(cropped), dim=1)
        predictions.append(pred)

    # Average probabilities (ensemble of N predictions)
    mean_pred = torch.stack(predictions).mean(dim=0)
    return mean_pred.argmax(dim=1)  # predicted class

import albumentations as A
from albumentations.pytorch import ToTensorV2

def get_medical_transforms(img_size: int = 512, modality: str = 'xray'):
    """
    Modality-specific augmentation for medical images.
    Different strategies for X-ray, histology, MRI, ultrasound.
    """
    common = [A.Resize(img_size, img_size)]

    if modality == 'xray':
        augmentations = [
            # Horizontal flip only (anatomically valid for chest)
            A.HorizontalFlip(p=0.5),
            # Slight rotation (patients are not always perfectly aligned)
            A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05,
                               rotate_limit=10, p=0.5),
            # CLAHE improves contrast on X-rays
            A.CLAHE(clip_limit=2.0, tile_grid_size=(8, 8), p=0.5),
            # Realistic sensor noise
            A.GaussNoise(var_limit=(5, 25), p=0.4),
            # Slight contrast variation
            A.RandomGamma(gamma_limit=(80, 120), p=0.5),
            # DO NOT use: vertical flip, color jitter, hue shift
        ]

    elif modality == 'histology':
        augmentations = [
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),    # OK for histology (no fixed orientation)
            A.RandomRotate90(p=0.5),
            # Color variation important for histology (different labs/staining)
            A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=20,
                                val_shift_limit=20, p=0.7),
            A.RandomBrightnessContrast(brightness_limit=0.2,
                                      contrast_limit=0.2, p=0.7),
            A.ElasticTransform(alpha=120, sigma=120 * 0.05,
                               alpha_affine=120 * 0.03, p=0.3),
        ]

    elif modality == 'mri':
        augmentations = [
            A.HorizontalFlip(p=0.5),
            A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1,
                               rotate_limit=15, p=0.5),
            # MRI: intensity variations between different scanners
            A.RandomBrightnessContrast(brightness_limit=0.3,
                                      contrast_limit=0.3, p=0.7),
            # Simulate MRI motion artifacts
            A.GaussianBlur(blur_limit=3, p=0.3),
            # Realistic anatomical elastic deformation
            A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=0.2),
        ]

    else:
        augmentations = []

    norm = [
        A.Normalize(mean=[0.5], std=[0.5]),  # Grayscale normalization
        # For RGB: A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]

    return A.Compose(common + augmentations + norm + [ToTensorV2()])

import albumentations as A
from albumentations.pytorch import ToTensorV2

def get_industrial_transforms(img_size: int = 256, is_train: bool = True):
    """
    Pipeline for industrial visual inspection.
    Goal: robustness to lighting variation, partial rotation,
    sensor noise, and small mechanical deformations.
    """
    if is_train:
        return A.Compose([
            A.RandomResizedCrop(img_size, img_size, scale=(0.8, 1.0)),

            # Rotation: products on conveyor belts have variable orientations
            A.RandomRotate90(p=0.5),
            A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
                               rotate_limit=30, border_mode=0, p=0.7),

            # Illumination: variations from industrial lighting (LED, fluorescent)
            A.OneOf([
                A.RandomBrightnessContrast(brightness_limit=0.4, contrast_limit=0.4),
                A.CLAHE(clip_limit=4.0),
                A.RandomGamma(gamma_limit=(70, 130))
            ], p=0.8),

            # Sensor: noise and blur from industrial acquisition systems
            A.OneOf([
                A.GaussNoise(var_limit=(10, 60)),
                A.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5)),
                A.MultiplicativeNoise(multiplier=[0.9, 1.1]),
            ], p=0.4),
            A.OneOf([
                A.GaussianBlur(blur_limit=(3, 5)),
                A.Defocus(radius=(1, 3)),
                A.MotionBlur(blur_limit=5)
            ], p=0.3),

            # Reflection/shadow artifacts
            A.RandomShadow(shadow_roi=(0, 0, 1, 1), p=0.2),
            A.Downscale(scale_min=0.7, scale_max=0.9, p=0.2),  # simulate low resolution

            # CoarseDropout simulates partial occlusion of the part
            A.CoarseDropout(max_holes=3, max_height=32, max_width=32, p=0.2),

            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])
    else:
        return A.Compose([
            A.Resize(img_size, img_size),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])

import torch
import pandas as pd
import albumentations as A
from albumentations.pytorch import ToTensorV2

def run_augmentation_ablation(
    model_class,
    dataset_path: str,
    augmentation_configs: dict,
    n_epochs: int = 30,
    device: str = 'cuda'
) -> pd.DataFrame:
    """
    Systematic ablation study over augmentation configurations.
    All configurations use identical training conditions.

    augmentation_configs = {
        'baseline':   A.Compose([A.Resize(224,224), A.Normalize(...), ToTensorV2()]),
        'flip+crop':  A.Compose([A.RandomResizedCrop(224,224), A.HorizontalFlip(0.5), ...]),
        'full':       get_classification_transforms(is_train=True),
    }
    """
    results = []

    for config_name, transform in augmentation_configs.items():
        print(f"\n=== Ablation: {config_name} ===")

        # Same random seed for every config to ensure fair comparison
        torch.manual_seed(42)
        model = model_class(num_classes=10).to(device)
        optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=n_epochs)
        criterion = torch.nn.CrossEntropyLoss()

        best_val_acc = 0.0
        for epoch in range(n_epochs):
            model.train()
            # ... training loop ...
            scheduler.step()

            # Validation
            model.eval()
            correct = total = 0
            with torch.no_grad():
                pass  # ... validation loop ...
            val_acc = 100.0 * correct / max(total, 1)
            best_val_acc = max(best_val_acc, val_acc)

        results.append({
            'config': config_name,
            'best_val_acc': round(best_val_acc, 2)
        })
        print(f"Best val accuracy: {best_val_acc:.2f}%")

    df = pd.DataFrame(results).sort_values('best_val_acc', ascending=False)
    print("\n=== Ablation Study Results ===")
    print(df.to_string(index=False))
    return df