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

from ultralytics import YOLO

file_urls = [
    'https://www.dropbox.com/s/b5g97xo901zb3ds/pothole_example.jpg?dl=1',
    'https://www.dropbox.com/s/86uxlxxlm1iaexa/pothole_screenshot.png?dl=1'
]

def download_file(url, save_name):
    url = url
    if not os.path.exists(save_name):
        file = requests.get(url)
        open(save_name, 'wb').write(file.content)

for i, url in enumerate(file_urls):    
    download_file(
        file_urls[i],
        f"image_{i}.jpg"
    )

model = YOLO('best.pt')
path  = [['image_0.jpg'], ['image_1.jpg']]

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=False,
    # live=True,
)

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