FairMOT在线跟踪框架及代码学习记录

在线跟踪框架及代码学习记录"/>

FairMOT在线跟踪框架及代码学习记录

在线跟踪匹配流程图

在线跟踪匹配源码

参考： .py

def update(self, im_blob, img0):self.frame_id += 1activated_starcks = []refind_stracks = []lost_stracks = []removed_stracks = []width = img0.shape[1]height = img0.shape[0]inp_height = im_blob.shape[2]inp_width = im_blob.shape[3]c = np.array([width / 2., height / 2.], dtype=np.float32)s = max(float(inp_width) / float(inp_height) * height, width) * 1.0meta = {'c': c, 's': s,'out_height': inp_height // self.opt.down_ratio,'out_width': inp_width // self.opt.down_ratio}''' Step 1: Network forward, get detections & embeddings'''with torch.no_grad():output = self.model(im_blob)[-1]hm = output['hm'].sigmoid_()wh = output['wh']id_feature = output['id']id_feature = F.normalize(id_feature, dim=1)reg = output['reg'] if self.opt.reg_offset else Nonedets, inds = mot_decode(hm, wh, reg=reg, ltrb=self.opt.ltrb, K=self.opt.K)id_feature = _tranpose_and_gather_feat(id_feature, inds)id_feature = id_feature.squeeze(0)id_feature = id_feature.cpu().numpy()dets = self.post_process(dets, meta)dets = self.merge_outputs([dets])[1]remain_inds = dets[:, 4] > self.opt.conf_thresdets = 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 = []tracked_stracks = []for track in self.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, self.lost_stracks)# Predict the current location with KFSTrack.multi_predict(strack_pool)dists = matching.embedding_distance(strack_pool, detections)dists = matching.fuse_motion(self.kalman_filter, dists, strack_pool, detections)matches, u_track, u_detection = matching.linear_assignment(dists, thresh=0.4)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: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.5)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:continuetrack.activate(self.kalman_filter, self.frame_id)activated_starcks.append(track)""" Step 5: Update state"""for track in self.lost_stracks:if self.frame_id - track.end_frame > self.max_time_lost:track.mark_removed()removed_stracks.append(track)