yolov5 deepsort Pedestrian tracking and counting system based on yolov5 and deepsort [yolov5] [deepsort] [pedestrian detection] [tracking count]

introduce

This project uses yolov5 as a detector and deepsort as a tracker to track and count the number of pedestrians walking in front of the camera.
The setting in the code is to only count the number of pedestrians. If you want to count other things such as vehicles, it is also very simple. You only need to modify the code a little bit. In , filter the categories to be detected by yolov5 through classes, and 0 means to detect parser.add_argument('--classes', default=0, type=int, help='filter by class: --class 0, or --class 0 1 2 3')people , 1 means bicycle, and there are other categories.

environment

Project test environment: win10, pycharm, python3.6+
Main required packages: pytorch 1.7.1 + cuda110, opencv

Project source code (github)

https://github.com/MichistaLin/yolov5-deepsort-pedestraintracking

part of the code

Guide package

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)

Detection and counting: This function uses yolov5 for detection and deepsort for tracking

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 function: initialize the model, load videos or pictures, process and detect and track people in videos and pictures

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')))

Project entry: the main external input parameters, default is the default parameter

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)

Effect: Sometimes the dummy in the clothing store on the street will also be recognized as a person, this yolo has no choice, after all, it is too similar

virtual environment

I added a virtual environment for the project, but the virtual environment is too big so I didn’t upload it to github, but put it in the 123 network disk, and the link is in the file of github. If you don’t want to download pytorch yourself, you can use README.mdit The virtual environment I provided runs the project directly, but transplanting the virtual environment requires some changes, as shown in the figure below:

Find the pyvenv.cfg file in the virtual environment, open it in Pycharm or Notepad, and change the home path to your own system In the directory where the Python interpreter is located, and then select the python interpreter in the virtual environment in the interpreter settings.

The pytorch version in the virtual environment is torch 1.7.1 + cu110
If the virtual environment does not work, you can clear the cache and restart the IDE to try