import gradio as gr
import torch
from PIL import ImageDraw
import pathlib

# Model  
model_path = 'model_torch.pt'
model = torch.hub.load('Ultralytics/yolov5', 'custom', model_path, verbose = False)
model.eval()

labels = model.names

colors = ["red", "blue", "green", "yellow"]

def detect_objects(image):

    draw = ImageDraw.Draw(image)

    detections = model(image)

    probabilities = {}

    for detection in detections.xyxy[0]:
        
        x1, y1, x2, y2, p, category_id = detection
        x1, y1, x2, y2, category_id = int(x1), int(y1), int(x2), int(y2), int(category_id)
        draw.rectangle((x1, y1, x2, y2), outline=colors[category_id], width=4)
        draw.text((x1, y1), labels[category_id], colors[category_id])
        probabilities[labels[category_id]] = float(p)

    return [image, probabilities]


demo = gr.Blocks()#(css=css)

title = '# 3D print failures detection App'
description = 'App for detect errors in the 3D printing'

with demo:
    gr.Markdown(title)
    gr.Markdown(description)
    
    with gr.Tabs():

        #Image static
        with gr.TabItem('Image Upload'):
            with gr.Row():
                
                with gr.Column():
                    img_input = gr.Image(type='pil')
                    examples_images2 = gr.Examples(examples = [[path.as_posix()] for path in sorted(pathlib.Path('images').rglob('*.jpg'))],
                                               inputs=img_input)
                    labels_bars = gr.Label(label = "Categories")
                    print(labels_bars)
                
                with gr.Column():
                    img_output= gr.Image()

            img_button = gr.Button('Detect')

        img_button.click(detect_objects,inputs=img_input,outputs=[img_output, labels_bars])
        
        #Image static
        with gr.TabItem('Webcam'):
            with gr.Row():
                with gr.Column():
                    webcam_input = gr.Image(shape=(320,240), source='webcam', type="pil", )
                    webcam_output = gr.Image()
                with gr.Column():
                    labels_bars2 = gr.Label(label = "Categories")
                
            webcam_button = gr.Button('Detect')

        webcam_button.click(detect_objects,inputs=webcam_input,outputs=[webcam_output, labels_bars2])
    
    

if __name__ == "__main__":
    demo.launch()