-
Notifications
You must be signed in to change notification settings - Fork 5
/
update_ava_kinetics_csv.py
137 lines (108 loc) · 5.21 KB
/
update_ava_kinetics_csv.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
import os, math
import shutil, argparse
def make_box_anno(llist):
box = [llist[2], llist[3], llist[4], llist[5]]
return [float(b) for b in box]
def read_kinetics_annotations(anno_file):
lines = open(anno_file, 'r').readlines()
annotations = {}
is_test = anno_file.find('test')>-1
for line in lines:
line = line.rstrip('\n')
line_list = line.split(',')
video_name = line_list[0]
time_stamp = float(line_list[1])
if len(line_list)>2:
box = make_box_anno(line_list)
label = int(line_list[-1])
if video_name not in annotations:
annotations[video_name] = [[time_stamp, box, label]]
else:
annotations[video_name] += [[time_stamp, box, label]]
# elif is_test:
else:
if video_name not in annotations:
annotations[line_list[0]] = [[time_stamp, None, None]]
else:
annotations[line_list[0]] = [[time_stamp, None, None]]
return annotations
def update_csvs(frames_dir, anno_files, anno_dir):
new_anno_dir = 'ava_kinetics_updated_csv/'
if not os.path.isdir(new_anno_dir):
os.makedirs(new_anno_dir)
trim_format='%06d'
frames_dir = args.frames_dir
for anno_name in anno_files:
anno_file = os.path.join(anno_dir, anno_name)
annotations = read_kinetics_annotations(anno_file)
new_csv = open(os.path.join(new_anno_dir, anno_name), 'w')
total_count = 1
found_count = 0
num_frames = 1
for ii, video in enumerate(annotations):
integer_time_stamp = int(math.floor(annotations[video][0][0]))
time_stamp = annotations[video][0][0]
if integer_time_stamp<= 3:
new_time_stamp = time_stamp
else:
new_time_stamp = 3.0 + time_stamp - integer_time_stamp
for anno in annotations[video]:
total_count += 1
video_name = '{:s}_{:s}'.format(video, trim_format % int(math.floor(time_stamp)))
video_frames_dir = os.path.join(frames_dir, video_name)
imglist = []
if os.path.isdir(os.path.join(frames_dir, video_name)):
imglist = os.listdir(video_frames_dir)
imglist = [img for img in imglist if img.endswith('.jpg')]
num_f= len(imglist)
min_frames = int(new_time_stamp*30+20)
if len(imglist)>=min_frames:
found_count += 1
if found_count%10000 == 0:
print(ii, video, len(imglist))
num_frames += num_f
wrtie_str = '{:s},{:f}'.format(video_name, new_time_stamp, )
if anno[1] is not None:
for b in anno[1]:
wrtie_str += ',{:f}'.format(b)
wrtie_str += ',{:d},{:d}'.format(anno[2],num_f)
wrtie_str += '\n'
new_csv.write(wrtie_str)
else:
shutil.rmtree(video_frames_dir)
new_csv.close()
print('{:s} found {:d}/{:d} average number of frames {:d}'.format(anno_name, found_count, total_count, int(num_frames/found_count)))
def move_dirs(frames_dir, anno_files, anno_dir):
trim_format='%06d'
frames_dir = args.frames_dir
to_move_dir = os.path.join('/'.join(frames_dir.split('/')[:-2]), 'test-images/')
print('TO MOVE DIR', to_move_dir)
if not os.path.isdir(to_move_dir):
os.makedirs(to_move_dir)
anno_name = anno_files[-1]
anno_file = os.path.join(anno_dir, anno_name)
annotations = read_kinetics_annotations(anno_file)
total_count = 0
for ii, video_name in enumerate(annotations):
total_count += 1
time_stamp = annotations[video_name][0][0]
# video_name = '{:s}_{:s}'.format(video, trim_format % int(math.floor(time_stamp)))
src_frames_dir = os.path.join(frames_dir, video_name)
# dst_frames_dir = os.path.join(frames_dir, video_name)
# print('move',video_name, src_frames_dir)
if os.path.isdir(src_frames_dir):
print(src_frames_dir, to_move_dir)
shutil.move(src_frames_dir, to_move_dir)
print('Moved ', total_count, ' dirs')
if __name__ == '__main__':
description = 'Helper script for updating cvs of kinetics dataset after video trimming.'
p = argparse.ArgumentParser(description=description)
p.add_argument('--frames_dir', type=str, default='/raid/susaha/datasets/ava-kinetics/kinetics/train-images/',
help='Video directory where videos are saved.')
args = p.parse_args()
anno_files = ['kinetics_train_v1.0.csv', 'kinetics_val_v1.0.csv', 'kinetics_test_v1.0.csv']
anno_files = ['kinetics_train.csv', 'kinetics_val.csv', 'kinetics_test.csv']
anno_dir = 'ava_kinetics_csv/'
anno_dir = '/raid/susaha/datasets/ava-kinetics/kinetics/annotations/'
# update_csvs(args.frames_dir, anno_files, anno_dir)
move_dirs(args.frames_dir, anno_files, anno_dir)