import gradio as gr
import cv2
import requests
import os
from ultralytics import YOLO

# ------
file_urls = [
    'https://i.imgur.com/wgZzKIk.jpg',
    'https://i.imgur.com/TvIo0Nq.jpg'
]

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):
    if 'mp4' in file_urls[i]:
        download_file(
            file_urls[i],
            f"video.mp4"
        )
    else:
        download_file(
            file_urls[i],
            f"image_{i}.jpg"
        )
# ------
model = YOLO('crack.pt')
path  = [['image_0.jpg'], ['image_1.jpg']]
video_path = [['video.mp4']]
# ------
def show_preds_image(source):
    global model
    res = model(source, conf=.5, iou=.5)
    res_plotted = res[0].plot()
    # converting BGR to RGB
    result = cv2.cvtColor(res_plotted, cv2.COLOR_BGR2RGB)
    return result
    
# ------
inputs_image = [
    gr.components.Image(type="filepath", label="Input Image"),
]
outputs_image = [
    gr.components.Image(type="numpy", label="Output Image"),
]
interface_image = gr.Interface(
    fn=show_preds_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Roof Crack detector app",
    examples=path,
    cache_examples=False,
)
# ------

gr.TabbedInterface(
    [interface_image],
    tab_names=['Image inference']
).queue().launch()