import gradio as gr
import torch
from PIL import Image, ImageFilter, ImageDraw
from ultralytics import YOLO
import spaces

MODEL_PATH = "./models/face_yolov8n_v2.pt"
MODEL = YOLO(MODEL_PATH)

def create_rounded_rectangle_mask(image:Image.Image, radius, alpha=255):
    size = image.size
    factor = 5  # Factor to increase the image size that I can later antialiaze the corners
    radius = radius * factor
    image = Image.new('RGBA', (size[0] * factor, size[1] * factor), (0, 0, 0, 0))

    # create corner
    corner = Image.new('RGBA', (radius, radius), (0, 0, 0, 0))
    draw = ImageDraw.Draw(corner)
    # added the fill = .. you only drew a line, no fill
    draw.pieslice((0, 0, radius * 2, radius * 2), 180, 270, fill=(255, 255, 255, alpha))

    # max_x, max_y
    mx, my = (size[0] * factor, size[1] * factor)

    # paste corner rotated as needed
    # use corners alpha channel as mask
    image.paste(corner, (0, 0), corner)
    image.paste(corner.rotate(90), (0, my - radius), corner.rotate(90))
    image.paste(corner.rotate(180), (mx - radius, my - radius), corner.rotate(180))
    image.paste(corner.rotate(270), (mx - radius, 0), corner.rotate(270))

    # draw both inner rects
    draw = ImageDraw.Draw(image)
    draw.rectangle([(radius, 0), (mx - radius, my)], fill=(255, 255, 255, alpha))
    draw.rectangle([(0, radius), (mx, my - radius)], fill=(255, 255, 255, alpha))
    return image.resize(size, Image.LANCZOS)  # Smooth the corners

def max_rounding_radius(rect_coords):
    # Extract the coordinates of the rectangle [x0, y0, x1, y1]
    rect_coords = [coord
-1 for coord in rect_coords]
    x0, y0, x1, y1 = rect_coords

    # Calculate the width and height of the rectangle
    width = x1 - x0
    height = y1 - y0

    # Determine the smallest dimension (width or height)
    min_dimension = min(width, height)

    # Calculate the maximum radius as half of the smallest dimension
    return min_dimension // 2

@spaces.GPU(enable_queue=True)
def generate_image(source_image:Image.Image, confidence=0.3, radius=50, blur_amount=10, margin=0):
    if source_image is None:
        return source_image

    device = "cuda" if torch.cuda.is_available() else "cpu"
    pred = MODEL(source_image, conf=confidence, device=device)
    bboxes = pred[0].boxes.xyxy.cpu().numpy()
    if bboxes.size == 0:
        result = None
    else:
        bboxes = bboxes.tolist()
        result = source_image.copy()

        for bboxe in bboxes:
            bboxe = [round(coord) for coord in bboxe]

            mean_margin = margin // 2
            new_bboxe = bboxe[0]-mean_margin, bboxe[1]-mean_margin, bboxe[2]+mean_margin, bboxe[3]+mean_margin

            new_x0, new_y0, new_x1, new_y1 = new_bboxe
            new_width = new_x1 - new_x0
            new_height = new_y1 - new_y0

            if not (new_width > source_image.width or new_height > source_image.height):
                bboxe = new_bboxe

            # Crop the region of interest
            region = result.crop(bboxe)

            final_radius = round(max_rounding_radius(bboxe)* (radius/100))
            mask = create_rounded_rectangle_mask(region, final_radius)

            # Apply blur filter to the cropped region
            blurred_region = region.filter(ImageFilter.GaussianBlur(radius=blur_amount))

            # Paste the blurred region back onto the original image
            result.paste(blurred_region, bboxe, mask)

    return result

css = """
img {
    max-height: 500px;
    object-fit: contain;
}
"""

with gr.Blocks(css=css) as FACE_2_BLUR:
    with gr.Row():
        with gr.Column():
            image = gr.Image(label="Upload your image", type="pil")
            generate_btn = gr.Button("Generate")
            confidence = gr.Slider(label="Detection model confidence threshold.", value=0.3, minimum=0.0, maximum=1.0, step=0.01)
            radius = gr.Slider(label="Edges radius in percentage.", value=50, minimum=0, maximum=100, step=1)
            blur_amount = gr.Number(label="Controls the strength of the blur effect.", value=10, minimum=1, maximum=20, step=1)
            margin = gr.Slider(label="Margin to add to the blurred box.", value=0, minimum=0, maximum=200)

        with gr.Column():
            image_out = gr.Image(label="Blurred face image", type="pil")

    generate_btn.click(
        fn=generate_image,
        inputs=[image, confidence, radius, blur_amount, margin],
        outputs=image_out,
        api_name="generate_blurred"
    )

FACE_2_BLUR.launch()