app.py

import gradio as gr
import torch
from PIL import ImageDraw

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

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

    return image


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

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

urls = ["https://c8.alamy.com/comp/J2AB4K/the-new-york-stock-exchange-on-the-wall-street-in-new-york-J2AB4K.jpg"]

with demo:
    gr.Markdown(title)
    gr.Markdown(description)
    
    with gr.Tabs():
        with gr.TabItem('Image Upload'):
            with gr.Row():
                img_input = gr.Image(type='pil')
                img_output= gr.Image()
                
            #with gr.Row(): 
            #    example_images = gr.Dataset(components=[img_input],
            #                                samples=[[path.as_posix()]
            #                                         for path in sorted(pathlib.Path('images').rglob('*.JPG'))])
                
            img_button = gr.Button('Detect')
        

    img_button.click(detect_objects,inputs=img_input,outputs=img_output)
    

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