import gradio as gr
import numpy as np
import uuid
from PIL import Image
import os


def inference(img):
    name_tag = uuid.uuid4()
    script_command = f"python yolov9/detect_dual.py --source {img} --img 640 --device cpu --weights yolov9/runs/train/exp2/weights/best.pt --name {name_tag}"
    os.system(script_command)
    return f"yolov9/runs/detect/{name_tag}/{img.split('/')[-1]}"


title = "Sketch GUI Element Detection"

description = "This is a demo for detecting GUI elements in a sketch image. Upload a sketch image and the model will detect the GUI elements in the image."

img_input = gr.Image(
    type="filepath", label="Upload a sketch image", width=300, height=300
)
prediction_output = gr.Image(label="Output Image", width=640, height=640)

example_lst = [
    ["test_images/Shipping-1.png"],
]

demo = gr.Interface(
    fn=inference,
    inputs=img_input,
    outputs=prediction_output,
    title=title,
    description=description,
    examples=example_lst,
)

demo.launch(share=True)