Robustness/local_corruptions.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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import cv2
import numpy as np
import random
from scipy.stats import norm


class LocalCorruptions:
    def __init__(self, severity, mode='light', gauss_position=None):
        # mode is either light or blur
        self.severity = severity
        self.mode = mode
        self.gauss_position = gauss_position

    def normalize(self, inp: np.ndarray):
        """Squash image input to the value range [0, 1] (no clipping)"""
        inp_out = (inp - np.min(inp)) / np.ptp(inp)
        return inp_out

    def _decay_value_radically_norm_in_matrix(self, mask_size, centers, max_value, min_value, dev):
        """
        _decay_value_radically_norm function in matrix format
        """
        center_prob = norm.pdf(0, 0, dev)
        x_value_rate = np.zeros((mask_size[1], mask_size[0]))
        for center in centers:
            coord_x = np.arange(mask_size[0])
            coord_y = np.arange(mask_size[1])
            xv, yv = np.meshgrid(coord_x, coord_y)
            dist_x = xv - center[0]
            dist_y = yv - center[1]
            dist = np.sqrt(np.power(dist_x, 2) + np.power(dist_y, 2))
            x_value_rate += norm.pdf(dist, 0, dev) / center_prob
        mask = x_value_rate * (max_value - min_value) + min_value
        mask[mask > 255] = 255

        return mask

    def generate_gauss_mask(self,
                                 mask_size,
                                 position=None,
                                 max_brightness=255,
                                 min_brightness=0):
        """
        Generate decayed guassian mask generated by given position, direction. Multiple positions are accepted.
        Args:
            mask_size: tuple of integers (w, h) defining generated mask size
            position: list of tuple of integers (x, y) defining the center of spotlight light position,
                      which is the reference point during rotating
            max_brightness: integer that max brightness in the mask
            min_brightness: integer that min brightness in the mask
            mode: the way that brightness decay from max to min: linear or gaussian
            linear_decay_rate: only valid in linear_static mode. Suggested value is within [0.2, 2]
            speedup: use `shrinkage then expansion` strategy to speed up vale calculation
        Return:
            gaussian: ndarray in float type consisting value from max_brightness to min_brightness. If in 'linear' mode
                        minimum value could be smaller than given min_brightness.
        """
        if position is None:
            position = [(random.randint(0 + 200, mask_size[0] - 200), random.randint(0 + 200, mask_size[1]) - 200)]

        mask = np.zeros(shape=(mask_size[1], mask_size[0]), dtype=np.float32)
        mu = np.sqrt(mask.shape[0] ** 2 + mask.shape[1] ** 2)
        dev = mu / self.sigma
        mask = self._decay_value_radically_norm_in_matrix(mask_size, position, max_brightness, min_brightness, dev)

        mask = np.asarray(mask, dtype=np.uint8)
        # add median blur
        mask = cv2.medianBlur(mask, 5)
        mask = 255 - mask

        return mask

    def alphaBlend(self, img1, mask):
        """
        alphaBlend img and blurred images with weighted mask
        """
        blur = cv2.GaussianBlur(img1, (self.kernel_size, self.kernel_size), self.kernel_size)

        if mask.ndim == 3 and mask.shape[-1] == 3:
            alpha = mask / 255.0
        else:
            alpha = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) / 255.0
        blended = cv2.convertScaleAbs(img1 * (1 - alpha) + blur * alpha)

        return blended

    def add_corruption(self, frame, light_position=None, max_brightness=255, min_brightness=0):
        """
        Add gaussian mask to given image
        """

        if self.mode == "light":
            c = \
            [(12, 0.9), (12, 0.8), (10, 0.8), (8, 0.75), (8, 0.65), (8, 0.5), (6,0.3), (6,0.2), (4,0.1), (2,0.1)][self.severity]
        if self.mode == "blur":
            c = [(12, 0.4, 31), (12, 0.4, 41), (10, 0.3, 51), (8, 0.3, 61), (8, 0.2, 71), (8, 0.2, 81), (6, 0.2, 91),
                 (6, 0.2, 101), (4, 0.1, 111), (4, 0.1, 121)][
                self.severity]
            self.kernel_size = c[2]

        self.sigma = c[0]
        self.transparency = c[1]

        frame = np.array(frame)

        # frame = (frame*255).astype(np.uint8)
        height, width, _ = frame.shape
        mask = self.generate_gauss_mask(mask_size=(width, height),
                                             position=self.gauss_position,
                                             max_brightness=max_brightness,
                                             min_brightness=min_brightness
                                             )

        mask = -(mask - np.max(mask))

        if self.mode == "light":
            hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HLS).astype(np.uint16)
            hsv[:, :, 1] = hsv[:, :, 1] + mask * (1 - self.transparency)  # -(mask-np.max(mask))
            hsv[hsv > 255] = 255

            hsv = hsv.astype(np.uint8)
            frame = cv2.cvtColor(hsv, cv2.COLOR_HLS2RGB)

        if self.mode == "blur":
            frame = self.alphaBlend(frame, mask)

        return self.normalize(frame)

    def __call__(self, img):
        img_c = self.add_corruption(img)
        return img_c