yolov5 deepsort 基于yolov5和deepsort的行人跟踪计数系统【yolov5】【deepsort】【行人检测】【跟踪计数】_再游于北方知寒

大大的周 02-07 7421

行人跟踪计数系统介绍环境项目源码（github)部分代码

介绍

本项目使用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就没办法了，毕竟实在是太相似了