AI-Naga's picture
Update app.py
87df76d
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(["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"),
]
# gradio Output
outputs = gr.outputs.Image(type="filepath", label="Output Image")
outputs_crops = gr.Gallery(label="Object crop")
title = "Perimeter security - to monitor suspicious human gestures"
# gradio examples: "Image", "Model", "Image Size", "Confidence Threshold", "IOU Threshold"
examples = [['1.jpg', 'Crime_Y5.pt', 640, 0.35, 0.45]
,['2.jpg', 'Crime_Y5.pt', 640, 0.35, 0.45]
,['3.jpg', 'Crime_Y5.pt', 640, 0.35, 0.45]]
# gradio app launch
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)