coutant's picture
simple person detection with Yolov8.
0316c1c
raw
history blame
1.58 kB
import cv2
import gradio as gr
from ultralytics import YOLO
def inference(path:str, threshold:float=0.6):
print("trying inference with path", path)
if path is None:
return None,0
model = YOLO('yolov8m.pt')
model.classes = [0] # only considering class 'person' and not the 79 other classes...
image = cv2.imread(path)
outputs = model.predict(source=path, return_outputs=True)
for output in outputs: # mono item batch
detections = output['det']
counter=0
for detection in detections:
if detection[4]<threshold:
break
cv2.rectangle(
image,
(int(detection[0]), int(detection[1])),
(int(detection[2]), int(detection[3])),
color=(0, 0, 255),
thickness=2,
)
counter+=1
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), counter
gr.Interface(
fn = inference,
inputs = [ gr.components.Image(type="filepath", label="Input"), gr.Slider(minimum=0.5, maximum=0.9, step=0.05, value=0.7, label="Confidence threshold") ],
outputs = [ gr.components.Image(type="numpy", label="Output"), gr.Label(label="nb of persons detected for given confidence threshold") ],
title="Person detection with YOLO v8",
description="Person detection, you can tweak the corresponding confidence threshold. Good results even when face not visible.",
examples=[ ['data/businessmen-612.jpg'], ['data/businessmen-back.jpg']],
allow_flagging="never"
).launch(debug=True, enable_queue=True)