import cv2
import gradio as gr
from typing import Union, Tuple
from PIL import Image, ImageOps
import numpy as np
import torch

model = torch.jit.load('./model/model.pt').eval()

def resize_with_padding(img: Image.Image, expected_size: Tuple[int, int]) -> Image.Image:
    img.thumbnail((expected_size[0], expected_size[1]))
    delta_width = expected_size[0] - img.size[0]
    delta_height = expected_size[1] - img.size[1]
    pad_width = delta_width // 2
    pad_height = delta_height // 2
    padding = (pad_width, pad_height, delta_width - pad_width, delta_height - pad_height)
    return ImageOps.expand(img, padding), padding

def preprocess_image(img: Image.Image, size: int = 512) -> Tuple[Image.Image, torch.tensor, Tuple[int]]:    
    pil_img, padding = resize_with_padding(img, (size, size))
    
    img = (np.array(pil_img).astype(np.float32) / 255) - np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
    img = img / np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)
    img = np.transpose(img, (2, 0, 1))
    
    return pil_img, torch.tensor(img[None]), padding

def soft_blur_with_mask(image: Image.Image, mask: torch.tensor, padding: Tuple[int]) -> Image.Image:
    image = np.array(image)
    # Create a blurred copy of the original image.
    blurred_image = cv2.GaussianBlur(image, (221, 221), sigmaX=20, sigmaY=20)
    image_height, image_width = image.shape[:2]
    mask = cv2.resize(mask.astype(np.uint8), (image_width, image_height), interpolation=cv2.INTER_NEAREST)
    # Blurring the mask itself to get a softer mask with no firm edges
    mask = cv2.GaussianBlur(mask.astype(np.float32), (11, 11), 10, 10)[:, :, None]

    # Take the blurred image where the mask it positive, and the original image where the image is original
    image = (mask * blurred_image + (1.0 - mask) * image)
    pad_w, pad_h, _, _ = padding
    img_w, img_h, _ = image.shape
    image = image[(pad_h):(img_h-pad_h), (pad_w):(img_w-pad_w), :]
    return Image.fromarray(image.astype(np.uint8))

def run(image, size):
    pil_image, torch_image, padding = preprocess_image(image, size=size)

    with torch.inference_mode():
        mask = model(torch_image)
    mask = mask.argmax(dim=1).numpy().squeeze()

    return soft_blur_with_mask(pil_image, mask, padding)

content_image_input = gr.inputs.Image(label="Entrada", type="pil")
model_image_size = gr.inputs.Radio([256, 384, 512, 1024], type="value", default=512, label="Ajustar nivel de inferencia")

app_interface = gr.Interface(fn=run,
                             inputs=[content_image_input, model_image_size],
                             outputs="image")
app_interface.launch()