File size: 1,762 Bytes
a80d6bb
 
 
 
 
 
 
 
 
c74a070
 
 
 
 
 
 
 
 
 
 
 
a80d6bb
 
c74a070
a80d6bb
 
 
 
 
 
 
 
c74a070
 
 
 
 
 
 
 
 
 
 
a80d6bb
 
c74a070
a80d6bb
 
 
 
c74a070
 
 
 
a80d6bb
c74a070
a80d6bb
 
 
 
c74a070
a80d6bb
 
c74a070
 
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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import albumentations as A


class DarkAug(object):
    """
    Extreme dark augmentation aiming at Aachen Day-Night
    """

    def __init__(self) -> None:
        self.augmentor = A.Compose(
            [
                A.RandomBrightnessContrast(
                    p=0.75, brightness_limit=(-0.6, 0.0), contrast_limit=(-0.5, 0.3)
                ),
                A.Blur(p=0.1, blur_limit=(3, 9)),
                A.MotionBlur(p=0.2, blur_limit=(3, 25)),
                A.RandomGamma(p=0.1, gamma_limit=(15, 65)),
                A.HueSaturationValue(p=0.1, val_shift_limit=(-100, -40)),
            ],
            p=0.75,
        )

    def __call__(self, x):
        return self.augmentor(image=x)["image"]


class MobileAug(object):
    """
    Random augmentations aiming at images of mobile/handhold devices.
    """

    def __init__(self):
        self.augmentor = A.Compose(
            [
                A.MotionBlur(p=0.25),
                A.ColorJitter(p=0.5),
                A.RandomRain(p=0.1),  # random occlusion
                A.RandomSunFlare(p=0.1),
                A.JpegCompression(p=0.25),
                A.ISONoise(p=0.25),
            ],
            p=1.0,
        )

    def __call__(self, x):
        return self.augmentor(image=x)["image"]


def build_augmentor(method=None, **kwargs):
    if method is not None:
        raise NotImplementedError(
            "Using of augmentation functions are not supported yet!"
        )
    if method == "dark":
        return DarkAug()
    elif method == "mobile":
        return MobileAug()
    elif method is None:
        return None
    else:
        raise ValueError(f"Invalid augmentation method: {method}")


if __name__ == "__main__":
    augmentor = build_augmentor("FDA")