# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/10_image_generator.ipynb (unless otherwise specified).

__all__ = ['random_point_inside_rectangle', 'create_newpoints', 'img_b64_to_arr', 'get_image', 'encode_polygon',
           'segmentation_mask', 'mask_generator', 'blacked_out_image_except_mask', 'get_segment_crop',
           'generate_new_image', 'create_newannotation', 'generate_synthetic_data']

# Cell

from PIL import Image, ImageDraw, ImageFilter, ImageFont
import json
import glob
import random
from pycocotools import mask as CocoMask
import cv2
import json
import io
import base64
import numpy as np

# Cell

def random_point_inside_rectangle(points):
    """Returns a random point inside a rectangle
    """
    min_x, max_x = min([a[0] for a in points]), max([a[0] for a in points])
    min_y, max_y = min(b[1] for b in points), max(b[1] for b in points)
    random_point_x = random.uniform(min_x, max_x)
    random_point_y = random.uniform(min_y, max_y)
    return [int(random_point_x), int(random_point_y)]

def create_newpoints(old_points, new_point):
    """ Returns new co-ordinates of a sticky note given the relative point on new image.
        Ideally uses the relative point co-ordinates and generates new co-ordinates of the note on new image
    """
    #Find distance of all points from single reference edge of rectangle
    diff_list = [[a[0]-old_points[0][0], a[1]- old_points[0][1]] for a in old_points[1:]]

    relative_point_x_diff = old_points[3][0]- old_points[0][0]
    relative_point_x_diff = relative_point_x_diff if relative_point_x_diff<0 else 0

    relative_point_y_diff = old_points[1][1]- old_points[0][1]
    relative_point_y_diff = relative_point_y_diff if relative_point_y_diff<0 else 0


    new_point = [new_point[0] - relative_point_x_diff, new_point[1] - relative_point_y_diff]
    new_list = [[new_point[0]+ a[0], new_point[1]+ a[1]] for a in diff_list]

    new_list.insert(0, new_point)
    return new_list

# Cell

def img_b64_to_arr(img_b64: str):
    """Image b64 to array
    """
    f = io.BytesIO()
    f.write(base64.b64decode(img_b64))
    img_arr = np.array(Image.open(f))
    return img_arr


def get_image(imgpath, imagedata, num):
    image = Image.fromarray(img_b64_to_arr(imagedata))
    w, h = image.size
    img_dict = {
        "license": None,
        "file_name": imgpath.rsplit("/")[-1],
        "url": None,
        "height": h,
        "width": w,
        "data_captured": None,
        "id": num,
    }
    return img_dict, (w, h)


def encode_polygon(points, img_dim):
    """
    points: [[x1, y1], [x2, y2], ....[xn, yn]]
    """
    points = [x for sublist in points for x in sublist]
    img = Image.new("L", img_dim, 0)
    ImageDraw.Draw(img).polygon(points, outline=1, fill=1)
    img_array = np.asarray(img)
    img_uint8 = np.expand_dims(img_array, axis=2).astype("uint8")
    img_array = np.asfortranarray(img_uint8)
    rs = CocoMask.encode(img_array)
    return rs, img_uint8


def segmentation_mask(mask):
    """
    """
    contours, _ = cv2.findContours((mask).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    segmentation = []

    for contour in contours:
        contour = contour.flatten().tolist()
        segmentation.append(contour)
        if len(contour) > 4:
            segmentation.append(contour)
    return segmentation


def mask_generator(sticky_coordinates, img_path):
    """ Creates binary mask of the required class. Takes box co-ordinates of a class and original Image path

        Returns a np array with required class mask as 1 and other part of image as o's
    """

    img_dict, img_dim = get_image(img_path, img_path, 1)
    rs, mask = encode_polygon(sticky_coordinates, img_dim)
    box = CocoMask.toBbox(rs).tolist()
    assert len(box) == 1
    area = CocoMask.area(rs).tolist()
    assert len(area) == 1
    segm = segmentation_mask(np.squeeze(mask, 2))
    mask_new = np.squeeze(mask, axis=(2,))
    return mask_new

def blacked_out_image_except_mask(img, mask):
    """Creates a new image in which except for the given class mask, rest of the image will be black pixels.
       If class is stickynote. This function returns an image in which except for stickynote mask part, rest of
       image is black pixels.

       Takes input image and input mask in binary. 1s for class label and 0's for rest of image

    """
    img_array = np.array(img)
    src1_mask=cv2.cvtColor(mask*255,cv2.COLOR_GRAY2BGR)#change mask to a 3 channel image
    mask_out=cv2.subtract(src1_mask,img_array)
    mask_out=cv2.subtract(src1_mask,mask_out)
    blacked_image = Image.fromarray(mask_out)
    return blacked_image

def get_segment_crop(img,mask=None):
    """Returns the segmentation crop from the image of a particular class. Essentially, this returns just the
       cropped part of image with stickynote in it.

       Takes input image and mask of the image and returns cropped stickynote.

    """
    img = np.array(img)
    if mask is None:
        mask = img > 0
    return img[np.ix_(mask.any(1), mask.any(0))]


def generate_new_image(source_img, seg_crop_img, random_coordinates):

    """Returns a new generated image. Takes source image and segmentation crop of sticky note and co-ordinates
       where the crop should be stiched on the image.

    """
    source_array = np.array(source_img)
    x,y,z = source_array.shape
    black_img_array = np.zeros([x, y, z], dtype=np.uint8)
    black_img_array.fill(0)
    black_img = Image.fromarray(black_img_array)

    black_img.paste(seg_crop_img, (random_coordinates[0], random_coordinates[1]))
    black_img_array = np.array(black_img)

    r,c,m = np.where(black_img_array==0)
    black_img_array[(r,c,m)] = source_array[(r,c,m)]
    generated_image = Image.fromarray(black_img_array)
    return generated_image

# Cell

def create_newannotation(old_annot_dict, new_points, result_folder, name_prefix):
    """ Writes the generated  annotation file of the generated image in the result folder
    """

    new_values = {'label': 'sticky_note',
                   'line_color': [0, 255, 0, 128],
                   'fill_color': [255, 0, 0, 128],
                   'shape_type': 'polygon',
                   'flags': {}}
    new_values['points'] = new_points

    old_annot_dict['shapes'].append(new_values)
    old_annot_dict["imagePath"] = name_prefix + old_annot_dict["imagePath"]
    with open('{}{}.json'.format(result_folder, old_annot_dict["imagePath"].split(".")[0]), "w") as outfile:
        outfile.write(json.dumps(old_annot_dict))

# Cell

def generate_synthetic_data(sticky_note_folder, normal_folder, result_folder, name_prefix="generated", num_images=100):

    """Takes input folder to pick sticky notes from
             source folder to pick normal tv images from
             number of samples to generate
             name of the images to generate
             result folder in which images need to be sent to

    """
    annot_files = glob.glob(sticky_note_folder + "*.json")
    normal_annot_files = glob.glob(normal_folder + "*.json")

    for i in tqdm(range(num_images)):
        try:
            # take a. random image with sticky note
            annot_random_file = random.randint(1,len(annot_files) - 1)
            with open(annot_files[annot_random_file]) as f:

                annot = json.load(f)
                labels = [a["label"] for a in annot["shapes"]]
                index = labels.index("sticky_note")
                sticky_coordinates = annot["shapes"][index]["points"]
                annot_img = Image.open(sticky_note_folder + annot["imagePath"])

            # below lines to generate a segmenation crop
            mask_original_image = mask_generator(sticky_coordinates, annot["imagePath"])
            mask_blacked_out_image = blacked_out_image_except_mask(annot_img, mask_original_image)

            sticky_note_crop_array = get_segment_crop(mask_blacked_out_image, mask_original_image)
            sticky_note_crop_img = Image.fromarray(sticky_note_crop_array)

            # take a random image to which sticky note needs to be sticked to
            normal_random_file = random.randint(1, len(normal_annot_files) -1)

            with open(normal_annot_files[normal_random_file]) as f:
                normal_annot = json.load(f)
                normal_labels = [a["label"] for a in normal_annot["shapes"]]

                normal_index = normal_labels.index("outer")
                outer_coordinates = normal_annot["shapes"][normal_index]["points"]

            normal_image_file = Image.open(normal_folder + normal_annot["imagePath"])

            random_point = random_point_inside_rectangle(outer_coordinates)
            # generate a new image
            generated_image = generate_new_image(normal_image_file, sticky_note_crop_img, random_point)
            generated_image.save(result_folder + name_prefix + normal_annot["imagePath"])
            create_newannotation(normal_annot,create_newpoints(sticky_coordinates, random_point), result_folder)

        except:
            pass
