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.val_libs import voc_eval
19 from libs.box_utils import draw_box_in_img
21 from help_utils import tools
24 def eval_with_plac(det_net, real_test_imgname_list, img_root, draw_imgs=False):
27 img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not BGR
28 img_batch = tf.cast(img_plac, tf.float32)
30 img_batch = short_side_resize_for_inference_data(img_tensor=img_batch,
31 target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
32 length_limitation=cfgs.IMG_MAX_LENGTH)
33 img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)
34 img_batch = tf.expand_dims(img_batch, axis=0)
36 detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
37 input_img_batch=img_batch,
41 tf.global_variables_initializer(),
42 tf.local_variables_initializer()
45 restorer, restore_ckpt = det_net.get_restorer()
47 config = tf.ConfigProto()
48 config.gpu_options.allow_growth = True
53 with tf.Session(config=config) as sess:
55 if not restorer is None:
56 restorer.restore(sess, restore_ckpt)
57 print('restore model')
60 for i, a_img_name in enumerate(real_test_imgname_list):
62 raw_img = cv2.imread(os.path.join(img_root, a_img_name))
63 raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
66 resized_img, detected_boxes, detected_scores, detected_categories = \
68 [img_batch, detection_boxes, detection_scores, detection_category],
69 feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB
72 compute_time = compute_time + (end - start)
73 compute_imgnum = compute_imgnum + 1
74 # print("{} cost time : {} ".format(img_name, (end - start)))
76 show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
77 show_scores = detected_scores[show_indices]
78 show_boxes = detected_boxes[show_indices]
79 show_categories = detected_categories[show_indices]
80 final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(np.squeeze(resized_img, 0),
82 labels=show_categories,
84 if not os.path.exists(cfgs.TEST_SAVE_PATH):
85 os.makedirs(cfgs.TEST_SAVE_PATH)
87 cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + a_img_name + '.jpg',
88 final_detections[:, :, ::-1])
90 xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
91 detected_boxes[:, 2], detected_boxes[:, 3]
93 resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
95 xmin = xmin * raw_w / resized_w
96 xmax = xmax * raw_w / resized_w
98 ymin = ymin * raw_h / resized_h
99 ymax = ymax * raw_h / resized_h
101 boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
102 dets = np.hstack((detected_categories.reshape(-1, 1),
103 detected_scores.reshape(-1, 1),
105 all_boxes.append(dets)
107 tools.view_bar('{} image cost {}s'.format(a_img_name, (end - start)), i + 1, len(real_test_imgname_list))
109 # save_dir = os.path.join(cfgs.EVALUATE_DIR, cfgs.VERSION)
110 # if not os.path.exists(save_dir):
111 # os.makedirs(save_dir)
112 # fw1 = open(os.path.join(save_dir, 'detections.pkl'), 'wb')
113 # pickle.dump(all_boxes, fw1)
114 print('\n average_training_time_per_image is' + str(compute_time / compute_imgnum))
118 def eval(num_imgs, eval_dir, annotation_dir, showbox):
120 # with open('/home/yjr/DataSet/VOC/VOC_test/VOC2007/ImageSets/Main/aeroplane_test.txt') as f:
121 # all_lines = f.readlines()
122 # test_imgname_list = [a_line.split()[0].strip() for a_line in all_lines]
124 test_imgname_list = [item for item in os.listdir(eval_dir)
125 if item.endswith(('.jpg', 'jpeg', '.png', '.tif', '.tiff'))]
126 if num_imgs == np.inf:
127 real_test_imgname_list = test_imgname_list
129 real_test_imgname_list = test_imgname_list[: num_imgs]
131 faster_rcnn = build_whole_network.DetectionNetwork(base_network_name=cfgs.NET_NAME,
133 all_boxes = eval_with_plac(det_net=faster_rcnn, real_test_imgname_list=real_test_imgname_list,
137 # save_dir = os.path.join(cfgs.EVALUATE_DIR, cfgs.VERSION)
138 # if not os.path.exists(save_dir):
139 # os.makedirs(save_dir)
140 # with open(os.path.join(save_dir, 'detections.pkl'), 'rb') as f:
141 # all_boxes = pickle.load(f)
143 # print(len(all_boxes))
145 voc_eval.voc_evaluate_detections(all_boxes=all_boxes,
146 test_annotation_path=annotation_dir,
147 test_imgid_list=real_test_imgname_list)
151 parser = argparse.ArgumentParser('evaluate the result with Pascal2007 stdand')
153 parser.add_argument('--eval_imgs', dest='eval_imgs',
154 help='evaluate imgs dir ',
155 default='../data/pcb_test/JPEGImages', type=str)
156 parser.add_argument('--annotation_dir', dest='test_annotation_dir',
157 help='the dir save annotations',
158 default='../data/pcb_test/Annotations', type=str)
159 parser.add_argument('--showbox', dest='showbox',
160 help='whether show detecion results when evaluation',
161 default=False, type=bool)
162 parser.add_argument('--GPU', dest='GPU',
164 default='2', type=str)
165 #parser.add_argument('--eval_num', dest='eval_num',
166 # help='the num of eval imgs',
167 # default=np.inf, type=int)
168 parser.add_argument('--eval_num', dest='eval_num',
169 help='the num of eval imgs',
170 default=100, type=int)
171 args = parser.parse_args()
175 if __name__ == '__main__':
181 os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU
182 eval(np.inf, # use np.inf to test all the imgs. use 10 to test 10 imgs.
183 eval_dir=args.eval_imgs,
184 annotation_dir=args.test_annotation_dir,
185 showbox=args.showbox)