Tiler augmenter proposal

[chamecall]

I needed to create tile image from different objects of an image for object detection task. I couldn't find anything similar so I implemented it with your iaa.meta.Augmenter class.

For those who are looking for such functionality, I attach the Tiler class, client code with sequence of augmenters including the Tiler augmenter and an example of launch.

Required imports:

from imgaug import augmenters as iaa 
import random
from random import randint
from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage
from imgaug.augmentables.batches import _BatchInAugmentation

Tiler class:

class Tiler(iaa.meta.Augmenter):

    seq = iaa.Sequential([
        iaa.GammaContrast(),
        iaa.Multiply(mul=(0.6, 1.2)),
        iaa.Affine(translate_percent={"x": (-0.1, 0.1), 'y': (-0.1, 0.1)}, scale=(0.7, 1.5), rotate=(-180, 180), fit_output=True),
        iaa.PerspectiveTransform(scale=(0, 0.1), keep_size=False),
        iaa.AdditiveGaussianNoise()
    ], random_order=True)

    min_patch_size_frac = 0.005

    def generate_patches(self, img, bbs):
        image_aug, bbs_aug = Tiler.seq(image=img, bounding_boxes=bbs)
        bbs_aug = bbs_aug.remove_out_of_image_fraction(0.3)

        patches = [(bb.label, bb.extract_from_image(image_aug)) for bb in bbs_aug]
        return random.sample(patches, len(patches))

    def get_parameters(self):
        pass

    def _augment_batch_(self, batch, random_state, parents, hooks):
        images = batch.images
        bbs = batch.bounding_boxes

        assert len(images) == 1
        assert len(bbs) == 1

        image = images[0]
        bounding_boxes = bbs[0]

        src_img_h, src_img_w = image.shape[:2]
        src_img_area = src_img_h * src_img_w

        patch_min_size = src_img_area * Tiler.min_patch_size_frac

        bounding_boxes = BoundingBoxesOnImage([bb for bb in bounding_boxes if patch_min_size <= bb.area], bounding_boxes.shape)
        if len(bounding_boxes) == 0:
            return batch

        blur_img = iaa.GaussianBlur(sigma=(21, 21))(image=image)
        blur_img = iaa.Affine(translate_percent=0, scale=(0.8, 1.2), fit_output=True)(image=blur_img)

        img_h, img_w = blur_img.shape[:2]
        img_area = img_h * img_w

        patches_area = 0
        patches = []

        while True:
            new_patches = self.generate_patches(image, bounding_boxes)
            for label, patch in new_patches:
                h, w = patch.shape[:2]

                h_ratio = h / img_h
                w_ratio = w / img_w

                if h_ratio > 0.3 or w_ratio > 0.3:
                    new_ratio = random.uniform(0.1, 0.3)
                    if h_ratio > w_ratio:
                        new_h = int(img_h * new_ratio)
                        patch = resize(patch, height=new_h)
                    else:
                        new_w = int(img_w * new_ratio)
                        patch = resize(patch, width=new_w)

                h, w = patch.shape[:2]
                patch_area = h * w

                if patches_area + patch_area > img_area:
                    break
                else:
                    patches.append((patch_area, (label, patch)))
                    patches_area += patch_area
            else:
                continue
            break

        patches = sorted(patches, reverse=False, key=lambda patch: patch[0])

        if randint(0, 1):
            random.shuffle(patches)

        margin_range = (5, 30)

        top_y = bottom_y = 0
        cur_x = randint(*margin_range)

        boxes = []

        for _, (label, patch) in patches:
            h_p, w_p = patch.shape[:2]

            if cur_x + w_p > img_w:
                top_y = bottom_y
                cur_x = randint(*margin_range)
                if cur_x + w_p > img_w:
                    cur_x = 0

            y_margin = randint(*margin_range)

            if y_margin + top_y + h_p > img_h:
                y_margin = 0

                if top_y + h_p > img_h:
                    break

            x1, y1 = cur_x, y_margin+top_y
            x2, y2 = cur_x+w_p, y_margin+top_y+h_p

            try:
                blur_img[y1:y2, x1:x2] = patch
            except Exception:
                print(blur_img.shape, img_h, img_w)
                print(patch.shape)
                print((x1, y1), (x2, y2))

            boxes.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, label=label))

            cur_x += w_p + randint(*margin_range)
            if y_margin + h_p > bottom_y - top_y:
                bottom_y = y_margin + top_y + h_p

        bbs = BoundingBoxesOnImage(boxes, shape=blur_img.shape)
        return _BatchInAugmentation(images=[blur_img], bounding_boxes=[bbs])

Client code:

img = cv2.imread('your_image.jpg')

bbs =' yourBoundingBoxesOnImage object'

transformer = \
iaa.Sequential([
    iaa.Grayscale(),
    iaa.Sometimes(
        p=0.1, 
        then_list=Tiler(),
        else_list=iaa.Sequential([
            iaa.GammaContrast(),
            iaa.Multiply(mul=(0.7, 1.3)),
            iaa.Affine(translate_percent={"x": (-0.1, 0.1), 'y': (-0.1, 0.1)}, scale=(0.5, 1.5), rotate=(-40, 40), fit_output=True),
            iaa.PerspectiveTransform(scale=(0, 0.1), keep_size=False),
            iaa.Sometimes(0.3, iaa.AdditiveGaussianNoise(3))
        ], random_order=True)
    )]
)

img, bbs = transformer(image=img, bounding_boxes=bbs)
ia.imshow(bbs.draw_on_image(img))

Source img:

Example of execution ordinary sequence of augmentations:

Examples of execution Tiler augmentator: