deflinear_assignment(cost_matrix:np.ndarray,thresh:float,use_lap:bool=True)->tuple:""" Perform linear assignment using scipy or lap.lapjv. Args: cost_matrix (np.ndarray): The matrix containing cost values for assignments. thresh (float): Threshold for considering an assignment valid. use_lap (bool, optional): Whether to use lap.lapjv. Defaults to True. Returns: Tuple with: - matched indices - unmatched indices from 'a' - unmatched indices from 'b' """ifcost_matrix.size==0:returnnp.empty((0,2),dtype=int),tuple(range(cost_matrix.shape[0])),tuple(range(cost_matrix.shape[1]))ifuse_lap:# Use lap.lapjv# https://github.com/gatagat/lap_,x,y=lap.lapjv(cost_matrix,extend_cost=True,cost_limit=thresh)matches=[[ix,mx]forix,mxinenumerate(x)ifmx>=0]unmatched_a=np.where(x<0)[0]unmatched_b=np.where(y<0)[0]else:# Use scipy.optimize.linear_sum_assignment# https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.linear_sum_assignment.htmlx,y=scipy.optimize.linear_sum_assignment(cost_matrix)# row x, col ymatches=np.asarray([[x[i],y[i]]foriinrange(len(x))ifcost_matrix[x[i],y[i]]<=thresh])iflen(matches)==0:unmatched_a=list(np.arange(cost_matrix.shape[0]))unmatched_b=list(np.arange(cost_matrix.shape[1]))else:unmatched_a=list(set(np.arange(cost_matrix.shape[0]))-set(matches[:,0]))unmatched_b=list(set(np.arange(cost_matrix.shape[1]))-set(matches[:,1]))returnmatches,unmatched_a,unmatched_b
defiou_distance(atracks:list,btracks:list)->np.ndarray:""" Compute cost based on Intersection over Union (IoU) between tracks. Args: atracks (list[STrack] | list[np.ndarray]): List of tracks 'a' or bounding boxes. btracks (list[STrack] | list[np.ndarray]): List of tracks 'b' or bounding boxes. Returns: (np.ndarray): Cost matrix computed based on IoU. """ifatracksandisinstance(atracks[0],np.ndarray)orbtracksandisinstance(btracks[0],np.ndarray):atlbrs=atracksbtlbrs=btrackselse:atlbrs=[track.xywhaiftrack.angleisnotNoneelsetrack.xyxyfortrackinatracks]btlbrs=[track.xywhaiftrack.angleisnotNoneelsetrack.xyxyfortrackinbtracks]ious=np.zeros((len(atlbrs),len(btlbrs)),dtype=np.float32)iflen(atlbrs)andlen(btlbrs):iflen(atlbrs[0])==5andlen(btlbrs[0])==5:ious=batch_probiou(np.ascontiguousarray(atlbrs,dtype=np.float32),np.ascontiguousarray(btlbrs,dtype=np.float32),).numpy()else:ious=bbox_ioa(np.ascontiguousarray(atlbrs,dtype=np.float32),np.ascontiguousarray(btlbrs,dtype=np.float32),iou=True,)return1-ious# cost matrix
defembedding_distance(tracks:list,detections:list,metric:str="cosine")->np.ndarray:""" Compute distance between tracks and detections based on embeddings. Args: tracks (list[STrack]): List of tracks. detections (list[BaseTrack]): List of detections. metric (str, optional): Metric for distance computation. Defaults to 'cosine'. Returns: (np.ndarray): Cost matrix computed based on embeddings. """cost_matrix=np.zeros((len(tracks),len(detections)),dtype=np.float32)ifcost_matrix.size==0:returncost_matrixdet_features=np.asarray([track.curr_featfortrackindetections],dtype=np.float32)# for i, track in enumerate(tracks):# cost_matrix[i, :] = np.maximum(0.0, cdist(track.smooth_feat.reshape(1,-1), det_features, metric))track_features=np.asarray([track.smooth_featfortrackintracks],dtype=np.float32)cost_matrix=np.maximum(0.0,cdist(track_features,det_features,metric))# Normalized featuresreturncost_matrix
deffuse_score(cost_matrix:np.ndarray,detections:list)->np.ndarray:""" Fuses cost matrix with detection scores to produce a single similarity matrix. Args: cost_matrix (np.ndarray): The matrix containing cost values for assignments. detections (list[BaseTrack]): List of detections with scores. Returns: (np.ndarray): Fused similarity matrix. """ifcost_matrix.size==0:returncost_matrixiou_sim=1-cost_matrixdet_scores=np.array([det.scorefordetindetections])det_scores=np.expand_dims(det_scores,axis=0).repeat(cost_matrix.shape[0],axis=0)fuse_sim=iou_sim*det_scoresreturn1-fuse_sim# fuse_cost