jde_tracker.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """
  15. This code is borrow from https://github.com/Zhongdao/Towards-Realtime-MOT/blob/master/tracker/multitracker.py
  16. """
  17. import paddle
  18. from ..matching import jde_matching as matching
  19. from .base_jde_tracker import TrackState, BaseTrack, STrack
  20. from .base_jde_tracker import joint_stracks, sub_stracks, remove_duplicate_stracks
  21. from paddlex.ppdet.core.workspace import register, serializable
  22. from paddlex.ppdet.utils.logger import setup_logger
  23. logger = setup_logger(__name__)
  24. __all__ = ['JDETracker']
  25. @register
  26. @serializable
  27. class JDETracker(object):
  28. __inject__ = ['motion']
  29. """
  30. JDE tracker
  31. Args:
  32. det_thresh (float): threshold of detection score
  33. track_buffer (int): buffer for tracker
  34. min_box_area (int): min box area to filter out low quality boxes
  35. tracked_thresh (float): linear assignment threshold of tracked
  36. stracks and detections
  37. r_tracked_thresh (float): linear assignment threshold of
  38. tracked stracks and unmatched detections
  39. unconfirmed_thresh (float): linear assignment threshold of
  40. unconfirmed stracks and unmatched detections
  41. motion (object): KalmanFilter instance
  42. conf_thres (float): confidence threshold for tracking
  43. metric_type (str): either "euclidean" or "cosine", the distance metric
  44. used for measurement to track association.
  45. """
  46. def __init__(self,
  47. det_thresh=0.3,
  48. track_buffer=30,
  49. min_box_area=200,
  50. tracked_thresh=0.7,
  51. r_tracked_thresh=0.5,
  52. unconfirmed_thresh=0.7,
  53. motion='KalmanFilter',
  54. conf_thres=0,
  55. metric_type='euclidean'):
  56. self.det_thresh = det_thresh
  57. self.track_buffer = track_buffer
  58. self.min_box_area = min_box_area
  59. self.tracked_thresh = tracked_thresh
  60. self.r_tracked_thresh = r_tracked_thresh
  61. self.unconfirmed_thresh = unconfirmed_thresh
  62. self.motion = motion
  63. self.conf_thres = conf_thres
  64. self.metric_type = metric_type
  65. self.frame_id = 0
  66. self.tracked_stracks = []
  67. self.lost_stracks = []
  68. self.removed_stracks = []
  69. self.max_time_lost = 0
  70. # max_time_lost will be calculated: int(frame_rate / 30.0 * track_buffer)
  71. def update(self, pred_dets, pred_embs):
  72. """
  73. Processes the image frame and finds bounding box(detections).
  74. Associates the detection with corresponding tracklets and also handles
  75. lost, removed, refound and active tracklets.
  76. Args:
  77. pred_dets (Tensor): Detection results of the image, shape is [N, 5].
  78. pred_embs (Tensor): Embedding results of the image, shape is [N, 512].
  79. Return:
  80. output_stracks (list): The list contains information regarding the
  81. online_tracklets for the recieved image tensor.
  82. """
  83. self.frame_id += 1
  84. activated_starcks = []
  85. # for storing active tracks, for the current frame
  86. refind_stracks = []
  87. # Lost Tracks whose detections are obtained in the current frame
  88. lost_stracks = []
  89. # The tracks which are not obtained in the current frame but are not
  90. # removed. (Lost for some time lesser than the threshold for removing)
  91. removed_stracks = []
  92. remain_inds = paddle.nonzero(pred_dets[:, 4] > self.conf_thres)
  93. if remain_inds.shape[0] == 0:
  94. pred_dets = paddle.zeros([0, 1])
  95. pred_embs = paddle.zeros([0, 1])
  96. else:
  97. pred_dets = paddle.gather(pred_dets, remain_inds)
  98. pred_embs = paddle.gather(pred_embs, remain_inds)
  99. # Filter out the image with box_num = 0. pred_dets = [[0.0, 0.0, 0.0 ,0.0]]
  100. empty_pred = True if len(pred_dets) == 1 and paddle.sum(
  101. pred_dets) == 0.0 else False
  102. """ Step 1: Network forward, get detections & embeddings"""
  103. if len(pred_dets) > 0 and not empty_pred:
  104. pred_dets = pred_dets.numpy()
  105. pred_embs = pred_embs.numpy()
  106. detections = [
  107. STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f, 30)
  108. for (tlbrs, f) in zip(pred_dets, pred_embs)
  109. ]
  110. else:
  111. detections = []
  112. ''' Add newly detected tracklets to tracked_stracks'''
  113. unconfirmed = []
  114. tracked_stracks = [] # type: list[STrack]
  115. for track in self.tracked_stracks:
  116. if not track.is_activated:
  117. # previous tracks which are not active in the current frame are added in unconfirmed list
  118. unconfirmed.append(track)
  119. else:
  120. # Active tracks are added to the local list 'tracked_stracks'
  121. tracked_stracks.append(track)
  122. """ Step 2: First association, with embedding"""
  123. # Combining currently tracked_stracks and lost_stracks
  124. strack_pool = joint_stracks(tracked_stracks, self.lost_stracks)
  125. # Predict the current location with KF
  126. STrack.multi_predict(strack_pool, self.motion)
  127. dists = matching.embedding_distance(
  128. strack_pool, detections, metric=self.metric_type)
  129. dists = matching.fuse_motion(self.motion, dists, strack_pool,
  130. detections)
  131. # The dists is the list of distances of the detection with the tracks in strack_pool
  132. matches, u_track, u_detection = matching.linear_assignment(
  133. dists, thresh=self.tracked_thresh)
  134. # The matches is the array for corresponding matches of the detection with the corresponding strack_pool
  135. for itracked, idet in matches:
  136. # itracked is the id of the track and idet is the detection
  137. track = strack_pool[itracked]
  138. det = detections[idet]
  139. if track.state == TrackState.Tracked:
  140. # If the track is active, add the detection to the track
  141. track.update(detections[idet], self.frame_id)
  142. activated_starcks.append(track)
  143. else:
  144. # We have obtained a detection from a track which is not active,
  145. # hence put the track in refind_stracks list
  146. track.re_activate(det, self.frame_id, new_id=False)
  147. refind_stracks.append(track)
  148. # None of the steps below happen if there are no undetected tracks.
  149. """ Step 3: Second association, with IOU"""
  150. detections = [detections[i] for i in u_detection]
  151. # detections is now a list of the unmatched detections
  152. r_tracked_stracks = []
  153. # This is container for stracks which were tracked till the previous
  154. # frame but no detection was found for it in the current frame.
  155. for i in u_track:
  156. if strack_pool[i].state == TrackState.Tracked:
  157. r_tracked_stracks.append(strack_pool[i])
  158. dists = matching.iou_distance(r_tracked_stracks, detections)
  159. matches, u_track, u_detection = matching.linear_assignment(
  160. dists, thresh=self.r_tracked_thresh)
  161. # matches is the list of detections which matched with corresponding
  162. # tracks by IOU distance method.
  163. for itracked, idet in matches:
  164. track = r_tracked_stracks[itracked]
  165. det = detections[idet]
  166. if track.state == TrackState.Tracked:
  167. track.update(det, self.frame_id)
  168. activated_starcks.append(track)
  169. else:
  170. track.re_activate(det, self.frame_id, new_id=False)
  171. refind_stracks.append(track)
  172. # Same process done for some unmatched detections, but now considering IOU_distance as measure
  173. for it in u_track:
  174. track = r_tracked_stracks[it]
  175. if not track.state == TrackState.Lost:
  176. track.mark_lost()
  177. lost_stracks.append(track)
  178. # If no detections are obtained for tracks (u_track), the tracks are added to lost_tracks list and are marked lost
  179. '''Deal with unconfirmed tracks, usually tracks with only one beginning frame'''
  180. detections = [detections[i] for i in u_detection]
  181. dists = matching.iou_distance(unconfirmed, detections)
  182. matches, u_unconfirmed, u_detection = matching.linear_assignment(
  183. dists, thresh=self.unconfirmed_thresh)
  184. for itracked, idet in matches:
  185. unconfirmed[itracked].update(detections[idet], self.frame_id)
  186. activated_starcks.append(unconfirmed[itracked])
  187. # The tracks which are yet not matched
  188. for it in u_unconfirmed:
  189. track = unconfirmed[it]
  190. track.mark_removed()
  191. removed_stracks.append(track)
  192. # after all these confirmation steps, if a new detection is found, it is initialized for a new track
  193. """ Step 4: Init new stracks"""
  194. for inew in u_detection:
  195. track = detections[inew]
  196. if track.score < self.det_thresh:
  197. continue
  198. track.activate(self.motion, self.frame_id)
  199. activated_starcks.append(track)
  200. """ Step 5: Update state"""
  201. # If the tracks are lost for more frames than the threshold number, the tracks are removed.
  202. for track in self.lost_stracks:
  203. if self.frame_id - track.end_frame > self.max_time_lost:
  204. track.mark_removed()
  205. removed_stracks.append(track)
  206. # Update the self.tracked_stracks and self.lost_stracks using the updates in this step.
  207. self.tracked_stracks = [
  208. t for t in self.tracked_stracks if t.state == TrackState.Tracked
  209. ]
  210. self.tracked_stracks = joint_stracks(self.tracked_stracks,
  211. activated_starcks)
  212. self.tracked_stracks = joint_stracks(self.tracked_stracks,
  213. refind_stracks)
  214. self.lost_stracks = sub_stracks(self.lost_stracks,
  215. self.tracked_stracks)
  216. self.lost_stracks.extend(lost_stracks)
  217. self.lost_stracks = sub_stracks(self.lost_stracks,
  218. self.removed_stracks)
  219. self.removed_stracks.extend(removed_stracks)
  220. self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(
  221. self.tracked_stracks, self.lost_stracks)
  222. # get scores of lost tracks
  223. output_stracks = [
  224. track for track in self.tracked_stracks if track.is_activated
  225. ]
  226. logger.debug('===========Frame {}=========='.format(self.frame_id))
  227. logger.debug('Activated: {}'.format(
  228. [track.track_id for track in activated_starcks]))
  229. logger.debug('Refind: {}'.format(
  230. [track.track_id for track in refind_stracks]))
  231. logger.debug('Lost: {}'.format(
  232. [track.track_id for track in lost_stracks]))
  233. logger.debug('Removed: {}'.format(
  234. [track.track_id for track in removed_stracks]))
  235. return output_stracks