import cv2 as cv
import gradio as gr
from dexined import Dexined
from huggingface_hub import hf_hub_download

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

# Initialize model
model = Dexined(modelPath=model_path)

def detect_edges(input_image):
    input_image = cv.cvtColor(input_image, cv.COLOR_RGB2BGR)
    result = model.infer(input_image)
    result = cv.cvtColor(result, cv.COLOR_BGR2RGB)
    return result

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

    gr.Markdown("### Edge Detection DexiNed (OpenCV DNN)")
    gr.Markdown("Upload an image to detect edges using OpenCV's ONNX-based edge detection using DexiNed model.")

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

    # 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_edges, 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/baboon.jpg"],
            ["examples/chicky_512.png"],
            ["examples/lena.jpg"],
            ["examples/messi5.jpg"]
        ],
        inputs=input_image
    )

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