security_vision/sort
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

WIP: update linear assignment.

Browse source 
- Fix for Scipy has a DeprecationWarning for linear_assignment_ module.
- LAP scales better than Scipy for frames with 100s-1000s of objects.
This commit is contained in:
Alex Bewley 2020-01-05 22:12:18 +01:00
parent 54e63a7e43
commit b69f2fc279
1 changed files with 77 additions and 56 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

69
sort.py
View file

@ -18,17 +18,22 @@
from __future__ import print_function from __future__ import print_function
from numba import jit from numba import jit
import os.path import os.path as osp
import numpy as np import numpy as np
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import matplotlib.patches as patches import matplotlib.patches as patches
from skimage import io from skimage import io
from sklearn.utils.linear_assignment_ import linear_assignment #from sklearn.utils.linear_assignment_ import linear_assignment
import lap
import glob import glob
import time import time
import argparse import argparse
from filterpy.kalman import KalmanFilter from filterpy.kalman import KalmanFilter
np.random.seed(0)
@jit @jit
def iou(bb_test, bb_gt): def iou(bb_test, bb_gt):
""" """
@ -74,7 +79,7 @@ def convert_x_to_bbox(x,score=None):
class KalmanBoxTracker(object): class KalmanBoxTracker(object):
""" """
This class represents the internel state of individual tracked objects observed as bbox. This class represents the internal state of individual tracked objects observed as bbox.
""" """
count = 0 count = 0
def __init__(self,bbox): def __init__(self,bbox):
@ -144,7 +149,20 @@ def associate_detections_to_trackers(detections,trackers,iou_threshold = 0.3):
for d,det in enumerate(detections): for d,det in enumerate(detections):
for t,trk in enumerate(trackers): for t,trk in enumerate(trackers):
iou_matrix[d,t] = iou(det,trk) iou_matrix[d,t] = iou(det,trk)
matched_indices = linear_assignment(-iou_matrix)
# TODO (bewley): remove rows and cols iou.max() < threshold
if min(iou_matrix.shape) > 0:
a = (iou_matrix > iou_threshold).astype(np.int32)
if a.sum(1).max() == 1 and a.sum(0).max() == 1:
#matrix is #TODO(this doesnt provide much gains)
matched_indices = np.stack(np.where(a), axis=1)
else:
_, x, y = lap.lapjv(-iou_matrix, extend_cost=True)
matched_indices = np.array([[y[i],i] for i in x if i >= 0]) #
else:
matched_indices = np.empty(shape=(0,2))
#matched_indices = linear_assignment(-iou_matrix)
unmatched_detections = [] unmatched_detections = []
for d, det in enumerate(detections): for d, det in enumerate(detections):
@ -171,7 +189,6 @@ def associate_detections_to_trackers(detections,trackers,iou_threshold = 0.3):
return matches, np.array(unmatched_detections), np.array(unmatched_trackers) return matches, np.array(unmatched_detections), np.array(unmatched_trackers)
class Sort(object): class Sort(object):
def __init__(self, max_age=1, min_hits=3): def __init__(self, max_age=1, min_hits=3):
""" """
@ -182,11 +199,11 @@ class Sort(object):
self.trackers = [] self.trackers = []
self.frame_count = 0 self.frame_count = 0
def update(self,dets): def update(self, dets=np.empty((0, 5))):
""" """
Params: Params:
dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...] dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]
Requires: this method must be called once for each frame even with empty detections. Requires: this method must be called once for each frame even with empty detections (use np.empty((0, 5)) for frames without detections).
Returns the a similar array, where the last column is the object ID. Returns the a similar array, where the last column is the object ID.
NOTE: The number of objects returned may differ from the number of detections provided. NOTE: The number of objects returned may differ from the number of detections provided.
@ -199,7 +216,7 @@ class Sort(object):
for t, trk in enumerate(trks): for t, trk in enumerate(trks):
pos = self.trackers[t].predict()[0] pos = self.trackers[t].predict()[0]
trk[:] = [pos[0], pos[1], pos[2], pos[3], 0] trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
if(np.any(np.isnan(pos))): if np.any(np.isnan(pos)):
to_del.append(t) to_del.append(t)
trks = np.ma.compress_rows(np.ma.masked_invalid(trks)) trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
for t in reversed(to_del): for t in reversed(to_del):
@ -207,10 +224,12 @@ class Sort(object):
matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(dets,trks) matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(dets,trks)
# update matched trackers with assigned detections # update matched trackers with assigned detections
for t,trk in enumerate(self.trackers): for m in matched:
if(t not in unmatched_trks): self.trackers[m[1]].update(dets[m[0], :])
d = matched[np.where(matched[:,1]==t)[0],0] #for t, trk in enumerate(self.trackers):
trk.update(dets[d,:][0]) # if(t not in unmatched_trks):
# d = matched[np.where(matched[:,1]==t)[0],0]
# trk.update(dets[d,:][0])
# create and initialise new trackers for unmatched detections # create and initialise new trackers for unmatched detections
for i in unmatched_dets: for i in unmatched_dets:
@ -219,7 +238,7 @@ class Sort(object):
i = len(self.trackers) i = len(self.trackers)
for trk in reversed(self.trackers): for trk in reversed(self.trackers):
d = trk.get_state()[0] d = trk.get_state()[0]
if((trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits)): if (trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits):
ret.append(np.concatenate((d,[trk.id+1])).reshape(1,-1)) # +1 as MOT benchmark requires positive ret.append(np.concatenate((d,[trk.id+1])).reshape(1,-1)) # +1 as MOT benchmark requires positive
i -= 1 i -= 1
# remove dead tracklet # remove dead tracklet
@ -233,31 +252,35 @@ def parse_args():
"""Parse input arguments.""" """Parse input arguments."""
parser = argparse.ArgumentParser(description='SORT demo') parser = argparse.ArgumentParser(description='SORT demo')
parser.add_argument('--display', dest='display', help='Display online tracker output (slow) [False]',action='store_true') parser.add_argument('--display', dest='display', help='Display online tracker output (slow) [False]',action='store_true')
parser.add_argument("--seq_path", help="Path to detections.", type=str, default='mot_benchmark')
parser.add_argument("--phase", help="Subdirectory in seq_path.", type=str, default='train')
args = parser.parse_args() args = parser.parse_args()
return args return args
if __name__ == '__main__': if __name__ == '__main__':
# all train # all train
sequences = ['PETS09-S2L1','TUD-Campus','TUD-Stadtmitte','ETH-Bahnhof','ETH-Sunnyday','ETH-Pedcross2','KITTI-13','KITTI-17','ADL-Rundle-6','ADL-Rundle-8','Venice-2']
args = parse_args() args = parse_args()
display = args.display display = args.display
phase = 'train' phase = args.phase
total_time = 0.0 total_time = 0.0
total_frames = 0 total_frames = 0
colours = np.random.rand(32, 3) #used only for display colours = np.random.rand(32, 3) #used only for display
if(display): if(display):
if not os.path.exists('mot_benchmark'): if not osp.exists('mot_benchmark'):
print('\n\tERROR: mot_benchmark link not found!\n\n Create a symbolic link to the MOT benchmark\n (https://motchallenge.net/data/2D_MOT_2015/#download). E.g.:\n\n $ ln -s /path/to/MOT2015_challenge/2DMOT2015 mot_benchmark\n\n') print('\n\tERROR: mot_benchmark link not found!\n\n Create a symbolic link to the MOT benchmark\n (https://motchallenge.net/data/2D_MOT_2015/#download). E.g.:\n\n $ ln -s /path/to/MOT2015_challenge/2DMOT2015 mot_benchmark\n\n')
exit() exit()
plt.ion() plt.ion()
fig = plt.figure() fig = plt.figure()
if not os.path.exists('output'): if not osp.exists('output'):
os.makedirs('output') os.makedirs('output')
pattern = osp.join(args.seq_path, phase, '*', 'det', 'det.txt')
for seq in sequences: for seq_dets_fn in glob.glob(pattern):
mot_tracker = Sort() #create instance of the SORT tracker mot_tracker = Sort() #create instance of the SORT tracker
seq_dets = np.loadtxt('data/%s/det.txt'%(seq),delimiter=',') #load detections print(seq_dets_fn)
seq_dets = np.loadtxt(seq_dets_fn, delimiter=',')
seq = seq_dets_fn[pattern.find('*'):].split('/')[0]
with open('output/%s.txt'%(seq),'w') as out_file: with open('output/%s.txt'%(seq),'w') as out_file:
print("Processing %s."%(seq)) print("Processing %s."%(seq))
for frame in range(int(seq_dets[:,0].max())): for frame in range(int(seq_dets[:,0].max())):
@ -283,7 +306,7 @@ if __name__ == '__main__':
if(display): if(display):
d = d.astype(np.int32) d = d.astype(np.int32)
ax1.add_patch(patches.Rectangle((d[0],d[1]),d[2]-d[0],d[3]-d[1],fill=False,lw=3,ec=colours[d[4]%32,:])) ax1.add_patch(patches.Rectangle((d[0],d[1]),d[2]-d[0],d[3]-d[1],fill=False,lw=3,ec=colours[d[4]%32,:]))
ax1.set_adjustable('box-forced') #ax1.set_adjustable('box-forced')
if(display): if(display):
fig.canvas.flush_events() fig.canvas.flush_events()
@ -291,8 +314,6 @@ if __name__ == '__main__':
ax1.cla() ax1.cla()
print("Total Tracking took: %.3f for %d frames or %.1f FPS" % (total_time, total_frames, total_frames / total_time)) print("Total Tracking took: %.3f for %d frames or %.1f FPS" % (total_time, total_frames, total_frames / total_time))
if(display): if(display):
print("Note: to get real runtime results run without the option: --display") print("Note: to get real runtime results run without the option: --display")