Spaces:

coutant
/

yolo-person

Running

App Files Files and versions Community

yolo-person / app.py

coutant

removed share, not supported in space.

4916e5f 5 months ago

raw history blame contribute delete

No virus

1.72 kB

	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 ] # only considering class 'person' and not the 79 other classes...
	model.conf = 0.6 # only considering detection above the threshold.

	def inference(img:PIL.Image.Image, threshold:float=0.6):
	if img is None:
	return None,0
	images:List[PIL.Image.Image] = [ img ] # inference operates on a list of images
	model.conf = threshold
	detections = model(images, size=640)
	predictions:torch.Tensor = detections.pred[0] # the predictions for our single image
	result_image = detections.ims[0] if hasattr(detections, "ims") else detections.imgs[0] # either model.common.Detections or torchvision.Detections
	detections.render() # bounding boxes and labels added into image
	return result_image, predictions.size(dim=0) # image and number of detections

	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)