Apply image augmentors on several components, with shared augmentation parameters. Example: .. code-block:: python ds = MyDataFlow() # produce [image(HWC), segmask(HW), keypoint(Nx2)] ds = AugmentImageComponents( ds, augs,
| 168 | |
| 169 | |
| 170 | class AugmentImageComponents(MapData): |
| 171 | """ |
| 172 | Apply image augmentors on several components, with shared augmentation parameters. |
| 173 | |
| 174 | Example: |
| 175 | |
| 176 | .. code-block:: python |
| 177 | |
| 178 | ds = MyDataFlow() # produce [image(HWC), segmask(HW), keypoint(Nx2)] |
| 179 | ds = AugmentImageComponents( |
| 180 | ds, augs, |
| 181 | index=(0,1), coords_index=(2,)) |
| 182 | |
| 183 | """ |
| 184 | |
| 185 | def __init__(self, ds, augmentors, index=(0, 1), coords_index=(), copy=True, catch_exceptions=False): |
| 186 | """ |
| 187 | Args: |
| 188 | ds (DataFlow): input DataFlow. |
| 189 | augmentors (AugmentorList): a list of :class:`imgaug.ImageAugmentor` instance to be applied in order. |
| 190 | index: tuple of indices of the image components. |
| 191 | coords_index: tuple of indices of the coordinates components. |
| 192 | copy, catch_exceptions: same as in :class:`AugmentImageComponent` |
| 193 | """ |
| 194 | if isinstance(augmentors, AugmentorList): |
| 195 | self.augs = augmentors |
| 196 | else: |
| 197 | self.augs = AugmentorList(augmentors) |
| 198 | self.ds = ds |
| 199 | self._exception_handler = ExceptionHandler(catch_exceptions) |
| 200 | self._copy = copy |
| 201 | self._index = index |
| 202 | self._coords_index = coords_index |
| 203 | |
| 204 | super(AugmentImageComponents, self).__init__(ds, self._aug_mapper) |
| 205 | |
| 206 | def reset_state(self): |
| 207 | self.ds.reset_state() |
| 208 | self.augs.reset_state() |
| 209 | |
| 210 | def _aug_mapper(self, dp): |
| 211 | dp = copy_mod.copy(dp) # always do a shallow copy, make sure the list is intact |
| 212 | copy_func = copy_mod.deepcopy if self._copy else lambda x: x # noqa |
| 213 | with self._exception_handler.catch(): |
| 214 | major_image = self._index[0] # image to be used to get params. TODO better design? |
| 215 | im = copy_func(dp[major_image]) |
| 216 | check_dtype(im) |
| 217 | tfms = self.augs.get_transform(im) |
| 218 | dp[major_image] = tfms.apply_image(im) |
| 219 | for idx in self._index[1:]: |
| 220 | check_dtype(dp[idx]) |
| 221 | dp[idx] = tfms.apply_image(copy_func(dp[idx])) |
| 222 | for idx in self._coords_index: |
| 223 | coords = copy_func(dp[idx]) |
| 224 | validate_coords(coords) |
| 225 | dp[idx] = tfms.apply_coords(coords) |
| 226 | return dp |
| 227 |
no outgoing calls
no test coverage detected