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 UI
demo = gr.Interface(
    fn=detect_edges,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Output"),
    title="Edge Detection DexiNed (OpenCV DNN)",
    allow_flagging="never",
    description="Upload an image to detect edges using OpenCV's ONNX-based edge detection using DexiNed model."
)

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