linear_assignment.py 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. # vim: expandtab:ts=4:sw=4
  2. from __future__ import absolute_import
  3. import numpy as np
  4. from scipy.optimize import linear_sum_assignment
  5. from . import kalman_filter
  6. INFTY_COST = 1e+5
  7. def min_cost_matching(
  8. distance_metric, max_distance, tracks, detections, track_indices=None,
  9. detection_indices=None):
  10. """Solve linear assignment problem.
  11. Parameters
  12. ----------
  13. distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray
  14. The distance metric is given a list of tracks and detections as well as
  15. a list of N track indices and M detection indices. The metric should
  16. return the NxM dimensional cost matrix, where element (i, j) is the
  17. association cost between the i-th track in the given track indices and
  18. the j-th detection in the given detection_indices.
  19. max_distance : float
  20. Gating threshold. Associations with cost larger than this value are
  21. disregarded.
  22. tracks : List[track.Track]
  23. A list of predicted tracks at the current time step.
  24. detections : List[detection.Detection]
  25. A list of detections at the current time step.
  26. track_indices : List[int]
  27. List of track indices that maps rows in `cost_matrix` to tracks in
  28. `tracks` (see description above).
  29. detection_indices : List[int]
  30. List of detection indices that maps columns in `cost_matrix` to
  31. detections in `detections` (see description above).
  32. Returns
  33. -------
  34. (List[(int, int)], List[int], List[int])
  35. Returns a tuple with the following three entries:
  36. * A list of matched track and detection indices.
  37. * A list of unmatched track indices.
  38. * A list of unmatched detection indices.
  39. """
  40. if track_indices is None:
  41. track_indices = np.arange(len(tracks))
  42. if detection_indices is None:
  43. detection_indices = np.arange(len(detections))
  44. if len(detection_indices) == 0 or len(track_indices) == 0:
  45. return [], track_indices, detection_indices # Nothing to match.
  46. cost_matrix = distance_metric(
  47. tracks, detections, track_indices, detection_indices)
  48. cost_matrix[cost_matrix > max_distance] = max_distance + 1e-5
  49. indices = linear_sum_assignment(cost_matrix)
  50. indices = np.asarray(indices)
  51. indices = np.transpose(indices)
  52. matches, unmatched_tracks, unmatched_detections = [], [], []
  53. for col, detection_idx in enumerate(detection_indices):
  54. if col not in indices[:, 1]:
  55. unmatched_detections.append(detection_idx)
  56. for row, track_idx in enumerate(track_indices):
  57. if row not in indices[:, 0]:
  58. unmatched_tracks.append(track_idx)
  59. for row, col in indices:
  60. track_idx = track_indices[row]
  61. detection_idx = detection_indices[col]
  62. if cost_matrix[row, col] > max_distance:
  63. unmatched_tracks.append(track_idx)
  64. unmatched_detections.append(detection_idx)
  65. else:
  66. matches.append((track_idx, detection_idx))
  67. return matches, unmatched_tracks, unmatched_detections
  68. def matching_cascade(
  69. distance_metric, max_distance, cascade_depth, tracks, detections,
  70. track_indices=None, detection_indices=None):
  71. """Run matching cascade.
  72. Parameters
  73. ----------
  74. distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray
  75. The distance metric is given a list of tracks and detections as well as
  76. a list of N track indices and M detection indices. The metric should
  77. return the NxM dimensional cost matrix, where element (i, j) is the
  78. association cost between the i-th track in the given track indices and
  79. the j-th detection in the given detection indices.
  80. max_distance : float
  81. Gating threshold. Associations with cost larger than this value are
  82. disregarded.
  83. cascade_depth: int
  84. The cascade depth, should be se to the maximum track age.
  85. tracks : List[track.Track]
  86. A list of predicted tracks at the current time step.
  87. detections : List[detection.Detection]
  88. A list of detections at the current time step.
  89. track_indices : Optional[List[int]]
  90. List of track indices that maps rows in `cost_matrix` to tracks in
  91. `tracks` (see description above). Defaults to all tracks.
  92. detection_indices : Optional[List[int]]
  93. List of detection indices that maps columns in `cost_matrix` to
  94. detections in `detections` (see description above). Defaults to all
  95. detections.
  96. Returns
  97. -------
  98. (List[(int, int)], List[int], List[int])
  99. Returns a tuple with the following three entries:
  100. * A list of matched track and detection indices.
  101. * A list of unmatched track indices.
  102. * A list of unmatched detection indices.
  103. """
  104. if track_indices is None:
  105. track_indices = list(range(len(tracks)))
  106. if detection_indices is None:
  107. detection_indices = list(range(len(detections)))
  108. unmatched_detections = detection_indices
  109. matches = []
  110. for level in range(cascade_depth):
  111. if len(unmatched_detections) == 0: # No detections left
  112. break
  113. track_indices_l = [
  114. k for k in track_indices
  115. if tracks[k].time_since_update == 1 + level
  116. ]
  117. if len(track_indices_l) == 0: # Nothing to match at this level
  118. continue
  119. matches_l, _, unmatched_detections = \
  120. min_cost_matching(
  121. distance_metric, max_distance, tracks, detections,
  122. track_indices_l, unmatched_detections)
  123. matches += matches_l
  124. unmatched_tracks = list(set(track_indices) - set(k for k, _ in matches))
  125. return matches, unmatched_tracks, unmatched_detections
  126. def gate_cost_matrix(
  127. kf, cost_matrix, tracks, detections, track_indices, detection_indices,
  128. gated_cost=INFTY_COST, only_position=False):
  129. """Invalidate infeasible entries in cost matrix based on the state
  130. distributions obtained by Kalman filtering.
  131. Parameters
  132. ----------
  133. kf : The Kalman filter.
  134. cost_matrix : ndarray
  135. The NxM dimensional cost matrix, where N is the number of track indices
  136. and M is the number of detection indices, such that entry (i, j) is the
  137. association cost between `tracks[track_indices[i]]` and
  138. `detections[detection_indices[j]]`.
  139. tracks : List[track.Track]
  140. A list of predicted tracks at the current time step.
  141. detections : List[detection.Detection]
  142. A list of detections at the current time step.
  143. track_indices : List[int]
  144. List of track indices that maps rows in `cost_matrix` to tracks in
  145. `tracks` (see description above).
  146. detection_indices : List[int]
  147. List of detection indices that maps columns in `cost_matrix` to
  148. detections in `detections` (see description above).
  149. gated_cost : Optional[float]
  150. Entries in the cost matrix corresponding to infeasible associations are
  151. set this value. Defaults to a very large value.
  152. only_position : Optional[bool]
  153. If True, only the x, y position of the state distribution is considered
  154. during gating. Defaults to False.
  155. Returns
  156. -------
  157. ndarray
  158. Returns the modified cost matrix.
  159. """
  160. gating_dim = 2 if only_position else 4
  161. gating_threshold = kalman_filter.chi2inv95[gating_dim]
  162. measurements = np.asarray(
  163. [detections[i].to_xyah() for i in detection_indices])
  164. for row, track_idx in enumerate(track_indices):
  165. track = tracks[track_idx]
  166. gating_distance = kf.gating_distance(
  167. track.mean, track.covariance, measurements, only_position)
  168. cost_matrix[row, gating_distance > gating_threshold] = gated_cost
  169. return cost_matrix