import numpy as np
import cv2
import gradio as gr

from opencv_zoo.models.face_detection_yunet.yunet import YuNet
from opencv_zoo.models.license_plate_detection_yunet.lpd_yunet import LPD_YuNet

# Instantiate face detection YuNet
face_model = YuNet(
    modelPath="opencv_zoo/models/face_detection_yunet/face_detection_yunet_2023mar.onnx",
    inputSize=[320, 320],
    confThreshold=0.9,
    nmsThreshold=0.3,
    topK=5000,
    backendId=cv2.dnn.DNN_BACKEND_OPENCV,
    targetId=cv2.dnn.DNN_TARGET_CPU,
)

# Instantiate license plate detection YuNet
lpd_model = LPD_YuNet(
    modelPath="opencv_zoo/models/license_plate_detection_yunet/license_plate_detection_lpd_yunet_2023mar.onnx",
    confThreshold=0.9,
    nmsThreshold=0.3,
    topK=5000,
    keepTopK=750,
    backendId=cv2.dnn.DNN_BACKEND_OPENCV,
    targetId=cv2.dnn.DNN_TARGET_CPU,
)


def json_detections(face_results, lpd_results):
    json_result = {}
    json_result["faces"] = []
    json_result["license_plates"] = []

    for det in face_results if face_results is not None else []:
        bbox = det[0:4].astype(np.int32)
        json_result["faces"].append(
            {
                "xmin": int(bbox[0]),
                "ymin": int(bbox[1]),
                "xmax": int(bbox[0]) + int(bbox[2]),
                "ymax": int(bbox[1]) + int(bbox[3]),
            }
        )

    for det in lpd_results if lpd_results is not None else []:
        bbox = det[:-1].astype(np.int32)
        x1, y1, x2, y2, x3, y3, x4, y4 = bbox

        xmin = min(x1, x2, x3, x4)
        xmax = max(x1, x2, x3, x4)
        ymin = min(y1, y2, y3, y4)
        ymax = max(y1, y2, y3, y4)

        json_result["license_plates"].append(
            {
                "xmin": int(xmin),
                "ymin": int(ymin),
                "xmax": int(xmax),
                "ymax": int(ymax),
            }
        )

    return json_result


def overlay_results(image, face_results, lpd_results):
    # Draw face results on the input image
    for det in face_results if face_results is not None else []:
        bbox = det[0:4].astype(np.int32)
        cv2.rectangle(
            image,
            (bbox[0], bbox[1]),
            (bbox[0] + bbox[2], bbox[1] + bbox[3]),
            (0, 0, 0),
            -1,
        )

    # Draw lpd results on the input image
    for det in lpd_results:
        bbox = det[:-1].astype(np.int32)
        x1, y1, x2, y2, x3, y3, x4, y4 = bbox

        # The output of this is technically a parallelogram, but we will
        # just black out the rectangle
        xmin = min(x1, x2, x3, x4)
        xmax = max(x1, x2, x3, x4)
        ymin = min(y1, y2, y3, y4)
        ymax = max(y1, y2, y3, y4)

        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 0, 0), -1)

    return image


def predict(image):

    h, w, _ = image.shape

    # Inference
    face_model.setInputSize([w, h])
    lpd_model.setInputSize([w, h])
    infer_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    face_results = face_model.infer(infer_image)
    lpd_results = lpd_model.infer(infer_image)

    # Process output
    image = overlay_results(image, face_results, lpd_results)
    json = json_detections(face_results, lpd_results)

    return image, json


demo = gr.Interface(
    title="Face and License Plate Obfuscator - YuNet",
    fn=predict,
    inputs=gr.Image(type="numpy", label="Original Image"),
    outputs=[
        gr.Image(type="numpy", label="Output Image"),
        gr.JSON(visible=False),
    ],
)

demo.launch()