-
Notifications
You must be signed in to change notification settings - Fork 22
/
generate_tracks.py
executable file
·70 lines (53 loc) · 2.19 KB
/
generate_tracks.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
import os
import yaml
import tqdm
import glob
import pickle
from tracker.iou_tracker import track_iou
from detect_faces_on_videos import DETECTIONS_FILE_NAME, DETECTIONS_ROOT
SIGMA_L = 0.3
SIGMA_H = 0.9
SIGMA_IOU = 0.3
T_MIN = 1
TRACKS_FILE_NAME = 'tracks.pkl'
def get_tracks(detections):
if len(detections) == 0:
return []
converted_detections = []
for i, detections_per_frame in enumerate(detections):
converted_detections_per_frame = []
for j, (bbox, score) in enumerate(zip(detections_per_frame['boxes'], detections_per_frame['scores'])):
bbox = tuple(bbox.tolist())
converted_detections_per_frame.append({'bbox': bbox, 'score': score})
converted_detections.append(converted_detections_per_frame)
tracks = track_iou(converted_detections, SIGMA_L, SIGMA_H, SIGMA_IOU, T_MIN)
tracks_converted = []
for track in tracks:
track_converted = []
start_frame = track['start_frame'] - 1
for i, bbox in enumerate(track['bboxes']):
track_converted.append((start_frame + i, bbox))
tracks_converted.append(track_converted)
return tracks_converted
def main():
with open('config.yaml', 'r') as f:
config = yaml.load(f)
root_dir = os.path.join(config['ARTIFACTS_PATH'], DETECTIONS_ROOT)
detections_content = []
for path in glob.iglob(os.path.join(root_dir, '**', DETECTIONS_FILE_NAME), recursive=True):
rel_path = path[len(root_dir) + 1:]
detections_content.append(rel_path)
detections_content = sorted(detections_content)
print('Total number of videos: {}'.format(len(detections_content)))
video_to_tracks = {}
for rel_path in tqdm.tqdm(detections_content):
video = os.path.dirname(rel_path)
with open(os.path.join(root_dir, rel_path), 'rb') as f:
detections = pickle.load(f)
video_to_tracks[video] = get_tracks(detections)
track_count = sum([len(tracks) for tracks in video_to_tracks.values()])
print('Total number of tracks: {}'.format(track_count))
with open(os.path.join(config['ARTIFACTS_PATH'], TRACKS_FILE_NAME), 'wb') as f:
pickle.dump(video_to_tracks, f)
if __name__ == '__main__':
main()