import gradio as gr
import torch
from ultralyticsplus import YOLO, render_result

def yoloFunc(image: gr.inputs.Image = None,
             image_size: int = 640,
             conf_threshold: float = 0.4,
             iou_threshold: float = 0.5):
    model_path = 'best.pt'
    model = YOLO(model_path)
    
    results = model.predict(image,
                            image_size=image_size,
                            conf_threshold=conf_threshold,
                            iou_threshold=iou_threshold
                            )
    
    box = results[0].boxes
    
    render = render_result(model=model, image=image, results=results[0])
    return render

inputs = [
    gr.inputs.Image(type='filepath', label="Input Image"),
    gr.inputs.Slider(minimum=320, maximum=1024, default=640, step=32, label="Image Size"),
    gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.4, steps=0.05, label="Confidence Threshold"),
    gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.5, steps=0.05, label="IOU Threshold")
]

outputs = gr.outputs.Image(type='filepath', label="Output Image")

title = "Pothole Detection"

yolo_app = gr.Interface(
    fn=yoloFunc,
    inputs=inputs,
    outputs=outputs,
    title=title,   
)

yolo_app.launch(debug=True, enable_queue=True)