123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139 |
- # vim: expandtab:ts=4:sw=4
- from __future__ import absolute_import
- import numpy as np
- from . import kalman_filter
- from . import linear_assignment
- from . import iou_matching
- from .track import Track
- class Tracker:
- """
- This is the multi-target tracker.
- Parameters
- ----------
- metric : nn_matching.NearestNeighborDistanceMetric
- A distance metric for measurement-to-track association.
- max_age : int
- Maximum number of missed misses before a track is deleted.
- n_init : int
- Number of consecutive detections before the track is confirmed. The
- track state is set to `Deleted` if a miss occurs within the first
- `n_init` frames.
- Attributes
- ----------
- metric : nn_matching.NearestNeighborDistanceMetric
- The distance metric used for measurement to track association.
- max_age : int
- Maximum number of missed misses before a track is deleted.
- n_init : int
- Number of frames that a track remains in initialization phase.
- kf : kalman_filter.KalmanFilter
- A Kalman filter to filter target trajectories in image space.
- tracks : List[Track]
- The list of active tracks at the current time step.
- """
- def __init__(self, metric, max_iou_distance=0.7, max_age=75, n_init=3):
- self.metric = metric
- self.max_iou_distance = max_iou_distance
- self.max_age = max_age
- self.n_init = n_init
- self.kf = kalman_filter.KalmanFilter()
- self.tracks = []
- self._next_id = 1
- def predict(self):
- """Propagate track state distributions one time step forward.
- This function should be called once every time step, before `update`.
- """
- for track in self.tracks:
- track.predict(self.kf)
- def update(self, detections):
- """Perform measurement update and track management.
- Parameters
- ----------
- detections : List[deep_sort.detection.Detection]
- A list of detections at the current time step.
- """
- # Run matching cascade.
- matches, unmatched_tracks, unmatched_detections = \
- self._match(detections)
- # Update track set.
- for track_idx, detection_idx in matches:
- self.tracks[track_idx].update(
- self.kf, detections[detection_idx])
- for track_idx in unmatched_tracks:
- self.tracks[track_idx].mark_missed()
- for detection_idx in unmatched_detections:
- self._initiate_track(detections[detection_idx])
- self.tracks = [t for t in self.tracks if not t.is_deleted()]
- # Update distance metric.
- active_targets = [t.track_id for t in self.tracks if t.is_confirmed()]
- features, targets = [], []
- for track in self.tracks:
- if not track.is_confirmed():
- continue
- features += track.features
- targets += [track.track_id for _ in track.features]
- track.features = []
- self.metric.partial_fit(
- np.asarray(features), np.asarray(targets), active_targets)
- def _match(self, detections):
- def gated_metric(tracks, dets, track_indices, detection_indices):
- features = np.array([dets[i].feature for i in detection_indices])
- targets = np.array([tracks[i].track_id for i in track_indices])
- cost_matrix = self.metric.distance(features, targets)
- cost_matrix = linear_assignment.gate_cost_matrix(
- self.kf, cost_matrix, tracks, dets, track_indices,
- detection_indices)
- return cost_matrix
- # Split track set into confirmed and unconfirmed tracks.
- confirmed_tracks = [
- i for i, t in enumerate(self.tracks) if t.is_confirmed()]
- unconfirmed_tracks = [
- i for i, t in enumerate(self.tracks) if not t.is_confirmed()]
- # Associate confirmed tracks using appearance features.
- matches_a, unmatched_tracks_a, unmatched_detections = \
- linear_assignment.matching_cascade(
- gated_metric, self.metric.matching_threshold, self.max_age,
- self.tracks, detections, confirmed_tracks)
- # Associate remaining tracks together with unconfirmed tracks using IOU.
- iou_track_candidates = unconfirmed_tracks + [
- k for k in unmatched_tracks_a if
- self.tracks[k].time_since_update == 1]
- unmatched_tracks_a = [
- k for k in unmatched_tracks_a if
- self.tracks[k].time_since_update != 1]
- matches_b, unmatched_tracks_b, unmatched_detections = \
- linear_assignment.min_cost_matching(
- iou_matching.iou_cost, self.max_iou_distance, self.tracks,
- detections, iou_track_candidates, unmatched_detections)
- matches = matches_a + matches_b
- unmatched_tracks = list(set(unmatched_tracks_a + unmatched_tracks_b))
- return matches, unmatched_tracks, unmatched_detections
- def _initiate_track(self, detection):
- mean, covariance = self.kf.initiate(detection.to_xyah())
- class_name = detection.get_class()
- self.tracks.append(Track(
- mean, covariance, self._next_id, self.n_init, self.max_age,
- detection.feature, class_name))
- self._next_id += 1
|