tracker.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. # vim: expandtab:ts=4:sw=4
  2. from __future__ import absolute_import
  3. import numpy as np
  4. from . import kalman_filter
  5. from . import linear_assignment
  6. from . import iou_matching
  7. from .track import Track
  8. class Tracker:
  9. """
  10. This is the multi-target tracker.
  11. Parameters
  12. ----------
  13. metric : nn_matching.NearestNeighborDistanceMetric
  14. A distance metric for measurement-to-track association.
  15. max_age : int
  16. Maximum number of missed misses before a track is deleted.
  17. n_init : int
  18. Number of consecutive detections before the track is confirmed. The
  19. track state is set to `Deleted` if a miss occurs within the first
  20. `n_init` frames.
  21. Attributes
  22. ----------
  23. metric : nn_matching.NearestNeighborDistanceMetric
  24. The distance metric used for measurement to track association.
  25. max_age : int
  26. Maximum number of missed misses before a track is deleted.
  27. n_init : int
  28. Number of frames that a track remains in initialization phase.
  29. kf : kalman_filter.KalmanFilter
  30. A Kalman filter to filter target trajectories in image space.
  31. tracks : List[Track]
  32. The list of active tracks at the current time step.
  33. """
  34. def __init__(self, metric, max_iou_distance=0.7, max_age=75, n_init=3):
  35. self.metric = metric
  36. self.max_iou_distance = max_iou_distance
  37. self.max_age = max_age
  38. self.n_init = n_init
  39. self.kf = kalman_filter.KalmanFilter()
  40. self.tracks = []
  41. self._next_id = 1
  42. def predict(self):
  43. """Propagate track state distributions one time step forward.
  44. This function should be called once every time step, before `update`.
  45. """
  46. for track in self.tracks:
  47. track.predict(self.kf)
  48. def update(self, detections):
  49. """Perform measurement update and track management.
  50. Parameters
  51. ----------
  52. detections : List[deep_sort.detection.Detection]
  53. A list of detections at the current time step.
  54. """
  55. # Run matching cascade.
  56. matches, unmatched_tracks, unmatched_detections = \
  57. self._match(detections)
  58. # Update track set.
  59. for track_idx, detection_idx in matches:
  60. self.tracks[track_idx].update(
  61. self.kf, detections[detection_idx])
  62. for track_idx in unmatched_tracks:
  63. self.tracks[track_idx].mark_missed()
  64. for detection_idx in unmatched_detections:
  65. self._initiate_track(detections[detection_idx])
  66. self.tracks = [t for t in self.tracks if not t.is_deleted()]
  67. # Update distance metric.
  68. active_targets = [t.track_id for t in self.tracks if t.is_confirmed()]
  69. features, targets = [], []
  70. for track in self.tracks:
  71. if not track.is_confirmed():
  72. continue
  73. features += track.features
  74. targets += [track.track_id for _ in track.features]
  75. track.features = []
  76. self.metric.partial_fit(
  77. np.asarray(features), np.asarray(targets), active_targets)
  78. def _match(self, detections):
  79. def gated_metric(tracks, dets, track_indices, detection_indices):
  80. features = np.array([dets[i].feature for i in detection_indices])
  81. targets = np.array([tracks[i].track_id for i in track_indices])
  82. cost_matrix = self.metric.distance(features, targets)
  83. cost_matrix = linear_assignment.gate_cost_matrix(
  84. self.kf, cost_matrix, tracks, dets, track_indices,
  85. detection_indices)
  86. return cost_matrix
  87. # Split track set into confirmed and unconfirmed tracks.
  88. confirmed_tracks = [
  89. i for i, t in enumerate(self.tracks) if t.is_confirmed()]
  90. unconfirmed_tracks = [
  91. i for i, t in enumerate(self.tracks) if not t.is_confirmed()]
  92. # Associate confirmed tracks using appearance features.
  93. matches_a, unmatched_tracks_a, unmatched_detections = \
  94. linear_assignment.matching_cascade(
  95. gated_metric, self.metric.matching_threshold, self.max_age,
  96. self.tracks, detections, confirmed_tracks)
  97. # Associate remaining tracks together with unconfirmed tracks using IOU.
  98. iou_track_candidates = unconfirmed_tracks + [
  99. k for k in unmatched_tracks_a if
  100. self.tracks[k].time_since_update == 1]
  101. unmatched_tracks_a = [
  102. k for k in unmatched_tracks_a if
  103. self.tracks[k].time_since_update != 1]
  104. matches_b, unmatched_tracks_b, unmatched_detections = \
  105. linear_assignment.min_cost_matching(
  106. iou_matching.iou_cost, self.max_iou_distance, self.tracks,
  107. detections, iou_track_candidates, unmatched_detections)
  108. matches = matches_a + matches_b
  109. unmatched_tracks = list(set(unmatched_tracks_a + unmatched_tracks_b))
  110. return matches, unmatched_tracks, unmatched_detections
  111. def _initiate_track(self, detection):
  112. mean, covariance = self.kf.initiate(detection.to_xyah())
  113. class_name = detection.get_class()
  114. self.tracks.append(Track(
  115. mean, covariance, self._next_id, self.n_init, self.max_age,
  116. detection.feature, class_name))
  117. self._next_id += 1