import gradio as gr
import math
import os
import yolov5
import cv2
import torch
import numpy as np

# load model
model_face = torch.load('face_model.pt')
model_face.eval()
model_plate = yolov5.load('keremberke/yolov5m-license-plate', device="cpu")
# # set model parameters
model_plate.conf = 0.25  # NMS confidence threshold
model_plate.iou = 0.45  # NMS IoU threshold
model_plate.agnostic = False  # NMS class-agnostic
model_plate.multi_label = False  # NMS multiple labels per box
model_plate.max_det = 1000  # maximum number of detections per image


def blur_plates_image(image, plate_blur):

    results_plate = model_plate(image, augment=True)
    boxes_plate_list = results_plate.pred[0][:, :4].cpu().numpy().tolist()
    for box in boxes_plate_list:
        ROI = image[int(box[1]):int(box[3]), int(box[0]):int(box[2])]

        blur_value = (int(plate_blur) * 2 - 1)
        blur = cv2.GaussianBlur(ROI, (blur_value, blur_value), 20, cv2.BORDER_DEFAULT)
        # Insert ROI back into image
        image[int(box[1]):int(box[3]), int(box[0]):int(box[2])] = blur

    return image


def blur_faces_image(image, face_blur):

    results = model_face(np.array(image))

    # parse results
    boxes_face_list = results.pred[0][:, :4].cpu().numpy().tolist() # x1, y1, x2, y2

    for box in boxes_face_list:
        p1 = (int(box[0]), int(box[1]))
        p2 = (int(box[2]), int(box[3]))
        w = p2[0] - p1[0]
        h = p2[1] - p1[1]

        circle_center = ((p1[0] + p2[0]) // 2, (p1[1] + p2[1]) // 2)
        circle_radius = int(math.sqrt(w * w + h * h) // 2)
        ROI = np.zeros(image.shape, dtype='uint8')
        cv2.circle(ROI, circle_center, circle_radius, (255, 255, 255), -1)

        blur_value = (int(face_blur) * 2 - 1)
        blur = cv2.GaussianBlur(image, (blur_value, blur_value), cv2.BORDER_DEFAULT)

        image = np.where(ROI > 0, blur, image)

    return image


def blur(image, face_blur, plate_blur):

    image = blur_plates_image(image, plate_blur)
    image = blur_faces_image(image, face_blur)

    return image


gr.Interface(
    fn=blur,
    inputs=[
        gr.Image(source="upload", type="numpy", optional=False),
        gr.Slider(minimum=0, maximum=30, step=1, value=10, label="Face Blur Intensity"),
        gr.Slider(minimum=0, maximum=30, step=1, value=10, label="Plate Blur Intensity"),
    ],
    outputs="image",
    title="Plate Licenses and Faces Blur",
    enable_queue=False,
    sidebar="<img src='images/22.png' width='100%'/>",
    description="This app uses image processing techniques to automatically blur out license plates and faces in photos. Protect the privacy of individuals and comply with regulations by obscuring sensitive information in your images with ease.",
    examples=[["images/22.png", 5, 5], ["images/18.jpg", 5, 5]],
    css=".gradio-container {margin-top:80px !important;}"
).launch()