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

__all__ = ['encode_polygon', 'get_image', 'load_annotation_points', 'order_points', 'four_point_transform',
           'get_perspective_transform', 'get_details', 'transform_points', 'check_boundaries', 'transform_keypoints',
           'get_length_4_sides', 'Perspective_Transform']

# Cell

import json
import cv2
import glob
import numpy as np
import matplotlib.pyplot as plt

from pycocotools import mask as CocoMask

# Cell

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 get_image(imgpath, imagedata, num):
    w, h = Image.fromarray(img_b64_to_arr(imagedata)).size
    img_dict = {
        "license": None,
        "file_name": args.image_loc + "/" + imgpath.rsplit("/")[-1],
        "url": None,
        "height": h,
        "width": w,
        "data_captured": None,
        "id": num
    }
    print(args.save_loc.split("/")[-1] + "/" + imgpath.rsplit("/")[-1])
    return img_dict, (w, h)

def load_annotation_points(json_file):

    """
    loads labels with annotations in dictionary
    """

    js = load_json(json_file)

    boxes =  {}

    for annotation in js["shapes"]:
        label = annotation["label"]
        boxes[label] = np.array(annotation["points"])

    return boxes

# Cell

"""
Gently borrowwed from : https://www.pyimagesearch.com/2014/08/25/4-point-opencv-getperspective-transform-example/
"""
def order_points(pts):
    # initialzie a list of coordinates that will be ordered
    # such that the first entry in the list is the top-left,
    # the second entry is the top-right, the third is the
    # bottom-right, and the fourth is the bottom-left
    rect = np.zeros((4, 2), dtype = "float32")
    # the top-left point will have the smallest sum, whereas
    # the bottom-right point will have the largest sum
    s = pts.sum(axis = 1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    # now, compute the difference between the points, the
    # top-right point will have the smallest difference,
    # whereas the bottom-left will have the largest difference
    diff = np.diff(pts, axis = 1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    # return the ordered coordinates
    return rect

def four_point_transform(image, pts):
    # obtain a consistent order of the points and unpack them
    # individually
    rect = order_points(pts)
    (tl, tr, br, bl) = rect
    # compute the width of the new image, which will be the
    # maximum distance between bottom-right and bottom-left
    # x-coordiates or the top-right and top-left x-coordinates
    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    maxWidth = max(int(widthA), int(widthB))
    # compute the height of the new image, which will be the
    # maximum distance between the top-right and bottom-right
    # y-coordinates or the top-left and bottom-left y-coordinates
    heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
    heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
    maxHeight = max(int(heightA), int(heightB))
    # now that we have the dimensions of the new image, construct
    # the set of destination points to obtain a "birds eye view",
    # (i.e. top-down view) of the image, again specifying points
    # in the top-left, top-right, bottom-right, and bottom-left
    # order
    dst = np.array([
        [0, 0],
        [maxWidth - 1, 0],
        [maxWidth - 1, maxHeight - 1],
        [0, maxHeight - 1]], dtype = "float32")
    # compute the perspective transform matrix and then apply it
    M = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
    # return the warped image
    return warped

def get_perspective_transform(image, pts):
    # obtain a consistent order of the points and unpack them
    # individually
    rect = order_points(pts)
    (tl, tr, br, bl) = rect
    # compute the width of the new image, which will be the
    # maximum distance between bottom-right and bottom-left
    # x-coordiates or the top-right and top-left x-coordinates
    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    maxWidth = max(int(widthA), int(widthB))
    # compute the height of the new image, which will be the
    # maximum distance between the top-right and bottom-right
    # y-coordinates or the top-left and bottom-left y-coordinates
    heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
    heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
    maxHeight = max(int(heightA), int(heightB))
    # now that we have the dimensions of the new image, construct
    # the set of destination points to obtain a "birds eye view",
    # (i.e. top-down view) of the image, again specifying points
    # in the top-left, top-right, bottom-right, and bottom-left
    # order
    dst = np.array([
        [0, 0],
        [maxWidth - 1, 0],
        [maxWidth - 1, maxHeight - 1],
        [0, maxHeight - 1]], dtype = "float32")
    # compute the perspective transform matrix and then apply it
    M = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
    # return the warped image
    return M, (maxWidth, maxHeight), warped

def get_details(json_file):
    js = load_json(json_file)

    img_path = root + js["imagePath"]
    boxes = js["shapes"]

    points = {}

    for label in boxes:
        points[label["label"]] = label["points"]

    return img_path,points

def transform_points(M,i,j):
    res = np.dot(M,[i,j,1])
    i2,j2,_  = (res / res[2] + 0.5).astype(int)
    return i2,j2

def check_boundaries(edge_length, point):
    if point<0:
        return 0
    if point>=edge_length:
        return edge_length-1
    return point


def transform_keypoints(warped_image, M, original_kps):
    transform_keypoints = []
    for (x,y) in original_kps:
        x_t,y_t = transform_points(M,x,y)
        x_t = check_boundaries(warped_image.shape[1],x_t)
        y_t = check_boundaries(warped_image.shape[0],y_t)
        transform_keypoints.append([x_t,y_t])

    return transform_keypoints

def get_length_4_sides(p):

    return([
        math.sqrt((p[i][0]-p[i-1][0])**2 + (p[i][1]-p[i-1][1])**2) for i in [1,2,3,0]
    ])

# Cell

class Perspective_Transform():

    def __init__(self,annotations_path):

        self.annotations_path = annotations_path
        self.json_files = glob.glob(self.annotations_path + "*.json")
        self.num_annotations = len(self.json_files)

    def __len__(self):
        return (self.num_annotations)

    def __getitem__(self,index):

        json_file = self.json_files[index]
        img_file = json_file[:-5] + ".jpg"

        boxes = load_annotation_points(json_file)

        img = cv2.imread(img_file)

        warped_inner = four_point_transform(img, (boxes["inner"]))
        warped_outer = four_point_transform(img, (boxes["outer"]))

        return img,warped_inner,warped_outer

# Cell
try: from nbdev.imports import IN_NOTEBOOK
except: IN_NOTEBOOK=False

if __name__ == "__main__" and not IN_NOTEBOOK:
    annotations_path = "/mnt/nfshome1/FRACTAL/saksham.jindal/DLG/bezel_detection/LG_Phi_Pana/inner_outer/"
    perspective_transform = Perspective_Transform(annotations_path)

    ID = np.random.randint(0,len(perspective_transform)-1)
    img,inner_t,outer_t = perspective_transform[ID]

    plt.figure(figsize=(50,10))
    f,ax = plt.subplots(1,3)

    ax[0].imshow(img)
    ax[1].imshow(inner_t)
    ax[2].imshow(outer_t)

    plt.show()