| 23 | |
| 24 | |
| 25 | class Tracker: |
| 26 | def __init__(self, metric, max_iou_distance=0.7, max_age=70, n_init=3): |
| 27 | self.metric = metric |
| 28 | self.max_iou_distance = max_iou_distance |
| 29 | self.max_age = max_age |
| 30 | self.n_init = n_init |
| 31 | |
| 32 | self.kf = kalman_filter.KalmanFilter() |
| 33 | self.tracks = [] |
| 34 | self._next_id = 1 |
| 35 | |
| 36 | def predict(self): |
| 37 | """Propagate track state distributions one time step forward. |
| 38 | This function should be called once every time step, before `update`. |
| 39 | """ |
| 40 | for track in self.tracks: |
| 41 | track.predict(self.kf) |
| 42 | |
| 43 | def increment_ages(self): |
| 44 | for track in self.tracks: |
| 45 | track.increment_age() |
| 46 | track.mark_missed() |
| 47 | |
| 48 | def update(self, detections, classes): |
| 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 | |
| 59 | # Update track set. |
| 60 | for track_idx, detection_idx in matches: |
| 61 | self.tracks[track_idx].update( |
| 62 | self.kf, detections[detection_idx]) |
| 63 | for track_idx in unmatched_tracks: |
| 64 | self.tracks[track_idx].mark_missed() |
| 65 | for detection_idx in unmatched_detections: |
| 66 | self._initiate_track(detections[detection_idx], classes[detection_idx].item()) |
| 67 | self.tracks = [t for t in self.tracks if not t.is_deleted()] |
| 68 | |
| 69 | # Update distance metric. |
| 70 | active_targets = [t.track_id for t in self.tracks if t.is_confirmed()] |
| 71 | features, targets = [], [] |
| 72 | for track in self.tracks: |
| 73 | if not track.is_confirmed(): |
| 74 | continue |
| 75 | features += track.features |
| 76 | targets += [track.track_id for _ in track.features] |
| 77 | track.features = [] |
| 78 | self.metric.partial_fit( |
| 79 | np.asarray(features), np.asarray(targets), active_targets) |
| 80 | |
| 81 | def _match(self, detections): |
| 82 | |