in container-batch-inference/resources/FairMOT/multitracker.py [0:0]
def associate_tracker(self, dets, meta, id_feature, stream_id):
activated_starcks = []
refind_stracks = []
lost_stracks = []
removed_stracks = []
tracked_stracks = []
frame_tracker = self.frame_trackers[stream_id]
dets = self.post_process(dets, meta)
dets = self.merge_outputs([dets])[1]
remain_inds = dets[:, 4] > self.opt.conf_thres
dets = dets[remain_inds]
id_feature = id_feature[remain_inds]
if len(dets) > 0:
'''Detections'''
detections = [STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f, 30) for
(tlbrs, f) in zip(dets[:, :5], id_feature)]
else:
detections = []
''' Add newly detected tracklets to tracked_stracks'''
unconfirmed = []
for track in frame_tracker.tracked_stracks:
if not track.is_activated:
unconfirmed.append(track)
else:
tracked_stracks.append(track)
''' Step 2: First association, with embedding'''
strack_pool = joint_stracks(tracked_stracks, frame_tracker.lost_stracks)
# Predict the current location with KF
#for strack in strack_pool:
#strack.predict()
STrack.multi_predict(strack_pool)
dists = matching.embedding_distance(strack_pool, detections)
#dists = matching.iou_distance(strack_pool, detections)
dists = matching.fuse_motion(frame_tracker.kalman_filter, dists, strack_pool, detections)
matches, u_track, u_detection = matching.linear_assignment(dists, thresh=0.8)
for itracked, idet in matches:
track = strack_pool[itracked]
det = detections[idet]
if track.state == TrackState.Tracked:
track.update(detections[idet], self.frame_id)
activated_starcks.append(track)
else:
# Frame_id is same for each video stream
track.re_activate(det, self.frame_id, new_id=False)
refind_stracks.append(track)
''' Step 3: Second association, with IOU'''
detections = [detections[i] for i in u_detection]
r_tracked_stracks = [strack_pool[i] for i in u_track if strack_pool[i].state == TrackState.Tracked]
dists = matching.iou_distance(r_tracked_stracks, detections)
matches, u_track, u_detection = matching.linear_assignment(dists, thresh=0.7)
for itracked, idet in matches:
track = r_tracked_stracks[itracked]
det = detections[idet]
if track.state == TrackState.Tracked:
track.update(det, self.frame_id)
activated_starcks.append(track)
else:
track.re_activate(det, self.frame_id, new_id=False)
refind_stracks.append(track)
for it in u_track:
track = r_tracked_stracks[it]
if not track.state == TrackState.Lost:
track.mark_lost()
lost_stracks.append(track)
'''Deal with unconfirmed tracks, usually tracks with only one beginning frame'''
detections = [detections[i] for i in u_detection]
dists = matching.iou_distance(unconfirmed, detections)
matches, u_unconfirmed, u_detection = matching.linear_assignment(dists, thresh=0.7)
for itracked, idet in matches:
unconfirmed[itracked].update(detections[idet], self.frame_id)
activated_starcks.append(unconfirmed[itracked])
for it in u_unconfirmed:
track = unconfirmed[it]
track.mark_removed()
removed_stracks.append(track)
""" Step 4: Init new stracks"""
for inew in u_detection:
track = detections[inew]
if track.score < self.det_thresh:
continue
track.activate(frame_tracker.kalman_filter, self.frame_id)
activated_starcks.append(track)
""" Step 5: Update state"""
for track in frame_tracker.lost_stracks:
if self.frame_id - track.end_frame > self.max_time_lost:
track.mark_removed()
removed_stracks.append(track)
frame_tracker.tracked_stracks = [t for t in frame_tracker.tracked_stracks if t.state == TrackState.Tracked]
frame_tracker.tracked_stracks = joint_stracks(frame_tracker.tracked_stracks, activated_starcks)
frame_tracker.tracked_stracks = joint_stracks(frame_tracker.tracked_stracks, refind_stracks)
frame_tracker.lost_stracks = sub_stracks(frame_tracker.lost_stracks,
frame_tracker.tracked_stracks)
frame_tracker.lost_stracks.extend(lost_stracks)
frame_tracker.lost_stracks = sub_stracks(frame_tracker.lost_stracks,
frame_tracker.removed_stracks)
frame_tracker.removed_stracks.extend(removed_stracks)
frame_tracker.tracked_stracks, frame_tracker.lost_stracks = remove_duplicate_stracks(
frame_tracker.tracked_stracks, frame_tracker.lost_stracks)
# get scores of lost tracks
self.frame_trackers[stream_id] = frame_tracker
logging.info('tracker_stracker: {}, removed_stracks: {}, lost_stracks: {}'.format(
len(frame_tracker.tracked_stracks),
len(frame_tracker.removed_stracks),
len(frame_tracker.lost_stracks)))
output_stracks = [track for track in frame_tracker.tracked_stracks if track.is_activated]
logging.debug('===========Frame {}=========='.format(self.frame_id))
logging.debug('Activated: {}'.format([track.track_id for track in activated_starcks]))
logging.debug('Refind: {}'.format([track.track_id for track in refind_stracks]))
logging.debug('Lost: {}'.format([track.track_id for track in lost_stracks]))
logging.debug('Removed: {}'.format([track.track_id for track in removed_stracks]))
return output_stracks