augment.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

"""
3Augment implementation
Data-augmentation (DA) based on dino DA (https://github.com/facebookresearch/dino)
and timm DA(https://github.com/rwightman/pytorch-image-models)
"""

import random

import numpy as np
import torch
from PIL import ImageFilter, ImageOps
from timm.data.mixup import mixup_target
from timm.data.transforms import RandomResizedCropAndInterpolation
from torch.utils.data._utils.collate import default_collate_fn_map, collate
from torchvision import transforms
from math import sqrt

DEBUG = False

from patch_sampler import GridSampler, RandomUniformSampler, RandomMultiscaleSampler, RandomCropV2, \
    RandomDelegatedSampler, SmartSampler, CentralMultiscaleSampler


class GaussianBlur(object):
    """
    Apply Gaussian Blur to the PIL image.
    """

    def __init__(self, p=0.1, radius_min=0.1, radius_max=2.):
        self.prob = p
        self.radius_min = radius_min
        self.radius_max = radius_max

    def __call__(self, img):
        do_it = random.random() <= self.prob
        if not do_it:
            return img

        img = img.filter(
            ImageFilter.GaussianBlur(
                radius=random.uniform(self.radius_min, self.radius_max)
            )
        )
        return img


class Solarization(object):
    """
    Apply Solarization to the PIL image.
    """

    def __init__(self, p=0.2):
        self.p = p

    def __call__(self, img):
        if random.random() < self.p:
            return ImageOps.solarize(img)
        else:
            return img


class gray_scale(object):
    """
    Apply Solarization to the PIL image.
    """

    def __init__(self, p=0.2):
        self.p = p
        self.transf = transforms.Grayscale(3)

    def __call__(self, img):
        if random.random() < self.p:
            return self.transf(img)
        else:
            return img


class horizontal_flip(object):
    """
    Apply Solarization to the PIL image.
    """

    def __init__(self, p=0.2, activate_pred=False):
        self.p = p
        self.transf = transforms.RandomHorizontalFlip(p=1.0)

    def __call__(self, img):
        if random.random() < self.p:
            return self.transf(img)
        else:
            return img


def new_data_aug_generator(args=None):
    img_size = args.input_size
    remove_random_resized_crop = args.src
    mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
    primary_tfl = []
    scale = (0.08, 1.0)
    interpolation = 'bicubic'
    if remove_random_resized_crop:
        primary_tfl = [
            transforms.Resize(img_size, interpolation=3),
            transforms.RandomCrop(img_size, padding=4, padding_mode='reflect'),
            transforms.RandomHorizontalFlip()
        ]
    else:
        primary_tfl = [
            RandomResizedCropAndInterpolation(
                img_size, scale=scale, interpolation=interpolation),
            transforms.RandomHorizontalFlip()
        ]

    secondary_tfl = [transforms.RandomChoice([gray_scale(p=1.0),
                                              Solarization(p=1.0),
                                              GaussianBlur(p=1.0)])]

    if args.color_jitter is not None and not args.color_jitter == 0:
        secondary_tfl.append(transforms.ColorJitter(args.color_jitter, args.color_jitter, args.color_jitter))
    final_tfl = [
        transforms.ToTensor(),
        transforms.Normalize(
            mean=torch.tensor(mean),
            std=torch.tensor(std))
    ]
    return transforms.Compose(primary_tfl + secondary_tfl + final_tfl)


class AugCollatePatchSampleMixup:
    def __init__(self, crop_sizes=None, smart_patches=None, random_patches=None, random_min_size=8, random_max_size=48,
                 grid_patch_size=None, random_patch_min_scale=None, random_patch_max_scale=None,
                 grid_to_random_ratio=0.7, augment=True, mixup_alpha=1., cutmix_alpha=0., prob=1.0, switch_prob=0.5,
                 correct_lam=True, label_smoothing=0.1, num_classes=1000,
                 central_patch_sizes=None,
                 merge_patches = False,
                 hard_dropout=True,
                 patch_modifications = None) -> None:
        self.augment = augment
        self.mixup_alpha = mixup_alpha
        self.cutmix_alpha = cutmix_alpha
        self.mix_prob = prob
        self.switch_prob = switch_prob
        self.label_smoothing = label_smoothing
        self.num_classes = num_classes
        self.correct_lam = correct_lam

        if smart_patches is not None:
            self.patch_sampler = SmartSampler(smart_patches, random_min_size, random_max_size)
            print('using smart sampler')
        elif grid_patch_size is not None:
            self.patch_sampler = GridSampler(patch_size=(grid_patch_size, grid_patch_size), patch_modifications=patch_modifications, merge_patches=merge_patches, hard_dropout=hard_dropout)
            print('using grid sampler')
            if random_patches is not None:
                random_patch_sampler = RandomUniformSampler(random_patches, random_min_size, random_max_size)
                print('using random sampler')
                self.patch_sampler = RandomDelegatedSampler([
                    (self.patch_sampler, grid_to_random_ratio),
                    (random_patch_sampler, 1 - grid_to_random_ratio)
                ])
        elif central_patch_sizes is not None:
            print('using central multiscale sampler')
            self.patch_sampler = CentralMultiscaleSampler(patch_sizes=central_patch_sizes)
        elif random_patch_min_scale is not None and random_patch_max_scale is not None and random_patches is not None:
            print('using random sampler v2')
            self.patch_sampler = RandomCropV2(patches_num=random_patches, min_scale=random_patch_min_scale,
                                              max_scale=random_patch_max_scale)
        elif random_patches is not None:
            print('using random sampler')
            self.patch_sampler = RandomUniformSampler(random_patches, random_min_size, random_max_size)
        elif crop_sizes is not None:
            print('using random multiscale sampler')
            self.patch_sampler = RandomMultiscaleSampler(eval(crop_sizes))
        else:
            raise ValueError('no patch sampling parameters specified')

    def _params_per_batch(self):
        lam = 1.
        use_cutmix = False
        if np.random.rand() < self.mix_prob:
            if self.mixup_alpha > 0. and self.cutmix_alpha > 0.:
                use_cutmix = np.random.rand() < self.switch_prob
                lam_mix = np.random.beta(self.cutmix_alpha, self.cutmix_alpha) if use_cutmix else \
                    np.random.beta(self.mixup_alpha, self.mixup_alpha)
            elif self.mixup_alpha > 0.:
                lam_mix = np.random.beta(self.mixup_alpha, self.mixup_alpha)
            elif self.cutmix_alpha > 0.:
                use_cutmix = True
                lam_mix = np.random.beta(self.cutmix_alpha, self.cutmix_alpha)
            else:
                assert False, "One of mixup_alpha > 0., cutmix_alpha > 0., cutmix_minmax not None should be true."
            lam = float(lam_mix)
        return lam, use_cutmix

    def _mixup_batch(self, x, lam):
        if lam == 1.:
            return 1.
        x_flipped = x.flip(0).mul_(1. - lam)
        x.mul_(lam).add_(x_flipped)
        return lam

    def _patchmix_batch(self, x, coords, lam):
        if lam == 1.:
            return 1.
        p = int(lam * x.shape[1])  # patches to keep
        lam = p / x.shape[1]  # round lambda to patch
        x[:, p:, ...] = x.flip(0)[:, p:, ...]
        coords[:, p:, ...] = coords.flip(0)[:, p:, ...]
        return lam

    def on_epoch(self):
        if not self.augment and self.patch_sampler is not None:
            self.patch_sampler.reset_state()

    def __call__(self, batch):
        collated_batch = collate(batch, collate_fn_map=default_collate_fn_map)
        if isinstance(collated_batch, dict):
            images, targets = collated_batch["image"], collated_batch["label"]
        else:
            images, targets = collated_batch
        assert images.shape[0] % 2 == 0, 'Batch size should be even when using this'


        #import matplotlib.pyplot as plt
        #plt.imsave('debug_0.png', images[0].detach().permute(1,2,0).numpy())
        #assert False

        lam, use_cutmix = self._params_per_batch()

        if self.augment and not use_cutmix:
            lam = self._mixup_batch(images, lam)

        patches = []
        coords = []
        keeps = []
        for im in images:
            pt, cs, ks = self.patch_sampler(im)
            patches.append(pt)
            coords.append(cs)
            if ks is not None:
                keeps.append(ks)
        images = torch.stack(patches, dim=0)
        del patches
        coords = torch.stack(coords, dim=0)
        if keeps:
            keeps = torch.stack(keeps, dim=0)
        else:
            keeps = None
        if self.augment and use_cutmix:
            lam = self._patchmix_batch(images, coords, lam)

        if self.augment:
            targets = mixup_target(targets, self.num_classes, lam, self.label_smoothing)  # , images.device)
        return (images, coords, keeps), targets