forked from Tencent/ActionDetection-DBG
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
104 lines (85 loc) · 3.53 KB
/
test.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
import os
import numpy as np
import pandas as pd
import tensorflow as tf
import tqdm
import data_loader
import model
from config import DBGConfig
dbg_config = DBGConfig()
checkpoint_dir = dbg_config.checkpoint_dir
result_dir = dbg_config.result_dir
tscale = dbg_config.tscale
feature_dim = dbg_config.feature_dim
# test batch size
batch_size = dbg_config.test_batch_size
test_mode = dbg_config.test_mode
mask = np.zeros([tscale, tscale], np.float32)
for i in range(tscale):
for j in range(i, tscale):
mask[i, j] = 1
tf_mask = tf.convert_to_tensor(mask, tf.float32)
tf_mask = tf.reshape(tf_mask, [1, tscale, tscale, 1])
if __name__ == "__main__":
X_feature = tf.placeholder(tf.float32, shape=(batch_size, tscale, feature_dim))
scores, iou_mat, x1, x2, xc, prop_start, prop_end = model.model(X_feature, training=False)
prop_start = prop_start * tf_mask
prop_end = prop_end * tf_mask
# boundary map fusion
pstart = tf.reduce_sum(prop_start, 2) / tf.maximum(tf.reduce_sum(tf_mask, 2), 1)
pend = tf.reduce_sum(prop_end, 1) / tf.maximum(tf.reduce_sum(tf_mask, 1), 1)
model_saver = tf.train.Saver()
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
sess = tf.InteractiveSession(config=tf_config)
tf.global_variables_initializer().run()
model_saver.restore(sess, os.path.join(checkpoint_dir, 'dbg_model_best'))
train_dict, val_dict, test_dict = data_loader.getDatasetDict()
if test_mode == 'validation':
video_dict = val_dict
else:
video_dict = test_dict
batch_video_list = data_loader.getBatchListTest(video_dict, batch_size, shuffle=False)
batch_result_xmin = []
batch_result_xmax = []
batch_result_iou = []
batch_result_pstart = []
batch_result_pend = []
print('Runing DBG model ...')
for idx in tqdm.tqdm(range(len(batch_video_list))):
batch_anchor_xmin, batch_anchor_xmax, batch_anchor_feature = data_loader.getProposalDataTest(
batch_video_list[idx], video_dict)
out_iou, out_start, out_end = sess.run([iou_mat, pstart, pend],
feed_dict={X_feature: batch_anchor_feature})
batch_result_xmin.append(batch_anchor_xmin)
batch_result_xmax.append(batch_anchor_xmax)
batch_result_iou.append(out_iou[:, :, :, 0])
batch_result_pstart.append(out_start[:, :, 0])
batch_result_pend.append(out_end[:, :, 0])
print('Saving results ...')
columns = ["iou", "start", "end", "xmin", "xmax"]
for idx in tqdm.tqdm(range(len(batch_video_list))):
b_video = batch_video_list[idx]
b_xmin = batch_result_xmin[idx]
b_xmax = batch_result_xmax[idx]
b_iou = batch_result_iou[idx]
b_pstart = batch_result_pstart[idx]
b_pend = batch_result_pend[idx]
for j in range(len(b_video)):
tmp_video = b_video[j]
tmp_xmin = b_xmin[j]
tmp_xmax = b_xmax[j]
tmp_iou = b_iou[j]
tmp_pstart = b_pstart[j]
tmp_pend = b_pend[j]
res = []
# save all proposals result
for i in range(tscale):
for j in range(i, tscale):
start = tmp_pstart[i]
end = tmp_pend[j]
iou = tmp_iou[i, j]
res.append([iou, start, end, tmp_xmin[i], tmp_xmax[j]])
tmp_result = np.stack(res)
tmp_df = pd.DataFrame(tmp_result, columns=columns)
tmp_df.to_csv(os.path.join(result_dir, tmp_video + '.csv'), index=False)