base_jde_tracker.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267
  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 numpy as np
  18. from collections import deque, OrderedDict
  19. from ..matching import jde_matching as matching
  20. from paddlex.ppdet.core.workspace import register, serializable
  21. __all__ = [
  22. 'TrackState',
  23. 'BaseTrack',
  24. 'STrack',
  25. 'joint_stracks',
  26. 'sub_stracks',
  27. 'remove_duplicate_stracks',
  28. ]
  29. class TrackState(object):
  30. New = 0
  31. Tracked = 1
  32. Lost = 2
  33. Removed = 3
  34. @register
  35. @serializable
  36. class BaseTrack(object):
  37. _count = 0
  38. track_id = 0
  39. is_activated = False
  40. state = TrackState.New
  41. history = OrderedDict()
  42. features = []
  43. curr_feature = None
  44. score = 0
  45. start_frame = 0
  46. frame_id = 0
  47. time_since_update = 0
  48. # multi-camera
  49. location = (np.inf, np.inf)
  50. @property
  51. def end_frame(self):
  52. return self.frame_id
  53. @staticmethod
  54. def next_id():
  55. BaseTrack._count += 1
  56. return BaseTrack._count
  57. def activate(self, *args):
  58. raise NotImplementedError
  59. def predict(self):
  60. raise NotImplementedError
  61. def update(self, *args, **kwargs):
  62. raise NotImplementedError
  63. def mark_lost(self):
  64. self.state = TrackState.Lost
  65. def mark_removed(self):
  66. self.state = TrackState.Removed
  67. @register
  68. @serializable
  69. class STrack(BaseTrack):
  70. def __init__(self, tlwh, score, temp_feat, buffer_size=30):
  71. # wait activate
  72. self._tlwh = np.asarray(tlwh, dtype=np.float)
  73. self.kalman_filter = None
  74. self.mean, self.covariance = None, None
  75. self.is_activated = False
  76. self.score = score
  77. self.tracklet_len = 0
  78. self.smooth_feat = None
  79. self.update_features(temp_feat)
  80. self.features = deque([], maxlen=buffer_size)
  81. self.alpha = 0.9
  82. def update_features(self, feat):
  83. feat /= np.linalg.norm(feat)
  84. self.curr_feat = feat
  85. if self.smooth_feat is None:
  86. self.smooth_feat = feat
  87. else:
  88. self.smooth_feat = self.alpha * self.smooth_feat + (1 - self.alpha
  89. ) * feat
  90. self.features.append(feat)
  91. self.smooth_feat /= np.linalg.norm(self.smooth_feat)
  92. def predict(self):
  93. mean_state = self.mean.copy()
  94. if self.state != TrackState.Tracked:
  95. mean_state[7] = 0
  96. self.mean, self.covariance = self.kalman_filter.predict(
  97. mean_state, self.covariance)
  98. @staticmethod
  99. def multi_predict(stracks, kalman_filter):
  100. if len(stracks) > 0:
  101. multi_mean = np.asarray([st.mean.copy() for st in stracks])
  102. multi_covariance = np.asarray([st.covariance for st in stracks])
  103. for i, st in enumerate(stracks):
  104. if st.state != TrackState.Tracked:
  105. multi_mean[i][7] = 0
  106. multi_mean, multi_covariance = kalman_filter.multi_predict(
  107. multi_mean, multi_covariance)
  108. for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
  109. stracks[i].mean = mean
  110. stracks[i].covariance = cov
  111. def activate(self, kalman_filter, frame_id):
  112. """Start a new tracklet"""
  113. self.kalman_filter = kalman_filter
  114. self.track_id = self.next_id()
  115. self.mean, self.covariance = self.kalman_filter.initiate(
  116. self.tlwh_to_xyah(self._tlwh))
  117. self.tracklet_len = 0
  118. self.state = TrackState.Tracked
  119. if frame_id == 1:
  120. self.is_activated = True
  121. self.frame_id = frame_id
  122. self.start_frame = frame_id
  123. def re_activate(self, new_track, frame_id, new_id=False):
  124. self.mean, self.covariance = self.kalman_filter.update(
  125. self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh))
  126. self.update_features(new_track.curr_feat)
  127. self.tracklet_len = 0
  128. self.state = TrackState.Tracked
  129. self.is_activated = True
  130. self.frame_id = frame_id
  131. if new_id:
  132. self.track_id = self.next_id()
  133. def update(self, new_track, frame_id, update_feature=True):
  134. self.frame_id = frame_id
  135. self.tracklet_len += 1
  136. new_tlwh = new_track.tlwh
  137. self.mean, self.covariance = self.kalman_filter.update(
  138. self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh))
  139. self.state = TrackState.Tracked
  140. self.is_activated = True
  141. self.score = new_track.score
  142. if update_feature:
  143. self.update_features(new_track.curr_feat)
  144. @property
  145. def tlwh(self):
  146. """
  147. Get current position in bounding box format `(top left x, top left y,
  148. width, height)`.
  149. """
  150. if self.mean is None:
  151. return self._tlwh.copy()
  152. ret = self.mean[:4].copy()
  153. ret[2] *= ret[3]
  154. ret[:2] -= ret[2:] / 2
  155. return ret
  156. @property
  157. def tlbr(self):
  158. """
  159. Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
  160. `(top left, bottom right)`.
  161. """
  162. ret = self.tlwh.copy()
  163. ret[2:] += ret[:2]
  164. return ret
  165. @staticmethod
  166. def tlwh_to_xyah(tlwh):
  167. """
  168. Convert bounding box to format `(center x, center y, aspect ratio,
  169. height)`, where the aspect ratio is `width / height`.
  170. """
  171. ret = np.asarray(tlwh).copy()
  172. ret[:2] += ret[2:] / 2
  173. ret[2] /= ret[3]
  174. return ret
  175. def to_xyah(self):
  176. return self.tlwh_to_xyah(self.tlwh)
  177. @staticmethod
  178. def tlbr_to_tlwh(tlbr):
  179. ret = np.asarray(tlbr).copy()
  180. ret[2:] -= ret[:2]
  181. return ret
  182. @staticmethod
  183. def tlwh_to_tlbr(tlwh):
  184. ret = np.asarray(tlwh).copy()
  185. ret[2:] += ret[:2]
  186. return ret
  187. def __repr__(self):
  188. return 'OT_{}_({}-{})'.format(self.track_id, self.start_frame,
  189. self.end_frame)
  190. def joint_stracks(tlista, tlistb):
  191. exists = {}
  192. res = []
  193. for t in tlista:
  194. exists[t.track_id] = 1
  195. res.append(t)
  196. for t in tlistb:
  197. tid = t.track_id
  198. if not exists.get(tid, 0):
  199. exists[tid] = 1
  200. res.append(t)
  201. return res
  202. def sub_stracks(tlista, tlistb):
  203. stracks = {}
  204. for t in tlista:
  205. stracks[t.track_id] = t
  206. for t in tlistb:
  207. tid = t.track_id
  208. if stracks.get(tid, 0):
  209. del stracks[tid]
  210. return list(stracks.values())
  211. def remove_duplicate_stracks(stracksa, stracksb):
  212. pdist = matching.iou_distance(stracksa, stracksb)
  213. pairs = np.where(pdist < 0.15)
  214. dupa, dupb = list(), list()
  215. for p, q in zip(*pairs):
  216. timep = stracksa[p].frame_id - stracksa[p].start_frame
  217. timeq = stracksb[q].frame_id - stracksb[q].start_frame
  218. if timep > timeq:
  219. dupb.append(q)
  220. else:
  221. dupa.append(p)
  222. resa = [t for i, t in enumerate(stracksa) if not i in dupa]
  223. resb = [t for i, t in enumerate(stracksb) if not i in dupb]
  224. return resa, resb