-
Notifications
You must be signed in to change notification settings - Fork 21
/
Copy pathdataset.lua
142 lines (131 loc) · 4.7 KB
/
dataset.lua
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
print('prepare data')
assert(opt.dataset == 'activitynet' or opt.datset == 'kinetics',
'dataset must be activitynet or kinetics')
if opt.dataset == 'activitynet' then
dataset_utils = dofile('activitynet_utils.lua')
elseif opt.dataset == 'kinetics' then
dataset_utils = dofile('kinetics_utils.lua')
end
local data = dataset_utils.load_annotation_data(opt.annotation_path)
local training_video_names, training_annotations,
validation_video_names, validation_annotations,
test_video_names =
dataset_utils.get_data_names_and_annotations(data)
local class_labels_map = dataset_utils.get_class_labels(data)
class_names_map = {}
for name, label in pairs(class_labels_map) do
class_names_map[label] = name
end
function prepare_data(video_names, annotations)
local dataset = {}
for i = 1, #video_names do
local video_path = paths.concat(opt.video_path, video_names[i])
if paths.dirp(video_path) then
local fps_file_path = paths.concat(video_path, 'fps')
local fps = utils.load_fps_file(fps_file_path)
for _, one_annotation in pairs(annotations[i]) do
local begin_t = math.ceil(one_annotation.segment[1] * fps)
local end_t = math.ceil(one_annotation.segment[2] * fps)
if begin_t == 0 then
begin_t = 1
end
local sample = {
video = video_path,
segment = {begin_t, end_t},
label = class_labels_map[one_annotation.label],
fps = fps
}
table.insert(dataset, sample)
end
end
end
return dataset
end
function prepare_data_kinetics(video_names, annotations)
local dataset = {}
for i = 1, #video_names do
local video_path = paths.concat(opt.video_path, video_names[i])
if paths.dirp(video_path) then
local n_frames_file_path = paths.concat(video_path, 'n_frames')
local n_frames = utils.load_n_frames_file(n_frames_file_path)
if n_frames > 0 then
n_frames = math.floor(n_frames)
local begin_t = 1
local end_t = n_frames
local sample = {
video = video_path,
segment = {begin_t, end_t},
n_frames = n_frames,
video_id = string.sub(video_names[i], 1, -15)
}
if annotations ~= nil then
sample.label = class_labels_map[annotations[i].label]
end
table.insert(dataset, sample)
end
end
end
return dataset
end
function prepare_data_per_video(video_names, annotations)
local dataset = {}
for i = 1, #video_names do
local video_path = paths.concat(opt.video_path, video_names[i])
if paths.dirp(video_path) then
local n_frames = math.floor(utils.get_n_frames(video_path))
local fps_file_path = paths.concat(video_path, 'fps')
local fps = utils.load_fps_file(fps_file_path)
local video_data = {}
video_data.video = video_path
video_data.video_id = string.sub(video_names[i], 3, -1)
video_data.n_frames = n_frames
video_data.fps = fps
video_data.annotations = {}
if annotations ~= nil then
for _, one_annotation in pairs(annotations[i]) do
local begin_t = math.ceil(one_annotation.segment[1] * fps)
local end_t = math.ceil(one_annotation.segment[2] * fps)
if begin_t == 0 then
begin_t = 1
end
local sample = {
segment = {begin_t, end_t},
label = class_labels_map[one_annotation.label]
}
table.insert(video_data.annotations, sample)
end
end
table.insert(dataset, video_data)
end
end
return dataset
end
if not opt.no_train then
if opt.dataset == 'activitynet' then
training_data = prepare_data(training_video_names, training_annotations)
elseif opt.dataset == 'kinetics' then
training_data = prepare_data_kinetics(training_video_names, training_annotations)
end
end
if not opt.no_test then
if opt.dataset == 'activitynet' then
validation_data = prepare_data(validation_video_names, validation_annotations)
elseif opt.dataset == 'kinetics' then
validation_data = prepare_data_kinetics(validation_video_names, validation_annotations)
end
end
if opt.test_video then
if opt.dataset == 'activitynet' then
if opt.test_subset == 'val' then
video_test_data = prepare_data_per_video(validation_video_names, validation_annotations)
elseif opt.test_subset == 'test' then
video_test_data = prepare_data_per_video(test_video_names)
end
elseif opt.dataset == 'kinetics' then
if opt.test_subset == 'val' then
video_test_data = prepare_data_kinetics(validation_video_names, validation_annotations)
elseif opt.test_subset == 'test' then
video_test_data = prepare_data_kinetics(test_video_names)
end
end
end