Add motion scores
This commit is contained in:
parent
27809fd27b
commit
57533d82c9
3 changed files with 25 additions and 6 deletions
|
@ -195,9 +195,9 @@ class YOLOLayer(nn.Module):
|
||||||
|
|
||||||
else:
|
else:
|
||||||
p_conf = torch.softmax(p_conf, dim=1)[:,1,...].unsqueeze(-1)
|
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 = F.normalize(p_emb.unsqueeze(1).repeat(1,self.nA,1,1,1).contiguous(), dim=-1)
|
||||||
p_emb_up = shift_tensor_vertically(p_emb, -self.shift[self.yolo_layer])
|
#p_emb_up = F.normalize(shift_tensor_vertically(p_emb, -self.shift[self.layer]), dim=-1)
|
||||||
p_emb_down = shift_tensor_vertically(p_emb, self.shift[self.yolo_layer])
|
#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_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], dim=-1)
|
||||||
#p = torch.cat([p_box, p_conf, p_cls, p_emb, p_emb_up, p_emb_down], dim=-1)
|
#p = torch.cat([p_box, p_conf, p_cls, p_emb, p_emb_up, p_emb_down], dim=-1)
|
||||||
|
|
|
@ -120,3 +120,20 @@ def gate_cost_matrix(kf, cost_matrix, tracks, detections, only_position=False):
|
||||||
track.mean, track.covariance, measurements, only_position)
|
track.mean, track.covariance, measurements, only_position)
|
||||||
cost_matrix[row, gating_distance > gating_threshold] = np.inf
|
cost_matrix[row, gating_distance > gating_threshold] = np.inf
|
||||||
return cost_matrix
|
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
|
||||||
|
|
|
@ -34,6 +34,7 @@ class STrack(BaseTrack):
|
||||||
self.alpha = 0.9
|
self.alpha = 0.9
|
||||||
|
|
||||||
def update_features(self, feat):
|
def update_features(self, feat):
|
||||||
|
feat /= np.linalg.norm(feat)
|
||||||
self.curr_feat = feat
|
self.curr_feat = feat
|
||||||
if self.smooth_feat is None:
|
if self.smooth_feat is None:
|
||||||
self.smooth_feat = feat
|
self.smooth_feat = feat
|
||||||
|
@ -186,7 +187,7 @@ class JDETracker(object):
|
||||||
scale_coords(self.opt.img_size, dets[:, :4], img0.shape).round()
|
scale_coords(self.opt.img_size, dets[:, :4], img0.shape).round()
|
||||||
'''Detections'''
|
'''Detections'''
|
||||||
detections = [STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f.numpy(), 30) for
|
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:
|
else:
|
||||||
detections = []
|
detections = []
|
||||||
|
|
||||||
|
@ -209,7 +210,8 @@ class JDETracker(object):
|
||||||
strack.predict()
|
strack.predict()
|
||||||
|
|
||||||
dists = matching.embedding_distance(strack_pool, detections)
|
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)
|
matches, u_track, u_detection = matching.linear_assignment(dists, thresh=0.7)
|
||||||
|
|
||||||
for itracked, idet in matches:
|
for itracked, idet in matches:
|
||||||
|
|
Loading…
Reference in a new issue