TransformDefin.py

import torch
from torchvision import transforms
from PIL import Image
import torchvision.transforms.functional as F
from CustomTransforms import PILToFloat255, Clamp255, TenCropNormalizeStack


# 3) Lighting (PCA jitter) transform
class Lighting(object):
    """
    PCA-based color augmentation.
    α_i ~ N(0, alphastd)
    shift = eigvecs @ (α * eigvals)
    """
    def __init__(self, alphastd, eigvals, eigvecs):
        self.alphastd = alphastd
        # register as tensors for fast matmul
        self.eigvals = torch.as_tensor(eigvals, dtype=torch.float32)
        self.eigvecs = torch.as_tensor(eigvecs, dtype=torch.float32)


    def __call__(self, img):
        """
        img: Tensor shape (3,H,W), values in [0,1]
        """
        if self.alphastd == 0:
            return img
        
        if not torch.is_floating_point(img):
            img = img.float()

        # sample alpha
        alpha = img.new_empty(3).normal_(0, self.alphastd) #This simply creates the tensor of (3,) sampling a random number from normal distribution with mean 0 and stddev alphastd
        #We use img.new_empty to ensure the tensor is on the same device(gpu) as img


        # compute shift: (3×3) @ (3,) → (3,)
        shift = (self.eigvecs @ (self.eigvals * alpha)).view(3, 1, 1)
        # reshape to (3,1,1) and add
        return img + shift

def float255_to_float01(t):
    return t / 255.0

if __name__ == "__main__":
    # Load PCA values from file
    pca_values = torch.load("pca_values.pth",weights_only=False)
    mean = pca_values['mean']
    std = pca_values['std']
    eigvals = pca_values['eigvals']
    eigvecs = pca_values['eigvecs']

    # Define your training transforms
    train_transforms = transforms.Compose([
        transforms.RandomCrop(224),
        transforms.RandomHorizontalFlip(p=0.5),

        PILToFloat255(),                      # -> float tensor 0..255
        Lighting(alphastd=0.01, eigvals=eigvals, eigvecs=eigvecs),
        Clamp255(),

        transforms.Normalize(mean=mean.tolist(), std=std.tolist()),
        ])

    val_transforms = transforms.Compose([
        transforms.TenCrop(224),  # 4 corners + center + their horizontal flips = 10
        TenCropNormalizeStack(mean=mean.tolist(), std=std.tolist()),
    ])

    # Save both (picklable)
    torch.save(train_transforms, "train_transform.pth")
    torch.save(val_transforms, "val_transform.pth")

    print("Saved full transforms")