import cv2 as cv
import numpy as np
import gradio as gr
from yunet import YuNet
from huggingface_hub import hf_hub_download

# Download ONNX model from Hugging Face
model_path = hf_hub_download(repo_id="opencv/face_detection_yunet", filename="face_detection_yunet_2023mar.onnx")

# Initialize YuNet model
model = YuNet(
    modelPath=model_path,
    inputSize=[320, 320],
    confThreshold=0.9,
    nmsThreshold=0.3,
    topK=5000,
    backendId=cv.dnn.DNN_BACKEND_OPENCV,
    targetId=cv.dnn.DNN_TARGET_CPU
)

def visualize(image, results, box_color=(0, 255, 0), text_color=(0, 0, 255)):
    output = image.copy()
    landmark_color = [
        (255,   0,   0), # right eye
        (  0,   0, 255), # left eye
        (  0, 255,   0), # nose tip
        (255,   0, 255), # right mouth corner
        (  0, 255, 255)  # left mouth corner
    ]

    for det in results:
        bbox = det[0:4].astype(np.int32)
        cv.rectangle(output, (bbox[0], bbox[1]), (bbox[0]+bbox[2], bbox[1]+bbox[3]), box_color, 2)
        conf = det[-1]
        cv.putText(output, '{:.2f}'.format(conf), (bbox[0], bbox[1] + 12), cv.FONT_HERSHEY_SIMPLEX, 0.5, text_color)

        landmarks = det[4:14].astype(np.int32).reshape((5, 2))
        for idx, landmark in enumerate(landmarks):
            cv.circle(output, tuple(landmark), 2, landmark_color[idx], 2)

    return output

def detect_faces(input_image):
    input_image = cv.cvtColor(input_image, cv.COLOR_RGB2BGR)
    h, w, _ = input_image.shape
    model.setInputSize([w, h])
    results = model.infer(input_image)
    if results is None or len(results) == 0:
        input_image = cv.cvtColor(input_image, cv.COLOR_BGR2RGB)
        return input_image
    output = visualize(input_image, results)
    output = cv.cvtColor(output, cv.COLOR_BGR2RGB)
    return output

# Gradio Interface
# demo = gr.Interface(
#     fn=detect_faces,
#     inputs=gr.Image(type="numpy", label="Upload Image"),
#     outputs=gr.Image(type="numpy", label="Detected Faces"),
#     title="Face Detection YuNet (OpenCV DNN)",
#     allow_flagging="never",
#     description="Upload an image to detect faces using OpenCV's ONNX-based YuNet face detector."
# )

# Gradio Interface
with gr.Blocks(css='''.example * {
    font-style: italic;
    font-size: 18px !important;
    color: #0ea5e9 !important;
    }''') as demo:

    gr.Markdown("### Face Detection YuNet (OpenCV DNN)")
    gr.Markdown("Upload an image to detect faces using OpenCV's ONNX-based YuNet face detector.")

    with gr.Row():
        input_image = gr.Image(type="numpy", label="Upload Image")
        output_image = gr.Image(type="numpy", label="Detected Faces")

    # Clear output when new image is uploaded
    input_image.change(fn=lambda: (None), outputs=output_image)

    with gr.Row():
        submit_btn = gr.Button("Submit", variant="primary")
        clear_btn = gr.Button("Clear")

    submit_btn.click(fn=detect_faces, inputs=input_image, outputs=output_image)
    clear_btn.click(fn=lambda:(None, None), outputs=[input_image, output_image])

    gr.Markdown("Click on any example to try it.", elem_classes=["example"])

    gr.Examples(
        examples=[
            ["examples/selfie.jpg"],
            ["examples/lena.jpg"],
            ["examples/messi5.jpg"]
        ],
        inputs=input_image
    )

if __name__ == "__main__":
    demo.launch()