def merge_result(result_dict): final_bbox = [ result_dict['final_bbox_m1'], result_dict['final_bbox_m2'], result_dict['final_bbox_m2'] ] final_scores = [ result_dict['final_scores_m1'], result_dict['final_scores_m2'], result_dict['final_scores_m3'] ] final_category = [ result_dict['final_category_m1'], result_dict['final_category_m2'], result_dict['final_category_m3'] ] final_bbox = np.concatenate(final_bbox, axis=0) final_scores = np.concatenate(final_scores, axis=0) final_category = np.concatenate(final_category, axis=0) keep = nms.py_cpu_nms(final_bbox, final_scores, cfgs.FINAL_NMS_IOU_THRESHOLD) final_bbox = final_bbox[keep] final_scores = final_scores[keep] final_category = final_category[keep] return final_bbox, final_scores, final_category
def inference(det_net, file_paths, des_folder, h_len, w_len, h_overlap, w_overlap, save_res=False): if save_res: assert cfgs.SHOW_SCORE_THRSHOLD >= 0.5, \ 'please set score threshold (example: SHOW_SCORE_THRSHOLD = 0.5) in cfgs.py' else: assert cfgs.SHOW_SCORE_THRSHOLD < 0.005, \ 'please set score threshold (example: SHOW_SCORE_THRSHOLD = 0.00) in cfgs.py' # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) img_batch = tf.cast(img_plac, tf.float32) img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) img_batch = short_side_resize_for_inference_data( img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, is_resize=False) det_boxes_h, det_scores_h, det_category_h, \ det_boxes_r, det_scores_r, det_category_r = det_net.build_whole_detection_network(input_img_batch=img_batch, gtboxes_h_batch=None, gtboxes_r_batch=None) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = det_net.get_restorer() config = tf.ConfigProto() config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') for count, img_path in enumerate(file_paths): start = timer() img = cv2.imread(img_path) box_res = [] label_res = [] score_res = [] box_res_rotate = [] label_res_rotate = [] score_res_rotate = [] imgH = img.shape[0] imgW = img.shape[1] if imgH < h_len: temp = np.zeros([h_len, imgW, 3], np.float32) temp[0:imgH, :, :] = img img = temp imgH = h_len if imgW < w_len: temp = np.zeros([imgH, w_len, 3], np.float32) temp[:, 0:imgW, :] = img img = temp imgW = w_len for hh in range(0, imgH, h_len - h_overlap): if imgH - hh - 1 < h_len: hh_ = imgH - h_len else: hh_ = hh for ww in range(0, imgW, w_len - w_overlap): if imgW - ww - 1 < w_len: ww_ = imgW - w_len else: ww_ = ww src_img = img[hh_:(hh_ + h_len), ww_:(ww_ + w_len), :] det_boxes_h_, det_scores_h_, det_category_h_, \ det_boxes_r_, det_scores_r_, det_category_r_ = \ sess.run( [det_boxes_h, det_scores_h, det_category_h, det_boxes_r, det_scores_r, det_category_r], feed_dict={img_plac: src_img[:, :, ::-1]} ) if len(det_boxes_h_) > 0: for ii in range(len(det_boxes_h_)): box = det_boxes_h_[ii] box[0] = box[0] + ww_ box[1] = box[1] + hh_ box[2] = box[2] + ww_ box[3] = box[3] + hh_ box_res.append(box) label_res.append(det_category_h_[ii]) score_res.append(det_scores_h_[ii]) if len(det_boxes_r_) > 0: for ii in range(len(det_boxes_r_)): box_rotate = det_boxes_r_[ii] box_rotate[0] = box_rotate[0] + ww_ box_rotate[1] = box_rotate[1] + hh_ box_res_rotate.append(box_rotate) label_res_rotate.append(det_category_r_[ii]) score_res_rotate.append(det_scores_r_[ii]) box_res = np.array(box_res) label_res = np.array(label_res) score_res = np.array(score_res) box_res_rotate = np.array(box_res_rotate) label_res_rotate = np.array(label_res_rotate) score_res_rotate = np.array(score_res_rotate) box_res_rotate_, label_res_rotate_, score_res_rotate_ = [], [], [] box_res_, label_res_, score_res_ = [], [], [] r_threshold = { 'roundabout': 0.1, 'tennis-court': 0.3, 'swimming-pool': 0.1, 'storage-tank': 0.2, 'soccer-ball-field': 0.3, 'small-vehicle': 0.2, 'ship': 0.05, 'plane': 0.3, 'large-vehicle': 0.1, 'helicopter': 0.2, 'harbor': 0.0001, 'ground-track-field': 0.3, 'bridge': 0.0001, 'basketball-court': 0.3, 'baseball-diamond': 0.3 } h_threshold = { 'roundabout': 0.35, 'tennis-court': 0.35, 'swimming-pool': 0.4, 'storage-tank': 0.3, 'soccer-ball-field': 0.3, 'small-vehicle': 0.4, 'ship': 0.35, 'plane': 0.35, 'large-vehicle': 0.4, 'helicopter': 0.4, 'harbor': 0.3, 'ground-track-field': 0.4, 'bridge': 0.3, 'basketball-court': 0.4, 'baseball-diamond': 0.3 } for sub_class in range(1, cfgs.CLASS_NUM + 1): index = np.where(label_res_rotate == sub_class)[0] if len(index) == 0: continue tmp_boxes_r = box_res_rotate[index] tmp_label_r = label_res_rotate[index] tmp_score_r = score_res_rotate[index] tmp_boxes_r = np.array(tmp_boxes_r) tmp = np.zeros( [tmp_boxes_r.shape[0], tmp_boxes_r.shape[1] + 1]) tmp[:, 0:-1] = tmp_boxes_r tmp[:, -1] = np.array(tmp_score_r) try: inx = nms_rotate.nms_rotate_cpu( boxes=np.array(tmp_boxes_r), scores=np.array(tmp_score_r), iou_threshold=r_threshold[LABEl_NAME_MAP[sub_class]], max_output_size=500) except: # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms jitter = np.zeros( [tmp_boxes_r.shape[0], tmp_boxes_r.shape[1] + 1]) jitter[:, 0] += np.random.rand(tmp_boxes_r.shape[0], ) / 1000 inx = rotate_gpu_nms( np.array(tmp, np.float32) + np.array(jitter, np.float32), float(r_threshold[LABEl_NAME_MAP[sub_class]]), 0) box_res_rotate_.extend(np.array(tmp_boxes_r)[inx]) score_res_rotate_.extend(np.array(tmp_score_r)[inx]) label_res_rotate_.extend(np.array(tmp_label_r)[inx]) for sub_class in range(1, cfgs.CLASS_NUM + 1): index = np.where(label_res == sub_class)[0] if len(index) == 0: continue tmp_boxes_h = box_res[index] tmp_label_h = label_res[index] tmp_score_h = score_res[index] tmp_boxes_h = np.array(tmp_boxes_h) tmp = np.zeros( [tmp_boxes_h.shape[0], tmp_boxes_h.shape[1] + 1]) tmp[:, 0:-1] = tmp_boxes_h tmp[:, -1] = np.array(tmp_score_h) inx = nms.py_cpu_nms( dets=np.array(tmp, np.float32), thresh=h_threshold[LABEl_NAME_MAP[sub_class]], max_output_size=500) box_res_.extend(np.array(tmp_boxes_h)[inx]) score_res_.extend(np.array(tmp_score_h)[inx]) label_res_.extend(np.array(tmp_label_h)[inx]) time_elapsed = timer() - start if save_res: det_detections_h = draw_box_in_img.draw_box_cv( np.array(img, np.float32) - np.array(cfgs.PIXEL_MEAN), boxes=np.array(box_res_), labels=np.array(label_res_), scores=np.array(score_res_)) det_detections_r = draw_box_in_img.draw_rotate_box_cv( np.array(img, np.float32) - np.array(cfgs.PIXEL_MEAN), boxes=np.array(box_res_rotate_), labels=np.array(label_res_rotate_), scores=np.array(score_res_rotate_)) save_dir = os.path.join(des_folder, cfgs.VERSION) tools.mkdir(save_dir) cv2.imwrite( save_dir + '/' + img_path.split('/')[-1].split('.')[0] + '_h.jpg', det_detections_h) cv2.imwrite( save_dir + '/' + img_path.split('/')[-1].split('.')[0] + '_r.jpg', det_detections_r) view_bar( '{} cost {}s'.format( img_path.split('/')[-1].split('.')[0], time_elapsed), count + 1, len(file_paths)) else: # eval txt CLASS_DOTA = NAME_LABEL_MAP.keys() # Task1 write_handle_r = {} write_handle_h_ = {} txt_dir_r = os.path.join('txt_output', cfgs.VERSION + '_r') txt_dir_h_minAreaRect = os.path.join( 'txt_output', cfgs.VERSION + '_h_minAreaRect') tools.mkdir(txt_dir_r) tools.mkdir(txt_dir_h_minAreaRect) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle_r[sub_class] = open( os.path.join(txt_dir_r, 'Task1_%s.txt' % sub_class), 'a+') write_handle_h_[sub_class] = open( os.path.join(txt_dir_h_minAreaRect, 'Task2_%s.txt' % sub_class), 'a+') rboxes = coordinate_convert.forward_convert(box_res_rotate_, with_label=False) for i, rbox in enumerate(rboxes): command = '%s %.3f %.1f %.1f %.1f %.1f %.1f %.1f %.1f %.1f\n' % ( img_path.split('/')[-1].split('.')[0], score_res_rotate_[i], rbox[0], rbox[1], rbox[2], rbox[3], rbox[4], rbox[5], rbox[6], rbox[7], ) command_ = '%s %.3f %.1f %.1f %.1f %.1f\n' % ( img_path.split('/')[-1].split('.')[0], score_res_rotate_[i], min(rbox[::2]), min( rbox[1::2]), max(rbox[::2]), max(rbox[1::2])) write_handle_r[LABEl_NAME_MAP[label_res_rotate_[i]]].write( command) write_handle_h_[LABEl_NAME_MAP[ label_res_rotate_[i]]].write(command_) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle_r[sub_class].close() # Task2 write_handle_h = {} txt_dir_h = os.path.join('txt_output', cfgs.VERSION + '_h') tools.mkdir(txt_dir_h) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle_h[sub_class] = open( os.path.join(txt_dir_h, 'Task2_%s.txt' % sub_class), 'a+') for i, hbox in enumerate(box_res_): command = '%s %.3f %.1f %.1f %.1f %.1f\n' % ( img_path.split('/')[-1].split('.')[0], score_res_[i], hbox[0], hbox[1], hbox[2], hbox[3]) write_handle_h[LABEl_NAME_MAP[label_res_[i]]].write( command) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle_h[sub_class].close() view_bar( '{} cost {}s'.format( img_path.split('/')[-1].split('.')[0], time_elapsed), count + 1, len(file_paths))
def inference(det_net, file_paths, des_folder, h_len, w_len, h_overlap, w_overlap, save_res=False): if save_res: assert cfgs.SHOW_SCORE_THRSHOLD >= 0.5, \ 'please set score threshold (example: SHOW_SCORE_THRSHOLD = 0.5) in cfgs.py' else: assert cfgs.SHOW_SCORE_THRSHOLD <= 0.005, \ 'please set score threshold (example: SHOW_SCORE_THRSHOLD = 0.00) in cfgs.py' # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) img_batch = tf.cast(img_plac, tf.float32) if cfgs.NET_NAME in ['resnet101_v1d']: img_batch = (img_batch / 255 - tf.constant(cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) img_batch = short_side_resize_for_inference_data(img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, is_resize=False) det_boxes_h, det_scores_h, det_category_h = det_net.build_whole_detection_network(input_img_batch=img_batch, gtboxes_batch=None) init_op = tf.group( tf.global_variables_initializer(), tf.local_variables_initializer() ) restorer, restore_ckpt = det_net.get_restorer() config = tf.ConfigProto() config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') if not os.path.exists('./tmp.txt'): fw = open('./tmp.txt', 'w') fw.close() fr = open('./tmp.txt', 'r') pass_img = fr.readlines() fr.close() for count, img_path in enumerate(file_paths): fw = open('./tmp.txt', 'a+') if img_path + '\n' in pass_img: continue start = timer() img = cv2.imread(img_path) box_res = [] label_res = [] score_res = [] imgH = img.shape[0] imgW = img.shape[1] if imgH < h_len: temp = np.zeros([h_len, imgW, 3], np.float32) temp[0:imgH, :, :] = img img = temp imgH = h_len if imgW < w_len: temp = np.zeros([imgH, w_len, 3], np.float32) temp[:, 0:imgW, :] = img img = temp imgW = w_len for hh in range(0, imgH, h_len - h_overlap): if imgH - hh - 1 < h_len: hh_ = imgH - h_len else: hh_ = hh for ww in range(0, imgW, w_len - w_overlap): if imgW - ww - 1 < w_len: ww_ = imgW - w_len else: ww_ = ww src_img = img[hh_:(hh_ + h_len), ww_:(ww_ + w_len), :] det_boxes_h_, det_scores_h_, det_category_h_ = \ sess.run( [det_boxes_h, det_scores_h, det_category_h], feed_dict={img_plac: src_img[:, :, ::-1]} ) if len(det_boxes_h_) > 0: for ii in range(len(det_boxes_h_)): box = det_boxes_h_[ii] box[0] = box[0] + ww_ box[1] = box[1] + hh_ box[2] = box[2] + ww_ box[3] = box[3] + hh_ box_res.append(box) label_res.append(det_category_h_[ii]) score_res.append(det_scores_h_[ii]) box_res = np.array(box_res) label_res = np.array(label_res) score_res = np.array(score_res) box_res_, label_res_, score_res_ = [], [], [] h_threshold = {'roundabout': 0.35, 'tennis-court': 0.35, 'swimming-pool': 0.4, 'storage-tank': 0.3, 'soccer-ball-field': 0.3, 'small-vehicle': 0.4, 'ship': 0.35, 'plane': 0.35, 'large-vehicle': 0.4, 'helicopter': 0.4, 'harbor': 0.3, 'ground-track-field': 0.4, 'bridge': 0.3, 'basketball-court': 0.4, 'baseball-diamond': 0.3} for sub_class in range(1, cfgs.CLASS_NUM + 1): index = np.where(label_res == sub_class)[0] if len(index) == 0: continue tmp_boxes_h = box_res[index] tmp_label_h = label_res[index] tmp_score_h = score_res[index] tmp_boxes_h = np.array(tmp_boxes_h) tmp = np.zeros([tmp_boxes_h.shape[0], tmp_boxes_h.shape[1] + 1]) tmp[:, 0:-1] = tmp_boxes_h tmp[:, -1] = np.array(tmp_score_h) inx = nms.py_cpu_nms(dets=np.array(tmp, np.float32), thresh=h_threshold[LABEl_NAME_MAP[sub_class]], max_output_size=500) box_res_.extend(np.array(tmp_boxes_h)[inx]) score_res_.extend(np.array(tmp_score_h)[inx]) label_res_.extend(np.array(tmp_label_h)[inx]) time_elapsed = timer() - start if save_res: scores = np.array(score_res_) labels = np.array(label_res_) boxes = np.array(box_res_) valid_show = scores > cfgs.SHOW_SCORE_THRSHOLD scores = scores[valid_show] boxes = boxes[valid_show] labels = labels[valid_show] det_detections_h = draw_box_in_img.draw_boxes_with_label_and_scores(np.array(img, np.float32), boxes=np.array(boxes), labels=np.array(labels), scores=np.array(scores), in_graph=False) save_dir = os.path.join(des_folder, cfgs.VERSION) tools.mkdir(save_dir) cv2.imwrite(save_dir + '/' + img_path.split('/')[-1].split('.')[0] + '_h.jpg', det_detections_h) view_bar('{} cost {}s'.format(img_path.split('/')[-1].split('.')[0], time_elapsed), count + 1, len(file_paths)) else: # eval txt CLASS_DOTA = NAME_LABEL_MAP.keys() # Task2 write_handle_h = {} txt_dir_h = os.path.join('txt_output', cfgs.VERSION + '_h') tools.mkdir(txt_dir_h) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle_h[sub_class] = open(os.path.join(txt_dir_h, 'Task2_%s.txt' % sub_class), 'a+') for i, hbox in enumerate(box_res_): command = '%s %.3f %.1f %.1f %.1f %.1f\n' % (img_path.split('/')[-1].split('.')[0], score_res_[i], hbox[0], hbox[1], hbox[2], hbox[3]) write_handle_h[LABEl_NAME_MAP[label_res_[i]]].write(command) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle_h[sub_class].close() view_bar('%s cost %.3fs' % (img_path.split('/')[-1].split('.')[0], time_elapsed), count + 1, len(file_paths)) fw.write('{}\n'.format(img_path)) fw.close() os.remove('./tmp.txt')
def detect_img(file_paths, des_folder, det_th, h_len, w_len, h_overlap, w_overlap, show_res=False): with tf.Graph().as_default(): img_plac = tf.placeholder(shape=[None, None, 3], dtype=tf.uint8) img_tensor = tf.cast(img_plac, tf.float32) - tf.constant( [103.939, 116.779, 123.68]) img_batch = image_preprocess.short_side_resize_for_inference_data( img_tensor, target_shortside_len=cfgs.SHORT_SIDE_LEN, is_resize=False) # *********************************************************************************************** # * share net * # *********************************************************************************************** _, share_net = get_network_byname(net_name=cfgs.NET_NAME, inputs=img_batch, num_classes=None, is_training=True, output_stride=None, global_pool=False, spatial_squeeze=False) # *********************************************************************************************** # * RPN * # *********************************************************************************************** rpn = build_rpn.RPN( net_name=cfgs.NET_NAME, inputs=img_batch, gtboxes_and_label=None, is_training=False, share_head=cfgs.SHARE_HEAD, share_net=share_net, stride=cfgs.STRIDE, anchor_ratios=cfgs.ANCHOR_RATIOS, anchor_scales=cfgs.ANCHOR_SCALES, scale_factors=cfgs.SCALE_FACTORS, base_anchor_size_list=cfgs. BASE_ANCHOR_SIZE_LIST, # P2, P3, P4, P5, P6 level=cfgs.LEVEL, top_k_nms=cfgs.RPN_TOP_K_NMS, rpn_nms_iou_threshold=cfgs.RPN_NMS_IOU_THRESHOLD, max_proposals_num=cfgs.MAX_PROPOSAL_NUM, rpn_iou_positive_threshold=cfgs.RPN_IOU_POSITIVE_THRESHOLD, rpn_iou_negative_threshold=cfgs.RPN_IOU_NEGATIVE_THRESHOLD, rpn_mini_batch_size=cfgs.RPN_MINIBATCH_SIZE, rpn_positives_ratio=cfgs.RPN_POSITIVE_RATE, remove_outside_anchors=False, # whether remove anchors outside rpn_weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME]) # rpn predict proposals rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals( ) # rpn_score shape: [300, ] # *********************************************************************************************** # * Fast RCNN * # *********************************************************************************************** fast_rcnn = build_fast_rcnn.FastRCNN( img_batch=img_batch, feature_pyramid=rpn.feature_pyramid, rpn_proposals_boxes=rpn_proposals_boxes, rpn_proposals_scores=rpn_proposals_scores, img_shape=tf.shape(img_batch), roi_size=cfgs.ROI_SIZE, scale_factors=cfgs.SCALE_FACTORS, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, gtboxes_and_label=None, fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD, fast_rcnn_maximum_boxes_per_img=100, fast_rcnn_nms_max_boxes_per_class=cfgs. FAST_RCNN_NMS_MAX_BOXES_PER_CLASS, show_detections_score_threshold=det_th, # show detections which score >= 0.6 num_classes=cfgs.CLASS_NUM, fast_rcnn_minibatch_size=cfgs.FAST_RCNN_MINIBATCH_SIZE, fast_rcnn_positives_ratio=cfgs.FAST_RCNN_POSITIVE_RATE, fast_rcnn_positives_iou_threshold=cfgs. FAST_RCNN_IOU_POSITIVE_THRESHOLD, use_dropout=False, weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME], is_training=False, level=cfgs.LEVEL) fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category = \ fast_rcnn.fast_rcnn_predict() init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = restore_model.get_restorer() config = tf.ConfigProto() # config.gpu_options.per_process_gpu_memory_fraction = 0.5 config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) for img_path in file_paths: start = timer() img = cv2.imread(img_path) box_res = [] label_res = [] score_res = [] imgH = img.shape[0] imgW = img.shape[1] for hh in range(0, imgH, h_len - h_overlap): if imgH - hh - 1 < h_len: hh_ = imgH - h_len else: hh_ = hh for ww in range(0, imgW, w_len - w_overlap): if imgW - ww - 1 < w_len: ww_ = imgW - w_len else: ww_ = ww src_img = img[hh_:(hh_ + h_len), ww_:(ww_ + w_len), :] boxes, labels, scores = sess.run( [ fast_rcnn_decode_boxes, detection_category, fast_rcnn_score ], feed_dict={img_plac: src_img}) if show_res: visualize_detection(src_img, boxes, scores) if len(boxes) > 0: for ii in range(len(boxes)): box = boxes[ii] box[0] = box[0] + hh_ box[1] = box[1] + ww_ box[2] = box[2] + hh_ box[3] = box[3] + ww_ box_res.append(box) label_res.append(labels[ii]) score_res.append(scores[ii]) box_res = np.array(box_res) label_res = np.array(label_res) score_res = np.array(score_res) box_res_, label_res_, score_res_ = [], [], [] for sub_class in range(1, cfgs.CLASS_NUM + 1): index = np.where(label_res == sub_class)[0] if len(index) == 0: continue tmp_boxes_h = box_res[index] tmp_label_h = label_res[index] tmp_score_h = score_res[index] tmp_boxes_h = np.array(tmp_boxes_h) tmp = np.zeros( [tmp_boxes_h.shape[0], tmp_boxes_h.shape[1] + 1]) tmp[:, 0:-1] = tmp_boxes_h tmp[:, -1] = np.array(tmp_score_h) inx = nms.py_cpu_nms(dets=np.array(tmp, np.float32), thresh=0.7, max_output_size=500) box_res_.extend(np.array(tmp_boxes_h)[inx]) score_res_.extend(np.array(tmp_score_h)[inx]) label_res_.extend(np.array(tmp_label_h)[inx]) time_elapsed = timer() - start print("{} detection time : {:.4f} sec".format( img_path.split('/')[-1].split('.')[0], time_elapsed)) mkdir(des_folder) img_np = draw_box_cv(np.array(img, np.float32) - np.array([103.939, 116.779, 123.68]), boxes=np.array(box_res_), labels=np.array(label_res_), scores=np.array(score_res_)) cv2.imwrite( des_folder + '/{}_fpn.jpg'.format( img_path.split('/')[-1].split('.')[0]), img_np) coord.request_stop() coord.join(threads)
def eval_coco(det_net, real_test_img_list, draw_imgs=False): # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not BGR img_batch = tf.cast(img_plac, tf.float32) img_batch = short_side_resize_for_inference_data(img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, length_limitation=cfgs.IMG_MAX_LENGTH, is_resize=False) if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']: img_batch = (img_batch / 255 - tf.constant(cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) # img_batch = (img_batch - tf.constant(cfgs.PIXEL_MEAN)) / (tf.constant(cfgs.PIXEL_STD)*255) img_batch = tf.expand_dims(img_batch, axis=0) detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network( input_img_batch=img_batch, gtboxes_batch=None) init_op = tf.group( tf.global_variables_initializer(), tf.local_variables_initializer() ) restorer, restore_ckpt = det_net.get_restorer() config = tf.ConfigProto() config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') save_path = os.path.join('./eval_coco', cfgs.VERSION) tools.mkdir(save_path) fw_json_dt = open(os.path.join(save_path, 'coco_minival_ms.json'), 'w') coco_det = [] for i, a_img in enumerate(real_test_img_list): record = json.loads(a_img) raw_img = cv2.imread(record['fpath']) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] start = time.time() detected_scores_, detected_boxes_, detected_categories_ = [], [], [] for ss in [600, 800, 1000, 1200]: # cfgs.IMG_SHORT_SIDE_LEN: img_resize = cv2.resize(raw_img, (ss, ss)) resized_img, tmp_detected_boxes, tmp_detected_scores, tmp_detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: img_resize[:, :, ::-1]} # cv is BGR. But need RGB ) eval_indices = tmp_detected_scores >= 0.01 tmp_detected_scores = tmp_detected_scores[eval_indices] tmp_detected_boxes = tmp_detected_boxes[eval_indices] tmp_detected_categories = tmp_detected_categories[eval_indices] xmin, ymin, xmax, ymax = tmp_detected_boxes[:, 0], tmp_detected_boxes[:, 1], \ tmp_detected_boxes[:, 2], tmp_detected_boxes[:, 3] resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] xmin = xmin * raw_w / resized_w xmax = xmax * raw_w / resized_w ymin = ymin * raw_h / resized_h ymax = ymax * raw_h / resized_h resize_boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax])) detected_scores_.append(tmp_detected_scores) detected_boxes_.append(resize_boxes) detected_categories_.append(tmp_detected_categories) detected_scores_ = np.concatenate(detected_scores_) detected_boxes_ = np.concatenate(detected_boxes_) detected_categories_ = np.concatenate(detected_categories_) detected_scores, detected_boxes, detected_categories = [], [], [] for sub_class in range(1, cfgs.CLASS_NUM + 1): index = np.where(detected_categories_ == sub_class)[0] if len(index) == 0: continue tmp_boxes_h = detected_boxes_[index] tmp_label_h = detected_categories_[index] tmp_score_h = detected_scores_[index] tmp_boxes_h = np.array(tmp_boxes_h) tmp = np.zeros([tmp_boxes_h.shape[0], tmp_boxes_h.shape[1] + 1]) tmp[:, 0:-1] = tmp_boxes_h tmp[:, -1] = np.array(tmp_score_h) inx = nms.py_cpu_nms(dets=np.array(tmp, np.float32), thresh=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD, max_output_size=500) detected_boxes.extend(np.array(tmp_boxes_h)[inx]) detected_scores.extend(np.array(tmp_score_h)[inx]) detected_categories.extend(np.array(tmp_label_h)[inx]) detected_scores = np.array(detected_scores) detected_boxes = np.array(detected_boxes) detected_categories = np.array(detected_categories) # print("{} cost time : {} ".format(img_name, (end - start))) if draw_imgs: show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD show_scores = detected_scores[show_indices] show_boxes = detected_boxes[show_indices] show_categories = detected_categories[show_indices] # if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']: # draw_img = (raw_img * np.array(cfgs.PIXEL_STD) + np.array(cfgs.PIXEL_MEAN_)) * 255 # else: # draw_img = raw_img + np.array(cfgs.PIXEL_MEAN) # draw_img = draw_img * (np.array(cfgs.PIXEL_STD)*255) + np.array(cfgs.PIXEL_MEAN) raw_img = np.array(raw_img, np.float32) final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(raw_img, boxes=show_boxes, labels=show_categories, scores=show_scores, in_graph=False) if not os.path.exists(cfgs.TEST_SAVE_PATH): os.makedirs(cfgs.TEST_SAVE_PATH) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + record['ID'], final_detections) # cost much time for j, box in enumerate(detected_boxes): coco_det.append({'bbox': [float(box[0]), float(box[1]), float(box[2]-box[0]), float(box[3]-box[1])], 'score': float(detected_scores[j]), 'image_id': int(record['ID'].split('.jpg')[0].split('_000000')[-1]), 'category_id': int(classes_originID[LABEl_NAME_MAP[detected_categories[j]]])}) end = time.time() tools.view_bar('%s image cost %.3fs' % (record['ID'], (end - start)), i + 1, len(real_test_img_list)) json.dump(coco_det, fw_json_dt) fw_json_dt.close() return os.path.join(save_path, 'coco_minival_ms.json')