diff --git a/README.md b/README.md index fe5edde..6edf7cd 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ If you find this repo useful in your research, please consider citing: ### Dependencies: -To install required dependencies run: +To install required dependencies for SORT and the demo, run: ``` $ pip install -r requirements.txt ``` @@ -51,7 +51,7 @@ To run the tracker with the provided detections: ``` $ cd path/to/sort -$ python sort.py +$ python demo.py ``` To display the results you need to: @@ -63,7 +63,7 @@ To display the results you need to: ``` 0. Run the demo with the ```--display``` flag ``` - $ python sort.py --display + $ python demo.py --display ``` @@ -84,7 +84,9 @@ Using the [MOT challenge devkit](https://motchallenge.net/devkit/) the method pr ### Using SORT in your own project -Below is the gist of how to instantiate and update SORT. See the ['__main__'](https://github.com/abewley/sort/blob/master/sort.py#L239) section of [sort.py](https://github.com/abewley/sort/blob/master/sort.py#L239) for a complete example. +To use SORT in your project, you can either install it locally with `pip install .` or simply copy [sort.py](sort/sort.py). + +Below is the gist of how to instantiate and update SORT. See the [`__main__`](demo.py#L28) section of [demo.py](demo.py) for a complete example. from sort import * diff --git a/demo.py b/demo.py new file mode 100644 index 0000000..50036e1 --- /dev/null +++ b/demo.py @@ -0,0 +1,85 @@ +from __future__ import print_function + +import os +import numpy as np +import matplotlib +matplotlib.use('TkAgg') +import matplotlib.pyplot as plt +import matplotlib.patches as patches +from skimage import io + +import glob +import time +import argparse + +from sort import Sort + +np.random.seed(0) + +def parse_args(): + """Parse input arguments.""" + 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("--seq_path", help="Path to detections.", type=str, default='data') + parser.add_argument("--phase", help="Subdirectory in seq_path.", type=str, default='train') + args = parser.parse_args() + return args + +if __name__ == '__main__': + # all train + args = parse_args() + display = args.display + phase = args.phase + total_time = 0.0 + total_frames = 0 + colours = np.random.rand(32, 3) #used only for display + if(display): + if not os.path.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') + exit() + plt.ion() + fig = plt.figure() + ax1 = fig.add_subplot(111, aspect='equal') + + if not os.path.exists('output'): + os.makedirs('output') + pattern = os.path.join(args.seq_path, phase, '*', 'det', 'det.txt') + for seq_dets_fn in glob.glob(pattern): + mot_tracker = Sort() #create instance of the SORT tracker + 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: + print("Processing %s."%(seq)) + for frame in range(int(seq_dets[:,0].max())): + frame += 1 #detection and frame numbers begin at 1 + dets = seq_dets[seq_dets[:, 0]==frame, 2:7] + dets[:, 2:4] += dets[:, 0:2] #convert to [x1,y1,w,h] to [x1,y1,x2,y2] + total_frames += 1 + + if(display): + fn = 'mot_benchmark/%s/%s/img1/%06d.jpg'%(phase, seq, frame) + im = io.imread(fn) + ax1.imshow(im) + plt.title(seq + ' Tracked Targets') + + start_time = time.time() + trackers = mot_tracker.update(dets) + cycle_time = time.time() - start_time + total_time += cycle_time + + for d in trackers: + print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1'%(frame,d[4],d[0],d[1],d[2]-d[0],d[3]-d[1]),file=out_file) + if(display): + 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,:])) + + if(display): + fig.canvas.flush_events() + plt.draw() + ax1.cla() + + print("Total Tracking took: %.3f seconds for %d frames or %.1f FPS" % (total_time, total_frames, total_frames / total_time)) + + if(display): + print("Note: to get real runtime results run without the option: --display") diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..89fa0a6 --- /dev/null +++ b/setup.py @@ -0,0 +1,11 @@ +from setuptools import setup, find_packages + +setup( + name='sort', + version='1.0', + author='Alex Bewley', + packages=find_packages(), + install_requires=['filterpy', 'lap'], + python_requires='>=3.6' +) + diff --git a/sort/__init__.py b/sort/__init__.py new file mode 100644 index 0000000..9d7de44 --- /dev/null +++ b/sort/__init__.py @@ -0,0 +1 @@ +from .sort import Sort diff --git a/sort.py b/sort/sort.py similarity index 72% rename from sort.py rename to sort/sort.py index 651a7f7..5a3b3b0 100644 --- a/sort.py +++ b/sort/sort.py @@ -15,19 +15,8 @@ You should have received a copy of the GNU General Public License along with this program. If not, see . """ -from __future__ import print_function -import os import numpy as np -import matplotlib -matplotlib.use('TkAgg') -import matplotlib.pyplot as plt -import matplotlib.patches as patches -from skimage import io - -import glob -import time -import argparse from filterpy.kalman import KalmanFilter try: @@ -36,8 +25,6 @@ def jit(func): return func -np.random.seed(0) - def linear_assignment(cost_matrix): try: @@ -258,70 +245,3 @@ def update(self, dets=np.empty((0, 5))): return np.concatenate(ret) return np.empty((0,5)) -def parse_args(): - """Parse input arguments.""" - 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("--seq_path", help="Path to detections.", type=str, default='data') - parser.add_argument("--phase", help="Subdirectory in seq_path.", type=str, default='train') - args = parser.parse_args() - return args - -if __name__ == '__main__': - # all train - args = parse_args() - display = args.display - phase = args.phase - total_time = 0.0 - total_frames = 0 - colours = np.random.rand(32, 3) #used only for display - if(display): - if not os.path.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') - exit() - plt.ion() - fig = plt.figure() - ax1 = fig.add_subplot(111, aspect='equal') - - if not os.path.exists('output'): - os.makedirs('output') - pattern = os.path.join(args.seq_path, phase, '*', 'det', 'det.txt') - for seq_dets_fn in glob.glob(pattern): - mot_tracker = Sort() #create instance of the SORT tracker - 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: - print("Processing %s."%(seq)) - for frame in range(int(seq_dets[:,0].max())): - frame += 1 #detection and frame numbers begin at 1 - dets = seq_dets[seq_dets[:, 0]==frame, 2:7] - dets[:, 2:4] += dets[:, 0:2] #convert to [x1,y1,w,h] to [x1,y1,x2,y2] - total_frames += 1 - - if(display): - fn = 'mot_benchmark/%s/%s/img1/%06d.jpg'%(phase, seq, frame) - im =io.imread(fn) - ax1.imshow(im) - plt.title(seq + ' Tracked Targets') - - start_time = time.time() - trackers = mot_tracker.update(dets) - cycle_time = time.time() - start_time - total_time += cycle_time - - for d in trackers: - print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1'%(frame,d[4],d[0],d[1],d[2]-d[0],d[3]-d[1]),file=out_file) - if(display): - 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,:])) - - if(display): - fig.canvas.flush_events() - plt.draw() - ax1.cla() - - print("Total Tracking took: %.3f seconds for %d frames or %.1f FPS" % (total_time, total_frames, total_frames / total_time)) - - if(display): - print("Note: to get real runtime results run without the option: --display")