Augment: horizontal flips OR rotate (0, 90, 180, 270 degrees). We use vertical flip and transpose for rotation implementation. All the images in the list use the same augmentation. Args: imgs (list[ndarray] | ndarray): Images to be augmented. If the input is an ndar
(imgs, hflip=True, rotation=True, flows=None, return_status=False)
| 127 | |
| 128 | |
| 129 | def augment(imgs, hflip=True, rotation=True, flows=None, return_status=False): |
| 130 | """Augment: horizontal flips OR rotate (0, 90, 180, 270 degrees). |
| 131 | |
| 132 | We use vertical flip and transpose for rotation implementation. |
| 133 | All the images in the list use the same augmentation. |
| 134 | |
| 135 | Args: |
| 136 | imgs (list[ndarray] | ndarray): Images to be augmented. If the input |
| 137 | is an ndarray, it will be transformed to a list. |
| 138 | hflip (bool): Horizontal flip. Default: True. |
| 139 | rotation (bool): Ratotation. Default: True. |
| 140 | flows (list[ndarray]: Flows to be augmented. If the input is an |
| 141 | ndarray, it will be transformed to a list. |
| 142 | Dimension is (h, w, 2). Default: None. |
| 143 | return_status (bool): Return the status of flip and rotation. |
| 144 | Default: False. |
| 145 | |
| 146 | Returns: |
| 147 | list[ndarray] | ndarray: Augmented images and flows. If returned |
| 148 | results only have one element, just return ndarray. |
| 149 | |
| 150 | """ |
| 151 | hflip = hflip and random.random() < 0.5 |
| 152 | vflip = rotation and random.random() < 0.5 |
| 153 | rot90 = rotation and random.random() < 0.5 |
| 154 | |
| 155 | def _augment(img): |
| 156 | if hflip: # horizontal |
| 157 | cv2.flip(img, 1, img) |
| 158 | if vflip: # vertical |
| 159 | cv2.flip(img, 0, img) |
| 160 | if rot90: |
| 161 | img = img.transpose(1, 0, 2) |
| 162 | return img |
| 163 | |
| 164 | def _augment_flow(flow): |
| 165 | if hflip: # horizontal |
| 166 | cv2.flip(flow, 1, flow) |
| 167 | flow[:, :, 0] *= -1 |
| 168 | if vflip: # vertical |
| 169 | cv2.flip(flow, 0, flow) |
| 170 | flow[:, :, 1] *= -1 |
| 171 | if rot90: |
| 172 | flow = flow.transpose(1, 0, 2) |
| 173 | flow = flow[:, :, [1, 0]] |
| 174 | return flow |
| 175 | |
| 176 | if not isinstance(imgs, list): |
| 177 | imgs = [imgs] |
| 178 | imgs = [_augment(img) for img in imgs] |
| 179 | if len(imgs) == 1: |
| 180 | imgs = imgs[0] |
| 181 | |
| 182 | if flows is not None: |
| 183 | if not isinstance(flows, list): |
| 184 | flows = [flows] |
| 185 | flows = [_augment_flow(flow) for flow in flows] |
| 186 | if len(flows) == 1: |
no test coverage detected