|
from typing import List |
|
import PIL.Image |
|
import torch |
|
import gradio as gr |
|
|
|
article = "<p style='text-align: center'><a href='https://github.com/scoutant/yolo-person-gradio' target='_blank' class='footer'>Github Repo</a></p>" |
|
|
|
model = torch.hub.load('ultralytics/yolov5', 'yolov5l') |
|
model.classes = [ 0 ] |
|
model.conf = 0.6 |
|
|
|
def inference(img:PIL.Image.Image, threshold:float=0.6): |
|
if img is None: |
|
return None,0 |
|
images:List[PIL.Image.Image] = [ img ] |
|
model.conf = threshold |
|
detections = model(images, size=640) |
|
predictions:torch.Tensor = detections.pred[0] |
|
result_image = detections.ims[0] if hasattr(detections, "ims") else detections.imgs[0] |
|
detections.render() |
|
return result_image, predictions.size(dim=0) |
|
|
|
gr.Interface( |
|
fn = inference, |
|
inputs = [ gr.Image(type="pil", label="Input"), gr.Slider(minimum=0.5, maximum=0.9, step=0.05, value=0.7, label="Confidence threshold") ], |
|
outputs = [ gr.Image(type="pil", label="Output"), gr.Label(label="nb of persons detected for given confidence threshold") ], |
|
title="Person detection with YOLO v5", |
|
description="Person detection, you can twik the corresponding confidence threshold. Good results even when face not visible.", |
|
article=article, |
|
examples=[['data/businessmen-612.jpg'], ['data/businessmen-back.jpg']], |
|
allow_flagging="never" |
|
).launch(debug=True, enable_queue=True) |