data/transforms.py

import albumentations as A

def train_transforms(size=512):
    return A.Compose([
        A.Resize(size, size),
        A.RandomBrightnessContrast(p=0.5),
        A.GaussianBlur(p=0.2),
        A.Rotate(limit=15, p=0.4),
        A.RandomCrop(height=size, width=size, p=0.2),
        A.GaussNoise(p=0.2),
    ])

def val_transforms(size=512):
    return A.Compose([
        A.Resize(size, size),
    ])