Add motion scores

This commit is contained in:
Zhongdao 2020-01-10 19:37:29 +08:00
parent 27809fd27b
commit 57533d82c9
3 changed files with 25 additions and 6 deletions

View File

@ -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)

View File

@ -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

View File

@ -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: