File size: 2,188 Bytes
9aa19cc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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) |