import gradio as gr
from gradio.outputs import Label
import cv2

from ultralytics import YOLO

model = YOLO('best.pt')
path  = [['pothole_example.jpg'], ['pothole_screenshot.png']]

def show_preds(image_path):
    image = cv2.imread(image_path)
    outputs = model.predict(source=image_path, return_outputs=True)
    for image_id, result in enumerate(outputs):
        print(result['det'])
        for i, det in enumerate(result['det']):
            print(det)
            cv2.rectangle(
                image,
                (int(det[0]), int(det[1])),
                (int(det[2]), int(det[3])),
                color=(0, 0, 255),
                thickness=2,
                lineType=cv2.LINE_AA
            )
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

gr_interface = gr.Interface(
    fn=show_preds,
    inputs=gr.inputs.Image(type="filepath", label="Input Image"),
    outputs=gr.outputs.Image(type="numpy", label="Output Image"),
    title="Pothole detector",
    examples=path,
    cache_examples=True,
    live=True,
)

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