【深度学习|目标跟踪】SSD+Sort实现多目标跟踪(Multiple Objections Track)!

源码地址

⚡️⚡️⚡️Github: SSD+Sort⚡️⚡️⚡️

1、🙌🏻匈牙利匹配算法

   关于匈牙利匹配算法,网上有十分多的教程,讲的也都很细致,这里我精简一下描述。
   参考文献:
   匹配算法之 匈牙利算法详解
   yolo目标追踪:卡尔曼滤波 + 匈牙利匹配算法 + deepsort算法
   多目标跟踪算法中之图匹配——匈牙利算法和KM算法详解
   超详细!图解匈牙利匹配算法

1.1 什么是匈牙利匹配

  匈牙利匹配算法,就是图论中寻找最大匹配的算法,即要保证匹配到最多的边,并且保证这些边的加权和最小。
  匈牙利匹配算法,用于解决与二分图匹配有关的问题。
  其中,匈牙利匹配算法分为不带权的二分图匹配问题(Hungarian algorithm)和带权重的二分图匹配问题(Kuhn-Munkres algorithm)。
  后面会具体的讲解什么叫匹配到最多的边,加权和又是指什么,什么是二分图这些问题。

1.2 什么是二分图:

  二分图是一种特殊的图,它是由两个部分组成的,各部分内部没有直接关系并不相连,如下图:
在这里插入图片描述
  二分图中的匹配讲的就是两个分布之间的连线,并且该连线的两端不能拥有共同的顶点,如上图,这就不是一个二分图的最终匹配,因为其中有很多的连线具有公共的端点,如(X1,Y2)和(X1,Y4)就是一组冲突的连线,我们可以通过去除一些公共端点的连线来使其变成匹配的二分图。下图这样的二分图则可以被称为匹配:
在这里插入图片描述
  这样的即可以被称作是该图的一个匹配
匹配有以下属性:

  • 匹配是由一组没有公共断电的不是圈的边构成的集合。
  • 匹配是边的集合。
  • 在该集合中,任意两条边不能有共同的顶点。

1.3 最大匹配

  二分图中的某一个匹配中含有最多的边数,我们即称之为是最大的匹配。(注意,一个二分图的最大匹配不一定只有一个!)

1.4 最优匹配

  二分图的某一个匹配是最大匹配,并且匹配中的每条边的权重相加求和是其他最大匹配中的最小值,即称这个最大匹配是最优匹配。

1.5 最小点覆盖

  最小顶点覆盖是指最少的顶点数使得二分图中的每条边都至少与其中一个点相关联,二分图的最小顶点覆盖=二分图的最大匹配数。

1.6 交替路

  从未匹配的点出发,依次经过未匹配的边和已匹配的边,即为交替路。如下图从2出发:2-5-1-7-4-8。
在这里插入图片描述
这里橘色的线我们认为是已经匹配的,黑色的连线认为是还每匹配的,我们可以看到2是未匹配的点,出发之后(2,5)是一个未匹配的边,(5,1)是匹配的边,(1,7)是未匹配的边,(7,4)是匹配的边,(4,8)是未匹配的边。

1.7 增广路

  如果交替路经过除出发点外的另一个为匹配点,则这条交替路被成为增广路,根据fig.3来看,2是一个出发的未匹配点,经过的5174都是已经匹配的点,直到8是一个未匹配的点。
观察这中现象,我们可以总结出以下几条规律:

  • 增广路径长度一定是一个奇数,因为最后会达到一个未匹配的点。
  • 对于增广路径,所有的奇数边都不在匹配的集合中,偶数边在匹配的集合中。
  • 当且仅当不存在关于该二分图的增广路径时,这个二分图此时达到最大匹配。

  最后我们可以将匈牙利算法总结成不断寻找增广路径来增加匹配个数的算法。

1.8 匈牙利匹配具体流程以及实例

  匈牙利匹配分为广度优先深度优先两种匹配原则。具体这两者有什么优劣我还没有具体了解过,我感觉就是一种方法的选择吧。

1.9 广度优先匹配(Hungarian algorithm)

  广度优先匹配的原则即:为先匹配上的一对保留原先配对,后面的竞争者想匹配上之前的已经匹配好的对象是不可能的,只有找它其他的能匹配且还未匹配好的对象。(这是我自己总结的大白话)我们用下图这个例子来举例说明:
按照从上至下的顺序来匹配,首先是A,有两个可以匹配的对象a和b,我们先让A与a匹配上。然后是B,B可以匹配的对象有两个,分别是a和b,但是a已经被A匹配上了,因此我们只能让B去匹配其他的对象b,最后C的话直接和c匹配即可。

在这里插入图片描述

1.10 深度优先匹配(Hungarian algorithm)

  深度优先匹配得原则即:每个点从另一个集合里挑选对象,没冲突的话就先安排上,要是冲突了,就找增广路径重新匹配。重复上述操作,直到所有的点找到对象,或者找不到对象也没有增广路了。 用下图的这个例子来说明:

在这里插入图片描述
按照这个逻辑,最后的匹配结果(绿色的阴影):
在这里插入图片描述

1.11 KM算法示例(实际应用中基本都是带加权的匹配)

  假设我们这里以检测到的汽车为例,ABC表示前一帧检测到的,abc表示当前帧检测到的,每个数表示两帧目标之间的距离:
在这里插入图片描述
在这里插入图片描述

1.12 实际使用方法(python)

  在实际使用中,我们首先需要确定的是两个分布之间的代价矩阵cost matrix,也就是这里面的数是怎么来的,通常情况下会根据任务的不同定制不同的计算方式,像在目标检测中使用目标跟踪的话可以使用IOU的计算方式,也可以使用欧氏距离(检测框质心的距离)等。
  如果是使用python来完成匈牙利匹配的任务的话,我们可以通过直接调用第三方库的方式来完成,如scipy库中的scipy.optimize.linear_sum_assignment 函数即可实现。输入一个代价矩阵,返回的是两个分布的对应索引。
  c++的实现可以参考我上传的github项目中的Kuhn_Munkres文件夹。

2、🚀卡尔曼滤波

此部分参考我的博客:【深度学习|目标跟踪】快速入门卡尔曼滤波!

3、⚡️Sort算法流程

  在上文中我们已经学习了匈牙利匹配算法(无权以及带权的),带权的匈牙利匹配算法的代价矩阵需要根据不同场景来使用不同的计算方式。下面我们来描述一下Sort算法在与SSD结合时的使用流程:

  • 1、 使用卡尔曼滤波,根据当前帧的检测结果来预测下一帧的检测框的位置。
  • 2、 然后根据下一帧的检测结果作为x(或者y,二分图的一边),预测框的集合作为二分图的y,然后使用IOU的计算方式,分别求得x与y两两检测框之间的IOU作为代价矩阵。
  • 3、 然后使用匈牙利匹配算法来得到x,y中配对的索引列表。(在python中我们可以调scipy.optimize.linear_sum_assignment()来进行匹配)。
  • 4、 这样就可以得到一一对应的当前帧检测框和当前帧的预测框的匹配关系了。
  • 5、 然后遍历所有的当前帧检测框,如果遍历到的检测框索引不在匹配的索引列表中,那么就将这个检测结果加入到未匹配的结果列表中。
  • 6、 然后遍历所有的当前帧预测框,如果遍历到的预测框索引不在匹配的索引列表中,那么就将这个结果加入到未匹配的结果列表中。
  • 7、 最后将对应的框两两求解IOU,并判断是否小于IOU_threshold,小于的化就将检测框和预测框分别按照5,6的步骤将其放入对应的未匹配列表中。如果大于IOU_threshold,则放入匹配的列表中。
  • 8、 最后则是更新阶段:对于匹配到的检测,我们遍历并用对应的当前检测来矫正(卡尔曼滤波),并保持他们的索引。对于未匹配到的检测,遍历并用对应的当前检测来创建新的跟踪对象(新的索引)。

参考以下图解以及注释:
在这里插入图片描述
在这里插入图片描述

4、🌈Sort算法代码详解

4.1 iou计算

  这一块很简单,就是传入两个检测框,格式都为x1,y1,x2,y2。这里会从第二帧开始检测的时候,将kalman filter在上一帧产生的预测框和当前帧的检测框进行一个IOU计算来得到匈牙利匹配算法所需要的代价矩阵。

"""
计算两个框[x1,y1,x2,y2]之间的IOU值.
"""
def iou_batch(bb_test, bb_gt):
    """
    From SORT: Computes IOU between two bboxes in the form [x1,y1,x2,y2]
    """
    bb_gt = np.expand_dims(bb_gt, 0)
    bb_test = np.expand_dims(bb_test, 1)

    my_bb = bb_test[..., 1]
    my_gt = bb_gt[..., 1]

    xx1 = np.maximum(bb_test[..., 0], bb_gt[..., 0])
    yy1 = np.maximum(bb_test[..., 1], bb_gt[..., 1])
    xx2 = np.minimum(bb_test[..., 2], bb_gt[..., 2])
    yy2 = np.minimum(bb_test[..., 3], bb_gt[..., 3])
    w = np.maximum(0., xx2 - xx1)
    h = np.maximum(0., yy2 - yy1)
    wh = w * h
    o = wh / ((bb_test[..., 2] - bb_test[..., 0]) * (bb_test[..., 3] - bb_test[..., 1])
              + (bb_gt[..., 2] - bb_gt[..., 0]) * (bb_gt[..., 3] - bb_gt[..., 1]) - wh)
    return (o)

4.2 convert_bbox_to_z

  这一步操作主要是为了将目标检测的检测框的坐标x1, y1, x2, y2 转成kalman filter需要的状态变量,即x,y,s,r的格式,其中x,y对应检测框的中心坐标,s对应检测框的面积,r对应检测框的纵横比。

"""
将[x1,y1,x2,y2]转成[x,y,s,r]的形式,其中x,y是框中心点的坐标,s是框的面积,r是框的宽高比例。
"""
def convert_bbox_to_z(bbox):
    """
    Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form
      [x,y,s,r] where x,y is the centre of the box and s is the scale/area and r is
      the aspect ratio
    """
    w = bbox[2] - bbox[0]
    h = bbox[3] - bbox[1]
    x = bbox[0] + w / 2.
    y = bbox[1] + h / 2.
    s = w * h  # scale is just area
    r = w / float(h)
    return np.array([x, y, s, r]).reshape((4, 1))

4.3 convert_x_to_bbox

  和4.2的相对应,就是将kalman filter预测完的检测框转化到x1,y1,x2,y2的格式上,以便计算IOU等。

"""
将[x,y,s,r]转成[x1,y1,x2,y2]的形式
"""
def convert_x_to_bbox(x, score=None):
    """
    Takes a bounding box in the centre form [x,y,s,r] and returns it in the form
      [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right
    """
    w = np.sqrt(x[2] * x[3])
    h = x[2] / w
    if (score == None):
        return np.array([x[0] - w / 2., x[1] - h / 2., x[0] + w / 2., x[1] + h / 2.]).reshape((1, 4))
    else:
        return np.array([x[0] - w / 2., x[1] - h / 2., x[0] + w / 2., x[1] + h / 2., score]).reshape((1, 5))

4.4 kalman filter预测检测框

  主要分为了四个部分,首先就是初始化的部分,初始化阶段调用filter.kalman库中的kalmanFilter来实例化一个kf属性,然后给这个属性赋初值,这些初值是在KalmanFilter中的超参数,包括了状态转移矩阵,观测矩阵,状态量,状态噪声,过程噪声等。状态转移矩阵和观测矩阵可以通过对系统进行运动学建模来获得,其他的超参数可以参考kalman filter超参数调节的经验。然后就是更新的部分,核心就是将获得的检测框坐标通过convert_bbox_to_z转成kaiman filter能够进行运算的状态变量x,然后调用KalmanFIlter中的update函数来进行状态量的更新。获取上一帧估计,调用KalmanFilter的predict()函数来进行下一帧的状态估计,最后返回的是x1,y1,x2,y2这样的检测框的格式。获取当前帧估计,即返回当前帧的估计结果。

class KalmanBoxTracker(object):
    """
    This class represents the internal state of individual tracked objects observed as bbox.
    """
    count = 0

    def __init__(self, bbox):
        """
        Initialises a tracker using initial bounding box.
        """
        # define constant velocity model
        """创建卡尔曼滤波器是需要设置状态向量和观测向量的维度"""
        # dim_x: 状态值是一个7维向量[u, v, s, r, u', v', s']
        # dim_z: 测量值是一个4维向量[u, v, s, r]
        self.kf = KalmanFilter(dim_x=7, dim_z=4)
        """状态转移矩阵"""
        self.kf.F = np.array(
            [[1, 0, 0, 0, 1, 0, 0], [0, 1, 0, 0, 0, 1, 0], [0, 0, 1, 0, 0, 0, 1], [0, 0, 0, 1, 0, 0, 0],
             [0, 0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 0, 1]])
        """观测矩阵"""
        self.kf.H = np.array(
            [[1, 0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0, 0]])

        self.kf.R[2:, 2:] *= 10.
        self.kf.P[4:, 4:] *= 1000.  # give high uncertainty to the unobservable initial velocities
        self.kf.P *= 10.
        self.kf.Q[-1, -1] *= 0.01
        self.kf.Q[4:, 4:] *= 0.01

        self.kf.x[:4] = convert_bbox_to_z(bbox)
        self.time_since_update = 0
        self.id = KalmanBoxTracker.count
        KalmanBoxTracker.count += 1
        self.history = []
        self.hits = 0
        self.hit_streak = 0
        self.age = 0

    def update(self, bbox):
        """
        Updates the state vector with observed bbox.
        """
        self.time_since_update = 0
        self.history = []
        self.hits += 1
        self.hit_streak += 1
        self.kf.update(convert_bbox_to_z(bbox))

    def predict(self):
        """
        Advances the state vector and returns the predicted bounding box estimate.
        """
        if ((self.kf.x[6] + self.kf.x[2]) <= 0):
            self.kf.x[6] *= 0.0
        self.kf.predict()
        self.age += 1
        if (self.time_since_update > 0):
            self.hit_streak = 0
        self.time_since_update += 1
        self.history.append(convert_x_to_bbox(self.kf.x))
        return self.history[-1]

    def get_state(self):
        """
        Returns the current bounding box estimate.
        """
        return convert_x_to_bbox(self.kf.x)

4.5 associate_detections_to_trackers

  这一块代码就是将检测框和跟踪器进行一个关联,最终返回三个内容,匹配的索引,未匹配的检测结果,未匹配的跟踪目标(这里的跟踪目标都会以id来指代)。

def associate_detections_to_trackers(detections, trackers, iou_threshold=0.3):
    """
    Assigns detections to tracked object (both represented as bounding boxes)

    Returns 3 lists of matches, unmatched_detections and unmatched_trackers
    """
    if (len(trackers) == 0):
        return np.empty((0, 2), dtype=int), np.arange(len(detections)), np.empty((0, 5), dtype=int)

    iou_matrix = iou_batch(detections, trackers)

    if min(iou_matrix.shape) > 0:
        a = (iou_matrix > iou_threshold).astype(np.int32)
        test = a.sum(0).max()
        if a.sum(1).max() == 1 and a.sum(0).max() == 1:
            matched_indices = np.stack(np.where(a), axis=1)
        else:
            matched_indices = linear_assignment(-iou_matrix)
    else:
        matched_indices = np.empty(shape=(0, 2))

    unmatched_detections = []
    for d, det in enumerate(detections):
        if (d not in matched_indices[:, 0]):
            unmatched_detections.append(d)
    unmatched_trackers = []
    for t, trk in enumerate(trackers):
        if (t not in matched_indices[:, 1]):
            unmatched_trackers.append(t)

    # filter out matched with low IOU
    matches = []
    for m in matched_indices:
        if (iou_matrix[m[0], m[1]] < iou_threshold):
            unmatched_detections.append(m[0])
            unmatched_trackers.append(m[1])
        else:
            matches.append(m.reshape(1, 2))
    if (len(matches) == 0):
        matches = np.empty((0, 2), dtype=int)
    else:
        matches = np.concatenate(matches, axis=0)

    return matches, np.array(unmatched_detections), np.array(unmatched_trackers)

4.6 sort

  这一部分是sort类的初始化构造函数,主要参数包括了max_age,即最大失联数,跟踪上的目标在max_age帧没关联上之后就删除该跟踪对象。min_hits,即连续min_hits帧跟踪到该目标之后就为其创建一个跟踪对象。iou_threshold作为关联检测框和预测框的阈值。self.trackers是一个存放跟踪对象的数组。self.frame_count是当前跟踪任务下的帧数计数参数。

class Sort(object):
    def __init__(self, max_age=1, min_hits=3, iou_threshold=0.3):
        """
        Sets key parameters for SORT
        """
        #--------------------------
        #   max age表示的是最大失联数,即一个目标超过这么帧没有匹配到的话就移除跟踪
        #   min hit表示的是最小命中数,即为了确认一个目标是否应该作为一个新的跟踪对象时的最小命中数
        #--------------------------
        self.max_age = max_age
        self.min_hits = min_hits
        self.iou_threshold = iou_threshold
        self.trackers = []
        self.frame_count = 0

  sort的更新部分,具体说明都在下面的代码注释中:

    def update(self, dets=np.empty((0, 5))):
        """
        Params:
          dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]
        Requires: this method must be called once for each frame even with empty detections (use np.empty((0, 5)) for frames without detections).
        Returns the a similar array, where the last column is the object ID.

        NOTE: The number of objects returned may differ from the number of detections provided.
        """
        self.frame_count += 1
        # get predicted locations from existing trackers.
        trks = np.zeros((len(self.trackers), 5))
        to_del = []
        ret = []
        # ----------------------------------------
        #   在第一帧时 由于trks的长度为0,所以不会进入这个for循环(还没有开始跟踪)
        #   从第二帧开始跟踪,进入下面的循环,会在循环中对每一个检测到的目标(第一帧)或者上一帧保留的跟踪对象(从第二帧开始)进行预测,然后将预测的状态量给到trk中
        # ----------------------------------------
        for t, trk in enumerate(trks):
            pos = self.trackers[t].predict()[0]
            trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
            #----------------------------------------
            #   将pos中含有nan值的添加进delete数组中,准备后面删除
            #   np.isnan()逐元素测试数组中的值是否为 NaN(Not a Number),并返回一个布尔数组
            #   np.any()用于测试数组中的元素是否有true,有的话则返回一个bool值 true
            #----------------------------------------
            if np.any(np.isnan(pos)):
                to_del.append(t)
        #----------------------------------------
        #   np.ma.masked_invalid()用于将数组中的NaN和inf值屏蔽掉,并返回一个masked_array对象
        #   np.ma.compress_rows()用于压缩二维数组中包含掩码值得行,换句话说,它会删除包含掩码值的所有行
        #----------------------------------------
        trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
        for t in reversed(to_del):
            self.trackers.pop(t)
        #----------------------------------------
        #   输入当前帧的检测结果与预测结果
        #----------------------------------------
        matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(dets, trks, self.iou_threshold)

        # update matched trackers with assigned detections
        #----------------------------------------
        #   匹配上的结果会在这里update,此时self.time_since_update会被更新成0,意味着没有丢失跟踪
        #   而丢失跟踪的目标此时会由于上面代码的predict而导致self.time_since_update+=1,意味着丢失跟踪了一帧,累加起来的就是自上一次更新之后未匹配到的帧数
        #----------------------------------------
        for m in matched:
            self.trackers[m[1]].update(dets[m[0], :])

        # create and initialise new trackers for unmatched detections
        #----------------------------------------
        #   为没有匹配到的检测创建新的跟踪实例,添加到trackers数组中
        #----------------------------------------
        for i in unmatched_dets:
            trk = KalmanBoxTracker(dets[i, :])
            self.trackers.append(trk)
        i = len(self.trackers)
        for trk in reversed(self.trackers):
            d = trk.get_state()[0]
            if (trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits):
                # ----------------------------------------
                #   给检测结果(kalman结合观测与先验估计得到的)加上id索引
                # ----------------------------------------
                ret.append(np.concatenate((d, [trk.id + 1])).reshape(1, -1))  # +1 as MOT benchmark requires positive
            i -= 1
            # remove dead tracklet
            if (trk.time_since_update > self.max_age):
                self.trackers.pop(i)
        if (len(ret) > 0):
            return np.concatenate(ret)
        return np.empty((0, 5))

5、💫python+openvino推理代码

import numpy as np
import colorsys
import torch
import sort
import cv2
import time

from PIL import Image
from PIL import Image, ImageDraw, ImageFont
from utils.anchors import get_anchors
from utils.utils import (cvtColor, get_classes, preprocess_input, resize_image,
                         show_config)
from utils.utils_bbox import BBoxUtility
from openvino.runtime import Model, Core

class OV_Infer():
    def __init__(self, **kwargs):
        self._defaults = {
            #-------------------------------
            #   模型路径
            #-------------------------------
            "model_path" : "ssd.onnx",

            #-------------------------------
            #   模型类别
            #-------------------------------
            "classes_path" : "model_data/voc_classes.txt",

            #-------------------------------
            #   是否开启sort目标跟踪
            #-------------------------------
            "track" : True,

            # -------------------------------
            #   backbone的输入尺寸
            # -------------------------------
            "input_shape" : [300, 300],

            #-------------------------------
            #   先验框的尺寸
            #-------------------------------
            'anchors_size' : [30, 60, 111, 162, 213, 264, 315],
            # -------------------------------
            #   是否对输入图像进行分辨率的保持(不失真)
            # -------------------------------
            'letterbox' : False,

        }
        # -------------------------------
        #   遍历设置好的字典的属性给到类属性  self._defaults--->self
        # -------------------------------
        self.__dict__.update(self._defaults)
        for name, value in kwargs.items():
            setattr(self, name, value)

        # -------------------------------
        #   设置类别
        # -------------------------------
        self.class_names, self.class_nums = get_classes(self.classes_path)

        # -------------------------------
        #   设置anchors,backbone取决于训练时的选择,选择了vgg那就“vgg”,选择了resnet18那就”resnet18“
        #   解析推理的结果时要用,因为ssd的output是根据先验框的相对位移来解码的
        # -------------------------------
        self.anchors = torch.from_numpy(get_anchors(input_shape=self.input_shape, anchors_size=self.anchors_size, backbone="vgg")).type(torch.FloatTensor)

        # if self.cuda:
        #     self.anchors = self.anchors.cuda()
        self.class_nums = self.class_nums + 1

        # -------------------------------
        #   设置不同框的颜色
        # -------------------------------
        hsv_tuples  = [(x / self.class_nums, 1., 1.) for x in range(self.class_nums)]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))

        self.box_util      = BBoxUtility(self.class_nums)
        self.Initial_model()

    def Initial_model(self):
        # ------------------------------
        #   实例化一个core,用于openvino推理
        # ------------------------------
        self.core = Core()

        # ------------------------------
        #   实例化一个ov的model对象
        # ------------------------------
        self.model = self.core.read_model(self.model_path)

        # ------------------------------
        #   实例化一个ov的编译模型
        # ------------------------------
        self.compile_model = self.core.compile_model(self.model, 'CPU')

        # ------------------------------
        #   如果开启了sort跟踪算法,实例化一个track对象
        # ------------------------------
        if self.track:
            self.mot_tracker = sort.Sort(max_age=1, min_hits=3, iou_threshold=0.3)
        else:
            self.mot_tracker = None

    def preprocess(self, image):
        # ------------------------------
        #   获取图像的shape
        # ------------------------------
        self.image_shape = np.array(np.shape(image)[0:2])

        # ------------------------------
        #   将输入图像转成RGB格式
        # ------------------------------
        image_data = cvtColor(image)

        # ------------------------------
        #   将不同尺寸的输入图像resize到网络的输入尺寸,即300 * 300
        # ------------------------------
        image_data = resize_image(image_data, (self.input_shape[1], self.input_shape[0]), self.letterbox)

        # ------------------------------
        #   图像像素减去均值
        # ------------------------------
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2,0,1)), 0)

        return image_data

    def detect(self, image, crop = False, count = False):
        # ------------------------------
        #   前处理
        # ------------------------------
        preprocessed_input = self.preprocess(image)

        # ------------------------------
        #   转化成torch的形式
        # ------------------------------
        preprocessed_input = torch.from_numpy(preprocessed_input).type(torch.FloatTensor)
        # if self.cuda:
        #     preprocessed_input = preprocessed_input.cuda()

        # ------------------------------
        #   前向传播
        # ------------------------------
        ov_outputs = self.compile_model(preprocessed_input)
        keys = ov_outputs.keys()
        ov_det, ov_cls = 0, 0
        for i, key in enumerate(keys):
            if i == 0:
                ov_det     = torch.from_numpy(ov_outputs[key])
            else:
                ov_cls     = torch.from_numpy(ov_outputs[key])
        outputs = (ov_det, ov_cls)
        results = self.box_util.decode_box(outputs, self.anchors, self.image_shape, self.input_shape,
                                           self.letterbox, nms_iou=0.45, confidence=0.5)

        # ------------------------------
        #   如果没有检测到物体 返回原图
        # ------------------------------
        if len(results[0]) <= 0:
            return image

        top_label = np.array(results[0][:,4], dtype='int32')
        top_conf  = results[0][:, 5]
        top_boxes = results[0][:, :4]

        # ---------------------------------------------------------#
        #   sort需要:torch.tensor的切片操作:跳过中间的元素,即不要类别索引
        # ---------------------------------------------------------#
        indices = torch.tensor([0, 1, 2, 3, 5])
        top_sort = results[0][:, indices]

        # ---------------------------------------------------------#
        #   如果开启了sort跟踪
        # ---------------------------------------------------------#
        if self.mot_tracker != None and self.track == True:
            trackers         = self.mot_tracker.update(top_sort)

        #---------------------------------------------------------#
        #   设置字体与边框厚度
        #---------------------------------------------------------#
        font = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32'))
        thickness = max((np.shape(image)[0] + np.shape(image)[1]) // self.input_shape[0], 1)

        # ---------------------------------------------------------#
        #   图像绘制
        # ---------------------------------------------------------#
        for i, c in list(enumerate(top_label)):
            predicted_class = self.class_names[top_label[len(top_label) - i - 1]]
            if self.mot_tracker == None:
                box = top_boxes[i]
                score = top_conf[i]

            '''
            使用跟踪结果
            '''
            if self.mot_tracker != None:
                if i >= len(trackers):
                    continue
                box             = trackers[i][:4]
                score           = top_conf[i]

            top, left, bottom, right = box

            top = max(0, np.floor(top).astype('int32'))
            left = max(0, np.floor(left).astype('int32'))
            bottom = min(image.size[1], np.floor(bottom).astype('int32'))
            right = min(image.size[0], np.floor(right).astype('int32'))

            if self.mot_tracker != None:
                label = '{} {:.2f} {}'.format(predicted_class, score, trackers[i][4])
            else:
                label = '{} {:.2f}'.format(predicted_class, score)
            draw = ImageDraw.Draw(image)
            label_size = draw.textbbox((0, 0), label, font)
            label = label.encode('utf-8')
            print(label, top, left, bottom, right)

            if top - label_size[1] >= 0:
                text_origin = np.array([left, top - label_size[1]])
            else:
                text_origin = np.array([left, top + 1])

            for i in range(thickness):
                draw.rectangle([left + i, top + i, right - i, bottom - i], outline=self.colors[c])
            draw.rectangle([tuple(text_origin), tuple(text_origin)], fill=self.colors[c])
            draw.text(text_origin, str(label, 'UTF-8'), fill=(0, 0, 0), font=font)
            del draw

        return image

if __name__ == "__main__":
    video_path = 0
    capture = cv2.VideoCapture(video_path)
    ref, _ = capture.read()
    if not ref:
        raise ValueError("video path is error!")

    fps = 0.0
    # ------------------------------
    #   初始化ov推理
    # ------------------------------
    ov_infer = OV_Infer()

    while True:
        t1 = time.time()

        ref, frame = capture.read()
        if not ref:
            break

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        frame = Image.fromarray(np.uint8(frame))

        # ------------------------------
        #   开始检测,将PIL的格式转成np格式给cv展示
        # ------------------------------
        frame = np.array(ov_infer.detect(frame))
        # ------------------------------
        #   调整一下通道到bgr给cv
        # ------------------------------
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        # ------------------------------
        #   计算fps并且显示在窗口上
        # ------------------------------
        fps = (fps + (1. / (time.time() - t1))) / 2
        print("fps= %.2f" % (fps))
        frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        cv2.imshow("video", frame)
        '''
        & 0xff是做了一个二进制的与操作,0xff是一个十六进制的数,对应的二进制是11111111。
        这个操作的目的是只保留返回值的最后八位,因为在某些系统中,cv2.waitKey的返回值不止8位。
        '''
        c = cv2.waitKey(1) & 0xff

        '''
        ASCII码27对应的是ESC按键
        '''
        if c == 27:
            capture.release()
            break

    print("Video Detection Done!")
    capture.release()
    cv2.destroyAllWindows()
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

澄鑫

谢谢,将继续努力提供技术方案

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值