From 57533d82c917f0d8ab75d6cfe0a728f771996869 Mon Sep 17 00:00:00 2001 From: Zhongdao Date: Fri, 10 Jan 2020 19:37:29 +0800 Subject: [PATCH] Add motion scores --- models.py | 6 +++--- tracker/matching.py | 17 +++++++++++++++++ tracker/multitracker.py | 8 +++++--- 3 files changed, 25 insertions(+), 6 deletions(-) diff --git a/models.py b/models.py index acd8fc4..2537f66 100644 --- a/models.py +++ b/models.py @@ -195,9 +195,9 @@ class YOLOLayer(nn.Module): else: p_conf = torch.softmax(p_conf, dim=1)[:,1,...].unsqueeze(-1) - p_emb = p_emb.unsqueeze(1).repeat(1,self.nA,1,1,1).contiguous() - p_emb_up = shift_tensor_vertically(p_emb, -self.shift[self.yolo_layer]) - p_emb_down = shift_tensor_vertically(p_emb, self.shift[self.yolo_layer]) + p_emb = F.normalize(p_emb.unsqueeze(1).repeat(1,self.nA,1,1,1).contiguous(), dim=-1) + #p_emb_up = F.normalize(shift_tensor_vertically(p_emb, -self.shift[self.layer]), dim=-1) + #p_emb_down = F.normalize(shift_tensor_vertically(p_emb, self.shift[self.layer]), dim=-1) p_cls = torch.zeros(nB,self.nA,nGh,nGw,1).cuda() # Temp p = torch.cat([p_box, p_conf, p_cls, p_emb], dim=-1) #p = torch.cat([p_box, p_conf, p_cls, p_emb, p_emb_up, p_emb_down], dim=-1) diff --git a/tracker/matching.py b/tracker/matching.py index 7456c4f..d4b583a 100644 --- a/tracker/matching.py +++ b/tracker/matching.py @@ -120,3 +120,20 @@ def gate_cost_matrix(kf, cost_matrix, tracks, detections, only_position=False): track.mean, track.covariance, measurements, only_position) cost_matrix[row, gating_distance > gating_threshold] = np.inf return cost_matrix + + +def fuse_motion(kf, cost_matrix, tracks, detections, only_position=False, lambda_=0.98): + if cost_matrix.size == 0: + return cost_matrix + gating_dim = 2 if only_position else 4 + gating_threshold = kalman_filter.chi2inv95[gating_dim] + measurements = np.asarray([det.to_xyah() for det in detections]) + for row, track in enumerate(tracks): + gating_distance = kf.gating_distance( + track.mean, track.covariance, measurements, only_position) + cost_matrix[row, gating_distance > gating_threshold] = np.inf + #print(cost_matrix[row]) + #print(gating_distance) + #print('-'*90) + cost_matrix[row] = lambda_ * cost_matrix[row] + (1-lambda_)* gating_distance + return cost_matrix diff --git a/tracker/multitracker.py b/tracker/multitracker.py index 7bf2682..84a92ee 100644 --- a/tracker/multitracker.py +++ b/tracker/multitracker.py @@ -34,7 +34,8 @@ class STrack(BaseTrack): self.alpha = 0.9 def update_features(self, feat): - self.curr_feat = feat + feat /= np.linalg.norm(feat) + self.curr_feat = feat if self.smooth_feat is None: self.smooth_feat = feat else: @@ -186,7 +187,7 @@ class JDETracker(object): scale_coords(self.opt.img_size, dets[:, :4], img0.shape).round() '''Detections''' detections = [STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f.numpy(), 30) for - (tlbrs, f) in zip(dets[:, :5], dets[:, -self.model.emb_dim:])] + (tlbrs, f) in zip(dets[:, :5], dets[:, 6:])] else: detections = [] @@ -209,7 +210,8 @@ class JDETracker(object): strack.predict() dists = matching.embedding_distance(strack_pool, detections) - dists = matching.gate_cost_matrix(self.kalman_filter, dists, strack_pool, detections) + #dists = matching.gate_cost_matrix(self.kalman_filter, dists, 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.7) for itracked, idet in matches: