3 from __future__ import absolute_import
4 from __future__ import print_function
5 from __future__ import division
8 import tensorflow as tf
13 sys.path.append("../")
15 from data.io.image_preprocess import short_side_resize_for_inference_data
16 from libs.configs import cfgs
17 from libs.networks import build_whole_network
18 from libs.box_utils import draw_box_in_img
19 from help_utils import tools
22 def detect(det_net, inference_save_path, real_test_imgname_list):
25 img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not GBR
26 img_batch = tf.cast(img_plac, tf.float32)
27 img_batch = short_side_resize_for_inference_data(img_tensor=img_batch,
28 target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
29 length_limitation=cfgs.IMG_MAX_LENGTH)
30 img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)
31 img_batch = tf.expand_dims(img_batch, axis=0) # [1, None, None, 3]
33 detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
34 input_img_batch=img_batch,
38 tf.global_variables_initializer(),
39 tf.local_variables_initializer()
42 restorer, restore_ckpt = det_net.get_restorer()
44 config = tf.ConfigProto()
45 config.gpu_options.allow_growth = True
47 with tf.Session(config=config) as sess:
49 if not restorer is None:
50 restorer.restore(sess, restore_ckpt)
51 print('restore model')
53 for i, a_img_name in enumerate(real_test_imgname_list):
55 raw_img = cv2.imread(a_img_name)
57 resized_img, detected_boxes, detected_scores, detected_categories = \
59 [img_batch, detection_boxes, detection_scores, detection_category],
60 feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB
63 # print("{} cost time : {} ".format(img_name, (end - start)))
65 raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
67 xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
68 detected_boxes[:, 2], detected_boxes[:, 3]
70 resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
72 xmin = xmin * raw_w / resized_w
73 xmax = xmax * raw_w / resized_w
75 ymin = ymin * raw_h / resized_h
76 ymax = ymax * raw_h / resized_h
78 detected_boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
80 show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
81 show_scores = detected_scores[show_indices]
82 show_boxes = detected_boxes[show_indices]
83 show_categories = detected_categories[show_indices]
84 final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(raw_img - np.array(cfgs.PIXEL_MEAN),
86 labels=show_categories,
88 nake_name = a_img_name.split('/')[-1]
89 # print (inference_save_path + '/' + nake_name)
90 cv2.imwrite(inference_save_path + '/' + nake_name,
91 final_detections[:, :, ::-1])
93 tools.view_bar('{} image cost {}s'.format(a_img_name, (end - start)), i + 1, len(real_test_imgname_list))
96 def test(test_dir, inference_save_path):
98 test_imgname_list = [os.path.join(test_dir, img_name) for img_name in os.listdir(test_dir)
99 if img_name.endswith(('.jpg', '.png', '.jpeg', '.tif', '.tiff'))]
100 assert len(test_imgname_list) != 0, 'test_dir has no imgs there.' \
101 ' Note that, we only support img format of (.jpg, .png, and .tiff) '
103 faster_rcnn = build_whole_network.DetectionNetwork(base_network_name=cfgs.NET_NAME,
105 detect(det_net=faster_rcnn, inference_save_path=inference_save_path, real_test_imgname_list=test_imgname_list)
110 Parse input arguments
112 parser = argparse.ArgumentParser(description='TestImgs...U need provide the test dir')
113 parser.add_argument('--data_dir', dest='data_dir',
115 default='demos', type=str)
116 parser.add_argument('--save_dir', dest='save_dir',
117 help='demo imgs to save',
118 default='inference_results', type=str)
119 parser.add_argument('--GPU', dest='GPU',
121 default='0', type=str)
123 if len(sys.argv) == 1:
127 args = parser.parse_args()
130 if __name__ == '__main__':
133 print('Called with args:')
135 os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU
137 inference_save_path=args.save_dir)