import gradio as gr
import onnxruntime as ort
import numpy as np
from PIL import Image, ImageDraw
import cv2

image_size = 224

def normalize_image(image, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
    image = (image/255.0).astype("float32")
    
    image[:, :, 0] = (image[:, :, 0] - mean[0]) / std[0]
    image[:, :, 1] = (image[:, :, 1] - mean[1]) / std[1]
    image[:, :, 2] = (image[:, :, 2] - mean[2]) / std[2]
    
    return image

def resize_longest_max_size(image, max_size=224):
    
    height, width = image.shape[:2]
    
    if width > height:
        ratio = max_size / width
    else:
        ratio = max_size / height
    
    new_width = int(width * ratio)
    new_height = int(height * ratio)
    
    resized_image = cv2.resize(image, (new_width, new_height), cv2.INTER_LINEAR)
    
    return resized_image

def pad_if_needed(image, target_size):
    height, width, _ = image.shape
    
    y0 = abs((height-target_size)//2)
    x0 = abs((width-target_size)//2)
    
    background = np.zeros((target_size, target_size, 3), dtype="uint8")
    background[y0:(y0+height), x0:(x0+width), :] = image
    
    return(background)

def heatmap2keypoints(heatmap: np.ndarray, image_size: int = 224) -> list:
    "Function to convert heatmap to keypoint x, y tensor"

    indx = heatmap.reshape(-1, image_size*image_size).argmax(axis=1)
    row = indx // image_size
    col = indx % image_size

    keypoints_array = np.stack((col, row), axis=1)
    keypoints_list = keypoints_array.tolist()

    return keypoints_list

def centercrop_keypoints(keypoints, crop_height, crop_width, image_height, image_width):
    y_diff = (image_height-crop_height)//2
    x_diff = (image_width-crop_width)//2
    
    keypoints_crop = [[x-x_diff, y-y_diff] for x, y in keypoints]
    return(keypoints_crop)

def resize_keypoints(keypoints, current_height, current_width, target_height, target_width):
    keypoints_resize = []
    for x, y in keypoints:
        x_resize = (x/current_width)*target_width
        y_resize = (y/current_height)*target_height
        keypoints_resize.append([int(x_resize), int(y_resize)])
    return(keypoints_resize)

def draw_keypoints(image, keypoints):
    draw = ImageDraw.Draw(image)
    w, h = image.size
    for keypoint in keypoints:
        x, y = keypoint
        # Draw a small circle at each keypoint
        radius = int(min(w, h) * 0.01)
        draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill='red', outline='red')
    return image

def point_dist(p0, p1):
    x0, y0 = p0
    x1, y1 = p1
    
    dist = ((x0-x1)**2 + (y0-y1)**2)**0.5
    
    return dist

def receipt_asp_ratio(keypoints, mode = "mean"):
    
    h0 = point_dist(keypoints[0], keypoints[3])
    h1 = point_dist(keypoints[1], keypoints[2])
    
    w0 = point_dist(keypoints[0], keypoints[1])
    w1 = point_dist(keypoints[2], keypoints[3])

    if mode == "max":
        h = max(h0, h1)
        w = max(w0, w1)
    elif mode == "mean":
        h = (h0+h1)/2
        w = (w0+w1)/2
    else:
        return("UNKNOWN MODE")
    
    return w/h

# Load the ONNX model
session = ort.InferenceSession("models/timm-mobilenetv3_small_100.onnx")

input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# Main function to handle the image input, apply preprocessing, run the model, and apply postprocessing
def process_image(input_image):

    # Convert PIL image to OpenCV image
    image = np.array(input_image.convert("RGB"))

    h, w, _ = image.shape

    # Preprocess the image
    image_resize = resize_longest_max_size(image)

    h_small, w_small, _ = image_resize.shape

    image_pad = pad_if_needed(image_resize, target_size=image_size)

    image_norm = normalize_image(image_pad)

    image_array = np.transpose(image_norm, (2, 0, 1))

    image_array = np.expand_dims(image_array, axis=0)

    # Run model inference
    output = session.run([output_name], {input_name: image_array})

    output_keypoints = heatmap2keypoints(output[0].squeeze())
    crop_keypoints = centercrop_keypoints(output_keypoints, h_small, w_small, image_size, image_size)
    large_keypoints = resize_keypoints(crop_keypoints, h_small, w_small, h, w)

    # Draw keypoints on the image
    image_with_keypoints = draw_keypoints(input_image, large_keypoints)

    persp_h = 1024
    persp_asp = receipt_asp_ratio(large_keypoints, mode="max")
    persp_w = int(persp_asp*persp_h)

    origin_points = np.float32([[x, y] for x, y in large_keypoints])
    target_points = np.float32([[0, 0], [persp_w-1, 0], [persp_w-1, persp_h-1], [0, persp_h-1]])

    persp_matrix = cv2.getPerspectiveTransform(origin_points, target_points)
    persp_image = cv2.warpPerspective(image, persp_matrix, (persp_w, persp_h), cv2.INTER_LINEAR)
    
    output_image = Image.fromarray(persp_image)

    return image_with_keypoints, output_image

demo_images = [
    "demo_images/image_1.jpg",
    "demo_images/image_2.jpg",
    "demo_images/image_3.jpg",
    "demo_images/image_flux_1.png",
    "demo_images/image_flux_2.png",
]

# Create Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Document corner detection and perspective correction")
    gr.Markdown("Upload an image to detect the corners of a document and correct the perspective.\n\nUses a UNet model to detect corners and OpenCV to correct the perspective.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Image", show_label=True, scale=1)
        
        with gr.Column():
            output_image1 = gr.Image(type="pil", label="Image with predicted corners", show_label=True, scale=1)

        with gr.Column():
            output_image2 = gr.Image(type="pil", label="Image with perspective correction", show_label=True, scale=1)

    with gr.Row():
            examples = gr.Examples(demo_images, input_image, cache_examples=False, label="Exampled documents (CORD dataset and FLUX.1-schnell generated)")
    
    input_image.change(fn=process_image, inputs=input_image, outputs=[output_image1, output_image2])

    gr.Markdown("Created by Kenneth Thorø Martinsen (kenneth2810@gmail.com)")

iface.launch()