Robustness/utils.py at main · TimJaspers0801/Robustness · GitHub

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
67
68
69
from torchvision import transforms
import os
import local_corruptions as loc_cor
import global_corruptions as glob_cor

class Compose:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img

# Custom Resize class
class Resize:
    def __init__(self, target_size):
        self.target_size = target_size

    def __call__(self, img):
        img = transforms.functional.resize(img, size=self.target_size,
                                           interpolation=transforms.InterpolationMode.LANCZOS)

        return img

def combine_paths(dir, files):
    paths = []
    for i in range(len(files)):
        paths.append(os.path.join(dir, files[i]))

    return paths


# define curruptions
def define_corruption(corruption, factor):
    corruptions = []

    if corruption == 'Resolution':
        corruptions.append(glob_cor.ReduceResolution(factor))
        corruptions.append(glob_cor.ToPIL())
    if corruption == 'JPEG':
        cwd = os.getcwd()
        corruptions.append(glob_cor.JPGcompression(os.path.join(cwd, 'JPG_image/temporary_image.jpg'), factor))
    if corruption == 'JPEG2000':
        cwd = os.getcwd()
        corruptions.append(glob_cor.JPG2000compression(os.path.join(cwd, 'JPG_image/temporary_image.jpg'), factor))
    if corruption == 'Overexposure':
        position = [(200, 200)]
        corruptions.append(loc_cor.LocalCorruptions(factor, mode='light', gauss_position = position))
        corruptions.append(glob_cor.ToPIL())
    if corruption == 'defocus-blur':
        corruptions.append(loc_cor.LocalCorruptions(factor, mode='blur'))
        corruptions.append(glob_cor.ToPIL())
    if corruption == 'Contrast':
        corruptions.append(glob_cor.ColorJitter(factor, deg='contrast'))
    if corruption == 'Brightness':
        corruptions.append(glob_cor.ColorJitter(factor, deg='brightness'))
    if corruption == 'Saturation':
        corruptions.append(glob_cor.ColorJitter(factor, deg='saturation'))
    if corruption == 'Hue':
        corruptions.append(glob_cor.ColorJitter(factor, deg='hue'))
    if corruption == 'Motion-blur':
        corruptions.append(glob_cor.motion_blur(factor))
    if corruption == 'Sharpness':
        corruptions.append(glob_cor.Sharpness(factor))
    if corruption == 'Underexposure':
        corruptions.append(glob_cor.ColorJitter(factor, deg='underexposure'))

    return corruptions