File size: 1,137 Bytes
3440f83
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
# code in this file is adpated from the ALBEF repo (https://github.com/salesforce/ALBEF)

from torchvision import transforms
from .randaugment import RandomAugment
from PIL import Image


def square_transform(size=224):
    return transforms.Compose(
        [
            transforms.Resize((size, size), interpolation=Image.BICUBIC),
            transforms.ToTensor(),
        ]
    )


def square_transform_randaug(size=224):
    return transforms.Compose(
        [
            transforms.RandomResizedCrop(size, scale=(0.8, 1.0), interpolation=Image.BICUBIC),
            transforms.RandomHorizontalFlip(),
            RandomAugment(
                2,
                7,
                isPIL=True,
                augs=[
                    "Identity",
                    "AutoContrast",
                    "Equalize",
                    "Brightness",
                    "Sharpness",
                    "ShearX",
                    "ShearY",
                    "TranslateX",
                    "TranslateY",
                    "Rotate",
                ],
            ),
            transforms.ToTensor(),
        ]
    )