jde_matching.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  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/matching.py
  16. """
  17. import lap
  18. import scipy
  19. import numpy as np
  20. from scipy.spatial.distance import cdist
  21. from ..motion import kalman_filter
  22. from paddlex.ppdet.utils.logger import setup_logger
  23. logger = setup_logger(__name__)
  24. __all__ = [
  25. 'merge_matches',
  26. 'linear_assignment',
  27. 'cython_bbox_ious',
  28. 'iou_distance',
  29. 'embedding_distance',
  30. 'fuse_motion',
  31. ]
  32. def merge_matches(m1, m2, shape):
  33. O, P, Q = shape
  34. m1 = np.asarray(m1)
  35. m2 = np.asarray(m2)
  36. M1 = scipy.sparse.coo_matrix(
  37. (np.ones(len(m1)), (m1[:, 0], m1[:, 1])), shape=(O, P))
  38. M2 = scipy.sparse.coo_matrix(
  39. (np.ones(len(m2)), (m2[:, 0], m2[:, 1])), shape=(P, Q))
  40. mask = M1 * M2
  41. match = mask.nonzero()
  42. match = list(zip(match[0], match[1]))
  43. unmatched_O = tuple(set(range(O)) - set([i for i, j in match]))
  44. unmatched_Q = tuple(set(range(Q)) - set([j for i, j in match]))
  45. return match, unmatched_O, unmatched_Q
  46. def linear_assignment(cost_matrix, thresh):
  47. if cost_matrix.size == 0:
  48. return np.empty(
  49. (0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(
  50. range(cost_matrix.shape[1]))
  51. matches, unmatched_a, unmatched_b = [], [], []
  52. cost, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh)
  53. for ix, mx in enumerate(x):
  54. if mx >= 0:
  55. matches.append([ix, mx])
  56. unmatched_a = np.where(x < 0)[0]
  57. unmatched_b = np.where(y < 0)[0]
  58. matches = np.asarray(matches)
  59. return matches, unmatched_a, unmatched_b
  60. def cython_bbox_ious(atlbrs, btlbrs):
  61. ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float)
  62. if ious.size == 0:
  63. return ious
  64. try:
  65. import cython_bbox
  66. except Exception as e:
  67. logger.error('cython_bbox not found, please install cython_bbox.'
  68. 'for example: `pip install cython_bbox`.')
  69. raise e
  70. ious = cython_bbox.bbox_overlaps(
  71. np.ascontiguousarray(
  72. atlbrs, dtype=np.float),
  73. np.ascontiguousarray(
  74. btlbrs, dtype=np.float))
  75. return ious
  76. def iou_distance(atracks, btracks):
  77. """
  78. Compute cost based on IoU between two list[STrack].
  79. """
  80. if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or (
  81. len(btracks) > 0 and isinstance(btracks[0], np.ndarray)):
  82. atlbrs = atracks
  83. btlbrs = btracks
  84. else:
  85. atlbrs = [track.tlbr for track in atracks]
  86. btlbrs = [track.tlbr for track in btracks]
  87. _ious = cython_bbox_ious(atlbrs, btlbrs)
  88. cost_matrix = 1 - _ious
  89. return cost_matrix
  90. def embedding_distance(tracks, detections, metric='euclidean'):
  91. """
  92. Compute cost based on features between two list[STrack].
  93. """
  94. cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float)
  95. if cost_matrix.size == 0:
  96. return cost_matrix
  97. det_features = np.asarray(
  98. [track.curr_feat for track in detections], dtype=np.float)
  99. track_features = np.asarray(
  100. [track.smooth_feat for track in tracks], dtype=np.float)
  101. cost_matrix = np.maximum(0.0, cdist(track_features, det_features,
  102. metric)) # Nomalized features
  103. return cost_matrix
  104. def fuse_motion(kf,
  105. cost_matrix,
  106. tracks,
  107. detections,
  108. only_position=False,
  109. lambda_=0.98):
  110. if cost_matrix.size == 0:
  111. return cost_matrix
  112. gating_dim = 2 if only_position else 4
  113. gating_threshold = kalman_filter.chi2inv95[gating_dim]
  114. measurements = np.asarray([det.to_xyah() for det in detections])
  115. for row, track in enumerate(tracks):
  116. gating_distance = kf.gating_distance(
  117. track.mean,
  118. track.covariance,
  119. measurements,
  120. only_position,
  121. metric='maha')
  122. cost_matrix[row, gating_distance > gating_threshold] = np.inf
  123. cost_matrix[row] = lambda_ * cost_matrix[row] + (1 - lambda_
  124. ) * gating_distance
  125. return cost_matrix