import cv2 as cv
import numpy as np
import gradio as gr
from lpd_yunet import LPD_YuNet
from huggingface_hub import hf_hub_download

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

# Initialize LPD-YuNet model
model = LPD_YuNet(
    modelPath=model_path,
    confThreshold=0.9,
    nmsThreshold=0.3,
    topK=5000,
    keepTopK=750,
    backendId=cv.dnn.DNN_BACKEND_OPENCV,
    targetId=cv.dnn.DNN_TARGET_CPU
)

def visualize(image, dets, line_color=(0, 255, 0), text_color=(0, 0, 255)):
    output = image.copy()
    for det in dets:
        bbox = det[:-1].astype(np.int32)
        x1, y1, x2, y2, x3, y3, x4, y4 = bbox
        cv.line(output, (x1, y1), (x2, y2), line_color, 2)
        cv.line(output, (x2, y2), (x3, y3), line_color, 2)
        cv.line(output, (x3, y3), (x4, y4), line_color, 2)
        cv.line(output, (x4, y4), (x1, y1), line_color, 2)
    return output

def detect_license_plates(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:
        return cv.cvtColor(input_image, cv.COLOR_BGR2RGB)
    output = visualize(input_image, results)
    output = cv.cvtColor(output, cv.COLOR_BGR2RGB)
    return output

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

    gr.Markdown("### License Plate Detection (LPD-YuNet)")
    gr.Markdown("Upload a vehicle image to detect license plates using OpenCV's ONNX-based LPD-YuNet model.")

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

    # 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_license_plates, 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/licenseplate1.jpg"],
            ["examples/licenseplate2.jpg"]
        ],
        inputs=input_image
    )

    gr.Markdown("Example images credit: https://unsplash.com/")

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