custom_augmentation.py

from transformations import color, shadow
import _pickle as cp
import glob
import os
import random


import cv2
import h5py
import imgaug
import numpy as np
import pylab
import scipy.ndimage as ndi
import imgaug as ia
from imgaug import augmenters as iaa
from imgaug.augmentables.polys import Polygon, PolygonsOnImage
from scipy.stats import norm
import copy
import imutils
from PIL import Image


class ParallelLight(iaa.meta.Augmenter):
    def __init__(self, name=None, deterministic=False, random_state=None):
        super().__init__(name=name, deterministic=deterministic, random_state=random_state)

    def _decayed_value_in_norm(self, x, max_value, min_value, center, range):
        """
        decay from max value to min value following Gaussian/Normal distribution
        """
        radius = range / 3
        center_prob = norm.pdf(center, center, radius)
        x_prob = norm.pdf(x, center, radius)
        x_value = (x_prob / center_prob) * (max_value - min_value) + min_value
        return x_value

    def _decayed_value_in_linear(self, x, max_value, padding_center, decay_rate):
        """
        decay from max value to min value with static linear decay rate.
        """
        x_value = max_value - abs(padding_center - x) * decay_rate
        if x_value < 0:
            x_value = 1
        return x_value

    def generate_parallel_light_mask(self, mask_size,
                                     position=None,
                                     direction=None,
                                     max_brightness=255,
                                     min_brightness=0,
                                     mode="gaussian",
                                     linear_decay_rate=None):
        """
        Generate decayed light mask generated by light strip given its position, direction
        Args:
            mask_size: tuple of integers (w, h) defining generated mask size
            position: tuple of integers (x, y) defining the center of light strip position,
                    which is the reference point during rotating
            direction: integer from 0 to 360 to indicate the rotation degree of light strip
            max_brightness: integer that max brightness in the mask
            min_brightness: integer that min brightness in the mask
            mode: the way that brightness decay from max to min: linear or gaussian
            linear_decay_rate: only valid in linear_static mode. Suggested value is within [0.2, 2]
        Return:
            light_mask: ndarray in float type consisting value from 0 to strength
        """
        if position is None:
            pos_x = random.randint(0, mask_size[0])
            pos_y = random.randint(0, mask_size[1])
        else:
            pos_x = position[0]
            pos_y = position[1]
        if direction is None:
            direction = random.randint(0, 360)
            # print("Rotate degree: ", direction)
        if linear_decay_rate is None:
            if mode == "linear_static":
                linear_decay_rate = random.uniform(0.2, 2)
            if mode == "linear_dynamic":
                linear_decay_rate = (
                    max_brightness - min_brightness) / max(mask_size)
        assert mode in ["linear_dynamic", "linear_static", "gaussian"], \
            "mode must be linear_dynamic, linear_static or gaussian"
        padding = int(max(mask_size) * np.sqrt(2))
        # add padding to satisfy cropping after rotating
        canvas_x = padding * 2 + mask_size[0]
        canvas_y = padding * 2 + mask_size[1]
        mask = np.zeros(shape=(canvas_y, canvas_x), dtype=np.float32)
        # initial mask's up left corner and bottom right corner coordinate
        init_mask_ul = (int(padding), int(padding))
        init_mask_br = (int(padding+mask_size[0]), int(padding+mask_size[1]))
        init_light_pos = (padding + pos_x, padding + pos_y)
        # fill in mask row by row with value decayed from center
        for i in range(canvas_y):
            if mode == "linear":
                i_value = self._decayed_value_in_linear(
                    i, max_brightness, init_light_pos[1], linear_decay_rate)
            elif mode == "gaussian":
                i_value = self._decayed_value_in_norm(
                    i, max_brightness, min_brightness, init_light_pos[1], mask_size[1])
            else:
                i_value = 0
            mask[i] = i_value
        # rotate mask
        rotate_M = cv2.getRotationMatrix2D(init_light_pos, direction, 1)
        mask = cv2.warpAffine(mask, rotate_M, (canvas_x,  canvas_y))
        # crop
        mask = mask[init_mask_ul[1]:init_mask_br[1],
                    init_mask_ul[0]:init_mask_br[0]]
        mask = np.asarray(mask, dtype=np.uint8)
        # add median blur
        mask = cv2.medianBlur(mask, 9)
        mask = 255 - mask
        # cv2.circle(mask, init_light_pos, 1, (0, 0, 255))
        # cv2.imshow("crop", mask[init_mask_ul[1]:init_mask_br[1], init_mask_ul[0]:init_mask_br[0]])
        # cv2.imshow("all", mask)
        # cv2.waitKey(0)
        return mask

    def add_parallel_light(self, image, light_position=None, direction=None, max_brightness=255, min_brightness=0,
                           mode="gaussian", linear_decay_rate=None, transparency=None):
        """
        Add mask generated from parallel light to given image
        """
        if transparency is None:
            transparency = random.uniform(0.5, 0.85)
        height, width, _ = image.shape
        hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        mask = self.generate_parallel_light_mask(mask_size=(width, height),
                                                 position=light_position,
                                                 direction=direction,
                                                 max_brightness=max_brightness,
                                                 min_brightness=min_brightness,
                                                 mode=mode,
                                                 linear_decay_rate=linear_decay_rate)
        hsv[:, :, 2] = hsv[:, :, 2] * transparency + mask * (1 - transparency)
        image = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
        image[image > 255] = 255
        image = np.asarray(image, dtype=np.uint8)
        return image

    def augment_image(self, image, **kwargs):
        return self.add_parallel_light(image)

    def _augment_images(self, images, **kwargs):
        return [self.add_parallel_light(each_img) for each_img in images]

    def _augment_keypoints(self, **kwargs):
        return

    def _augment_heatmaps(self, **kwargs):
        return

    def get_parameters(self):
        return


class SpotLight(iaa.meta.Augmenter):
    def __init__(self, name=None, deterministic=False, random_state=None):
        super().__init__(name=name, deterministic=deterministic, random_state=random_state)

    def generate_spot_light_mask(self, mask_size,
                                 position=None,
                                 max_brightness=255,
                                 min_brightness=0,
                                 mode="gaussian",
                                 linear_decay_rate=None,
                                 speedup=False):
        """
        Generate decayed light mask generated by spot light given position, direction. Multiple spotlights are accepted.
        Args:
            mask_size: tuple of integers (w, h) defining generated mask size
            position: list of tuple of integers (x, y) defining the center of spotlight light position,
                    which is the reference point during rotating
            max_brightness: integer that max brightness in the mask
            min_brightness: integer that min brightness in the mask
            mode: the way that brightness decay from max to min: linear or gaussian
            linear_decay_rate: only valid in linear_static mode. Suggested value is within [0.2, 2]
            speedup: use `shrinkage then expansion` strategy to speed up vale calculation
        Return:
            light_mask: ndarray in float type consisting value from max_brightness to min_brightness. If in 'linear' mode
                        minimum value could be smaller than given min_brightness.
        """
        if position is None:
            position = [(random.randint(0, mask_size[0]),
                         random.randint(0, mask_size[1]))]
        if linear_decay_rate is None:
            if mode == "linear_static":
                linear_decay_rate = random.uniform(0.25, 1)
        assert mode in ["linear", "gaussian"], \
            "mode must be linear_dynamic, linear_static or gaussian"
        mask = np.zeros(shape=(mask_size[1], mask_size[0]), dtype=np.float32)
        if mode == "gaussian":
            mu = np.sqrt(mask.shape[0]**2+mask.shape[1]**2)
            dev = mu / 3.5
            mask = self._decay_value_radically_norm_in_matrix(
                mask_size, position, max_brightness, min_brightness, dev)
        mask = np.asarray(mask, dtype=np.uint8)
        # add median blur
        mask = cv2.medianBlur(mask, 5)
        mask = 255 - mask
        # cv2.imshow("mask", mask)
        # cv2.waitKey(0)
        return mask

    def _decay_value_radically_norm_in_matrix(self, mask_size, centers, max_value, min_value, dev):
        """
        _decay_value_radically_norm function in matrix format
        """
        center_prob = norm.pdf(0, 0, dev)
        x_value_rate = np.zeros((mask_size[1], mask_size[0]))
        for center in centers:
            coord_x = np.arange(mask_size[0])
            coord_y = np.arange(mask_size[1])
            xv, yv = np.meshgrid(coord_x, coord_y)
            dist_x = xv - center[0]
            dist_y = yv - center[1]
            dist = np.sqrt(np.power(dist_x, 2) + np.power(dist_y, 2))
            x_value_rate += norm.pdf(dist, 0, dev) / center_prob
        mask = x_value_rate * (max_value - min_value) + min_value
        mask[mask > 255] = 255
        return mask

    def _decay_value_radically_norm(self, x, centers, max_value, min_value, dev):
        """
        Calculate point value decayed from center following Gaussian decay. If multiple centers are given, value
        from each center sums up while limiting the accumulated value into [0, 255]
        NOTE: assuming light at each center is identical: same brightness and same decay rate
        """
        center_prob = norm.pdf(0, 0, dev)
        x_value_rate = 0
        for center in centers:
            distance = np.sqrt((center[0]-x[0])**2 + (center[1]-x[1])**2)
            x_value_rate += norm.pdf(distance, 0, dev) / center_prob
        x_value = x_value_rate * (max_value - min_value) + min_value
        x_value = 255 if x_value > 255 else x_value
        return x_value

    def add_spot_light(self, image, light_position=None, max_brightness=255, min_brightness=0,
                       mode='gaussian', linear_decay_rate=None, transparency=None):
        """
        Add mask generated from spot light to given image
        """
        if transparency is None:
            transparency = random.uniform(0.5, 0.85)
        height, width, _ = image.shape
        hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        mask = self.generate_spot_light_mask(mask_size=(width, height),
                                             position=light_position,
                                             max_brightness=max_brightness,
                                             min_brightness=min_brightness,
                                             mode=mode,
                                             linear_decay_rate=linear_decay_rate)
        hsv[:, :, 2] = hsv[:, :, 2] * transparency + mask * (1 - transparency)
        image = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
        image[image > 255] = 255
        image = np.asarray(image, dtype=np.uint8)
        return image

    def augment_image(self, image, **kwargs):
        return self.add_spot_light(image)

    def _augment_images(self, images, **kwargs):
        return [self.add_spot_light(each_img) for each_img in images]

    def _augment_keypoints(self, **kwargs):
        return

    def _augment_heatmaps(self, **kwargs):
        return

    def get_parameters(self):
        return


def imgaug_to_normal(augmented_anno, raw_anno):
    all_anno = []
    for index, polygon in enumerate(augmented_anno.polygons):
        polygon = polygon.exterior
        all_anno.append(
            {
                "polygon": [(int(point[0]), int(point[1])) for point in polygon],#+[((int(polygon[0][0]), int(polygon[0][1])))],
                "text": raw_anno[index]['text'],
                "key": raw_anno[index]['key']
            }
        )
    return all_anno


def imgaug_to_normal_polygon(augmented_anno):
    all_anno = []
    for index, polygon in enumerate(augmented_anno.polygons):
        polygon = polygon.exterior
        all_anno.append([(int(point[0]), int(point[1])) for point in polygon]+[((int(polygon[0][0]), int(polygon[0][1])))])
    return all_anno


class OverLayImage():

    def __init__(self, background_path='./data/SynthText/'):
        self.background_path = background_path
        self.backgrounds = self._get_usable_background()
        if len(self.backgrounds) == 0:
            raise FileNotFoundError(
                "Can't file any background image in", os.path.join(self.background_path, 'bg_img', ", please refer to the README to download some"))

    def _get_usable_background(self):
        with open(os.path.join(self.background_path, 'imnames.cp'), 'rb') as f:
            filtered_imnames = set(cp.load(f))
        usable = []
        for imname in os.listdir(os.path.join(self.background_path, 'bg_img')):
            # ignore if not in filetered list:
            if imname not in filtered_imnames:
                continue
            else:
                usable.append(os.path.join(
                    self.background_path, 'bg_img', imname))
        return usable

    def _resize_background(self, image, id_size):
        p = random.uniform(1.35, 1.65)
        bg_size = image.shape
        width = int(id_size[1]*p)
        height = int(id_size[0]*p)
        if bg_size[0] >= id_size[0] and bg_size[1] >= id_size[1]:
            resizer = iaa.OneOf([
                iaa.CropToFixedSize(width, height),
                iaa.Resize({"height": height, "width": width})
            ])
        else:
            resizer = iaa.Resize({"height": height, "width": width})
        return resizer.augment_image(image)

    def _resize_id(self, image, anno, raw_anno):
        p = random.uniform(1, 2.08)
        width = int(image.shape[1]/p)
        height = int(image.shape[0]/p)
        resizer = iaa.Resize({"height": height, "width": width})
        psoi = imgaug.PolygonsOnImage(
            [Polygon(each['polygon']) for each in anno], shape=image.shape)
        image, anno = resizer(image=image, polygons=psoi)
        anno = imgaug_to_normal(anno, raw_anno)
        return image, anno

    def put4ChannelImageOn4ChannelImage(self, back, fore, x, y):
        back = cv2.cvtColor(back, cv2.COLOR_BGR2RGBA)
        rows, cols, channels = fore.shape
        trans_indices = fore[..., 3] != 0  # Where not transparent
        overlay_copy = back[y:y+rows, x:x+cols]
        overlay_copy[trans_indices] = fore[trans_indices]
        back[y:y+rows, x:x+cols] = overlay_copy
        return back

    def _overlay_id_on_background(self, id, anno, background):
        bg_h, bg_w, _ = background.shape
        id_h, id_w, _ = id.shape

        # print('BG shape:', background.shape)
        # print('ID shape: ', id.shape)

        x = random.randint(0, (bg_w - id_w))
        y = random.randint(0, (bg_h - id_h))

        # Calculate new coords
        for index, line in enumerate(anno):
            # import ipdb; ipdb.set_trace()
            new_coords = [[point[0]+x, point[1]+y]
                          for point in line['polygon']]
            anno[index]['polygon'] = new_coords

        augmented_img = self.put4ChannelImageOn4ChannelImage(
            background, id, x, y)
        return augmented_img, anno

    def augment_image(self, images, annos, all_raw_anno, is_scan=False):
        # Random resize the images
        resized = [self._resize_id(image, anno, raw_anno) for image, anno, raw_anno in zip(images, annos, all_raw_anno)]
        id_size_w = 0
        id_size_h = 0
        for img, anno in resized:
            h, w = img.shape[:2]
            id_size_h += h
            id_size_w += w
        id_size = (id_size_h, id_size_w)
        # print('Calling augment with input:', id.shape)
        if not is_scan:
            background = cv2.imread(random.choice(self.backgrounds))
        else:
            p = random.uniform(1.35, 1.65)
            width = int(id_size[1]*p)
            height = int(id_size[0]*p)
            background = 255 * np.ones(shape=[height, width, 3], dtype=np.uint8)
        # print('Choose background with shape:', background.shape)
        background = self._resize_background(background, id_size)
        # print('BG Shape:', background.shape)
        return self._overlay_id_on_background(resized, all_raw_anno, background)


class OverLayImageMask():

    def __init__(self, background_path='./data/SynthText/'):
        self.background_path = background_path
        self.backgrounds = self._get_usable_background()
        if len(self.backgrounds) == 0:
            raise FileNotFoundError(
                "Can't file any background image in", os.path.join(self.background_path, 'bg_img', ", please refer to the README to download some"))

    def _get_usable_background(self):
        with open(os.path.join(self.background_path, 'imnames.cp'), 'rb') as f:
            filtered_imnames = set(cp.load(f))
        usable = []
        for imname in os.listdir(os.path.join(self.background_path, 'bg_img')) + os.listdir('data/document_background'):
            # ignore if not in filetered list:
            if imname not in filtered_imnames:
                continue
            else:
                if 'data/document_background' in imname:
                    usable.append(os.path.join(
                        'data/document_background', imname))
                else:
                    usable.append(os.path.join(
                        self.background_path, 'bg_img', imname))
        return usable


    def _resize_background(self, image, id_size, f_min=1.35, f_max=1.65):
        p = random.uniform(f_min, f_max)
        bg_size = image.shape
        width = int(id_size[1]*p)
        height = int(id_size[0]*p)
        if bg_size[0] >= id_size[0] and bg_size[1] >= id_size[1]:
            resizer = iaa.OneOf([
                iaa.CropToFixedSize(width, height),
                iaa.Resize({"height": height, "width": width})
            ])
        else:
            resizer = iaa.Resize({"height": height, "width": width})
        return resizer.augment_image(image)


    def put4ChannelImageOn4ChannelImage(self, back, fore, x, y, is_binary):
        if not is_binary:
            tmp = cv2.cvtColor(fore, cv2.COLOR_BGR2GRAY)
            _, alpha = cv2.threshold(tmp, 0, 255, cv2.THRESH_BINARY)
            b, g, r = cv2.split(fore)
            rgba = [b, g, r, alpha]
            fore = cv2.merge(rgba, 4)
        else:
            fore = cv2.cvtColor(fore, cv2.COLOR_BGR2RGBA)

        back = cv2.cvtColor(back, cv2.COLOR_BGR2RGBA)
        rows, cols, channels = fore.shape
        trans_indices = fore[..., 3] != 0  # Where not transparent
        overlay_copy = back[y:y+rows, x:x+cols]
        overlay_copy[trans_indices] = fore[trans_indices]
        back[y:y+rows, x:x+cols] = overlay_copy
        # alpha_s = fore[:, :, 3] / 255.0
        # alpha_l = 1.0 - alpha_s

        # for c in range(0, 3):
        #     back[y:y+rows, x:x+cols, c] = (alpha_s * fore[:, :, c] +
        #                               alpha_l * back[y:y+rows, x:x+cols, c])
        return back

    def augment_image(self, image, mask, background=None, is_binary=False, f_min=1.35, f_max=1.65):
        id_size = image.shape[:-1]
        # print('Calling augment with input:', id.shape)
        if background is None:
            background = cv2.imread(random.choice(self.backgrounds))
            # print('Choose background with shape:', background.shape)
            background = self._resize_background(
                background, id_size, f_min, f_max)

        bg_h, bg_w, _ = background.shape
        id_h, id_w, _ = image.shape
        x = random.randint(0, (bg_w - id_w))
        y = random.randint(0, (bg_h - id_h))

        new_mask = np.zeros(
            (background.shape[0], background.shape[1], mask.shape[2]))

        rows, cols, channels = mask.shape
        for i in range(mask.shape[2]):
            new_mask[y:y+rows, x:x+cols, i] = mask[:, :, i]

        return self.put4ChannelImageOn4ChannelImage(background, image, x, y, is_binary), new_mask

    def augment_image_ocr(self, image, background=None, is_binary=False):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        id_size = image.shape[:-1]
        # print('Calling augment with input:', id.shape)
        if background is None:
            background = cv2.imread(random.choice(self.backgrounds))
            # print('Choose background with shape:', background.shape)
            background = self._resize_background(
                background, id_size, f_min=1.03, f_max=1.04)
        bg_h, bg_w, _ = background.shape
        id_h, id_w, _ = image.shape
        x = random.randint(0, (bg_w - id_w))
        y = random.randint(0, (bg_h - id_h))
        return self.put4ChannelImageOn4ChannelImage(background, image, x, y, is_binary)


class OverLayImageVbpl():
    def __init__(self, background_path='./data/background/'):
        self.background_path = background_path
        self.backgrounds = self._get_usable_background()
        if len(self.backgrounds) == 0:
            raise FileNotFoundError(
                "Can't file any background image in", os.path.join(self.background_path, 'bg_img', ", please refer to the README to download some"))

    def _get_usable_background(self):
        usable = []
        for imname in os.listdir(os.path.join(self.background_path)):
            usable.append(os.path.join(self.background_path, imname))
        return usable

    def _resize_id(self, image):
        p = random.uniform(1, 2.08)
        width = int(image.shape[1]/p)
        height = int(image.shape[0]/p)
        resizer = iaa.Resize({"height": height, "width": width})
        # image = resizer(image=image)
        image = cv2.resize(image, (width, height))
        return image

    def _resize_background(self, image, id_size):
        p = random.uniform(1.35, 1.65)
        try:
            bg_size = image.shape
        except:
            return None
        width = int(id_size[1]*p)
        height = int(id_size[0]*p)
        if bg_size[0] >= id_size[0] and bg_size[1] >= id_size[1]:
            resizer = iaa.OneOf([
                iaa.CropToFixedSize(width, height),
                iaa.Resize({"height": height, "width": width})
            ])
        else:
            resizer = iaa.Resize({"height": height, "width": width})
        return resizer.augment_image(image)

    def put4ChannelImageOn4ChannelImage(self, back, fore, x, y):
        back = cv2.cvtColor(back, cv2.COLOR_BGR2RGBA)
        rows, cols, channels = fore.shape
        # print(fore.shape)
        # img_height, img_width, channels = fore.shape

        out = fore.copy()

        # # Test
        # cv2.imshow("Background", back)
        # cv2.imshow("Foreground", fore)

        fore[:, :, 3] = (
            255 * (fore[:, :, :3] != 255).any(axis=2)).astype(np.uint8)
        # fore[:, :, 3] = (255 - x[:, :, :3].mean(axis=2)).astype(np.uint8)

        # trans_indices = fore[:, :, 3] != 0  # Where not transparent
        trans_indices = fore[..., 3] != 0  # Where not transparent

        overlay_copy = back[y:y+rows, x:x+cols]

        from random import uniform
        alpha = uniform(0, 0.5)
        cv2.addWeighted(overlay_copy, alpha, fore, 1 - alpha,
                        0, fore)

        overlay_copy[trans_indices] = fore[trans_indices]

        # # Test
        # cv2.imshow("Overlay", overlay_copy)
        back[y:y+rows, x:x+cols] = overlay_copy

        # # Test
        # cv2.imshow("IMAGE COMPLETED", back)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        return back

    def _overlay_on_background(self, image, background):
        bg_h, bg_w, _ = background.shape
        id_h, id_w, _ = image.shape

        # print('BG shape:', background.shape)
        # print('image shape: ', image.shape)

        x = random.randint(0, (bg_w - id_w))
        y = random.randint(0, (bg_h - id_h))

        augmented_img = self.put4ChannelImageOn4ChannelImage(
            background, image, x, y)
        return augmented_img

    def augment_image(self, image):
        image = self._resize_id(image)
        image_size = image.shape[:-1]
        background = cv2.imread(random.choice(self.backgrounds))
        background = self._resize_background(background, image_size)
        while background is None:
            background = cv2.imread(random.choice(self.backgrounds))
            background = self._resize_background(background, image_size)
        return self._overlay_on_background(image, background)


class LightFlare(iaa.meta.Augmenter):
    def __init__(self, name=None, deterministic=False, random_state=None):
        super().__init__(name=name, deterministic=deterministic, random_state=random_state)

    def flare_source(self, image, point, radius, src_color):
        overlay = image.copy()
        output = image.copy()
        num_times = radius//10
        alpha = np.linspace(0.0, 1, num=num_times)
        rad = np.linspace(1, radius, num=num_times)
        for i in range(num_times):
            cv2.circle(overlay, point, int(rad[i]), src_color, -1)
            alp = alpha[num_times-i-1] * \
                alpha[num_times-i-1]*alpha[num_times-i-1]
            cv2.addWeighted(overlay, alp, output, 1 - alp, 0, output)
        return output

    def augment_image(self, image, hooks=None):
        radius = int(image.shape[1]*random.uniform(0.05, 0.2))
        x = random.randint(0, image.shape[1])
        y = random.randint(0, image.shape[0])
        color = (random.randint(230, 255), random.randint(
            230, 255), random.randint(230, 255))
        return self.flare_source(image, (x, y), radius, color)

    def _augment_images(self, images, **kwargs):
        return [self.augment_image(each_img) for each_img in images]

    def _augment_keypoints(self, **kwargs):
        return

    def _augment_heatmaps(self, **kwargs):
        return

    def get_parameters(self):
        return


class WarpTexture(iaa.meta.Augmenter):
    def __init__(self, texture_path='./data/texture/', name=None, deterministic=False, random_state=None):
        super().__init__(name=name, deterministic=deterministic, random_state=random_state)
        # self.texture_file = [cv2.cvtColor(cv2.imread(each), cv2.COLOR_BGR2RGB) for each in glob.glob(os.path.join(texture_path, '*.jpg'))]
        self.texture_file = glob.glob(os.path.join(texture_path, '*.jpg'))

    def augment_image(self, image, hooks=None):
        foreground = cv2.imread(random.choice(self.texture_file)) # Just BGR
        f_h, f_w, _ = foreground.shape
        foreground = cv2.cvtColor(foreground, cv2.COLOR_BGR2RGB)
        
        background = image
        b_h, b_w, _ = background.shape

        if f_h < b_h or f_w < b_w:
            new_fore_size = (b_w, b_h)
            foreground = cv2.resize(foreground, new_fore_size)
        
        # foreground = random.choice(self.texture_file)
        
        # Random crop
        x = random.randint(0, foreground.shape[0]-background.shape[0])
        y = random.randint(0, foreground.shape[1]-background.shape[1])

        foreground = foreground[x:x+background.shape[0], y:y+background.shape[1]]
        final = cv2.addWeighted(
            background, 1, foreground, random.uniform(0.5, 1), 0)
        return final

    def _augment_images(self, images, **kwargs):
        return [self.augment_image(each_img) for each_img in images]

    def _augment_keypoints(self, **kwargs):
        return

    def _augment_heatmaps(self, **kwargs):
        return

    def get_parameters(self):
        return


class AddCrapyStamp(iaa.meta.Augmenter):
    def __init__(self, stamp_path='./data/stamp/', name=None, deterministic=False, random_state=None):
        super().__init__(name=name, deterministic=deterministic, random_state=random_state)
        # self.texture_file = [cv2.cvtColor(cv2.imread(each), cv2.COLOR_BGR2RGB) for each in glob.glob(os.path.join(texture_path, '*.jpg'))]
        self.stamp_file = glob.glob(os.path.join(stamp_path, '*'))

    def augment_image(self, image, hooks=None):
        background = image.copy()
        image_h, image_w, image_c = background.shape
        stamp = cv2.imread(random.choice(self.stamp_file))
        curr_stamp_h, curr_stamp_w, curr_stamp_c = stamp.shape
        if random.uniform(0, 1) < 0.6:  # big
            ratio = random.uniform(3, 4)
        else:
            ratio = random.uniform(0.9, 2)
        new_stamp_h = int(ratio * image_h)
        new_stamp_w = int((new_stamp_h/curr_stamp_h) * curr_stamp_w)
        new_stamp = cv2.resize(stamp, (new_stamp_w, new_stamp_h))
        plate = np.zeros(background.shape, dtype=np.uint8)

        mode = random.randint(1, 3)
        if mode == 1:  # At the start of the line
            x = random.randint(-int(0.05*image_w), int(0.2*image_w))
            y = random.randint(-int(0.05*image_h), int(0.2*image_h))
        elif mode == 2:  # At the end of the line
            x = random.randint(int(0.5*image_w),  int(0.7*image_w))
            y = random.randint(int(0.2*image_h),  int(0.3*image_h))
        else:  # Anywhere beetween
            x = random.randint(int(0.1*image_w), int(0.7*image_w))
            y = random.randint(int(0.2*image_h), int(0.3*image_h))
        # print(mode)
        # print(y, x)
        # print(plate.shape)
        # print(new_stamp.shape)

        if x in range(0, image_w) and y in range(0, image_h):
            if x + new_stamp_w <= image_w and y + new_stamp_h <= image_h:
                plate[y:y+new_stamp_h, x:x+new_stamp_w, :] = new_stamp
            elif x + new_stamp_w <= image_w and y + new_stamp_h > image_h:
                viewable_h = image_h - y
                plate[y:y+viewable_h, x:x+new_stamp_w,
                      :] = new_stamp[0:viewable_h, :, :]
            elif x + new_stamp_w > image_w and y + new_stamp_h <= image_h:
                viewable_w = image_w - x
                plate[y:y+new_stamp_h, x:x+viewable_w,
                      :] = new_stamp[:, 0:viewable_w, :]
            else:
                viewable_h = image_h - y
                viewable_w = image_w - x
                plate[y:y+viewable_h, x:x+viewable_w,
                      :] = new_stamp[0:viewable_h, 0:viewable_w, :]
        else:
            if y < 0 and x >= 0 and y + new_stamp_h <= image_h and x + new_stamp_w <= image_w:
                viewable_h = new_stamp_h - abs(y)
                plate[0:viewable_h, x:x+new_stamp_w,
                      :] = new_stamp[abs(y):, :, :]
            elif x < 0 and y >= 0 and x + new_stamp_w <= image_w and y + new_stamp_h <= image_h:
                viewable_w = new_stamp_w - abs(x)
                plate[y:y+new_stamp_h, 0:viewable_w,
                      :] = new_stamp[:, abs(x):, :]
            elif x < 0 and y >= 0 and x + new_stamp_w <= image_w and y + new_stamp_h >= image_h:
                viewable_w = new_stamp_w - abs(x)
                viewable_h = image_h - y
                plate[y:image_h, 0:viewable_w,
                      :] = new_stamp[:viewable_h, abs(x):, :]
            elif x < 0 and y < 0 and x + new_stamp_w <= image_w and y + new_stamp_h <= image_h:
                viewable_w = new_stamp_w - abs(x)
                viewable_h = new_stamp_h - abs(y)
                plate[0:viewable_h, 0:viewable_w,
                      :] = new_stamp[abs(y):, abs(x):, :]
            elif y < 0 and x >= 0 and x + new_stamp_w >= image_w and y + new_stamp_h <= image_h:
                viewable_h = new_stamp_h - abs(y)
                plate[0:viewable_h, x:, :] = new_stamp[abs(y):, x:image_w, :]

        data = plate
        r1, g1, b1 = 255, 255, 255  # Original value
        r2, g2, b2 = 0, 0, 0  # Value that we want to replace it with
        red, green, blue = data[:, :, 0], data[:, :, 1], data[:, :, 2]
        mask = (red == r1) & (green == g1) & (blue == b1)
        data[:, :, :3][mask] = [r2, g2, b2]
        plate = data

        background_alpha = random.uniform(0.4, 0.6)
        final = cv2.addWeighted(
            background, background_alpha, plate, 1.5-background_alpha, 0)
        return final

    def _augment_images(self, images, **kwargs):
        return [self.augment_image(each_img) for each_img in images]

    def _augment_keypoints(self, **kwargs):
        return

    def _augment_heatmaps(self, **kwargs):
        return

    def get_parameters(self):
        return


class AddCrapyStampCRAFT(iaa.meta.Augmenter):
    def __init__(self, stamp_path='./data/stamp/', name=None, deterministic=False, random_state=None):
        super().__init__(name=name, deterministic=deterministic, random_state=random_state)
        # self.texture_file = [cv2.cvtColor(cv2.imread(each), cv2.COLOR_BGR2RGB) for each in glob.glob(os.path.join(texture_path, '*.jpg'))]
        self.stamp_file = glob.glob(os.path.join(stamp_path, '*'))

    def brightness_augment(self, img, factor=0.5): 
        hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV) #convert to hsv
        hsv = np.array(hsv, dtype=np.float64)
        hsv[:, :, 2] = hsv[:, :, 2] * (factor + np.random.uniform()) #scale channel V uniformly
        hsv[:, :, 2][hsv[:, :, 2] > 255] = 255 #reset out of range values
        rgb = cv2.cvtColor(np.array(hsv, dtype=np.uint8), cv2.COLOR_HSV2RGB)
        return rgb

    def augment_image(self, image, hooks=None):
        background = image.copy()
        plate = np.zeros(background.shape, dtype=np.uint8)
        for _ in range(random.randint(1, 7)):
            image_h, image_w, image_c = background.shape
            stamp = cv2.imread(random.choice(self.stamp_file))
            curr_stamp_h, curr_stamp_w, curr_stamp_c = stamp.shape
            ratio = random.uniform(0.1, 0.125)
            new_stamp_h = int(ratio * image_h)
            new_stamp_w = int((new_stamp_h/curr_stamp_h) * curr_stamp_w)
            # new_stamp_w = int(random.uniform(0.05, 0.1)*image_w)
            # new_stamp_h = int(random.uniform(0.05, 0.1)*image_h)
            new_stamp = cv2.resize(stamp, (new_stamp_w, new_stamp_h))
            x = random.randint(int(0.1*image_w), int(0.95*image_w))
            y = random.randint(int(0.1*image_h), int(0.95*image_h))
            if x in range(0, image_w) and y in range(0, image_h):
                if x + new_stamp_w <= image_w and y + new_stamp_h <= image_h:
                    plate[y:y+new_stamp_h, x:x+new_stamp_w, :] = new_stamp
                elif x + new_stamp_w <= image_w and y + new_stamp_h > image_h:
                    viewable_h = image_h - y
                    plate[y:y+viewable_h, x:x+new_stamp_w,
                        :] = new_stamp[0:viewable_h, :, :]
                elif x + new_stamp_w > image_w and y + new_stamp_h <= image_h:
                    viewable_w = image_w - x
                    plate[y:y+new_stamp_h, x:x+viewable_w,
                        :] = new_stamp[:, 0:viewable_w, :]
                else:
                    viewable_h = image_h - y
                    viewable_w = image_w - x
                    plate[y:y+viewable_h, x:x+viewable_w,
                        :] = new_stamp[0:viewable_h, 0:viewable_w, :]
            else:
                if y < 0 and x >= 0 and y + new_stamp_h <= image_h and x + new_stamp_w <= image_w:
                    viewable_h = new_stamp_h - abs(y)
                    plate[0:viewable_h, x:x+new_stamp_w,
                        :] = new_stamp[abs(y):, :, :]
                elif x < 0 and y >= 0 and x + new_stamp_w <= image_w and y + new_stamp_h <= image_h:
                    viewable_w = new_stamp_w - abs(x)
                    plate[y:y+new_stamp_h, 0:viewable_w,
                        :] = new_stamp[:, abs(x):, :]
                elif x < 0 and y >= 0 and x + new_stamp_w <= image_w and y + new_stamp_h >= image_h:
                    viewable_w = new_stamp_w - abs(x)
                    viewable_h = image_h - y
                    plate[y:image_h, 0:viewable_w,
                        :] = new_stamp[:viewable_h, abs(x):, :]
                elif x < 0 and y < 0 and x + new_stamp_w <= image_w and y + new_stamp_h <= image_h:
                    viewable_w = new_stamp_w - abs(x)
                    viewable_h = new_stamp_h - abs(y)
                    plate[0:viewable_h, 0:viewable_w,
                        :] = new_stamp[abs(y):, abs(x):, :]
                elif y < 0 and x >= 0 and x + new_stamp_w >= image_w and y + new_stamp_h <= image_h:
                    viewable_h = new_stamp_h - abs(y)
                    plate[0:viewable_h, x:, :] = new_stamp[abs(y):, x:image_w, :]

            data = plate
            r1, g1, b1 = 255, 255, 255  # Original value
            r2, g2, b2 = 0, 0, 0  # Value that we want to replace it with
            red, green, blue = data[:, :, 0], data[:, :, 1], data[:, :, 2]
            mask = (red == r1) & (green == g1) & (blue == b1)
            data[:, :, :3][mask] = [r2, g2, b2]
            plate = data

        background_alpha = random.uniform(0.6, 0.75)
        final = cv2.addWeighted(
            background, background_alpha, plate, 1.5-background_alpha, 0)
        return self.brightness_augment(final, factor=1.5)

    def _augment_images(self, images, **kwargs):
        return [self.augment_image(each_img) for each_img in images]

    def _augment_keypoints(self, **kwargs):
        return

    def _augment_heatmaps(self, **kwargs):
        return

    def get_parameters(self):
        return


class BlobAndShadow(iaa.meta.Augmenter):
    def __init__(self, name=None, deterministic=False, random_state=None):
        super().__init__(name=name, deterministic=deterministic, random_state=random_state)

    def augment_image(self, image, hooks=None):
        image = augment_blobs(image)
        image = shadow.add_n_random_shadows(image)
        return image
    
    def _augment_images(self, images, **kwargs):
        return [self.augment_image(each_img) for each_img in images]

    def _augment_keypoints(self, **kwargs):
        return

    def _augment_heatmaps(self, **kwargs):
        return

    def get_parameters(self):
        return



class SimpleOverlayImage():
    def __init__(self, background_path='./data/SynthText/'):
        self.background_path = background_path
        self.backgrounds = self._get_usable_background()
        if len(self.backgrounds) == 0:
            raise FileNotFoundError(
                "Can't file any background image in", os.path.join(self.background_path, 'bg_img', ", please refer to the README to download some"))

    def _get_usable_background(self):
        with open(os.path.join(self.background_path, 'imnames.cp'), 'rb') as f:
            filtered_imnames = set(cp.load(f))
        usable = []
        for imname in os.listdir(os.path.join(self.background_path, 'bg_img')):
            # ignore if not in filetered list:
            if imname not in filtered_imnames:
                continue
            else:
                usable.append(os.path.join(
                    self.background_path, 'bg_img', imname))
        return usable

    def _resize_background(self, image, id_size):
        if random.uniform(0, 1) <= 0.95:
            p = random.uniform(1.05, 1.25)
        else:
            p = random.uniform(2, 3)
        bg_size = image.shape
        width = int(id_size[1]*p)
        height = int(id_size[0]*p)
        if bg_size[0] >= id_size[0] and bg_size[1] >= id_size[1]:
            resizer = iaa.OneOf([
                iaa.CropToFixedSize(width, height),
                iaa.Resize({"height": height, "width": width})
            ])
        else:
            resizer = iaa.Resize({"height": height, "width": width})
        return resizer.augment_image(image)

    def _resize_id(self, image, anno):
        p = random.uniform(0.5, 2.08)
        width = int(image.shape[1]/p)
        height = int(image.shape[0]/p)
        img = cv2.resize(image, (height, width))
        new_anno = []
        for each in anno:
            new_anno.append([[int(e[0]/p), int(e[1]/p)] for e in each])
        return image, anno

    def put4ChannelImageOn4ChannelImage(self, back, fore, x, y):
        back = cv2.cvtColor(back, cv2.COLOR_BGR2RGBA)
        rows, cols, channels = fore.shape
        trans_indices = fore[..., 3] != 0  # Where not transparent
        overlay_copy = back[y:y+rows, x:x+cols]
        overlay_copy[trans_indices] = fore[trans_indices]
        back[y:y+rows, x:x+cols] = overlay_copy
        return back

    def _overlay_id_on_background(self, id, background, anno):
        bg_h, bg_w, _ = background.shape
        id_h, id_w, _ = id.shape

        # print('BG shape:', background.shape)
        # print('ID shape: ', id.shape)

        x = random.randint(0, (bg_w - id_w))
        y = random.randint(0, (bg_h - id_h))

        # Calculate new coords
        for index, line in enumerate(anno):
            new_coords = [[point[0]+x, point[1]+y]
                          for point in line]
            anno[index] = new_coords

        augmented_img = self.put4ChannelImageOn4ChannelImage(
            background, id, x, y)
        return augmented_img, anno

    def augment_image(self, image, anno, is_scan=False):
        # Random resize the images
        id, anno = self._resize_id(image, anno)
        id_size = image.shape[:-1]
        # print('Calling augment with input:', id.shape)
        if not is_scan:
            background = cv2.imread(random.choice(self.backgrounds))
            background = cv2.cvtColor(background, cv2.COLOR_BGR2RGB)
        else:
            background = 255 * np.ones(shape=[5000, 5000, 3], dtype=np.uint8)
        # print('Choose background with shape:', background.shape)
        background = self._resize_background(background, id_size)
        return self._overlay_id_on_background(id, background, anno)


class RandomLine(iaa.meta.Augmenter):
    def __init__(self, is_binary=False, name=None, deterministic=False, random_state=None):
        self.middleline_thickness = [1, 2]
        self.middleline_thickness_p = [0.8, 0.2]
        self.is_binary = is_binary
        super().__init__(name=name, deterministic=deterministic, random_state=random_state)

    def augment_image(self, image, **kwargs):
        h, w, = image.shape[:2]
        for _ in range(random.randint(0, 5)):
            thickness = np.random.choice(
                self.middleline_thickness, p=self.middleline_thickness_p)
            image = cv2.line(image,
                             (random.randint(0, w), random.randint(0, h)),
                             (random.randint(0, w), random.randint(0, h)),
                             color=(random.randint(0, 255), random.randint(
                                 0, 255), random.randint(0, 255)) if not self.is_binary else random.choice([(0, 0, 0), (255, 255, 255)]),
                             thickness=thickness,
                             lineType=cv2.LINE_AA)
        return image

    def _augment_images(self, images, **kwargs):
        return [self.augment_image(each_img) for each_img in images]

    def _augment_keypoints(self, **kwargs):
        return

    def _augment_heatmaps(self, **kwargs):
        return

    def get_parameters(self):
        return


class MimicRotateBound(iaa.meta.Augmenter):
    def __init__(self, name=None, deterministic=False, random_state=None):
        super().__init__(name=name, deterministic=deterministic, random_state=random_state)

    def augment_image(self, image, hooks=None):
        if random.randint(0, 1) == 0:
            angle = random.uniform(-1, 1)
        else:
            angle = random.uniform(-5, 5)
        image = imutils.rotate_bound(image, angle)
            # Now try to remove black part because of background and rotate bound
        # y_nonzero, x_nonzero, _ = np.nonzero(image)
        # image = image[np.min(y_nonzero):np.max(y_nonzero), np.min(x_nonzero):np.max(x_nonzero)]
        top_row = 0
        for index, row in enumerate(image):
            if np.sum(row) != 0:
                break
            else:
                top_row += 1
        h, w, _ = image.shape
        bot_row = h
        for index, row in enumerate(image[::-1]):
            if np.sum(row) != 0:
                break
            else:
                bot_row -= 1
        left_col = 0
        for index, col in enumerate(image.T):
            if np.sum(col) != 0:
                break
            else:
                left_col += 1
            left_col = 0
        right_col = w
        for index, col in enumerate(image.T[::-1]):
            if np.sum(col) != 0:
                break
            else:
                right_col -= 1

        image = image[top_row:bot_row, left_col:right_col, :]

        return image

    def _augment_images(self, images, **kwargs):
        return [self.augment_image(each_img) for each_img in images]

    def _augment_keypoints(self, **kwargs):
        return

    def _augment_heatmaps(self, **kwargs):
        return

    def get_parameters(self):
        return


def random_blobs(shape, blobdensity, size, roughness=2.0):
    h, w = shape[0], shape[1]
    numblobs = int(blobdensity * w * h)
    mask = np.zeros((h, w), 'i')
    for i in range(numblobs):
        mask[random.randint(0, h-1), random.randint(0, w-1)] = 1
    dt = ndi.distance_transform_edt(1-mask)
    mask = np.array(dt < size, 'f')
    mask = ndi.gaussian_filter(mask, size/(2*roughness))
    mask -= np.amin(mask)
    mask /= np.amax(mask)
    noise = pylab.rand(h, w)
    noise = ndi.gaussian_filter(noise, size/(2*roughness))
    noise -= np.amin(noise)
    noise /= np.amax(noise)
    return np.array(mask * noise > 0.5, 'int')


def augment_blobs(image):
    bg = random_blobs(image.shape, 2e-4, size=10 + random.randint(-3, 10))
    redImg = copy.deepcopy(image)
    redImg[bg == 1] = (random.randint(0, 255), random.randint(
        0, 255), random.randint(0, 255))

    alpha = random.uniform(0.3, 0.8)
    cv2.addWeighted(redImg, alpha, image, 1-alpha, 0, image)

    return image


class OverLayMultipleImage():

    def __init__(self, background_path='./data/SynthText/'):
        self.background_path = background_path
        self.backgrounds = self._get_usable_background()
        if len(self.backgrounds) == 0:
            raise FileNotFoundError(
                "Can't file any background image in", os.path.join(self.background_path, 'bg_img', ", please refer to the README to download some"))

    def _get_usable_background(self):
        with open(os.path.join(self.background_path, 'imnames.cp'), 'rb') as f:
            filtered_imnames = set(cp.load(f))
        usable = []
        for imname in os.listdir(os.path.join(self.background_path, 'bg_img')):
            # ignore if not in filetered list:
            if imname not in filtered_imnames:
                continue
            else:
                usable.append(os.path.join(
                    self.background_path, 'bg_img', imname))
        return usable

    def _resize_background(self, image, id_size):
        p = random.uniform(1.1, 1.7)
        bg_size = image.shape
        width = int(id_size[1]*p)
        height = int(id_size[0]*p)
        if bg_size[0] >= id_size[0] and bg_size[1] >= id_size[1]:
            resizer = iaa.OneOf([
                iaa.CropToFixedSize(width, height),
                iaa.Resize({"height": height, "width": width})
            ])
        else:
            resizer = iaa.Resize({"height": height, "width": width})
        return resizer.augment_image(image)

    def _resize_id(self, image, anno, raw_anno, p=None):
        if p is None:
            p = random.uniform(1.3, 2)
            width = int(image.shape[1]/p)
            height = int(image.shape[0]/p)
        else:
            width, height = p
        resizer = iaa.Resize({"height": height, "width": width})
        psoi = imgaug.PolygonsOnImage(
            [Polygon(each['polygon']) for each in anno], shape=image.shape)
        image, anno = resizer(image=image, polygons=psoi)
        anno = imgaug_to_normal(anno, raw_anno)
        return image, anno

    def put4ChannelImageOn4ChannelImage(self, back, fore, x, y):
        if back.shape[-1] != 4:
            back = cv2.cvtColor(back, cv2.COLOR_BGR2RGBA)
        # print(x, y)
        # print('Back', back.shape)
        # print('Fore', fore.shape)
        rows, cols, channels = fore.shape
        trans_indices = fore[..., 3] != 0  # Where not transparent
        overlay_copy = back[y:y+rows, x:x+cols]
        # print('Overlay copy', overlay_copy.shape)
        overlay_copy[trans_indices] = fore[trans_indices]
        back[y:y+rows, x:x+cols] = overlay_copy
        # print('-')
        return back

    def get_random_location(self, fore_imgs, background):
        HEI, WID, _ = background.shape
        old_crd = []
        all_loc = []
        for (img, anno) in fore_imgs:
            hei, wid, _ = img.shape
            count = 0
            while True:
                x1 = random.randint(0, WID-wid)
                y1 = random.randint(0, HEI-hei)
                
                bbox = np.array([x1,y1, wid, hei])
                overlapped = self.check_overlap(bbox, np.asarray(old_crd))
                if not overlapped:
                    break

                count += 1
                if count > 100:
                    print("overlapped")
                    break

            old_crd.append((x1, y1, wid, hei))
            all_loc.append((x1,y1, overlapped))

        return all_loc

    def check_overlap(self, bbox, candidates):
        if len(candidates)==0:
            return False
        bbox_tl, bbox_br = bbox[:2], bbox[:2] + bbox[2:]
        candidates_tl = candidates[:, :2]
        candidates_br = candidates[:, :2] + candidates[:, 2:]

        tl = np.c_[np.maximum(bbox_tl[0], candidates_tl[:, 0])[:, np.newaxis],
                np.maximum(bbox_tl[1], candidates_tl[:, 1])[:, np.newaxis]]
        br = np.c_[np.minimum(bbox_br[0], candidates_br[:, 0])[:, np.newaxis],
                np.minimum(bbox_br[1], candidates_br[:, 1])[:, np.newaxis]]
        wh = np.maximum(0., br - tl)

        area_intersection = wh.prod(axis=1)
        area_bbox = bbox[2:].prod()
        area_candidates = candidates[:, 2:].prod(axis=1)
        
        IOU =  area_intersection / (area_bbox + area_candidates - area_intersection)
        
        if np.mean(IOU)== 0:
            return False
        else: 
            return True
            

    def _overlay_id_on_background(self, all_resized, all_raw_anno, background):
        bg_h, bg_w, _ = background.shape
        # all_loc = get_random_location(num_rects=len(all_resized), region_shape=(bg_w, bg_h), factor=100)
        # import ipdb; ipdb.set_trace()
        all_loc = self.get_random_location(all_resized, background)
        augmented_img = background
        augmented_anno = []


        for (id, anno), loc, raw_anno in zip(all_resized, all_loc, all_raw_anno):
            id_h, id_w, _ = id.shape
            # x = loc.min.x
            # y = loc.min.y

            #genos 
            x = loc[0]
            y = loc[1]
            overlapped = loc[2]
            if overlapped:
                continue

            x2, y2 = x+id_w, y+id_h

            # print(loc)
            # Calculate new coords
            org_anno = copy.deepcopy(anno)
            for index, line in enumerate(anno):
                new_coords = [[point[0]+x, point[1]+y]
                            for point in line['polygon']]
                anno[index]['polygon'] = new_coords
            try:
                augmented_img = self.put4ChannelImageOn4ChannelImage(
                    augmented_img, id, x=x, y=y)
            except IndexError:
                # print('SCALING DOWN ID')
                # Image need to be resize to a smaller size
                target_w = loc.max.x - loc.min.x
                target_h = loc.max.y - loc.min.y
                # print('Before scale shape', id.shape)
                # print(target_w, target_h)
                id, anno = self._resize_id(id, org_anno, raw_anno, p=(target_w, target_h))
                for index, line in enumerate(anno):
                    new_coords = [[point[0]+x, point[1]+y]
                                for point in line['polygon']]
                    anno[index]['polygon'] = new_coords
                # print('After scale shape', id.shape)
                augmented_img = self.put4ChannelImageOn4ChannelImage(
                    augmented_img, id, x=x, y=y)
            augmented_anno.append(anno)
        # print('------')
        return augmented_img, augmented_anno

    def augment_image(self, imgs, annos, all_raw_anno, is_scan=False):
        # Random resize the images
        resized = [self._resize_id(image, anno, raw_anno) for image, anno, raw_anno in zip(imgs, annos, all_raw_anno)]
        id_size_w = 0
        id_size_h = 0
        for img, anno in resized:
            h, w = img.shape[:2]
            id_size_h += h
            id_size_w += w
        id_size = (id_size_h, id_size_w)
        # print('Calling augment with input:', id.shape)
        if not is_scan:
            background = cv2.imread(random.choice(self.backgrounds))
        else:
            p = random.uniform(1.35, 1.65)
            width = int(id_size[1]*p)
            height = int(id_size[0]*p)
            background = 255 * np.ones(shape=[height, width, 3], dtype=np.uint8)
        # print('Choose background with shape:', background.shape)
        background = self._resize_background(background, id_size)
        # print('BG Shape:', background.shape)
        return self._overlay_id_on_background(resized, all_raw_anno, background)


class OverlayImages:
    def __init__(self, background_path='./data/SynthText/'):
        self.background_path = background_path
        self.backgrounds = self._get_usable_background()
        if len(self.backgrounds) == 0:
            raise FileNotFoundError(
                "Can't file any background image in", os.path.join(self.background_path, 'bg_img', ", please refer to the README to download some"))
    

    def _get_usable_background(self):
        '''SynthText dataset has some unusable images.
        '''
        with open(os.path.join(self.background_path, 'imnames.cp'), 'rb') as f:
            filtered_imnames = set(cp.load(f))
        usable = []
        for imname in os.listdir(os.path.join(self.background_path, 'bg_img')):
            # ignore if not in filetered list:
            if imname not in filtered_imnames:
                continue
            else:
                usable.append(os.path.join(
                    self.background_path, 'bg_img', imname))
        return usable

    def _resize_foreground(self, image, anno, all_raw_anno):
        pass


    def _resize_background(self, background, imgs_size, max_ratio_bg):
        p = random.uniform(1.1, max_ratio_bg)
        w_bg, h_bg = background.size
        bg = np.array(background)
        w_imgs, h_imgs = imgs_size
        width, height = int(w_imgs*p), int(h_imgs*p)

        if w_bg > w_imgs and h_bg > h_imgs:
            resizer = iaa.OneOf([
                iaa.CropToFixedSize(width, height),
                iaa.Resize({"height": height, "width": width})
            ])
        else:
            resizer = iaa.Resize({"height": height, "width": width})
        
        background = resizer.augment_image(background)
        background = Image.from_array(background)
        return background


    def augment_images(self, images, annos, all_raw_anno, ratio, is_scan=False):
        '''Overlay images into background

        Params
        ------
        :images: list of foreground (augmented) images
        :annos: list of (augmented) annotations of images 
        :all_raw_anno: list of raw annotations of images
        :ratio: ratio of foreground with background
        :is_scan: (bool) overlay scanned images in white ground

        Returns
        -------
        :augmented_image: 
        :anno: annotation of augmented_image
        '''
        resized_imgs = [self._resize_foreground(image, anno, all_raw_anno)
                        for image, anno, raw_anno in zip(images, annos, all_raw_anno)] 
        imgs_size_w = 0
        imgs_size_h = 0
        for img, _ in resized_imgs:
            h, w, _ = img.shape 
            imgs_size_h += h
            imgs_size_w += w
        imgs_size = (imgs_size_w, imgs_size_h)

        if not is_scan:
            background_path = random.choice(self.backgrounds)
            background = Image.open(background_path)
        else:
            background = Image.new("RGB", imgs_size, color=(255, 255, 255))

        max_ratio_bg = 1/ratio  
        background = self._resize_background(background, imgs_size, max_ratio_bg)
        augmented_image, anno = self._overlay_img_on_background()
        return augmented_image, anno

     

if __name__ == "__main__":
    aug = AddCrapyStampCRAFT()
#     raw_img = cv2.imread('./result/6.png')
#     raw_img = cv2.cvtColor(raw_img, cv2.COLOR_BGR2RGB)
#     augted_img = aug.augment_images([raw_img]*5)
#     for each in augted_img:
#         augted_img = cv2.cvtColor(each, cv2.COLOR_RGB2BGR)
#         cv2.imshow('test1', augted_img)
#         # cv2.imshow('test2', raw_img_2)
#         cv2.waitKey(0)
#     cv2.destroyAllWindows()
#     # a = OverLayImage()
#     a = OverLayImageVbpl()
    # aug = WarpTexture()
    raw_img = cv2.imread('/media/geneous/01D62877FB2A4900/Techainer/OCR/Fake-Data-Generator/result_data/result_template_rcj/rcj_0.png')
    raw_img = cv2.cvtColor(raw_img, cv2.COLOR_BGR2RGB)
    final = aug.augment_image(raw_img)
    final = cv2.cvtColor(final, cv2.COLOR_RGB2BGR)
    # cv2.imwrite('test.jpg', final)
    cv2.namedWindow("", cv2.WINDOW_NORMAL)
    cv2.imshow("", final)
    cv2.waitKey(0)