• 应用第三方ByteTrack实现目标跟踪

    在上一篇博文，我们实现了应用官网ByteTrack实现的目标跟踪。但吹毛求疵地说，官网的ByteTrack有一些不足：1、为了通用性及科研的要求，代码过于冗长，但这不利于集成到自己的程序中；2、目标跟踪结果没有目标类别的信息，需要自己额外添加；3、再次运行算法时，ID不会复位，它会接着上次的ID继续排序。因此，我们完全有理由对其进行优化。

    在github中，有很多优化后的ByteTrack，所以我们没有必要自己去优化。在这里我们选择下面这个ByteTrack：

    https://github.com/jahongir7174/ByteTrack/

    为此，我们把核心函数整合到一个名为track.py的文件内，并把其放入到mytrack的文件夹内，然后把该文件夹复制到当前目录下。我把track.py文件内的完整内容复制到本博文的后面。

    为使用该文件，先要导入：

    from mytrack import track

    然后实例化ByteTrack：

    bytetracker = track.BYTETracker(fps)

    BYTETracker只需要一个输入参数，其表示视频的帧率，默认为30。

    使用BYTETracker也很简单：

    tracks = bytetracker.update(boxes, scores, object_classes)

    三个输入参数都为目标检测得到的每一帧的检测结果，数据类型为numpy.array。boxes为目标边框的左上角和右下角坐标，scores为目标置信值，object_classes为目标类型。

    输出参数为该帧的跟踪结果，前4个元素为目标边框的左上角和右下角坐标（其与boxes有所不同），第4个元素为跟踪ID，第5个元素为目标得分值，第6个元素为目标类型，第7个元素为boxes所对应的索引值。

    可以看出该第三方ByteTrack与官网的ByteTrack使用方法区别不大。

    为使用该ByteTrack，只需要安装以下软件包即可：

    1. conda install python=3.8
    2. pip install scipy -i https://pypi.tuna.tsinghua.edu.cn/simple
    3. conda install -c conda-forge lap

    从中可以看出，该第三方ByteTrack所需的软件包很少，安装起来也不会出现各类莫名其妙的问题。

    下面就可以编写目标跟踪代码了。在这里，我们只把if outputs is not None:内的部分展示出来，而其他部分与上一篇博文的代码是一样的：

    1. for output in outputs:
    2.     x1, y1, x2, y2 = list(map(int, output[:4]))
    3.     boxes.append([x1, y1, x2, y2])
    4.     confidences.append(output[4])
    5.     object_classes.append(output[5])
    6.          
    7.     tracks = bytetracker.update(np.array(boxes), np.array(confidences), np.array(object_classes))
    8.     
    9.     if len(tracks) > 0:
    10.         identities = tracks[:, 4]
    11.         object_classes = tracks[:, 6]
    12.         idxs = tracks[:, 7]
    13.         for i, identity in enumerate(identities):
    14.             if object_classes[i] == 2:
    15.                 box_label(frame, outputs[int(idxs[i]),:4], '#'+str(int(identity))+' car' , (167, 146, 11))
    16.             elif object_classes[i] == 5:
    17.                 box_label(frame, outputs[int(idxs[i]),:4], '#'+str(int(identity))+' bus', (186, 55, 2))
    18.             elif object_classes[i] == 7:
    19.                 box_label(frame, outputs[int(idxs[i]),:4], '#'+str(int(identity))+' truck', (19, 222, 24))

    下面是完整的track.py内容：

    1. import numpy
    2. import lap
    3. import scipy
    4.  
    5.  
    6. def linear_assignment(cost_matrix, thresh):
    7.     # Linear assignment implementations with scipy and lap.lapjv
    8.     if cost_matrix.size == 0:
    9.         matches = numpy.empty((0, 2), dtype=int)
    10.         unmatched_a = tuple(range(cost_matrix.shape[0]))
    11.         unmatched_b = tuple(range(cost_matrix.shape[1]))
    12.         return matches, unmatched_a, unmatched_b
    13.  
    14.     _, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh)
    15.     matches = [[ix, mx] for ix, mx in enumerate(x) if mx >= 0]
    16.     unmatched_a = numpy.where(x < 0)[0]
    17.     unmatched_b = numpy.where(y < 0)[0]
    18.    
    19.     return matches, unmatched_a, unmatched_b
    20.  
    21.  
    22. def compute_iou(a_boxes, b_boxes):
    23.     """
    24.     Compute cost based on IoU
    25.     :type a_boxes: list[tlbr] | np.ndarray
    26.     :type b_boxes: list[tlbr] | np.ndarray
    28.     :rtype iou | np.ndarray
    29.     """
    30.     iou = numpy.zeros((len(a_boxes), len(b_boxes)), dtype=numpy.float32)
    31.     if iou.size == 0:
    32.         return iou
    33.     a_boxes = numpy.ascontiguousarray(a_boxes, dtype=numpy.float32)
    34.     b_boxes = numpy.ascontiguousarray(b_boxes, dtype=numpy.float32)
    35.     # Get the coordinates of bounding boxes
    36.     b1_x1, b1_y1, b1_x2, b1_y2 = a_boxes.T
    37.     b2_x1, b2_y1, b2_x2, b2_y2 = b_boxes.T
    38.  
    39.     # Intersection area
    40.     inter_area = (numpy.minimum(b1_x2[:, None], b2_x2) - numpy.maximum(b1_x1[:, None], b2_x1)).clip(0) * \
    41.                  (numpy.minimum(b1_y2[:, None], b2_y2) - numpy.maximum(b1_y1[:, None], b2_y1)).clip(0)
    42.  
    43.     # box2 area
    44.     box1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    45.     box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
    46.     return inter_area / (box2_area + box1_area[:, None] - inter_area + 1E-7)
    47.  
    48.  
    49. def iou_distance(a_tracks, b_tracks):
    50.     """
    51.     Compute cost based on IoU
    52.     :type a_tracks: list[STrack]
    53.     :type b_tracks: list[STrack]
    55.     :rtype cost_matrix np.ndarray
    56.     """
    57.  
    58.     if (len(a_tracks) > 0 and isinstance(a_tracks[0], numpy.ndarray)) \
    59.             or (len(b_tracks) > 0 and isinstance(b_tracks[0], numpy.ndarray)):
    60.         a_boxes = a_tracks
    61.         b_boxes = b_tracks
    62.     else:
    63.         a_boxes = [track.tlbr for track in a_tracks]
    64.         b_boxes = [track.tlbr for track in b_tracks]
    65.     return 1 - compute_iou(a_boxes, b_boxes)  # cost matrix
    66.  
    67.  
    68. def fuse_score(cost_matrix, detections):
    69.     if cost_matrix.size == 0:
    70.         return cost_matrix
    71.     iou_sim = 1 - cost_matrix
    72.     det_scores = numpy.array([det.score for det in detections])
    73.     det_scores = numpy.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
    74.     fuse_sim = iou_sim * det_scores
    75.     return 1 - fuse_sim  # fuse_cost
    76.  
    77.  
    78.  
    79.  
    80. class KalmanFilterXYAH:
    81.     """
    82.     A Kalman filter for tracking bounding boxes in image space.
    84.     The 8-dimensional state space
    86.         x, y, a, h, vx, vy, va, vh
    88.     contains the bounding box center position (x, y), aspect ratio a, height h,
    89.     and their respective velocities.
    91.     Object motion follows a constant velocity model. The bounding box location
    92.     (x, y, a, h) is taken as direct observation of the state space (linear
    93.     observation model).
    95.     """
    96.  
    97.     def __init__(self):
    98.         ndim, dt = 4, 1.
    99.  
    100.         # Create Kalman filter model matrices.
    101.         self._motion_mat = numpy.eye(2 * ndim, 2 * ndim)
    102.         for i in range(ndim):
    103.             self._motion_mat[i, ndim + i] = dt
    104.         self._update_mat = numpy.eye(ndim, 2 * ndim)
    105.  
    106.         # Motion and observation uncertainty are chosen relative to the current
    107.         # state estimate. These weights control the amount of uncertainty in
    108.         # the model. This is a bit hacky.
    109.         self._std_weight_position = 1. / 20
    110.         self._std_weight_velocity = 1. / 160
    111.  
    112.     def initiate(self, measurement):
    113.         """
    114.         Create track from unassociated measurement.
    116.         Parameters
    117.         ----------
    118.         measurement : ndarray
    119.             Bounding box coordinates (x, y, a, h) with center position (x, y),
    120.             aspect ratio a, and height h.
    122.         Returns
    123.         -------
    124.         (ndarray, ndarray)
    125.             Returns the mean vector (8 dimensional) and covariance matrix (8x8
    126.             dimensional) of the new track. Unobserved velocities are initialized
    127.             to 0 mean.
    129.         """
    130.         mean_pos = measurement
    131.         mean_vel = numpy.zeros_like(mean_pos)
    132.         mean = numpy.r_[mean_pos, mean_vel]
    133.  
    134.         std = [2 * self._std_weight_position * measurement[3],
    135.                2 * self._std_weight_position * measurement[3],
    136.                1e-2,
    137.                2 * self._std_weight_position * measurement[3],
    138.                10 * self._std_weight_velocity * measurement[3],
    139.                10 * self._std_weight_velocity * measurement[3],
    140.                1e-5,
    141.                10 * self._std_weight_velocity * measurement[3]]
    142.         covariance = numpy.diag(numpy.square(std))
    143.         return mean, covariance
    144.  
    145.     def predict(self, mean, covariance):
    146.         """
    147.         Run Kalman filter prediction step.
    149.         Parameters
    150.         ----------
    151.         mean : ndarray
    152.             The 8 dimensional mean vector of the object state at the previous
    153.             time step.
    154.         covariance : ndarray
    155.             The 8x8 dimensional covariance matrix of the object state at the
    156.             previous time step.
    158.         Returns
    159.         -------
    160.         (ndarray, ndarray)
    161.             Returns the mean vector and covariance matrix of the predicted
    162.             state. Unobserved velocities are initialized to 0 mean.
    164.         """
    165.         std_pos = [self._std_weight_position * mean[3],
    166.                    self._std_weight_position * mean[3],
    167.                    1e-2,
    168.                    self._std_weight_position * mean[3]]
    169.         std_vel = [self._std_weight_velocity * mean[3],
    170.                    self._std_weight_velocity * mean[3],
    171.                    1e-5,
    172.                    self._std_weight_velocity * mean[3]]
    173.         motion_cov = numpy.diag(numpy.square(numpy.r_[std_pos, std_vel]))
    174.  
    175.         # mean = np.dot(self._motion_mat, mean)
    176.         mean = numpy.dot(mean, self._motion_mat.T)
    177.         covariance = numpy.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T)) + motion_cov
    178.  
    179.         return mean, covariance
    180.  
    181.     def project(self, mean, covariance):
    182.         """
    183.         Project state distribution to measurement space.
    185.         Parameters
    186.         ----------
    187.         mean : ndarray
    188.             The state's mean vector (8 dimensional array).
    189.         covariance : ndarray
    190.             The state's covariance matrix (8x8 dimensional).
    192.         Returns
    193.         -------
    194.         (ndarray, ndarray)
    195.             Returns the projected mean and covariance matrix of the given state
    196.             estimate.
    198.         """
    199.         std = [self._std_weight_position * mean[3],
    200.                self._std_weight_position * mean[3],
    201.                1e-1,
    202.                self._std_weight_position * mean[3]]
    203.         innovation_cov = numpy.diag(numpy.square(std))
    204.  
    205.         mean = numpy.dot(self._update_mat, mean)
    206.         covariance = numpy.linalg.multi_dot((self._update_mat, covariance, self._update_mat.T))
    207.         return mean, covariance + innovation_cov
    208.  
    209.     def multi_predict(self, mean, covariance):
    210.         """
    211.         Run Kalman filter prediction step (Vectorized version).
    212.         Parameters
    213.         ----------
    214.         mean : ndarray
    215.             The Nx8 dimensional mean matrix of the object states at the previous
    216.             time step.
    217.         covariance : ndarray
    218.             The Nx8x8 dimensional covariance matrix of the object states at the
    219.             previous time step.
    220.         Returns
    221.         -------
    222.         (ndarray, ndarray)
    223.             Returns the mean vector and covariance matrix of the predicted
    224.             state. Unobserved velocities are initialized to 0 mean.
    225.         """
    226.         std_pos = [self._std_weight_position * mean[:, 3],
    227.                    self._std_weight_position * mean[:, 3],
    228.                    1e-2 * numpy.ones_like(mean[:, 3]),
    229.                    self._std_weight_position * mean[:, 3]]
    230.         std_vel = [self._std_weight_velocity * mean[:, 3],
    231.                    self._std_weight_velocity * mean[:, 3],
    232.                    1e-5 * numpy.ones_like(mean[:, 3]),
    233.                    self._std_weight_velocity * mean[:, 3]]
    234.         sqr = numpy.square(numpy.r_[std_pos, std_vel]).T
    235.  
    236.  
    237.         motion_cov = [numpy.diag(sqr[i]) for i in range(len(mean))]
    238.         motion_cov = numpy.asarray(motion_cov)
    239.  
    240.         #print(mean)
    241.         #print('eee')
    242.         #print(self._motion_mat.T)
    243.         mean = numpy.dot(mean, self._motion_mat.T)
    244.         #print('fff')
    245.         left = numpy.dot(self._motion_mat, covariance).transpose((1, 0, 2))
    246.         covariance = numpy.dot(left, self._motion_mat.T) + motion_cov
    247.     
    248.         return mean, covariance
    249.  
    250.     def update(self, mean, covariance, measurement):
    251.         """
    252.         Run Kalman filter correction step.
    254.         Parameters
    255.         ----------
    256.         mean : ndarray
    257.             The predicted state's mean vector (8 dimensional).
    258.         covariance : ndarray
    259.             The state's covariance matrix (8x8 dimensional).
    260.         measurement : ndarray
    261.             The 4 dimensional measurement vector (x, y, a, h), where (x, y)
    262.             is the center position, a the aspect ratio, and h the height of the
    263.             bounding box.
    265.         Returns
    266.         -------
    267.         (ndarray, ndarray)
    268.             Returns the measurement-corrected state distribution.
    270.         """
    271.         projected_mean, projected_cov = self.project(mean, covariance)
    272.  
    273.         chol_factor, lower = scipy.linalg.cho_factor(projected_cov, lower=True, check_finite=False)
    274.         kalman_gain = scipy.linalg.cho_solve((chol_factor, lower),
    275.                                              numpy.dot(covariance, self._update_mat.T).T,
    276.                                              check_finite=False).T
    277.         innovation = measurement - projected_mean
    278.  
    279.         new_mean = mean + numpy.dot(innovation, kalman_gain.T)
    280.         new_covariance = covariance - numpy.linalg.multi_dot((kalman_gain, projected_cov, kalman_gain.T))
    281.         return new_mean, new_covariance
    282.  
    283.     def gating_distance(self, mean, covariance, measurements, only_position=False, metric='maha'):
    284.         """
    285.         Compute gating distance between state distribution and measurements.
    286.         A suitable distance threshold can be obtained from `chi2inv95`. If
    287.         `only_position` is False, the chi-square distribution has 4 degrees of
    288.         freedom, otherwise 2.
    289.         Parameters
    290.         ----------
    291.         mean : ndarray
    292.             Mean vector over the state distribution (8 dimensional).
    293.         covariance : ndarray
    294.             Covariance of the state distribution (8x8 dimensional).
    295.         measurements : ndarray
    296.             An Nx4 dimensional matrix of N measurements, each in
    297.             format (x, y, a, h) where (x, y) is the bounding box center
    298.             position, a the aspect ratio, and h the height.
    299.         only_position : Optional[bool]
    300.             If True, distance computation is done with respect to the bounding
    301.             box center position only.
    302.         metric : str
    303.             Distance metric.
    304.         Returns
    305.         -------
    306.         ndarray
    307.             Returns an array of length N, where the i-th element contains the
    308.             squared Mahalanobis distance between (mean, covariance) and
    309.             `measurements[i]`.
    310.         """
    311.         mean, covariance = self.project(mean, covariance)
    312.         if only_position:
    313.             mean, covariance = mean[:2], covariance[:2, :2]
    314.             measurements = measurements[:, :2]
    315.  
    316.         d = measurements - mean
    317.         if metric == 'gaussian':
    318.             return numpy.sum(d * d, axis=1)
    319.         elif metric == 'maha':
    320.             factor = numpy.linalg.cholesky(covariance)
    321.             z = scipy.linalg.solve_triangular(factor, d.T, lower=True, check_finite=False, overwrite_b=True)
    322.             return numpy.sum(z * z, axis=0)  # square maha
    323.         else:
    324.             raise ValueError('invalid distance metric')
    325.  
    326.  
    327. class State:
    328.     New = 0
    329.     Tracked = 1
    330.     Lost = 2
    331.     Removed = 3
    332.  
    333.  
    334. class Track:
    335.     count = 0
    336.     shared_kalman = KalmanFilterXYAH()
    337.  
    338.     def __init__(self, tlwh, score, cls):
    339.  
    340.         # wait activate
    341.         self._tlwh = numpy.asarray(self.tlbr_to_tlwh(tlwh[:-1]), dtype=numpy.float32)
    342.         self.kalman_filter = None
    343.         self.mean, self.covariance = None, None
    344.         self.is_activated = False
    345.  
    346.         self.score = score
    347.         self.tracklet_len = 0
    348.         self.cls = cls
    349.         self.idx = tlwh[-1]
    350.  
    351.     def predict(self):
    352.         mean_state = self.mean.copy()
    353.         if self.state != State.Tracked:
    354.             mean_state[7] = 0
    355.         self.mean, self.covariance = self.kalman_filter.predict(mean_state, self.covariance)
    356.  
    357.     @staticmethod
    358.     def multi_predict(tracks):
    359.         if len(tracks) <= 0:
    360.             return
    361.         multi_mean = numpy.asarray([st.mean.copy() for st in tracks])
    362.         multi_covariance = numpy.asarray([st.covariance for st in tracks])
    363.         for i, st in enumerate(tracks):
    364.             if st.state != State.Tracked:
    365.                 multi_mean[i][7] = 0
    366.  
    367.         multi_mean, multi_covariance = Track.shared_kalman.multi_predict(multi_mean, multi_covariance)
    368.         #print('eee')
    369.  
    370.         for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
    371.             tracks[i].mean = mean
    372.             tracks[i].covariance = cov
    373.  
    374.  
    375.     def activate(self, kalman_filter, frame_id):
    376.         """Start a new tracklet"""
    377.         self.kalman_filter = kalman_filter
    378.         self.track_id = self.next_id()
    379.         self.mean, self.covariance = self.kalman_filter.initiate(self.convert_coords(self._tlwh))
    380.  
    381.         self.tracklet_len = 0
    382.         self.state = State.Tracked
    383.         if frame_id == 1:
    384.             self.is_activated = True
    385.         self.frame_id = frame_id
    386.         self.start_frame = frame_id
    387.  
    388.     def re_activate(self, new_track, frame_id, new_id=False):
    389.         self.mean, self.covariance = self.kalman_filter.update(self.mean, self.covariance,
    390.                                                                self.convert_coords(new_track.tlwh))
    391.         self.tracklet_len = 0
    392.         self.state = State.Tracked
    393.         self.is_activated = True
    394.         self.frame_id = frame_id
    395.         if new_id:
    396.             self.track_id = self.next_id()
    397.         self.score = new_track.score
    398.         self.cls = new_track.cls
    399.         self.idx = new_track.idx
    400.  
    401.     def update(self, new_track, frame_id):
    402.         """
    403.         Update a matched track
    404.         :type new_track: Track
    405.         :type frame_id: int
    406.         :return:
    407.         """
    408.         self.frame_id = frame_id
    409.         self.tracklet_len += 1
    410.  
    411.         new_tlwh = new_track.tlwh
    412.         self.mean, self.covariance = self.kalman_filter.update(self.mean, self.covariance,
    413.                                                                self.convert_coords(new_tlwh))
    414.         self.state = State.Tracked
    415.         self.is_activated = True
    416.  
    417.         self.score = new_track.score
    418.         self.cls = new_track.cls
    419.         self.idx = new_track.idx
    420.  
    421.     def convert_coords(self, tlwh):
    422.         return self.tlwh_to_xyah(tlwh)
    423.  
    424.     def mark_lost(self):
    425.         self.state = State.Lost
    426.  
    427.     def mark_removed(self):
    428.         self.state = State.Removed
    429.  
    430.     @property
    431.     def end_frame(self):
    432.         return self.frame_id
    433.  
    434.     @staticmethod
    435.     def next_id():
    436.         Track.count += 1
    437.         return Track.count
    438.  
    439.     @property
    440.     def tlwh(self):
    441.         """Get current position in bounding box format `(top left x, top left y,
    442.         width, height)`.
    443.         """
    444.         if self.mean is None:
    445.             return self._tlwh.copy()
    446.         ret = self.mean[:4].copy()
    447.         ret[2] *= ret[3]
    448.         ret[:2] -= ret[2:] / 2
    449.         return ret
    450.  
    451.     @property
    452.     def tlbr(self):
    453.         """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
    454.         `(top left, bottom right)`.
    455.         """
    456.         ret = self.tlwh.copy()
    457.         ret[2:] += ret[:2]
    458.         return ret
    459.  
    460.     @staticmethod
    461.     def reset_id():
    462.         Track.count = 0
    463.  
    464.     @staticmethod
    465.     def tlwh_to_xyah(tlwh):
    466.         """Convert bounding box to format `(center x, center y, aspect ratio,
    467.         height)`, where the aspect ratio is `width / height`.
    468.         """
    469.         ret = numpy.asarray(tlwh).copy()
    470.         ret[:2] += ret[2:] / 2
    471.         ret[2] /= ret[3]
    472.         return ret
    473.  
    474.     @staticmethod
    475.     def tlbr_to_tlwh(tlbr):
    476.         ret = numpy.asarray(tlbr).copy()
    477.         ret[2:] -= ret[:2]
    478.         return ret
    479.  
    480.     @staticmethod
    481.     def tlwh_to_tlbr(tlwh):
    482.         ret = numpy.asarray(tlwh).copy()
    483.         ret[2:] += ret[:2]
    484.         return ret
    485.  
    486.     def __repr__(self):
    487.         return f'OT_{self.track_id}_({self.start_frame}-{self.end_frame})'
    488.  
    489.  
    490. class BYTETracker:
    491.     def __init__(self, frame_rate=30):
    492.         self.tracked_tracks = []
    493.         self.lost_tracks = []
    494.         self.removed_tracks = []
    495.  
    496.         self.frame_id = 0
    497.         self.max_time_lost = int(frame_rate)
    498.         self.kalman_filter = KalmanFilterXYAH()
    499.         self.reset_id()
    500.  
    501.     def update(self, boxes, scores, object_classes):
    502.         self.frame_id += 1
    503.         activated_tracks = []
    504.         re_find_tracks = []
    505.         lost_tracks = []
    506.         removed_tracks = []
    507.  
    508.         # add index
    509.         boxes = numpy.concatenate([boxes, numpy.arange(len(boxes)).reshape(-1, 1)], axis=-1)
    510.  
    511.         indices_low = scores > 0.1
    512.         indices_high = scores < 0.5
    513.         indices_remain = scores > 0.5
    514.  
    515.         indices_second = numpy.logical_and(indices_low, indices_high)
    516.         boxes_second = boxes[indices_second]
    517.         boxes = boxes[indices_remain]
    518.         scores_keep = scores[indices_remain]
    519.         scores_second = scores[indices_second]
    520.         cls_keep = object_classes[indices_remain]
    521.         cls_second = object_classes[indices_second]
    522.  
    523.         detections = self.init_track(boxes, scores_keep, cls_keep)
    524.         """ Add newly detected tracklets to tracked_stracks"""
    525.         unconfirmed = []
    526.         tracked_stracks = []
    527.         for track in self.tracked_tracks:
    528.             if not track.is_activated:
    529.                 unconfirmed.append(track)
    530.             else:
    531.                 tracked_stracks.append(track)
    532.         """ Step 2: First association, with high score detection boxes"""
    533.         track_pool = self.joint_stracks(tracked_stracks, self.lost_tracks)
    534.         # Predict the current location with KF
    535.         self.multi_predict(track_pool)
    536.  
    537.  
    538.         #print('ddd')
    539.  
    540.         dists = self.get_dists(track_pool, detections)
    541.         matches, u_track, u_detection = linear_assignment(dists, thresh=0.8)
    542.         for tracked_i, box_i in matches:
    543.             track = track_pool[tracked_i]
    544.             det = detections[box_i]
    545.             if track.state == State.Tracked:
    546.                 track.update(det, self.frame_id)
    547.                 activated_tracks.append(track)
    548.             else:
    549.                 track.re_activate(det, self.frame_id, new_id=False)
    550.                 re_find_tracks.append(track)
    551.         """ Step 3: Second association, with low score detection boxes"""
    552.         # association the untrack to the low score detections
    553.         detections_second = self.init_track(boxes_second, scores_second, cls_second)
    554.         r_tracked_tracks = [track_pool[i] for i in u_track if track_pool[i].state == State.Tracked]
    555.         dists = iou_distance(r_tracked_tracks, detections_second)
    556.         matches, u_track, u_detection_second = linear_assignment(dists, thresh=0.5)
    557.         for tracked_i, box_i in matches:
    558.             track = r_tracked_tracks[tracked_i]
    559.             det = detections_second[box_i]
    560.             if track.state == State.Tracked:
    561.                 track.update(det, self.frame_id)
    562.                 activated_tracks.append(track)
    563.             else:
    564.                 track.re_activate(det, self.frame_id, new_id=False)
    565.                 re_find_tracks.append(track)
    566.  
    567.         for it in u_track:
    568.             track = r_tracked_tracks[it]
    569.             if track.state != State.Lost:
    570.                 track.mark_lost()
    571.                 lost_tracks.append(track)
    572.         """Deal with unconfirmed tracks, usually tracks with only one beginning frame"""
    573.         detections = [detections[i] for i in u_detection]
    574.         dists = self.get_dists(unconfirmed, detections)
    575.         matches, u_unconfirmed, u_detection = linear_assignment(dists, thresh=0.7)
    576.         for tracked_i, box_i in matches:
    577.             unconfirmed[tracked_i].update(detections[box_i], self.frame_id)
    578.             activated_tracks.append(unconfirmed[tracked_i])
    579.         for it in u_unconfirmed:
    580.             track = unconfirmed[it]
    581.             track.mark_removed()
    582.             removed_tracks.append(track)
    583.         """ Step 4: Init new stracks"""
    584.         for new_i in u_detection:
    585.             track = detections[new_i]
    586.             if track.score < 0.6:
    587.                 continue
    588.             track.activate(self.kalman_filter, self.frame_id)
    589.             activated_tracks.append(track)
    590.         """ Step 5: Update state"""
    591.         for track in self.lost_tracks:
    592.             if self.frame_id - track.end_frame > self.max_time_lost:
    593.                 track.mark_removed()
    594.                 removed_tracks.append(track)
    595.  
    596.  
    597.         #print('ccc')
    598.  
    599.  
    600.         self.tracked_tracks = [t for t in self.tracked_tracks if t.state == State.Tracked]
    601.         self.tracked_tracks = self.joint_stracks(self.tracked_tracks, activated_tracks)
    602.         self.tracked_tracks = self.joint_stracks(self.tracked_tracks, re_find_tracks)
    603.         self.lost_tracks = self.sub_stracks(self.lost_tracks, self.tracked_tracks)
    604.         self.lost_tracks.extend(lost_tracks)
    605.         self.lost_tracks = self.sub_stracks(self.lost_tracks, self.removed_tracks)
    606.         self.removed_tracks.extend(removed_tracks)
    607.         self.tracked_tracks, self.lost_tracks = self.remove_duplicate_stracks(self.tracked_tracks, self.lost_tracks)
    608.         output = [track.tlbr.tolist() + [track.track_id,
    609.                                          track.score,
    610.                                          track.cls,
    611.                                          track.idx] for track in self.tracked_tracks if track.is_activated]
    612.         return numpy.asarray(output, dtype=numpy.float32)
    613.  
    614.     @staticmethod
    615.     def init_track(boxes, scores, cls):
    616.         return [Track(box, s, c) for (box, s, c) in zip(boxes, scores, cls)] if len(boxes) else []  # detections
    617.  
    618.     @staticmethod
    619.     def get_dists(tracks, detections):
    620.         dists = iou_distance(tracks, detections)
    621.         dists = fuse_score(dists, detections)
    622.         return dists
    623.  
    624.     @staticmethod
    625.     def multi_predict(tracks):
    626.         Track.multi_predict(tracks)
    627.  
    628.     @staticmethod
    629.     def reset_id():
    630.         Track.reset_id()
    631.  
    632.     @staticmethod
    633.     def joint_stracks(tlista, tlistb):
    634.         exists = {}
    635.         res = []
    636.         for t in tlista:
    637.             exists[t.track_id] = 1
    638.             res.append(t)
    639.         for t in tlistb:
    640.             tid = t.track_id
    641.             if not exists.get(tid, 0):
    642.                 exists[tid] = 1
    643.                 res.append(t)
    644.         return res
    645.  
    646.     @staticmethod
    647.     def sub_stracks(tlista, tlistb):
    648.         stracks = {t.track_id: t for t in tlista}
    649.         for t in tlistb:
    650.             tid = t.track_id
    651.             if stracks.get(tid, 0):
    652.                 del stracks[tid]
    653.         return list(stracks.values())
    654.  
    655.     @staticmethod
    656.     def remove_duplicate_stracks(stracksa, stracksb):
    657.         pdist = iou_distance(stracksa, stracksb)
    658.         pairs = numpy.where(pdist < 0.15)
    659.         dupa, dupb = [], []
    660.         for p, q in zip(*pairs):
    661.             timep = stracksa[p].frame_id - stracksa[p].start_frame
    662.             timeq = stracksb[q].frame_id - stracksb[q].start_frame
    663.             if timep > timeq:
    664.                 dupb.append(q)
    665.             else:
    666.                 dupa.append(p)
    667.         resa = [t for i, t in enumerate(stracksa) if i not in dupa]
    668.         resb = [t for i, t in enumerate(stracksb) if i not in dupb]
    669.         return resa, resb

  • 相关阅读:
    「UI开发」DevExpress WPF Pivot Grid组件可轻松实现多维数据分析！（一）
    【小沐学前端】Node.js实现基于Protobuf协议的WebSocket通信
    R语言使用dev.print函数将当前最近的可视化结果保存为指定格式、dev.print函数不打开图像设备
    2011年09月01日 Go生态洞察：Go语言词法扫描与App Engine演示
    王学岗音视频开发(一)—————配置NDK开发环境
    关于mysql数据文件损坏导致的mysql无法启动的问题
    提高倾斜摄影三维模型顶层合并构建效率的技术方法初探
    单机高性能网络模型
    API是什么&API管理存在哪些问题，如何解决？
    STM32F103 CubeMX 使用USB生成鼠标设备
  • 原文地址：https://blog.csdn.net/zhaocj/article/details/133612221