Spaces:
Build error
Build error
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) |