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

file_urls = [
    'https://c8.alamy.com/zooms/9/382c1e254fe14207998df9ec56354291/wce5tj.jpg',
    'https://c8.alamy.com/comp/W20YYR/freight-container-on-the-back-of-a-truck-stuck-in-traffic-on-the-interstate-in-georgia-usa-W20YYR.jpg',
    'https://www.shutterstock.com/shutterstock/photos/318604739/display_1500/stock-photo-highway-and-container-truck-at-china-318604739.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):    
    download_file(
        file_urls[i],
        f"image_{i}.jpg"
    )

# model_path = 'Container_YOLOV5'

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 ):


    model = yolov5.load(model_path, device="cpu")
    model.conf = conf_threshold
    model.iou = iou_threshold
    results = model([image], size=image_size)
    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
        

inputs = [
    gr.inputs.Image(type="pil", label="Input Image"),
    gr.inputs.Dropdown(["Crime_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"),
]

outputs = gr.outputs.Image(type="filepath", label="Output Image")
outputs_crops = gr.Gallery(label="Object crop")
title = "Container code detection - YOLO V5"
description = "YOLOv5 is a family of object detection models pretrained on COCO dataset. This model is a pip implementation of the original YOLOv5 model."

examples = [['1.jpg', 'Crime_Y5.pt', 640, 0.35, 0.45]
          ,['2.jpg', 'Crime_Y5.pt', 640, 0.35, 0.45]
          ,['4.jpg', 'Crime_Y5.pt', 640, 0.35, 0.45]]

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)