| 13 | from . import flow_transforms |
| 14 | |
| 15 | class FlowAugmentor: |
| 16 | def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=True, pwc_aug=False): |
| 17 | |
| 18 | # spatial augmentation params |
| 19 | self.crop_size = crop_size |
| 20 | self.min_scale = min_scale |
| 21 | self.max_scale = max_scale |
| 22 | self.spatial_aug_prob = 0.8 |
| 23 | self.stretch_prob = 0.8 |
| 24 | self.max_stretch = 0.2 |
| 25 | |
| 26 | # flip augmentation params |
| 27 | self.do_flip = do_flip |
| 28 | self.h_flip_prob = 0.5 |
| 29 | self.v_flip_prob = 0.1 |
| 30 | |
| 31 | # photometric augmentation params |
| 32 | self.photo_aug = ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.5/3.14) |
| 33 | self.asymmetric_color_aug_prob = 0.2 |
| 34 | self.eraser_aug_prob = 0.5 |
| 35 | self.pwc_aug = pwc_aug |
| 36 | if self.pwc_aug: |
| 37 | print("[Using pwc-style spatial augmentation]") |
| 38 | |
| 39 | def color_transform(self, img1, img2): |
| 40 | """ Photometric augmentation """ |
| 41 | |
| 42 | # asymmetric |
| 43 | if np.random.rand() < self.asymmetric_color_aug_prob: |
| 44 | img1 = np.array(self.photo_aug(Image.fromarray(img1)), dtype=np.uint8) |
| 45 | img2 = np.array(self.photo_aug(Image.fromarray(img2)), dtype=np.uint8) |
| 46 | |
| 47 | # symmetric |
| 48 | else: |
| 49 | image_stack = np.concatenate([img1, img2], axis=0) |
| 50 | image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8) |
| 51 | img1, img2 = np.split(image_stack, 2, axis=0) |
| 52 | |
| 53 | return img1, img2 |
| 54 | |
| 55 | def eraser_transform(self, img1, img2, bounds=[50, 100]): |
| 56 | """ Occlusion augmentation """ |
| 57 | |
| 58 | ht, wd = img1.shape[:2] |
| 59 | if np.random.rand() < self.eraser_aug_prob: |
| 60 | mean_color = np.mean(img2.reshape(-1, 3), axis=0) |
| 61 | for _ in range(np.random.randint(1, 3)): |
| 62 | x0 = np.random.randint(0, wd) |
| 63 | y0 = np.random.randint(0, ht) |
| 64 | dx = np.random.randint(bounds[0], bounds[1]) |
| 65 | dy = np.random.randint(bounds[0], bounds[1]) |
| 66 | img2[y0:y0+dy, x0:x0+dx, :] = mean_color |
| 67 | |
| 68 | return img1, img2 |
| 69 | |
| 70 | def spatial_transform(self, img1, img2, flow): |
| 71 | # randomly sample scale |
| 72 | ht, wd = img1.shape[:2] |