-
Notifications
You must be signed in to change notification settings - Fork 16
/
predict.py
176 lines (135 loc) · 6.5 KB
/
predict.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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
"""
Copyright (c) College of Mechatronics and Control Engineering, Shenzhen University.
All rights reserved.
Description :
predition
Author:Team Li
"""
import tensorflow as tf
import numpy as np
import cv2
import os
from time import time
from model.factory import model_factory
from dataset.hazy_person import provider
import utils.test_tools as test_tools
from utils.logging import logger
import config
FLAGS = tf.app.flags.FLAGS
slim = tf.contrib.slim
tf.app.flags.DEFINE_string(
'model_name', None,
'The name of the architecture to train.')
tf.app.flags.DEFINE_string(
'attention_module', None,
'The name of attention module to apply.')
tf.app.flags.DEFINE_string(
'checkpoint_dir', '',
'The full file name to a checkpoint from which to fine-tune.')
tf.app.flags.DEFINE_float(
'select_threshold', 0.3, 'obj score less than it would be filter')
tf.app.flags.DEFINE_float(
'nms_threshold', 0.6, 'nms threshold')
tf.app.flags.DEFINE_integer(
'keep_top_k', 30, 'maximun num of obj after nms')
tf.app.flags.DEFINE_integer(
'vis_img_height', 512, 'the img height when visulize')
tf.app.flags.DEFINE_integer(
'vis_img_width', 512, 'the img width when visulize')
#### config only for prioriboxes_mbn ####
tf.app.flags.DEFINE_string(
'backbone_name', None,
'support mobilenet_v1 and mobilenet_v2')
tf.app.flags.DEFINE_boolean(
'multiscale_feats', True,
'whether merge different scale features')
tf.app.flags.DEFINE_boolean(
'whether_aug', True,
'whether use augmentation to prediction')
## define placeholder ##
inputs = tf.placeholder(tf.float32,
shape=(None, config.img_size[0], config.img_size[1], 3))
def build_graph(model_name, attention_module, config_dict, is_training):
"""build tf graph for predict
Args:
model_name: choose a model to build
attention_module: must be "se_block" or "cbam_block"
config_dict: some config for building net
is_training: whether to train or test, here must be False
Return:
det_loss: a tensor with a shape [bs, priori_boxes_num, 4]
clf_loss: a tensor with a shape [bs, priori_boxes_num, 2]
"""
assert is_training == False
net = model_factory(inputs=inputs, model_name=model_name,
attention_module=attention_module, is_training=is_training,
config_dict=config_dict)
corner_bboxes, clf_pred = net.get_output_for_test()
score, bboxes = test_tools.bboxes_select(clf_pred, corner_bboxes,
select_threshold= FLAGS.select_threshold)
score, bboxes = test_tools.bboxes_sort(score, bboxes)
rscores, rbboxes = test_tools.bboxes_nms_batch(score, bboxes,
nms_threshold=FLAGS.nms_threshold,
keep_top_k=FLAGS.keep_top_k)
return rscores, rbboxes
def main(_):
config_dict = {'multiscale_feats': FLAGS.multiscale_feats,
'backbone': FLAGS.backbone_name}
scores, bboxes = build_graph(FLAGS.model_name, FLAGS.attention_module,
config_dict=config_dict, is_training=False)
saver = tf.train.Saver(tf.global_variables())
ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_dir)
configuretion = tf.ConfigProto()
configuretion.gpu_options.allow_growth = True
with tf.Session(config=configuretion) as sess:
if ckpt:
logger.info('loading %s...' % str(ckpt.model_checkpoint_path))
saver.restore(sess, ckpt.model_checkpoint_path)
logger.info('Load checkpoint success...')
else:
raise ValueError("can not find checkpoint, pls check checkpoint_dir")
pd = provider(batch_size=1, for_what="predict", whether_aug=FLAGS.whether_aug)
logger.info("Please press any key to skip picture...")
while (True):
# start = time()
# norm_imgs, labels, corner_bboxes_gt = pd.load_batch()
norm_imgs, corner_bboxes_gt = pd.load_batch()
#print(corner_bboxes_gt)
imgs = np.uint8((norm_imgs[0] + 1.)*255 / 2)
imgs_for_gt = cv2.resize(imgs, dsize=(FLAGS.vis_img_height, FLAGS.vis_img_width))
imgs_for_pred = imgs_for_gt.copy()
corner_bboxes_gt = corner_bboxes_gt[0]
corner_bboxes_gt[:, 0] = corner_bboxes_gt[:, 0] * FLAGS.vis_img_height
corner_bboxes_gt[:, 1] = corner_bboxes_gt[:, 1] * FLAGS.vis_img_width
corner_bboxes_gt[:, 2] = corner_bboxes_gt[:, 2] * FLAGS.vis_img_height
corner_bboxes_gt[:, 3] = corner_bboxes_gt[:, 3] * FLAGS.vis_img_width
corner_bboxes_gt = np.int32(corner_bboxes_gt)
scores_pred, bboxes_pred = sess.run([scores, bboxes], feed_dict={inputs:np.array(norm_imgs)})
scores_pred = list(scores_pred.values())
bboxes_pred = list(bboxes_pred.values())
scores_pred = scores_pred[0][0]
bboxes_pred = bboxes_pred[0][0]
bboxes_pred[:, 0] = bboxes_pred[:, 0] * FLAGS.vis_img_height
bboxes_pred[:, 1] = bboxes_pred[:, 1] * FLAGS.vis_img_width
bboxes_pred[:, 2] = bboxes_pred[:, 2] * FLAGS.vis_img_height
bboxes_pred[:, 3] = bboxes_pred[:, 3] * FLAGS.vis_img_width
bboxes_pred = np.int32(bboxes_pred)
## vis ##
imgs_for_gt = cv2.cvtColor(imgs_for_gt, cv2.COLOR_BGR2RGB)
imgs_for_pred = cv2.cvtColor(imgs_for_pred, cv2.COLOR_BGR2RGB)
label = np.ones(corner_bboxes_gt.shape[0], dtype=np.int32)
imgs_for_gt = test_tools.visualize_boxes_and_labels_on_image_array(imgs_for_gt, corner_bboxes_gt, label,
None, config.category_index, skip_labels=False)
label = np.ones(bboxes_pred.shape[0], dtype=np.int32)
imgs_for_pred = test_tools.visualize_boxes_and_labels_on_image_array(imgs_for_pred, bboxes_pred, label,
scores_pred, config.category_index,
skip_labels=False)
imgs_for_gt = cv2.cvtColor(imgs_for_gt, cv2.COLOR_RGB2BGR)
imgs_for_pred = cv2.cvtColor(imgs_for_pred, cv2.COLOR_RGB2BGR)
cv2.imshow("ground-truth", imgs_for_gt)
cv2.imshow("prediction", imgs_for_pred)
cv2.waitKey(0)
cv2.destroyAllWindows()
pass
if __name__ == '__main__':
tf.app.run()