| 92 | |
| 93 | |
| 94 | class ReidAugmentationCV2(object): |
| 95 | def __init__(self, height, width, re,bri,contrast,brightness_delta=16,contrast_range=(0.8, 1.2), vit=False): |
| 96 | self.normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], |
| 97 | std=[0.229, 0.224, 0.225]) |
| 98 | self.height = height |
| 99 | self.width = width |
| 100 | self.padding_size = 10 |
| 101 | self.brightness_delta = brightness_delta |
| 102 | self.contrast_lower, self.contrast_upper = contrast_range |
| 103 | self.re = re |
| 104 | self.bri = bri |
| 105 | self.contrast = contrast |
| 106 | self.vit = vit |
| 107 | |
| 108 | def pad_images(self, img): |
| 109 | h, w = img.shape[:2] |
| 110 | width_max = w + 2 * self.padding_size |
| 111 | height_max = h + 2 * self.padding_size |
| 112 | |
| 113 | diff_vert = height_max - h |
| 114 | pad_top = diff_vert//2 |
| 115 | pad_bottom = diff_vert - pad_top |
| 116 | diff_hori = width_max - w |
| 117 | pad_left = diff_hori//2 |
| 118 | pad_right = diff_hori - pad_left |
| 119 | img_padded = cv2.copyMakeBorder(img, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=0) |
| 120 | assert img_padded.shape[:2] == (height_max, width_max) |
| 121 | return img_padded |
| 122 | |
| 123 | def random_crop(self, img): |
| 124 | h, w = img.shape[:2] |
| 125 | max_x = w - self.width |
| 126 | max_y = h - self.height |
| 127 | |
| 128 | x = np.random.randint(0, max_x) |
| 129 | y = np.random.randint(0, max_y) |
| 130 | |
| 131 | crop = img[y:y+self.height,x:x+self.width] |
| 132 | return crop |
| 133 | |
| 134 | def random_brightness(self, img, binary_mask=None): |
| 135 | if np.random.rand() <= 0.5: |
| 136 | img=img.astype(np.float32) |
| 137 | delta = random.uniform(-self.brightness_delta, |
| 138 | self.brightness_delta) |
| 139 | if binary_mask is not None: |
| 140 | delta = delta * binary_mask |
| 141 | img += delta |
| 142 | return img |
| 143 | |
| 144 | def random_contrast(self, img, binary_mask=None): |
| 145 | if np.random.rand() <= 0.5: |
| 146 | img=img.astype(np.float32) |
| 147 | alpha = random.uniform(self.contrast_lower, |
| 148 | self.contrast_upper) |
| 149 | if binary_mask is not None: |
| 150 | alpha = alpha * binary_mask |
| 151 | img *= alpha |