def process_augmentations(aug_seq, image, bbs, threshold):
"""
Applying data augmentation. Removing the bbox when it is outside the
image or the fraction of bbox area outside of the image plane > threshold
Parameters
----------
aug_seq: imgaug.augmenters.Sequential
image: image numpy array (H,W,C)
bbs: imgaug.augmentables.bbs,BoundingBoxesOnImage
threshold: float
Deleting the bbox when the fraction of bbox area outside of the image
plane > threshold
Returns image, bbs
-------
"""
if isinstance(aug_seq, (iaa.Sequential, iaa.Augmenter)):
image_aug, bbs_aug = aug_seq(image=image, bounding_boxes=bbs)
bbs_fraction_outside = [
x.compute_out_of_image_fraction(image) for x in bbs_aug
]
if_reserved = np.less(bbs_fraction_outside, threshold)
bbs_aug = BoundingBoxesOnImage(
[bbs_aug[i] for i in range(len(if_reserved)) if if_reserved[i]],
shape=image_aug.shape)
bbs_aug = bbs_aug.remove_out_of_image().clip_out_of_image()
return image_aug, bbs_aug
else:
raise TypeError(
'Aug_seq must be an instance of iaa.Sequential or iaa.Augmenter')
def apply_augmentation(
self,
policy: List[POLICY_TUPLE],
image: np.array,
bounding_boxes: List[List[int]],
labels: Iterable[int],
) -> Tuple[np.array, np.array]:
"""
Applies the augmentations to the image.
:type policy: List[POLICY_TUPLE]
:param policy: Augmentation policy to apply to the image
:type image: np.array
:param image: Image to augment
:type bounding_boxes: List[List[int]]
:param bounding_boxes: Bounding boxes for the image in the format:
[x_min, y_min, x_max, y_max]
:type labels: Iterable[int]
:param labels: Iterable containing class labels as integers
:rtype: Tuple[np.array, np.array]
:return: Tuple containing the augmented image and bounding boxes
"""
bbs = BoundingBoxesOnImage(
[
BoundingBox(*bb, label=label)
for bb, label in zip(bounding_boxes, labels)
],
image.shape
)
for i in policy:
if i.name.endswith('BBox'):
new_bbs = []
for box in bbs:
if i.probability > np.random.random():
if i.name == 'Cutout_BBox':
kwargs = self._cutout_kwargs(image.shape)
aug = self[i.name](i.magnitude, **kwargs)
else:
aug = self[i.name](i.magnitude)
image, box_aug = aug(
image=image,
bounding_boxes=BoundingBoxesOnImage([box], image.shape) # noqa: E501
)
new_bbs.append(box_aug[0])
else:
new_bbs.append(box)
bbs = BoundingBoxesOnImage(new_bbs, image.shape)
elif i.probability > np.random.random():
if i.name == 'Cutout':
kwargs = self._cutout_kwargs(image.shape)
aug = self[i.name](i.magnitude, **kwargs)
elif i.name == 'Cutout_Fraction':
if len(bbs) == 0:
aug = self[i.name](i.magnitude)
else:
random_bb = np.random.choice(bbs)
kwargs = {
'height_bbox': random_bb.height,
'width_bbox': random_bb.width,
'height': image.shape[0],
'width': image.shape[1]
}
aug = self[i.name](i.magnitude, **kwargs)
else:
aug = self[i.name](i.magnitude)
image, bbs = aug(image=image, bounding_boxes=bbs)
bbs = bbs.remove_out_of_image().clip_out_of_image()
if self.return_yolo:
bbs = self._bbs_to_percent(bbs, image.shape[0], image.shape[1])
else:
bbs = self._bbs_to_pixel(bbs)
return image, bbs
