【目标跟踪】pytorch实现DeepSORT+YOLOV5 YOLOFastestv2 含代码

目录

系列文章

一、非常简短的介绍

二、极其方便的上手

1.项目结构

2.执行demo

3.修改前置物体检测算法和特征提取模型 

4.修改deep_sort相关配置

三、明了清晰的代码

1.物体检测

2.提取特征

3.卡尔曼滤波predict

4.执行Matching

5.卡尔曼滤波update

系列文章

【目标跟踪】卡尔曼滤波器(Kalman Filter) 含源码

【目标跟踪】一图看懂DeepSORT大流程

【目标跟踪】pytorch YOLOV5 YOLOFastestv2 DeepSORT

一、非常简短的介绍

        项目链接：GitHub - oaifaye/dcmtracking

        效果演示：【目标跟踪】Pytorch实现YOLOV5+DeepSORT】 

        dcmtracking（dreams create miracles），中文：大聪明跟踪工具包。该项目构建的目的是集成当今SOTA的Tracking算法，提供算法工具箱，给出各种算法的实验数据，给算法落地带来便利。项目本着方便开发者的目的，开箱即用，直接将dcmtracking目录考到项目中，实例化一个类，然后调用即可。

        该项目现在实现了基于pytorch的YOLOV5+DeepSORT和YOLOFastestv2+DeepSORT，将持续更新，欢迎关注。

二、极其方便的上手

1.项目结构

dcmtracking                                    项目主目录

        — dcmtracking                       实现该项目所有核心功能，移植的时候直接考这个目录就可以

                — deep_sort                   deep_sort 的主目录

                        + deep                     图像提取特征的功能

                        + model_data          存放模型文件

                        + sort                      卡尔曼滤波等算法

                        + tracker                  暴露一些接口和实现类，供开发者使用

                        deep_sort.py           实现了deep_sort

                        deep_sort.yaml        deep_sort的配置文件

                — detection                     存放物体检测算法

                        yolo_fastestv2        yolo_fastestv2的主目录   

                        yolov5                     yolov5的注目录

                — utils                             工具

                demo.py                           提供可直接执行的demo

2.执行demo

        先下载一些模型问价和测试视频，网盘地址如下，目录结构已经排好，下载之后直接覆盖到项目根目录即可：

        链接：https://pan.baidu.com/s/1PjkiM2HNV20gQtCtT3oikg?pwd=902r 
        提取码：902r 

         然后，直接执行python demo.py

        如果想检测自己的视频，修改输入输出路径即可。

if __name__ == '__main__':
    # 执行yolov5s+deepsort
    demo_yolov5_deep_sort_tracker('data/test5.mp4', 'data/out5.flv')
    # 执行yolovfastestv2+deepsort
    demo_yolo_fastestv2_deep_sort_tracker('data/test3.mp4', 'data/out3_f.flv')

3.修改前置物体检测算法和特征提取模型 

        要实现跟踪功能需要集成dcmtracking.deep_sort.tracker.base_tracker.BaseTracker.py，并实现init_extractor()和detect()方法，项目中已经提供了yolo_fastestv2_deep_sort_tracker.py和yolov5_deep_sort_tracker.py两种实现，如果不能满足要求，可以自行添加新的方法。

        detect()方法：前置的物体检测方法。一般情况下需要在实现类init方法中初始化检测模型实例。项目中默认实现了yolo_fastestv2和yolov5两种物体检测算法，并提供了基于coco数据集的预训模型，如果是检测人、车等常规任务，可以直接使用。

        init_extractor()：初始化特征提取器，需要返回一个特征提取模型实例，项目中的使用一个简单的10层卷积的分类模型，推力时只取backbone，返回512的特征向量。项目提供了基于Market1501数据集的预训练模型，如果是执行人物跟踪人物，可以直接使用。

       示例代码如下：

# coding=utf-8
# ================================================================
#
#   File name   : yolov5_deep_sort_tracker.py
#   Author      : Faye
#   Created date: 2022/10/19 16:18 
#   Description : Yolov5s+deepsort
#
# ================================================================
from dcmtracking.deep_sort.tracker.base_tracker import BaseTracker
from dcmtracking.detection.yolov5.yolo import YOLO
from PIL import Image
import cv2
import torch
from dcmtracking.deep_sort.deep.feature_extractor import Extractor


class Yolov5DeepSortTracker(BaseTracker):
    def __init__(self, need_speed=False, need_angle=False):
        # 执行父类的init方法
        BaseTracker.__init__(self, need_speed=need_speed, need_angle=need_angle)
        # 初始化目标检测类
        self.yolo = YOLO()

    def init_extractor(self):
        """
        实现父类的init_extractor方法，初始化特征提取器
        Parameters
        ----------
        im

        Returns
        -------

        """
        model_path = "dcmtracking/deep_sort/deep/checkpoint/ckpt.t7"
        return Extractor(model_path, use_cuda=torch.cuda.is_available())

    def detect(self, im):
        """
        实现父类的detect方法
        Parameters
        ----------
        im

        Returns
        -------

        """
        im_h, im_w, _ = im.shape
        im_pil = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
        im_pil = Image.fromarray(im_pil)
        pred_boxes = []
        top_label, top_boxes, top_conf = self.yolo.detect_image(im_pil)
        if top_label is not None:
            for (y1, x1, y2, x2), lbl, conf in zip(top_boxes, top_label, top_conf):
                if lbl != 0:
                    continue
                pred_boxes.append(
                    (int(x1), int(y1), int(x2), int(y2), lbl, conf))
        return im, pred_boxes

4.修改deep_sort相关配置

        deep_sort的配置文件是dcmtracking/deep_sort/deep_sort.yaml，里面的配置如下：

DEEPSORT:
  # 物体匹配的阈值。距离较大的样本被认为是无效匹配。
  MAX_DIST: 0.5
  # 最小置信度，小于这个值，认为是无效物体
  MIN_CONFIDENCE: 0.3
  # 执行nms时，最大重叠占比，两个bbox的iou大于这个值，将认为是同一物体
  NMS_MAX_OVERLAP: 0.5
  # 执行IOU匹配时，大于此值的关联被忽略。
  MAX_IOU_DISTANCE: 0.7
  # 在删除track之前的最大miss数。
  MAX_AGE: 70
  # 在一个track被确认之前的连续探测次数。如果在第一个n_init帧内发生miss，则track状态被设置为' Deleted '。
  N_INIT: 3
  # 是否需要在原图上画框
  NEED_DRAW_BBOXES: False
  # 是否需要标注速度，速度单位pix/s，只有need_draw_bboxes=True时起作用
  NEED_SPEED: True
  # 是否需要标注运动方向，只有need_draw_bboxes=True时起作用
  NEED_ANGLE: True

三、明了清晰的代码

为了更好地理解代码，推荐大家先看一下下面两篇文章：

文章一：卡尔曼滤波相关： 【五分钟会，半小时懂】卡尔曼滤波器(Kalman Filter)—目标跟踪（含源码）_小殊小殊的博客

文章二：DeepSORT整个大流程：

一图看懂DeepSORT整个大流程，多目标跟踪_小殊小殊的博客

为了方便我把文章一中的公式和文章二中的流程图直接拿过来：

DeepSORT流程图

1.物体检测

        对视频每一帧执行物体检测，对应流程图中步骤3。

        代码位置：dcmtracking/deep_sort/tracker/base_tracker.py的deal_one_frame方法中,self.detect()返回图像本身和bbox信息，该方法需要在base_tracker.py的实现类中实现具体检测算法，代码如下：

def deal_one_frame(self, image, speed_skip, need_detect=True):
        """
        处理视频中的一帧
        Parameters
        ----------
        image:cv2读取的视频帧
        speed_skip:用于计算速度，一般传fps
        need_detect:知否需要执行目标检测，如果传False将使用上一次检测的结果，该参数主要用于加速

        Returns
        -------
        im:     返回画好框的图片
        ids:     这一帧出现的目标ids
        bboxes:  这一帧出现的目标框坐标,左上和右下

        """
        self.frames_count += 1
        if self.last_deepsort_outputs is None or need_detect:
            t1 = time.time()
            # 1.执行目标检测
            _, bboxes = self.detect(image)
        ......

2.提取特征

        提取bbox中图像的特征向量，并新建Detections。

        代码位置：dcmtracking/deep_sort/deep_sort.py的update方法：

def update(self, bbox_xywh, confidences, ori_img):
        """
        根据目标检测的结果，执行跟踪、更新DeepSort历史状态
        Parameters
        ----------
        extractor:图像特征提取器
        bbox_xywh:目标框的中心点和宽高
        confidences：置信度
        ori_img：图片

        Returns
        -------

        """
        self.height, self.width = ori_img.shape[:2]
        # 将图片按照bbox切割 每块生成特征向量（特征向量默认长度512）
        features = self._get_features(bbox_xywh, ori_img)
        # 将左上右下的四个坐标 转换成中心点和宽高
        bbox_tlwh = self._xywh_to_tlwh(bbox_xywh)
        # 根据features和bbox_tlwh生成detections 每个detection有features/tlwh/confidence 三个属性
        detections = [Detection(bbox_tlwh[i], conf, features[i]) for i,conf in enumerate(confidences) if conf > self.min_confidence]

        # 执行nms 去掉重复的detection 其实在目标检测阶段已经做了nms 这里不做也行
        boxes = np.array([d.tlwh for d in detections])
        scores = np.array([d.confidence for d in detections])
        indices = non_max_suppression(boxes, self.nms_max_overlap, scores)
        detections = [detections[i] for i in indices]
    ......

3.卡尔曼滤波predict

def predict(self, mean, covariance):
        """执行卡尔曼滤波的predict步骤.

        Parameters
        ----------
        mean : ndarray
            前一个轮次的物体状态的8维向量的期望（均值）。
        covariance : ndarray
            前一个轮次的物体状态的8x8维协方差矩阵

        Returns
        -------
        (ndarray, ndarray)
            返回预测状态的平均向量和协方差矩阵。未观测到的速度初始化为平均值0。

        """
        std_pos = [
            self._std_weight_position * mean[3],
            self._std_weight_position * mean[3],
            1e-2,
            self._std_weight_position * mean[3]]
        std_vel = [
            self._std_weight_velocity * mean[3],
            self._std_weight_velocity * mean[3],
            1e-5,
            self._std_weight_velocity * mean[3]]
        motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))

        mean = np.dot(self._motion_mat, mean)
        covariance = np.linalg.multi_dot((
            self._motion_mat, covariance, self._motion_mat.T)) + motion_cov

        return mean, covariance

4.执行Matching

        Matching分为两步，一个matching_cascade和iou_matching（两种匹配的具体解释请看文章），对应图中步骤4567。

        代码位置：dcmtracking/deep_sort/sort/tracker.py的_match()方法:

def _match(self, detections):

        def gated_metric(tracks, dets, track_indices, detection_indices):
            """
            基于外观信息和马氏距离，计算卡尔曼滤波预测的tracks和当前时刻检测到的detections的代价矩阵
            Parameters
            ----------
            tracks
            dets
            track_indices
            detection_indices

            Returns
            -------
            cost_matrix 代价矩阵

            """
            features = np.array([dets[i].feature for i in detection_indices])
            targets = np.array([tracks[i].track_id for i in track_indices])
            # 基于外观的特征向量，计算tracks和detections的余弦距离代价矩阵
            cost_matrix = self.metric.distance(features, targets)
            # 基于马氏距离，过滤掉代价矩阵中一些不合适的项 (将其设置为一个较大的值)
            cost_matrix = linear_assignment.gate_cost_matrix(
                self.kf, cost_matrix, tracks, dets, track_indices,
                detection_indices)

            return cost_matrix

        # 将已经存在的tracks分成已确定和未确定，感觉这里可以优化
        confirmed_tracks = [
            i for i, t in enumerate(self.tracks) if t.is_confirmed()]
        unconfirmed_tracks = [
            i for i, t in enumerate(self.tracks) if not t.is_confirmed()]

        # 对confirmd tracks进行级联匹配
        matches_a, unmatched_tracks_a, unmatched_detections = \
            linear_assignment.matching_cascade(
                gated_metric, self.metric.matching_threshold, self.max_age,
                self.tracks, detections, confirmed_tracks)

        # 对级联匹配中未匹配的tracks和unconfirmed tracks中time_since_update为1的tracks进行IOU匹配
        iou_track_candidates = unconfirmed_tracks + [
            k for k in unmatched_tracks_a if
            self.tracks[k].time_since_update == 1]
        unmatched_tracks_a = [
            k for k in unmatched_tracks_a if
            self.tracks[k].time_since_update != 1]
        matches_b, unmatched_tracks_b, unmatched_detections = \
            linear_assignment.min_cost_matching(
                iou_matching.iou_cost, self.max_iou_distance, self.tracks,
                detections, iou_track_candidates, unmatched_detections)

        # 整合所有的匹配对和未匹配的tracks
        matches = matches_a + matches_b
        unmatched_tracks = list(set(unmatched_tracks_a + unmatched_tracks_b))
        return matches, unmatched_tracks, unmatched_detections

5.卡尔曼滤波update

         执行卡尔曼滤波的update步骤，对观测值进行校正，实现公式345，分别对应下方代码注释中的123；对应流程图中的步骤11、12、13。

        代码位置：dcmtracking/deep_sort/sort/kalman_filter.py的update方法：

def update(self, mean, covariance, measurement):
        """卡尔曼滤波的update步骤，对观测值进行校正。

        Parameters
        ----------
        mean : ndarray
            预测状态的平均向量(8维)。
        covariance : ndarray
            状态的协方差矩阵(8x8维)。
        measurement : ndarray
            4维测量向量(x, y, a, h)，其中(x, y)是中心位置，a是纵横比，h是包围框的高度。

        Returns
        -------
        (ndarray, ndarray)
            返回经过测量校正的状态分布。

        """
        # 1.计算卡尔曼增益K
        projected_mean, projected_cov = self.project(mean, covariance)
        chol_factor, lower = scipy.linalg.cho_factor(
            projected_cov, lower=True, check_finite=False)
        kalman_gain = scipy.linalg.cho_solve(
            (chol_factor, lower), np.dot(covariance, self._update_mat.T).T,
            check_finite=False).T
        innovation = measurement - projected_mean
        # 2.计算当前步最优估计
        new_mean = mean + np.dot(innovation, kalman_gain.T)
        # 3.更新过程噪声协方差矩阵
        new_covariance = covariance - np.linalg.multi_dot((
            kalman_gain, projected_cov, kalman_gain.T))
        return new_mean, new_covariance

        dcmtracking项目就简单介绍到这里，该项目会持续更新，相信会越来越完善。

 关注订阅号了解更多精品文章