def eval_ship(img_num, mode): with tf.Graph().as_default(): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=False) gtboxes_and_label = tf.py_func( back_forward_convert, inp=[tf.squeeze(gtboxes_and_label_batch, 0)], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape( gtboxes_and_label_minAreaRectangle, [-1, 5]) # *********************************************************************************************** # * 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=False, 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_rcnn1.FastRCNN( 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, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=None, gtboxes_and_label_minAreaRectangle= gtboxes_and_label_minAreaRectangle, 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=cfgs.FINAL_SCORE_THRESHOLD, # 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, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, 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_decode_boxes_rotate, fast_rcnn_score_rotate, num_of_objects_rotate, detection_category_rotate = \ fast_rcnn.fast_rcnn_predict() if mode == 0: fast_rcnn_decode_boxes_rotate = get_horizen_minAreaRectangle( fast_rcnn_decode_boxes_rotate, False) # train init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = restore_model.get_restorer() with tf.Session() 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) gtboxes_horizontal_dict = {} predict_horizontal_dict = {} gtboxes_rotate_dict = {} predict_rotate_dict = {} for i in range(img_num): start = time.time() _img_name_batch, _img_batch, _gtboxes_and_label, _gtboxes_and_label_minAreaRectangle, \ _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category, _fast_rcnn_decode_boxes_rotate, \ _fast_rcnn_score_rotate, _detection_category_rotate \ = sess.run([img_name_batch, img_batch, gtboxes_and_label, gtboxes_and_label_minAreaRectangle, fast_rcnn_decode_boxes, fast_rcnn_score, detection_category, fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, detection_category_rotate]) end = time.time() # gtboxes convert dict gtboxes_horizontal_dict[str(_img_name_batch[0])] = [] predict_horizontal_dict[str(_img_name_batch[0])] = [] gtboxes_rotate_dict[str(_img_name_batch[0])] = [] predict_rotate_dict[str(_img_name_batch[0])] = [] gtbox_horizontal_list, predict_horizontal_list = \ make_dict_packle(_gtboxes_and_label_minAreaRectangle, _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category) if mode == 0: gtbox_rotate_list, predict_rotate_list = \ make_dict_packle(_gtboxes_and_label_minAreaRectangle, _fast_rcnn_decode_boxes_rotate, _fast_rcnn_score_rotate, _detection_category_rotate) else: gtbox_rotate_list, predict_rotate_list = \ make_dict_packle(_gtboxes_and_label, _fast_rcnn_decode_boxes_rotate, _fast_rcnn_score_rotate, _detection_category_rotate) gtboxes_horizontal_dict[str( _img_name_batch[0])].extend(gtbox_horizontal_list) predict_horizontal_dict[str( _img_name_batch[0])].extend(predict_horizontal_list) gtboxes_rotate_dict[str( _img_name_batch[0])].extend(gtbox_rotate_list) predict_rotate_dict[str( _img_name_batch[0])].extend(predict_rotate_list) view_bar( '{} image cost {}s'.format(str(_img_name_batch[0]), (end - start)), i + 1, img_num) fw1 = open('gtboxes_horizontal_dict.pkl', 'w') fw2 = open('predict_horizontal_dict.pkl', 'w') fw3 = open('gtboxes_rotate_dict.pkl', 'w') fw4 = open('predict_rotate_dict.pkl', 'w') pickle.dump(gtboxes_horizontal_dict, fw1) pickle.dump(predict_horizontal_dict, fw2) pickle.dump(gtboxes_rotate_dict, fw3) pickle.dump(predict_rotate_dict, fw4) fw1.close() fw2.close() fw3.close() fw4.close() coord.request_stop() coord.join(threads)
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_rcnn1.FastRCNN(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, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=None, gtboxes_and_label_minAreaRectangle=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, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, 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_decode_boxes_rotate, fast_rcnn_score_rotate, num_of_objects_rotate, detection_category_rotate = \ 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 = [] box_res_rotate = [] label_res_rotate = [] score_res_rotate = [] imgH = img.shape[0] imgW = img.shape[1] for hh in range(0, imgH, h_len - h_overlap): h_size = min(h_len, imgH - hh) if h_size < 10: break for ww in range(0, imgW, w_len - w_overlap): w_size = min(w_len, imgW - ww) if w_size < 10: break src_img = img[hh:(hh + h_size), ww:(ww + w_size), :] boxes, labels, scores = sess.run([fast_rcnn_decode_boxes, detection_category, fast_rcnn_score], feed_dict={img_plac: src_img}) boxes_rotate, labels_rotate, scores_rotate = sess.run([fast_rcnn_decode_boxes_rotate, detection_category_rotate, fast_rcnn_score_rotate], feed_dict={img_plac: src_img}) 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]) if len(boxes_rotate) > 0: for ii in range(len(boxes_rotate)): box_rotate = boxes_rotate[ii] box_rotate[0] = box_rotate[0] + hh box_rotate[1] = box_rotate[1] + ww box_res_rotate.append(box_rotate) label_res_rotate.append(labels_rotate[ii]) score_res_rotate.append(scores_rotate[ii]) # inx = nms_rotate.nms_rotate_cpu(boxes=np.array(box_res_rotate), scores=np.array(score_res_rotate), # iou_threshold=0.5, max_output_size=100) # box_res_rotate = np.array(box_res_rotate)[inx] # score_res_rotate = np.array(score_res_rotate)[inx] # label_res_rotate = np.array(label_res_rotate)[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)) img_np_rotate = draw_rotate_box_cv(np.array(img, np.float32) - np.array([103.939, 116.779, 123.68]), boxes=np.array(box_res_rotate), labels=np.array(label_res_rotate), scores=np.array(score_res_rotate)) cv2.imwrite(des_folder + '/{}_horizontal_fpn.jpg'.format(img_path.split('/')[-1].split('.')[0]), img_np) cv2.imwrite(des_folder + '/{}_rotate_fpn.jpg'.format(img_path.split('/')[-1].split('.')[0]), img_np_rotate) coord.request_stop() coord.join(threads)
def inference(): 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) # *********************************************************************************************** # * 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_rcnn1.FastRCNN(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, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=None, gtboxes_and_label_minAreaRectangle=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=cfgs.FINAL_SCORE_THRESHOLD, # 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, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, 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_decode_boxes_rotate, fast_rcnn_score_rotate, num_of_objects_rotate, detection_category_rotate = \ 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) imgs, img_names = get_imgs() for i, img in enumerate(imgs): start = time.time() _img_batch, _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category, \ _fast_rcnn_decode_boxes_rotate, _fast_rcnn_score_rotate, _detection_category_rotate = \ sess.run([img_batch, fast_rcnn_decode_boxes, fast_rcnn_score, detection_category, fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, detection_category_rotate], feed_dict={img_plac: img}) end = time.time() img_np = np.squeeze(_img_batch, axis=0) img_horizontal_np = draw_box_cv(img_np, boxes=_fast_rcnn_decode_boxes, labels=_detection_category, scores=_fast_rcnn_score) img_rotate_np = draw_rotate_box_cv(img_np, boxes=_fast_rcnn_decode_boxes_rotate, labels=_detection_category_rotate, scores=_fast_rcnn_score_rotate) mkdir(cfgs.INFERENCE_SAVE_PATH) cv2.imwrite(cfgs.INFERENCE_SAVE_PATH + '/{}_horizontal_fpn.jpg'.format(img_names[i]), img_horizontal_np) cv2.imwrite(cfgs.INFERENCE_SAVE_PATH + '/{}_rotate_fpn.jpg'.format(img_names[i]), img_rotate_np) view_bar('{} cost {}s'.format(img_names[i], (end - start)), i + 1, len(imgs)) coord.request_stop() coord.join(threads)
def test(img_num): with tf.Graph().as_default(): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=False) gtboxes_and_label = tf.py_func(back_forward_convert, inp=[tf.squeeze(gtboxes_and_label_batch, 0)], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle(gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape(gtboxes_and_label_minAreaRectangle, [-1, 5]) # *********************************************************************************************** # * 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_rcnn1.FastRCNN(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, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=None, gtboxes_and_label_minAreaRectangle=gtboxes_and_label_minAreaRectangle, 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=cfgs.FINAL_SCORE_THRESHOLD, # 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, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, 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_decode_boxes_rotate, fast_rcnn_score_rotate, num_of_objects_rotate, detection_category_rotate = \ fast_rcnn.fast_rcnn_predict() # train 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 i in range(img_num): start = time.time() _img_name_batch, _img_batch, _gtboxes_and_label, _gtboxes_and_label_minAreaRectangle, \ _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category, _fast_rcnn_decode_boxes_rotate, \ _fast_rcnn_score_rotate, _detection_category_rotate \ = sess.run([img_name_batch, img_batch, gtboxes_and_label, gtboxes_and_label_minAreaRectangle, fast_rcnn_decode_boxes, fast_rcnn_score, detection_category, fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, detection_category_rotate]) end = time.time() _img_batch = np.squeeze(_img_batch, axis=0) _img_batch_fpn_horizonal = help_utils.draw_box_cv(_img_batch, boxes=_fast_rcnn_decode_boxes, labels=_detection_category, scores=_fast_rcnn_score) _img_batch_fpn_rotate = help_utils.draw_rotate_box_cv(_img_batch, boxes=_fast_rcnn_decode_boxes_rotate, labels=_detection_category_rotate, scores=_fast_rcnn_score_rotate) mkdir(cfgs.TEST_SAVE_PATH) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/{}_horizontal_fpn.jpg'.format(str(_img_name_batch[0])), _img_batch_fpn_horizonal) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/{}_rotate_fpn.jpg'.format(str(_img_name_batch[0])), _img_batch_fpn_rotate) temp_label_horizontal = np.reshape(_gtboxes_and_label[:, -1:], [-1, ]).astype(np.int64) temp_label_rotate = np.reshape(_gtboxes_and_label[:, -1:], [-1, ]).astype(np.int64) _img_batch_gt_horizontal = help_utils.draw_box_cv(_img_batch, boxes=_gtboxes_and_label_minAreaRectangle[:, :-1], labels=temp_label_horizontal, scores=None) _img_batch_gt_rotate = help_utils.draw_rotate_box_cv(_img_batch, boxes=_gtboxes_and_label[:, :-1], labels=temp_label_rotate, scores=None) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/{}_horizontal_gt.jpg'.format(str(_img_name_batch[0])), _img_batch_gt_horizontal) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/{}_rotate_gt.jpg'.format(str(_img_name_batch[0])), _img_batch_gt_rotate) #view_bar('{} image cost {}s'.format(str(_img_name_batch[0]), (end - start)), i + 1, img_num) print('{} image cost {}s'.format(str(_img_name_batch[0]), (end - start)) + str(i + 1)+'/'+str(img_num)) coord.request_stop() coord.join(threads)
def train(): with tf.Graph().as_default(): with tf.name_scope('get_batch'): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=True) gtboxes_and_label = tf.py_func(back_forward_convert, inp=[tf.squeeze(gtboxes_and_label_batch, 0)], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle(gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape(gtboxes_and_label_minAreaRectangle, [-1, 5]) with tf.name_scope('draw_gtboxes'): gtboxes_in_img = draw_box_with_color(img_batch, tf.reshape(gtboxes_and_label_minAreaRectangle, [-1, 5])[:, :-1], text=tf.shape(gtboxes_and_label_minAreaRectangle)[0]) gtboxes_rotate_in_img = draw_box_with_color_rotate(img_batch, tf.reshape(gtboxes_and_label, [-1, 6])[:, :-1], text=tf.shape(gtboxes_and_label)[0]) # *********************************************************************************************** # * 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=gtboxes_and_label_minAreaRectangle, is_training=True, share_head=False, 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, # iou>=0.7 is positive box, iou< 0.3 is negative 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_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals() # rpn_score shape: [300, ] rpn_location_loss, rpn_classification_loss = rpn.rpn_losses() rpn_total_loss = rpn_classification_loss + rpn_location_loss with tf.name_scope('draw_proposals'): # score > 0.5 is object rpn_object_boxes_indices = tf.reshape(tf.where(tf.greater(rpn_proposals_scores, 0.5)), [-1]) rpn_object_boxes = tf.gather(rpn_proposals_boxes, rpn_object_boxes_indices) rpn_proposals_objcet_boxes_in_img = draw_box_with_color(img_batch, rpn_object_boxes, text=tf.shape(rpn_object_boxes)[0]) rpn_proposals_boxes_in_img = draw_box_with_color(img_batch, rpn_proposals_boxes, text=tf.shape(rpn_proposals_boxes)[0]) # *********************************************************************************************** # * Fast RCNN * # *********************************************************************************************** fast_rcnn = build_fast_rcnn1.FastRCNN(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, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=gtboxes_and_label, gtboxes_and_label_minAreaRectangle=gtboxes_and_label_minAreaRectangle, 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=cfgs.FINAL_SCORE_THRESHOLD, # 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, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME], is_training=True, level=cfgs.LEVEL) fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category, \ fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, num_of_objects_rotate, detection_category_rotate = \ fast_rcnn.fast_rcnn_predict() fast_rcnn_location_loss, fast_rcnn_classification_loss, \ fast_rcnn_location_rotate_loss, fast_rcnn_classification_rotate_loss = fast_rcnn.fast_rcnn_loss() fast_rcnn_total_loss = fast_rcnn_location_loss + fast_rcnn_classification_loss + \ fast_rcnn_location_rotate_loss + fast_rcnn_classification_rotate_loss with tf.name_scope('draw_boxes_with_categories'): fast_rcnn_predict_boxes_in_imgs = draw_boxes_with_categories(img_batch=img_batch, boxes=fast_rcnn_decode_boxes, labels=detection_category, scores=fast_rcnn_score) fast_rcnn_predict_rotate_boxes_in_imgs = draw_boxes_with_categories_rotate(img_batch=img_batch, boxes=fast_rcnn_decode_boxes_rotate, labels=detection_category_rotate, scores=fast_rcnn_score_rotate) # train total_loss = slim.losses.get_total_loss() global_step = slim.get_or_create_global_step() lr = tf.train.piecewise_constant(global_step, boundaries=[np.int64(20000), np.int64(40000)], values=[0.001, 0.0001, 0.00001]) tf.summary.scalar('lr', lr) optimizer = tf.train.MomentumOptimizer(lr, momentum=cfgs.MOMENTUM) train_op = slim.learning.create_train_op(total_loss, optimizer, global_step) # rpn_total_loss, # train_op = optimizer.minimize(second_classification_loss, global_step) # *********************************************************************************************** # * Summary * # *********************************************************************************************** # ground truth and predict tf.summary.image('img/gtboxes', gtboxes_in_img) tf.summary.image('img/gtboxes_rotate', gtboxes_rotate_in_img) tf.summary.image('img/faster_rcnn_predict', fast_rcnn_predict_boxes_in_imgs) tf.summary.image('img/faster_rcnn_predict_rotate', fast_rcnn_predict_rotate_boxes_in_imgs) # rpn loss and image tf.summary.scalar('rpn/rpn_location_loss', rpn_location_loss) tf.summary.scalar('rpn/rpn_classification_loss', rpn_classification_loss) tf.summary.scalar('rpn/rpn_total_loss', rpn_total_loss) tf.summary.scalar('fast_rcnn/fast_rcnn_location_loss', fast_rcnn_location_loss) tf.summary.scalar('fast_rcnn/fast_rcnn_classification_loss', fast_rcnn_classification_loss) tf.summary.scalar('fast_rcnn/fast_rcnn_location_rotate_loss', fast_rcnn_location_rotate_loss) tf.summary.scalar('fast_rcnn/fast_rcnn_classification_rotate_loss', fast_rcnn_classification_rotate_loss) tf.summary.scalar('fast_rcnn/fast_rcnn_total_loss', fast_rcnn_total_loss) tf.summary.scalar('loss/total_loss', total_loss) tf.summary.image('rpn/rpn_all_boxes', rpn_proposals_boxes_in_img) tf.summary.image('rpn/rpn_object_boxes', rpn_proposals_objcet_boxes_in_img) # learning_rate tf.summary.scalar('learning_rate', lr) summary_op = tf.summary.merge_all() init_op = tf.group( tf.global_variables_initializer(), tf.local_variables_initializer() ) restorer, restore_ckpt = restore_model.get_restorer() saver = tf.train.Saver(max_to_keep=10) with tf.Session() 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) summary_path = os.path.join(cfgs.SUMMARY_PATH, cfgs.VERSION) tools.mkdir(summary_path) summary_writer = tf.summary.FileWriter(summary_path, graph=sess.graph) for step in range(cfgs.MAX_ITERATION): training_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())) start = time.time() _global_step, _img_name_batch, _rpn_location_loss, _rpn_classification_loss, \ _rpn_total_loss, _fast_rcnn_location_loss, _fast_rcnn_classification_loss, \ _fast_rcnn_location_rotate_loss, _fast_rcnn_classification_rotate_loss, \ _fast_rcnn_total_loss, _total_loss, _ = \ sess.run([global_step, img_name_batch, rpn_location_loss, rpn_classification_loss, rpn_total_loss, fast_rcnn_location_loss, fast_rcnn_classification_loss, fast_rcnn_location_rotate_loss, fast_rcnn_classification_rotate_loss, fast_rcnn_total_loss, total_loss, train_op]) end = time.time() if step % 10 == 0: print(""" {}: step{} image_name:{} |\t rpn_loc_loss:{} |\t rpn_cla_loss:{} |\t rpn_total_loss:{} | fast_rcnn_loc_loss:{} |\t fast_rcnn_cla_loss:{} |\t fast_rcnn_loc_rotate_loss:{} |\t fast_rcnn_cla_rotate_loss:{} |\t fast_rcnn_total_loss:{} |\t total_loss:{} |\t pre_cost_time:{}s""" \ .format(training_time, _global_step, str(_img_name_batch[0]), _rpn_location_loss, _rpn_classification_loss, _rpn_total_loss, _fast_rcnn_location_loss, _fast_rcnn_classification_loss, _fast_rcnn_location_rotate_loss, _fast_rcnn_classification_rotate_loss, _fast_rcnn_total_loss, _total_loss, (end - start))) if step % 50 == 0: summary_str = sess.run(summary_op) summary_writer.add_summary(summary_str, _global_step) summary_writer.flush() if (step > 0 and step % 1000 == 0) or (step == cfgs.MAX_ITERATION - 1): save_dir = os.path.join(FLAGS.trained_checkpoint, cfgs.VERSION) if not os.path.exists(save_dir): os.mkdir(save_dir) save_ckpt = os.path.join(save_dir, 'voc_'+str(_global_step)+'model.ckpt') saver.save(sess, save_ckpt) print(' weights had been saved') coord.request_stop() coord.join(threads)