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


# 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(["PPE_Safety_Y5.pt"], label="Model", default = 'PPE_Safety_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 = "Identify violations of Personal Protective Equipment (PPE) protocols for improved safety"

# gradio examples: "Image", "Model", "Image Size", "Confidence Threshold", "IOU Threshold"
examples = [['image_1.jpg', 'PPE_Safety_Y5.pt', 640, 0.35, 0.45]
           ,['image_0.jpg', 'PPE_Safety_Y5.pt', 640, 0.35, 0.45]
           ,['image_2.jpg', 'PPE_Safety_Y5.pt', 640, 0.35, 0.45],
           ]

# gradio app launch
demo_app = gr.Interface(
    fn=yolov5_inference,
    inputs=inputs,
    outputs=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)