| 68 | return img1, img2 |
| 69 | |
| 70 | def spatial_transform(self, img1, img2, flow): |
| 71 | # randomly sample scale |
| 72 | ht, wd = img1.shape[:2] |
| 73 | min_scale = np.maximum( |
| 74 | (self.crop_size[0] + 8) / float(ht), |
| 75 | (self.crop_size[1] + 8) / float(wd)) |
| 76 | |
| 77 | scale = 2 ** np.random.uniform(self.min_scale, self.max_scale) |
| 78 | scale_x = scale |
| 79 | scale_y = scale |
| 80 | if np.random.rand() < self.stretch_prob: |
| 81 | scale_x *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch) |
| 82 | scale_y *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch) |
| 83 | scale_x = np.clip(scale_x, min_scale, None) |
| 84 | scale_y = np.clip(scale_y, min_scale, None) |
| 85 | |
| 86 | if np.random.rand() < self.spatial_aug_prob: |
| 87 | # rescale the images |
| 88 | img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR) |
| 89 | img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR) |
| 90 | flow = cv2.resize(flow, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR) |
| 91 | flow = flow * [scale_x, scale_y] |
| 92 | |
| 93 | if self.do_flip: |
| 94 | if np.random.rand() < self.h_flip_prob: # h-flip |
| 95 | img1 = img1[:, ::-1] |
| 96 | img2 = img2[:, ::-1] |
| 97 | flow = flow[:, ::-1] * [-1.0, 1.0] |
| 98 | |
| 99 | if np.random.rand() < self.v_flip_prob: # v-flip |
| 100 | img1 = img1[::-1, :] |
| 101 | img2 = img2[::-1, :] |
| 102 | flow = flow[::-1, :] * [1.0, -1.0] |
| 103 | |
| 104 | if img1.shape[0] == self.crop_size[0]: |
| 105 | y0 = 0 |
| 106 | else: |
| 107 | y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0]) |
| 108 | if img1.shape[1] == self.crop_size[1]: |
| 109 | x0 = 0 |
| 110 | else: |
| 111 | x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1]) |
| 112 | |
| 113 | img1 = img1[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]] |
| 114 | img2 = img2[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]] |
| 115 | flow = flow[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]] |
| 116 | |
| 117 | return img1, img2, flow |
| 118 | |
| 119 | def __call__(self, img1, img2, flow): |
| 120 | img1, img2 = self.color_transform(img1, img2) |