-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathdet_tracking.py
135 lines (114 loc) · 5.76 KB
/
det_tracking.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jan 10 18:42:29 2019
@author: rwzhang
"""
# =============================================================================
# Pre-Launch Check: Make sure Detections are done and the result is saved to the designated Place
# This is dedicated to run Tracking on pre-existing detection result
# If Detection result is saved in JSON format, please launch from ../yolo_video.py --tracking 1
# =============================================================================
from time import time
import argparse
from progressbar import ProgressBar
from util import load_mot, iou, save_to_csv
import pandas as pd
import os
def track_iou(detections, sigma_l, sigma_h, sigma_iou, t_min):
"""
Simple IOU based tracker.
See "High-Speed Tracking-by-Detection Without Using Image Information by E. Bochinski, V. Eiselein, T. Sikora" for
more information.
Args:
detections (list): list of detections per frame, usually generated by util.load_mot
sigma_l (float): low detection threshold.
sigma_h (float): high detection threshold.
sigma_iou (float): IOU threshold.
t_min (float): minimum track length in frames.
Returns:
list: list of tracks.
"""
tracks_active = []
tracks_finished = []
pbar = ProgressBar(maxval=len(detections))
pbar.start()
for frame_num, detections_frame in enumerate(detections):
# for frame_num in range(5):
# detections_frame = detections[frame_num]
# update the starting point to be 0 not 1
# apply low threshold to detections
dets = [det for det in detections_frame if det['score'] >= sigma_l]
updated_tracks = []
for track in tracks_active:
if len(dets) > 0:
# get det with highest iou
print("current_frame ", frame_num, track['bboxes'])
best_match = max(dets, key=lambda x: iou(track['bboxes'][-1], x['bbox']))
if iou(track['bboxes'][-1], best_match['bbox']) >= sigma_iou:
track['bboxes'].append(best_match['bbox'])
track['max_score'] = max(track['max_score'], best_match['score'])
#added class info:
track['class'] = track['class']
updated_tracks.append(track)
# remove from best matching detection from detections
del dets[dets.index(best_match)]
# if track was not updated
if len(updated_tracks) == 0 or track is not updated_tracks[-1]:
# finish track when the conditions are met
if track['max_score'] >= sigma_h and len(track['bboxes']) >= t_min:
tracks_finished.append(track)
# create new tracks
new_tracks = [{'bboxes': [det['bbox']], 'max_score': det['score'], 'start_frame': frame_num, 'class': det['class']} for det in dets]
tracks_active = updated_tracks + new_tracks
pbar.update(frame_num)
pbar.finish()
# finish all remaining active tracks
tracks_finished += [track for track in tracks_active
if track['max_score'] >= sigma_h and len(track['bboxes']) >= t_min]
return tracks_finished
def idorder_frameorder(orig_out_path, out_path):
'''
orig_out_path: Tracking result from iou_tracking ordred in ID order
out_path: Convert iou_tracking result to frame order
'''
temp_df = pd.read_table(orig_out_path, sep = ',', header=None, names = ['frame', 'id', 'x', 'y', 'w', 'h', 'score', 'wx', 'wy', 'wz', 'class'])
temp_df.sort_values(by=['frame', 'id'], inplace=True)
open(out_path, 'w').close()
temp_df.to_csv(out_path, sep = ',', header=None, index = False)
def main(args):
print('loading detections...')
det_path = os.path.expanduser('~/') + args.detection_path
detections = load_mot(det_path)
print('tracking...')
start = time()
tracks = track_iou(detections, args.sigma_l, args.sigma_h, args.sigma_iou, args.t_min)
end = time()
print("Finished Tracking", tracks)
num_frames = len(detections)
print("finished at " + str(int(num_frames / (end - start))) + " fps!")
orig_out_path = os.path.expanduser('~/') + args.output_path
save_to_csv(orig_out_path, tracks)
out_path = os.path.expanduser('~/') + args.output_path + 'ordredfrm'
if args.frameorder:
#reading in all the data again for re-formatting
idorder_frameorder(orig_out_path, out_path)
if __name__ == '__main__':
print('Usaage: python det_tracking.py --detection_path "2018AICity_TeamUW/data/Track3/Loc1_1/det.txt" --output "2018AICity_TeamUW/data/Track3/Loc1_1/res.txt"')
parser = argparse.ArgumentParser(description="IOU Tracker demo script")
parser.add_argument('-d', '--detection_path', type=str, required=True,
help="full path to CSV file containing the detections")
parser.add_argument('-o', '--output_path', type=str, required=True,
help="output path to store the tracking results (MOT challenge devkit compatible format)")
parser.add_argument('-sl', '--sigma_l', type=float, default=0,
help="low detection threshold")
parser.add_argument('-sh', '--sigma_h', type=float, default=0.5,
help="high detection threshold")
parser.add_argument('-si', '--sigma_iou', type=float, default=0.5,
help="intersection-over-union threshold")
parser.add_argument('-tm', '--t_min', type=float, default=2,
help="minimum track length")
parser.add_argument('-fo', '--frameorder', type=bool, default=False,
help="Option to convert output to be in frame order")
args = parser.parse_args()
main(args)