ai-art

 1import torch
 2import torch.nn as nn
 3
 4import kornia.augmentation as K
 5
 6CUTOUTS = {
 7    "Ji": K.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1, p=0.5),
 8    "Sh": K.RandomSharpness(sharpness=0.5, p=0.5),
 9    "Gn": K.RandomGaussianNoise(mean=0.0, std=1.0, p=0.5),
10    "Pe": K.RandomPerspective(distortion_scale=0.5, p=0.5),
11    "Ro": K.RandomRotation(degrees=15, p=0.5),
12    "Af": K.RandomAffine(
13        degrees=15, translate=0.1, shear=15, padding_mode="border", keepdim=True, p=0.5
14    ),
15    "Et": K.RandomElasticTransform(p=0.5),
16    "Hf": K.RandomHorizontalFlip(p=0.5),
17    "Ts": K.RandomThinPlateSpline(scale=0.2, same_on_batch=False, p=0.5),
18    "Er": K.RandomErasing(
19        scale=(0.02, 0.33), ratio=(0.3, 3.3), same_on_batch=False, p=0.5
20    ),
21}
22
23
24class MakeCutouts(nn.Module):
25    def __init__(self, augments, cut_size, cutn, cut_pow=1.0):
26        super().__init__()
27        self.cut_size = cut_size
28        self.cutn = cutn
29        self.cut_pow = cut_pow
30
31        augment_list = []
32        for item in augments:
33            if item == "Cr":
34                aug = K.RandomCrop(size=(self.cut_size, self.cut_size), p=0.5)
35            elif item == "Re":
36                aug = K.RandomResizedCrop(
37                    size=(self.cut_size, self.cut_size), cropping_mode="resample", p=0.5
38                )
39            else:
40                aug = CUTOUTS[item]
41            augment_list.append(aug)
42
43        print(f"Augmentations: {augment_list}")
44        self.augs = nn.Sequential(*augment_list)
45
46        self.noise_fac = 0.1
47
48        # Pooling
49        self.av_pool = nn.AdaptiveAvgPool2d((self.cut_size, self.cut_size))
50        self.max_pool = nn.AdaptiveMaxPool2d((self.cut_size, self.cut_size))
51
52    def forward(self, input):
53        cutouts = []
54
55        for _ in range(self.cutn):
56            # Use Pooling
57            cutout = (self.av_pool(input) + self.max_pool(input)) / 2
58            cutouts.append(cutout)
59
60        batch = self.augs(torch.cat(cutouts, dim=0))
61
62        if self.noise_fac:
63            facs = batch.new_empty([self.cutn, 1, 1, 1]).uniform_(0, self.noise_fac)
64            batch = batch + facs * torch.randn_like(batch)
65        return batch