trimming_synthentic.py

import glob
import random
import numpy as np
import cv2
import albumentations as A
import random
import cv2


class DrawHorizontal:
    def __call__(self, image):
        W, H, C = image.shape
        if random.random() > 0.8:
            color = np.zeros((3,)).astype(np.uint8)
        elif random.random() > 0.5:
            color = np.ones((3,)).astype(np.uint8)
        else:
            color = np.random.randint(0, 256, size=(3, )).astype(np.uint8)
        if random.random() > 0.9:
            D = int(H/2)
        elif random.random() > 0.5:
            D = H/24
        else:
            D = 10
        thickness = np.random.randint(1, D)
        r = random.random()
        r = max(r, 0.8)
        if random.random() > 0.5:
            color_bar = np.zeros_like(image[:, :thickness, :])
            color_bar += color
            image[:, :thickness, :] = image[:,
                                            :thickness, :]*(1-r) + r*color_bar
        else:
            color_bar = np.zeros_like(image[:, -thickness:, :])
            color_bar += color
            image[:, -thickness:, :] = image[:, -
                                             thickness:, :]*(1-r) + r*color_bar
        return image


class DrawVertical:
    def __call__(self, image):
        W, H, C = image.shape
        if random.random() > 0.8:
            color = np.zeros((3,)).astype(np.uint8)
        elif random.random() > 0.5:
            color = np.ones((3,)).astype(np.uint8)
        else:
            color = np.random.randint(0, 256, size=(3, )).astype(np.uint8)
        if random.random() > 0.9:
            D = int(H/2)
        elif random.random() > 0.5:
            D = H/24
        else:
            D = 10
        thickness = np.random.randint(1, D)
        r = random.random()
        r = max(r, 0.8)
        if random.random() > 0.5:
            color_bar = np.zeros_like(image[:thickness, :, :])
            color_bar += color
            image[:thickness, :, :] = image[:thickness, :, :] * \
                (1-r) + r*color_bar
        else:
            color_bar = np.zeros_like(image[-thickness:, :, :])
            color_bar += color
            image[-thickness:, :, :] = image[-thickness:, :, :] * \
                (1-r) + r*color_bar
        return image


class RandomBorderJitter(A.ImageOnlyTransform):
    def __init__(self, brightness=1, constrast=1, saturation=1, hue=1, always_apply=True, p=0.5):
        super(RandomBorderJitter, self).__init__(always_apply, p)
        self.random_jitter = A.ColorJitter(
            brightness, constrast, saturation, always_apply, p)
        self.p = p

    def __call__(self, image):
        image = image
        thickness = np.random.randint(1, 10)
        positions = ['top', 'left', 'bottom', 'right']
        pos = random.choice(positions)
        if pos == 'top':
            image[:, :thickness, :] = self.random_jitter(
                image=image[:, :thickness, :])['image']
        if pos == 'right':
            image[-thickness:, :,
                  :] = self.random_jitter(image=image[-thickness:, :, :])['image']
        if pos == 'bottom':
            image[:, -thickness:,
                  :] = self.random_jitter(image=image[:, -thickness:, :])['image']
        if pos == 'left':
            image[:thickness, :, :] = self.random_jitter(
                image=image[:thickness, :, :])['image']
        return image


class PolyTrimmed:
    def __call__(self, image):
        W, H, C = image.shape
        thickness = np.random.randint(1, 20)
        length = np.random.randint(W/16, W)
        if random.random() > 0.8:
            color = np.zeros((3,)).astype(np.uint8)
        elif random.random() > 0.5:
            color = np.ones((3,)).astype(np.uint8)
        else:
            color = np.random.randint(0, 256, size=(3, )).astype(np.uint8)
        color = (int(color[0]), int(color[1]), int(color[2]))
        p1 = [0, 0]
        p2 = [0, thickness]
        p3 = [length, 0]
        cv2.fillPoly(image, [np.array([p1, p2, p3])],
                     color=color, lineType=cv2.LINE_AA)
        return image


class AugmentImage():
    def __init__(self, draw_line=True, img_size=(384, 384), always_apply=True, p=1.):
        self.draw_line = draw_line
        self.img_size = img_size
        self.c_augment = A.Compose([
            A.RandomResizedCrop(height=img_size[0], width=img_size[1], scale=(
                0.75, 1), always_apply=True, interpolation=cv2.INTER_CUBIC)
        ])
        self.augment = [
            DrawHorizontal(),
            DrawVertical(),
            RandomBorderJitter(always_apply=always_apply),
            PolyTrimmed(),
        ]

    def __call__(self, image):
        '''
        image - a numpy array (H, W, C)
        '''
        image = self.c_augment(image=image)['image']
        image = random.choice(self.augment)(image=image)
        return image