代码收藏家 yolov5 deepsort 基于yolov5和deepsort的行人跟踪计数系统【yolov5】【deepsort】【行人检测】【跟踪计数】
行人跟踪计数系统
介绍
本项目使用yolov5作为检测器，使用deepsort作为跟踪器，跟踪并计数镜头前走过的行人数量。
代码中设置的是只计数行人的数量，如果要计数其他东西比如车辆之类的也非常简单，只需要稍微修改一下代码即可，在parser.add_argument('--classes', default=0, type=int, help='filter by class: --class 0, or --class 0 1 2 3')中，通过classes来过滤yolov5要检测类别，0表示检测人，1表示自行车，还有其他的类别。
环境
项目测试环境：win10，pycharm，python3.6+
主要需要的包：pytorch >= 1.7.0，opencv
项目源码（github)
https://github.com/MichistaLin/yolov5-deepsort-pedestraintracking
部分代码
导包
import torch, sys, argparse, cv2, os, time
from datetime import datetime
from self_utils.multi_tasks import Counting_Processing
from self_utils.overall_method import Object_Counter, Image_Capture
from deep_sort.configs.parser import get_config
from deep_sort.deep_sort import DeepSort
import imutils
if __name__ == "__main__":
torch.multiprocessing.set_start_method('spawn')
parser = argparse.ArgumentParser()
# 视频的路径，默认是本项目中的一个测试视频test.mp4，可自行更改
parser.add_argument('--input', type=str, default="./test.mp4",
help='test imgs folder or video or camera') # 输入'0'表示调用电脑默认摄像头
# 处理后视频的输出路径
parser.add_argument('--output', type=str, default="./output",
help='folder to save result imgs, can not use input folder')
parser.add_argument('--weights', type=str, default='weights/yolov5l.pt', help='model.pt path(s)')
parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf_thres', type=float, default=0.4, help='object confidence threshold')
parser.add_argument('--iou_thres', type=float, default=0.4, help='IOU threshold for NMS')
# GPU（0表示设备的默认的显卡）或CPU
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
# 通过classes来过滤yolo要检测类别，0表示检测人，1表示检测自行车，更多具体类别数字可以在19行附近打印出来
# ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', ...]
parser.add_argument('--classes', default=0, type=int, help='filter by class: --class 0, or --class 0 1 2 3')
yolo5_config = parser.parse_args()
print(yolo5_config)
main(yolo5_config)
print("结果保存在：", yolo5_config.output)
检测和计数:该函数中使用yolov5进行检测，使用deepsort进行跟踪
def Counting_Processing(input_img,yolo5_config,model,class_names,Tracker,Obj_Counter):
try:
# 把图片转换成tensor类型
tensor_img=img_preprocessing(input_img,yolo5_config.device,yolo5_config.img_size)
# 使用yolov5进行检测
pred=yolov5_prediction(model,tensor_img,yolo5_config.conf_thres, yolo5_config.iou_thres,yolo5_config.classes)
# 跟踪计数
result_img=count_post_processing(input_img,pred,class_names,tensor_img.shape,Tracker,Obj_Counter)
return result_img
except Exception as e:
print("Wrong:",e, 'multi_tasks.py--->Counting_Processing')
return e
main函数：初始化模型，加载视频或图片，处理并检测跟踪视频、图片中的人
def main(yolo5_config):
print("=> main task started: {}".format(datetime.now().strftime('%H:%M:%S')))
a = time.time()
class_names = []
# 加载模型
Model = torch.load(yolo5_config.weights, map_location=lambda storage, loc: storage.cuda(int(yolo5_config.device)))[
'model'].float().fuse().eval()
# 模型能检测的类别['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', ...]
classnames = Model.module.names if hasattr(Model, 'module') else Model.names
# print(classnames)
# 只检测人这一个类别，所以只需要在框上面标识出是person即可
# 如果想检测其他物体需要在100行附近更改'--classes'的数值，然后在这里把标签改为对应即可
class_names.append(classnames[0])
b = time.time()
print("==> class names: ", class_names)
print("=> load model, cost:{:.2f}s".format(b - a))
os.makedirs(yolo5_config.output, exist_ok=True)
c = time.time()
# 初始化追踪器deepssort_tracker
cfg = get_config()
cfg.merge_from_file("deep_sort/configs/deep_sort.yaml")
deepsort_tracker = DeepSort(cfg.DEEPSORT.REID_CKPT, max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE, max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True, use_appearence=True)
# 输入需要检测的图片或图片目录或视频并处理
mycap = Image_Capture(yolo5_config.input)
# 创建可视化计数器对象
Obj_Counter = Object_Counter(class_names)
# 总帧数
total_num = mycap.get_length()
videowriter = None
fps = int(mycap.get(5))
t = int(1000 / fps)
while mycap.ifcontinue():
ret, img = mycap.read()
if ret:
# 开始检测图片中的人
result_img = Counting_Processing(img, yolo5_config, Model, class_names, deepsort_tracker, Obj_Counter)
# print(result_img)
if videowriter is None:
fourcc = cv2.VideoWriter_fourcc(
'm', 'p', '4', 'v') # opencv3.0
videowriter = cv2.VideoWriter(
'./output/result.mp4', fourcc, fps, (result_img.shape[1], result_img.shape[0]))
videowriter.write(result_img)
result_img = imutils.resize(result_img, height=500)
cv2.imshow('video', result_img)
cv2.waitKey(t)
if cv2.getWindowProperty('video', cv2.WND_PROP_AUTOSIZE) < 1:
# 点x退出
break
sys.stdout.write("\r=> processing at %d; total: %d" % (mycap.get_index(), total_num))
sys.stdout.flush()
videowriter.release()
cv2.destroyAllWindows()
mycap.release()
print("\n=> process done {}/{} images, total cost: {:.2f}s [{:.2f} fps]".format(len(os.listdir(yolo5_config.output)),total_num, time.time() - c,len(os.listdir(yolo5_config.output)) / (time.time() - c)))
print("=> main task finished: {}".format(datetime.now().strftime('%H:%M:%S')))
项目入口：主要的外部输入的参数，defalt为默认参数
if __name__ == "__main__":
torch.multiprocessing.set_start_method('spawn')
parser = argparse.ArgumentParser()
# 视频的路径，默认是本项目中的一个测试视频test.mp4，可自行更改
parser.add_argument('--input', type=str, default="./test.mp4",
help='test imgs folder or video or camera') # 输入'0'表示调用电脑默认摄像头
# 处理后视频的输出路径
parser.add_argument('--output', type=str, default="./output",
help='folder to save result imgs, can not use input folder')
parser.add_argument('--weights', type=str, default='weights/yolov5l.pt', help='model.pt path(s)')
parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf_thres', type=float, default=0.4, help='object confidence threshold')
parser.add_argument('--iou_thres', type=float, default=0.4, help='IOU threshold for NMS')
# GPU（0表示设备的默认的显卡）或CPU
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
# 通过classes来过滤yolo要检测类别，0表示检测人，1表示自行车，更多具体类别数字可以在19行附近打印出来
parser.add_argument('--classes', default=0, type=int, help='filter by class: --class 0, or --class 0 1 2 3')
yolo5_config = parser.parse_args()
print(yolo5_config)
main(yolo5_config)
print("结果保存在：", yolo5_config.output)
效果：有的时候街道上服装店里的假人也会识别为一个人，这yolo就没办法了，毕竟实在是太相似了
