AK391
all files
7734d5b
raw
history blame
7.49 kB
"""
Copyright (c) https://github.com/xingyizhou/CenterTrack
Modified by Peize Sun, Rufeng Zhang
"""
# coding: utf-8
import torch
from scipy.optimize import linear_sum_assignment
from util import box_ops
import copy
class Tracker(object):
def __init__(self, score_thresh, max_age=32):
self.score_thresh = score_thresh
self.low_thresh = 0.2
self.high_thresh = score_thresh + 0.1
self.max_age = max_age
self.id_count = 0
self.tracks_dict = dict()
self.tracks = list()
self.unmatched_tracks = list()
self.reset_all()
def reset_all(self):
self.id_count = 0
self.tracks_dict = dict()
self.tracks = list()
self.unmatched_tracks = list()
def init_track(self, results):
scores = results["scores"]
classes = results["labels"]
bboxes = results["boxes"] # x1y1x2y2
ret = list()
ret_dict = dict()
for idx in range(scores.shape[0]):
if scores[idx] >= self.score_thresh:
self.id_count += 1
obj = dict()
obj["score"] = float(scores[idx])
obj["bbox"] = bboxes[idx, :].cpu().numpy().tolist()
obj["tracking_id"] = self.id_count
obj['active'] = 1
obj['age'] = 1
ret.append(obj)
ret_dict[idx] = obj
self.tracks = ret
self.tracks_dict = ret_dict
return copy.deepcopy(ret)
def step(self, output_results):
scores = output_results["scores"]
bboxes = output_results["boxes"] # x1y1x2y2
track_bboxes = output_results["track_boxes"] if "track_boxes" in output_results else None # x1y1x2y2
results = list()
results_dict = dict()
results_second = list()
tracks = list()
for idx in range(scores.shape[0]):
if idx in self.tracks_dict and track_bboxes is not None:
self.tracks_dict[idx]["bbox"] = track_bboxes[idx, :].cpu().numpy().tolist()
if scores[idx] >= self.score_thresh:
obj = dict()
obj["score"] = float(scores[idx])
obj["bbox"] = bboxes[idx, :].cpu().numpy().tolist()
results.append(obj)
results_dict[idx] = obj
elif scores[idx] >= self.low_thresh:
second_obj = dict()
second_obj["score"] = float(scores[idx])
second_obj["bbox"] = bboxes[idx, :].cpu().numpy().tolist()
results_second.append(second_obj)
results_dict[idx] = second_obj
tracks = [v for v in self.tracks_dict.values()] + self.unmatched_tracks
# for trackss in tracks:
# print(trackss.keys())
N = len(results)
M = len(tracks)
ret = list()
unmatched_tracks = [t for t in range(M)]
unmatched_dets = [d for d in range(N)]
if N > 0 and M > 0:
det_box = torch.stack([torch.tensor(obj['bbox']) for obj in results], dim=0) # N x 4
track_box = torch.stack([torch.tensor(obj['bbox']) for obj in tracks], dim=0) # M x 4
cost_bbox = 1.0 - box_ops.generalized_box_iou(det_box, track_box) # N x M
matched_indices = linear_sum_assignment(cost_bbox)
unmatched_dets = [d for d in range(N) if not (d in matched_indices[0])]
unmatched_tracks = [d for d in range(M) if not (d in matched_indices[1])]
matches = [[],[]]
for (m0, m1) in zip(matched_indices[0], matched_indices[1]):
if cost_bbox[m0, m1] > 1.2:
unmatched_dets.append(m0)
unmatched_tracks.append(m1)
else:
matches[0].append(m0)
matches[1].append(m1)
for (m0, m1) in zip(matches[0], matches[1]):
track = results[m0]
track['tracking_id'] = tracks[m1]['tracking_id']
track['age'] = 1
track['active'] = 1
ret.append(track)
# second association
N_second = len(results_second)
unmatched_tracks_obj = list()
for i in unmatched_tracks:
#print(tracks[i].keys())
track = tracks[i]
if track['active'] == 1:
unmatched_tracks_obj.append(track)
M_second = len(unmatched_tracks_obj)
unmatched_tracks_second = [t for t in range(M_second)]
if N_second > 0 and M_second > 0:
det_box_second = torch.stack([torch.tensor(obj['bbox']) for obj in results_second], dim=0) # N_second x 4
track_box_second = torch.stack([torch.tensor(obj['bbox']) for obj in unmatched_tracks_obj], dim=0) # M_second x 4
cost_bbox_second = 1.0 - box_ops.generalized_box_iou(det_box_second, track_box_second) # N_second x M_second
matched_indices_second = linear_sum_assignment(cost_bbox_second)
unmatched_tracks_second = [d for d in range(M_second) if not (d in matched_indices_second[1])]
matches_second = [[],[]]
for (m0, m1) in zip(matched_indices_second[0], matched_indices_second[1]):
if cost_bbox_second[m0, m1] > 0.8:
unmatched_tracks_second.append(m1)
else:
matches_second[0].append(m0)
matches_second[1].append(m1)
for (m0, m1) in zip(matches_second[0], matches_second[1]):
track = results_second[m0]
track['tracking_id'] = unmatched_tracks_obj[m1]['tracking_id']
track['age'] = 1
track['active'] = 1
ret.append(track)
for i in unmatched_dets:
trackd = results[i]
if trackd["score"] >= self.high_thresh:
self.id_count += 1
trackd['tracking_id'] = self.id_count
trackd['age'] = 1
trackd['active'] = 1
ret.append(trackd)
# ------------------------------------------------------ #
ret_unmatched_tracks = []
for j in unmatched_tracks:
track = tracks[j]
if track['active'] == 0 and track['age'] < self.max_age:
track['age'] += 1
track['active'] = 0
ret.append(track)
ret_unmatched_tracks.append(track)
for i in unmatched_tracks_second:
track = unmatched_tracks_obj[i]
if track['age'] < self.max_age:
track['age'] += 1
track['active'] = 0
ret.append(track)
ret_unmatched_tracks.append(track)
# for i in unmatched_tracks:
# track = tracks[i]
# if track['age'] < self.max_age:
# track['age'] += 1
# track['active'] = 0
# ret.append(track)
# ret_unmatched_tracks.append(track)
#print(len(ret_unmatched_tracks))
self.tracks = ret
self.tracks_dict = {red_ind:red for red_ind, red in results_dict.items() if 'tracking_id' in red}
self.unmatched_tracks = ret_unmatched_tracks
return copy.deepcopy(ret)