forked from bochinski/iou-tracker
-
Notifications
You must be signed in to change notification settings - Fork 2
/
util.py
147 lines (112 loc) · 4.9 KB
/
util.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
136
137
138
139
140
141
142
143
144
145
146
147
# ---------------------------------------------------------
# IOU Tracker
# Copyright (c) 2017 TU Berlin, Communication Systems Group
# Licensed under The MIT License [see LICENSE for details]
# Written by Erik Bochinski
# ---------------------------------------------------------
import numpy as np
import csv
def load_mot(detections):
"""
Loads detections stored in a mot-challenge like formatted CSV or numpy array (fieldNames = ['frame', 'id', 'x', 'y',
'w', 'h', 'score']).
Args:
detections
Returns:
list: list containing the detections for each frame.
"""
data = []
if type(detections) is str:
raw = np.genfromtxt(detections, delimiter=',', dtype=np.float32)
else:
# assume it is an array
assert isinstance(detections, np.ndarray), "only numpy arrays or *.csv paths are supported as detections."
raw = detections.astype(np.float32)
end_frame = int(np.max(raw[:, 0]))
for i in range(1, end_frame+1):
idx = raw[:, 0] == i
bbox = raw[idx, 2:6]
bbox[:, 2:4] += bbox[:, 0:2] # x1, y1, w, h -> x1, y1, x2, y2
scores = raw[idx, 6]
dets = []
for bb, s in zip(bbox, scores):
# Adapt to MRCNN format for bboxes: y0, x0, y1, x1
dets.append({'roi': [bb[1], bb[0], bb[3], bb[2]], 'score': s, 'centroid': [0.5*(bb[0] + bb[2]), 0.5*(bb[1] + bb[3])], 'frame': i})
data.append(dets)
return data
def save_to_csv(out_path, tracks):
"""
Saves tracks to a CSV file.
Args:
out_path (str): path to output csv file.
tracks (list): list of tracks to store.
"""
with open(out_path, "w") as ofile:
field_names = ['frame', 'id', 'x', 'y', 'w', 'h', 'score', 'wx', 'wy', 'wz']
odict = csv.DictWriter(ofile, field_names)
id_ = 1
for track in tracks:
for i, bbox in enumerate(track['bboxes']):
row = {'id': id_,
'frame': track['start_frame'] + i,
'y': bbox[1],
'x': bbox[0],
'h': bbox[3] - bbox[1],
'w': bbox[2] - bbox[0],
'score': track['max_score'],
'wx': -1,
'wy': -1,
'wz': -1}
odict.writerow(row)
id_ += 1
def iou(bbox1, bbox2):
"""
Calculates the intersection-over-union of two bounding boxes.
Args:
bbox1 (numpy.array, list of floats): bounding box in format x1,y1,x2,y2.
bbox2 (numpy.array, list of floats): bounding box in format x1,y1,x2,y2.
Returns:
int: intersection-over-onion of bbox1, bbox2
"""
bbox1 = [float(x) for x in bbox1]
bbox2 = [float(x) for x in bbox2]
# (x0_1, y0_1, x1_1, y1_1) = bbox1
# (x0_2, y0_2, x1_2, y1_2) = bbox2
(y0_1, x0_1, y1_1, x1_1) = bbox1
(y0_2, x0_2, y1_2, x1_2) = bbox2
# get the overlap rectangle
overlap_x0 = max(x0_1, x0_2)
overlap_y0 = max(y0_1, y0_2)
overlap_x1 = min(x1_1, x1_2)
overlap_y1 = min(y1_1, y1_2)
# check if there is an overlap
if overlap_x1 - overlap_x0 <= 0 or overlap_y1 - overlap_y0 <= 0:
return 0
# if yes, calculate the ratio of the overlap to each ROI size and the unified size
size_1 = (x1_1 - x0_1) * (y1_1 - y0_1)
size_2 = (x1_2 - x0_2) * (y1_2 - y0_2)
size_intersection = (overlap_x1 - overlap_x0) * (overlap_y1 - overlap_y0)
size_union = size_1 + size_2 - size_intersection
return size_intersection / size_union
def interp_tracks(tracks_finished):
"""
The Kalman-IOU tracker can skip frames, however the DETRAC toolkit takes off points for each frame missing,
and skipping 50% of the frames will effectively cap the maximum MOTA at 50%.
We perform a simple linear interpolation to fill in the gaps.
"""
furnished_tracks = []
for ftracks in tracks_finished:
starting_frame = ftracks[0]['frame']
ending_frame = ftracks[-1]['frame']
interp_track = np.zeros((ending_frame - starting_frame + 1, 4))
frames_present = []
for fframe in ftracks:
interp_track[fframe['frame'] - starting_frame, :] = fframe['roi']
frames_present.append(fframe['frame'])
frames_present_abs = (np.array(frames_present) - starting_frame).tolist()
frames_missing = [f for f in range(starting_frame, ending_frame + 1) if f not in frames_present]
frames_missing_abs = (np.array(frames_missing) - starting_frame).tolist()
for i in range(4):
interp_track[frames_missing_abs, i] = np.interp(frames_missing, frames_present, interp_track[frames_present_abs, i])
furnished_tracks.append([{'roi': interp_track[f, :].tolist(), 'frame': f + starting_frame} for f in range(ending_frame - starting_frame + 1)])
return furnished_tracks