yolov5 ディープソート yolov5 とディープソートに基づく歩行者追跡およびカウント システム [yolov5] [ディープソート] [歩行者検出] [追跡カウント]

導入

このプロジェクトでは、yolov5 を検出器として使用し、ディープソートをトラッカーとして使用して、カメラの前を歩く歩行者の数を追跡およびカウントします。
コード内の設定は歩行者の数のみをカウントするようにしています。車両など他のものをカウントしたい場合も非常に簡単です。コードを少し変更するだけです。 で、検出するカテゴリをフィルタリングします。 yolov5 によるクラスによるもので、0 は人を検出することを意味しparser.add_argument('--classes', default=0, type=int, help='filter by class: --class 0, or --class 0 1 2 3')、1 は自転車を意味し、その他のカテゴリがあります。

環境

プロジェクトのテスト環境：win10、pycharm、python3.6+
主な必須パッケージ：pytorch 1.7.1 + cuda110、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 を使用し、追跡にディープソートを使用します。

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

主な機能: モデルの初期化、ビデオまたは写真のロード、ビデオおよび写真内の人物の処理および検出および追跡

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

プロジェクトエントリ: メインの外部入力パラメータ、デフォルトはデフォルトパラメータです

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)

効果: 時々、路上の衣料品店のダミーも人として認識されます。このヨロは選択の余地がありません。結局のところ、それはあまりにも似ています

仮想環境

プロジェクトに仮想環境を追加しましたが、この仮想環境は大きすぎるため、github にはアップロードせず、123 ネットワーク ディスクに置き、リンクは github のファイル内にあります。 pytorch を自分でダウンロードしたい場合はREADME.md、それを使用できます。私が提供した仮想環境ではプロジェクトが直接実行されますが、仮想環境を移植するには、次の図に示すようにいくつかの変更が必要です。仮想環境で pyvenv.cfg ファイルを見つけて、それを開きます

。 Pycharm または Notepad を使用し、Python インタープリターが配置されているディレクトリでホーム パスを独自のシステムに変更し、インタープリター設定で仮想環境の Python インタープリターを選択します。

仮想環境の pytorch バージョンは次のとおりです。torch 1.7.1 + cu110
仮想環境が機能しない場合は、キャッシュをクリアして IDE を再起動して試す​​ことができます。