| 145 | return img1, img2, flow |
| 146 | |
| 147 | class SparseFlowAugmentor: |
| 148 | def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=False): |
| 149 | # spatial augmentation params |
| 150 | self.crop_size = crop_size |
| 151 | self.min_scale = min_scale |
| 152 | self.max_scale = max_scale |
| 153 | self.spatial_aug_prob = 0.8 |
| 154 | self.stretch_prob = 0.8 |
| 155 | self.max_stretch = 0.2 |
| 156 | |
| 157 | # flip augmentation params |
| 158 | self.do_flip = do_flip |
| 159 | self.h_flip_prob = 0.5 |
| 160 | self.v_flip_prob = 0.1 |
| 161 | |
| 162 | # photometric augmentation params |
| 163 | self.photo_aug = ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3/3.14) |
| 164 | self.asymmetric_color_aug_prob = 0.2 |
| 165 | self.eraser_aug_prob = 0.5 |
| 166 | |
| 167 | def color_transform(self, img1, img2): |
| 168 | image_stack = np.concatenate([img1, img2], axis=0) |
| 169 | image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8) |
| 170 | img1, img2 = np.split(image_stack, 2, axis=0) |
| 171 | return img1, img2 |
| 172 | |
| 173 | def eraser_transform(self, img1, img2): |
| 174 | ht, wd = img1.shape[:2] |
| 175 | if np.random.rand() < self.eraser_aug_prob: |
| 176 | mean_color = np.mean(img2.reshape(-1, 3), axis=0) |
| 177 | for _ in range(np.random.randint(1, 3)): |
| 178 | x0 = np.random.randint(0, wd) |
| 179 | y0 = np.random.randint(0, ht) |
| 180 | dx = np.random.randint(50, 100) |
| 181 | dy = np.random.randint(50, 100) |
| 182 | img2[y0:y0+dy, x0:x0+dx, :] = mean_color |
| 183 | |
| 184 | return img1, img2 |
| 185 | |
| 186 | def resize_sparse_flow_map(self, flow, valid, fx=1.0, fy=1.0): |
| 187 | ht, wd = flow.shape[:2] |
| 188 | coords = np.meshgrid(np.arange(wd), np.arange(ht)) |
| 189 | coords = np.stack(coords, axis=-1) |
| 190 | |
| 191 | coords = coords.reshape(-1, 2).astype(np.float32) |
| 192 | flow = flow.reshape(-1, 2).astype(np.float32) |
| 193 | valid = valid.reshape(-1).astype(np.float32) |
| 194 | |
| 195 | coords0 = coords[valid>=1] |
| 196 | flow0 = flow[valid>=1] |
| 197 | |
| 198 | ht1 = int(round(ht * fy)) |
| 199 | wd1 = int(round(wd * fx)) |
| 200 | |
| 201 | coords1 = coords0 * [fx, fy] |
| 202 | flow1 = flow0 * [fx, fy] |
| 203 | |
| 204 | xx = np.round(coords1[:,0]).astype(np.int32) |