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

from ultralytics import YOLO
import yolov5


# Image download
# file_urls = [
# ]

# 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"
#     )

# Function for inference
def yolov5_inference(
    image: gr.inputs.Image = None,
    model_path: gr.inputs.Dropdown = None,
    image_size: gr.inputs.Slider = 640,
    conf_threshold: gr.inputs.Slider = 0.25,
    iou_threshold: gr.inputs.Slider = 0.45 ):

    # Loading Yolo V5 model
    model = yolov5.load(model_path, device="cpu")

    # Setting model configuration 
    model.conf = conf_threshold
    model.iou = iou_threshold

    # Inference
    results = model([image], size=image_size)

    # Cropping the predictions    
    crops = results.crop(save=False)
    img_crops = []
    for i in range(len(crops)):
        img_crops.append(crops[i]["im"][..., ::-1])
    return results.render()[0], img_crops
        
# gradio Input
inputs = [
    gr.inputs.Image(type="pil", label="Input Image"),
    gr.inputs.Dropdown(["Damage_Vehicle_Y5.pt","yolov5s.pt", "yolov5m.pt", "yolov5l.pt", "yolov5x.pt"], label="Model", default = 'Crime_Y5.pt'),
    gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
]

# gradio Output
outputs = gr.outputs.Image(type="filepath", label="Output Image")
outputs_crops = gr.Gallery(label="Object crop")

title = "Vehicle damage detection"

# gradio examples: "Image", "Model", "Image Size", "Confidence Threshold", "IOU Threshold"
examples = [['1.jpg', 'Damage_Vehicle_Y5.pt', 640, 0.35, 0.45]
          ,['2.jpg', 'Damage_Vehicle_Y5.pt', 640, 0.35, 0.45]
          ,['3.jpg', 'Damage_Vehicle_Y5.pt', 640, 0.35, 0.45]]

# gradio app launch
demo_app = gr.Interface(
    fn=yolov5_inference,
    inputs=inputs,
    outputs=[outputs,outputs_crops],
    title=title,
    examples=examples,
    cache_examples=True,
    live=True,
    theme='huggingface',
)
demo_app.launch(debug=True, enable_queue=True, width=50, height=50)