基于YOLO和DeepSORT的行人检测跟踪源码

import numpy as np from scipy.optimize import linear_sum_assignment as linear_assignment chi2inv95 = { 1: 3.8415, 2: 5.9915, 3: 7.8147, 4: 9.4877, 5: 11.070, 6: 12.592, 7: 14.067, 8: 15.507, 9: 16.919 } def min_cost_matching( distance_metric, max_distance_threshold, tracks, detections, track_indices=None, detection_indices=None): """使用匈牙利算法解决线性分配问题 Returns ------- (List[(int, int)], List[int], List[int]) Returns a tuple with the following three entries: * A list of matched track and detection indices. * A list of unmatched track indices. * A list of unmatched detection indices. """ if track_indices is None: track_indices = np.arange(len(tracks)) if detection_indices is None: detection_indices = np.arange(len(detections)) if len(detection_indices) == 0 or len(track_indices) == 0: return [], track_indices, detection_indices # 无需匹配 cost_matrix = distance_metric( tracks, detections, track_indices, detection_indices ) cost_matrix[cost_matrix > max_distance_threshold] = max_distance_threshold + 1e-5 indices = linear_assignment(cost_matrix) indices = np.transpose(np.array(indices)) matches, unmatched_tracks, unmatched_detections = [], [], [] for col, detection_idx in enumerate(detection_indices): if col not in indices[:, 1]: unmatched_detections.append(detection_idx) for row, track_idx in enumerate(track_indices): if row not in indices[:, 0]: unmatched_tracks.append(track_idx) for row, col in indices: track_idx = track_indices[row] detection_idx = detection_indices[col] if cost_matrix[row, col] > max_distance_threshold: unmatched_tracks.append(track_idx) unmatched_detections.append(detection_idx) else: matches.append((track_idx, detection_idx)) return matches, unmatched_tracks, unmatched_detections def matching_cascade( distance_metric, max_distance, cascade_depth, tracks, detections, track_indices=None, detection_indices=None): """Run matching cascade. 级联匹配 Parameters ---------- distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray The distance metric is given a list of tracks and detections as well as a list of N track indices and M detection indices. The metric should return the NxM dimensional cost matrix, where element (i, j) is the association cost between the i-th track in the given track indices and the j-th detection in the given detection indices. max_distance : float Gating threshold. Associations with cost larger than this value are disregarded. cascade_depth: int The cascade depth, should be se to the maximum track age. tracks : List[track.Track] A list of predicted tracks at the current time step. detections : List[detection.Detection] A list of detections at the current time step. track_indices : Optional[List[int]] List of track indices that maps rows in `cost_matrix` to tracks in `tracks` (see description above). Defaults to all tracks. detection_indices : Optional[List[int]] List of detection indices that maps columns in `cost_matrix` to detections in `detections` (see description above). Defaults to all detections. Returns ------- (List[(int, int)], List[int], List[int]) Returns a tuple with the following three entries: * A list of matched track and detection indices. * A list of unmatched track indices. * A list of unmatched detection indices. """ if track_indices is None: track_indices = list(range(len(tracks))) if detection_indices is None: detection_indices = list(range(len(detections))) unmatched_detections = detection_indices matches = [] for level in range(cascade_depth): if len(unmatched_detections) == 0: # No detections left break track_indices_l = [ k for k in track_indices if tracks[k].time_since_update == 1 + level ] if len(track_indices_l) == 0: # Nothing to match at this level continue matches_l, _, unmatched_detections = \ min_cost_matching( distance_metric, max_distance, tracks, detections, track_indices_l, unmatched_detections) matches += matches_l unmatched_tracks = list(set(track_indices) - set(k for k, _ in matches)) return matches, unmatched_tracks, unmatched_detections def gate_cost_matrix( kf, cost_matrix, tracks, detections, track_indices, detection_indices, gated_cost=1e+5, only_position=False): """ 未理解 """ gating_dim = 2 if only_position else 4 gating_threshold = chi2inv95[gating_dim] measurements = np.asarray( [detections[i].to_xyah() for i in detection_indices] ) for row, track_idx in enumerate(track_indices): track = tracks[track_idx] gating_distance = kf.gating_distance( track.mean, track.covariance, measurements, only_position ) cost_matrix[row, gating_distance > gating_threshold] = gated_cost return cost_matrix