Exemplo n.º 1
0
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=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_rcnn.FastRCNN(
            feature_pyramid=rpn.feature_pyramid,
            rpn_proposals_boxes=rpn_proposals_boxes,
            rpn_proposals_scores=rpn_proposals_scores,
            img_shape=tf.shape(img_batch),
            img_batch=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,
            head_quadrant=None)

        fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category, \
        fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, fast_rcnn_head_quadrant, \
        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()  # 创建一个线程协调器,用来管理之后在Session中启动的所有线程
            threads = tf.train.start_queue_runners(
                sess,
                coord)  # 启动入队线程,把文件读入Filename Queue中(一般系统有多少个核,就会启动多少个入队线程)

            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, _fast_rcnn_head_quadrant, \
                _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, fast_rcnn_head_quadrant,
                              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,
                    head=np.argmax(_fast_rcnn_head_quadrant, axis=1))
                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)
Exemplo n.º 2
0
def eval_dict_convert(img_num, mode):
    with tf.Graph().as_default():

        # img = tf.placeholder(shape=[None, None, 3], dtype=tf.uint8)

        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=gtboxes_and_label,
            is_training=False,
            share_head=cfgs.SHARE_HEAD,
            share_net=share_net,
            anchor_ratios=cfgs.ANCHOR_RATIOS,
            anchor_scales=cfgs.ANCHOR_SCALES,
            anchor_angles=cfgs.ANCHOR_ANGLES,
            scale_factors=cfgs.SCALE_FACTORS,
            base_anchor_size_list=cfgs.
            BASE_ANCHOR_SIZE_LIST,  # P2, P3, P4, P5, P6
            level=cfgs.LEVEL,
            anchor_stride=cfgs.ANCHOR_STRIDE,
            pool_stride=cfgs.POOL_STRIDE,
            top_k_nms=cfgs.RPN_TOP_K_NMS,
            kernel_size=cfgs.KERNEL_SIZE,
            use_angles_condition=False,
            anchor_angle_threshold=cfgs.RPN_ANCHOR_ANGLES_THRESHOLD,
            nms_angle_threshold=cfgs.RPN_NMS_ANGLES_THRESHOLD,
            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=cfgs.
            IS_FILTER_OUTSIDE_BOXES,  # whether remove anchors outside
            rpn_weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME],
            scope='')

        rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals(
        )  # rpn_score shape: [300, ]
        _, _, rpn_predict_boxes, rpn_predict_scores = rpn.rpn_losses()

        # ***********************************************************************************************
        # *                                         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,
            stop_gradient_for_proposals=False,
            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,
            fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD,
            top_k_nms=cfgs.FAST_RCNN_TOP_K_NMS,
            nms_angle_threshold=cfgs.FAST_RCNN_NMS_ANGLES_THRESHOLD,
            use_angle_condition=False,
            level=cfgs.LEVEL,
            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,
            boxes_angle_threshold=cfgs.FAST_RCNN_BOXES_ANGLES_THRESHOLD,
            use_dropout=cfgs.USE_DROPOUT,
            weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME],
            is_training=False)

        fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category = \
            fast_rcnn.fast_rcnn_predict()

        ##############################################################################################
        if cfgs.NEED_AUXILIARY:
            predict_boxes = tf.concat(
                [fast_rcnn_decode_boxes, rpn_predict_boxes], axis=0)
            predict_scores = tf.concat(
                [fast_rcnn_score, rpn_predict_scores - 0.2], axis=0)
            rpn_predict_label = tf.ones([
                tf.shape(rpn_predict_scores)[0],
            ], tf.int64)
            labels = tf.concat([detection_category, rpn_predict_label], axis=0)

            # valid_indices = nms_rotate.nms_rotate(decode_boxes=predict_boxes,
            #                                       scores=predict_scores,
            #                                       iou_threshold=0.15,
            #                                       max_output_size=30,
            #                                       use_angle_condition=False,
            #                                       angle_threshold=15,
            #                                       use_gpu=True)
            valid_indices = tf.py_func(nms_rotate.nms_rotate_cpu,
                                       inp=[
                                           predict_boxes, predict_scores,
                                           tf.constant(0.15, tf.float32),
                                           tf.constant(30, tf.float32)
                                       ],
                                       Tout=tf.int64)

            fast_rcnn_decode_boxes = tf.gather(predict_boxes, valid_indices)
            fast_rcnn_score = tf.gather(predict_scores, valid_indices)
            detection_category = tf.gather(labels, valid_indices)

        ##############################################################################################
        if mode == 0:
            fast_rcnn_decode_boxes = get_horizen_minAreaRectangle(
                fast_rcnn_decode_boxes, False)

        # 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)

            gtboxes_dict = {}
            predict_dict = {}

            for i in range(img_num):

                start = time.time()

                _img_name_batch, _img_batch, _gtboxes_and_label, _fast_rcnn_decode_boxes, \
                _gtboxes_and_label_minAreaRectangle, _fast_rcnn_score, _detection_category \
                    = sess.run([img_name_batch, img_batch, gtboxes_and_label, fast_rcnn_decode_boxes,
                                gtboxes_and_label_minAreaRectangle, fast_rcnn_score, detection_category])
                end = time.time()

                # gtboxes convert dict
                gtboxes_dict[str(_img_name_batch[0])] = []
                predict_dict[str(_img_name_batch[0])] = []

                # for j, box in enumerate(_gtboxes_and_label):
                #     bbox_dict = {}
                #     bbox_dict['bbox'] = np.array(_gtboxes_and_label[j, :-1], np.float64)
                #     bbox_dict['name'] = LABEl_NAME_MAP[int(_gtboxes_and_label[j, -1])]
                #     gtbox_dict[str(_img_name_batch[0])].append(bbox_dict)
                #
                # for label in NAME_LABEL_MAP.keys():
                #     if label == 'back_ground':
                #         continue
                #     else:
                #         temp_dict = {}
                #         temp_dict['name'] = label
                #
                #         ind = np.where(_detection_category == NAME_LABEL_MAP[label])[0]
                #         temp_boxes = _fast_rcnn_decode_boxes[ind]
                #         temp_score = np.reshape(_fast_rcnn_score[ind], [-1, 1])
                #         temp_dict['bbox'] = np.array(np.concatenate([temp_boxes, temp_score], axis=1), np.float64)
                #         predict_dict[str(_img_name_batch[0])].append(temp_dict)

                if mode == 0:
                    gtboxes_list, predict_list = \
                        make_dict_packle(_gtboxes_and_label_minAreaRectangle, _fast_rcnn_decode_boxes,
                                         _fast_rcnn_score, _detection_category)
                else:
                    gtboxes_list, predict_list = \
                        make_dict_packle(_gtboxes_and_label, _fast_rcnn_decode_boxes,
                                         _fast_rcnn_score, _detection_category)

                gtboxes_dict[str(_img_name_batch[0])].extend(gtboxes_list)
                predict_dict[str(_img_name_batch[0])].extend(predict_list)

                view_bar(
                    '{} image cost {}s'.format(str(_img_name_batch[0]),
                                               (end - start)), i + 1, img_num)

            fw1 = open('gtboxes_dict.pkl', 'w')
            fw2 = open('predict_dict.pkl', 'w')
            pickle.dump(gtboxes_dict, fw1)
            pickle.dump(predict_dict, fw2)
            fw1.close()
            fw2.close()
            coord.request_stop()
            coord.join(threads)
Exemplo n.º 3
0
def test(img_num):
    with tf.Graph().as_default():

        # img = tf.placeholder(shape=[None, None, 3], dtype=tf.uint8)

        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])

        # ***********************************************************************************************
        # *                                         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,
            is_training=False,
            share_head=cfgs.SHARE_HEAD,
            share_net=share_net,
            anchor_ratios=cfgs.ANCHOR_RATIOS,
            anchor_scales=cfgs.ANCHOR_SCALES,
            anchor_angles=cfgs.ANCHOR_ANGLES,
            scale_factors=cfgs.SCALE_FACTORS,
            base_anchor_size_list=cfgs.
            BASE_ANCHOR_SIZE_LIST,  # P2, P3, P4, P5, P6
            level=cfgs.LEVEL,
            anchor_stride=cfgs.ANCHOR_STRIDE,
            pool_stride=cfgs.POOL_STRIDE,
            top_k_nms=cfgs.RPN_TOP_K_NMS,
            kernel_size=cfgs.KERNEL_SIZE,
            use_angles_condition=True,
            anchor_angle_threshold=cfgs.RPN_ANCHOR_ANGLES_THRESHOLD,
            nms_angle_threshold=cfgs.RPN_NMS_ANGLES_THRESHOLD,
            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],
            scope='')

        # rpn predict proposals
        rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals(
        )  # rpn_score shape: [300, ]
        _, _, rpn_top_k_boxes, rpn_top_k_scores = rpn.rpn_losses()

        # ***********************************************************************************************
        # *                                         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,
            stop_gradient_for_proposals=False,
            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,
            fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD,
            top_k_nms=cfgs.FAST_RCNN_TOP_K_NMS,
            nms_angle_threshold=cfgs.FAST_RCNN_NMS_ANGLES_THRESHOLD,
            use_angle_condition=False,
            level=cfgs.LEVEL,
            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,
            boxes_angle_threshold=cfgs.FAST_RCNN_BOXES_ANGLES_THRESHOLD,
            use_dropout=cfgs.USE_DROPOUT,
            weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME],
            is_training=False)

        fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category = \
            fast_rcnn.fast_rcnn_predict()

        ##############################################################################################
        if cfgs.NEED_AUXILIARY:
            predict_boxes = tf.concat(
                [fast_rcnn_decode_boxes, rpn_top_k_boxes], axis=0)
            predict_scores = tf.concat(
                [fast_rcnn_score, rpn_top_k_scores - 0.2], axis=0)
            rpn_top_k_label = tf.ones([
                tf.shape(rpn_top_k_scores)[0],
            ], tf.int64)
            labels = tf.concat([detection_category, rpn_top_k_label], axis=0)

            # valid_indices = nms_rotate.nms_rotate(decode_boxes=predict_boxes,
            #                                       scores=predict_scores,
            #                                       iou_threshold=0.15,
            #                                       max_output_size=30,
            #                                       use_angle_condition=False,
            #                                       angle_threshold=15,
            #                                       use_gpu=True)
            valid_indices = tf.py_func(nms_rotate.nms_rotate_cpu,
                                       inp=[
                                           predict_boxes, predict_scores,
                                           tf.constant(0.15, tf.float32),
                                           tf.constant(30, tf.float32)
                                       ],
                                       Tout=tf.int64)

            fast_rcnn_decode_boxes = tf.gather(predict_boxes, valid_indices)
            fast_rcnn_score = tf.gather(predict_scores, valid_indices)
            detection_category = tf.gather(labels, valid_indices)

            ##############################################################################################

        # 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, _fast_rcnn_decode_boxes, \
                _fast_rcnn_score, _detection_category \
                    = sess.run([img_name_batch, img_batch, gtboxes_and_label, fast_rcnn_decode_boxes,
                                fast_rcnn_score, detection_category])
                end = time.time()

                _img_batch = np.squeeze(_img_batch, axis=0)

                _img_batch_fpn = help_utils.draw_box_cv(
                    _img_batch,
                    boxes=_fast_rcnn_decode_boxes,
                    labels=_detection_category,
                    scores=_fast_rcnn_score)
                mkdir(cfgs.TEST_SAVE_PATH)
                cv2.imwrite(
                    cfgs.TEST_SAVE_PATH + '/{}_fpn.jpg'.format(
                        str(_img_name_batch[0]).split('.tif')[0]),
                    _img_batch_fpn)

                # _gtboxes_and_label_batch = np.squeeze(_gtboxes_and_label_batch, axis=0)

                temp_label = np.reshape(_gtboxes_and_label[:, -1:], [
                    -1,
                ]).astype(np.int64)
                _img_batch_gt = help_utils.draw_box_cv(
                    _img_batch,
                    boxes=_gtboxes_and_label[:, :-1],
                    labels=temp_label,
                    scores=None)

                cv2.imwrite(
                    cfgs.TEST_SAVE_PATH + '/{}_gt.jpg'.format(
                        str(_img_name_batch[0]).split('.tif')[0]),
                    _img_batch_gt)

                #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)
Exemplo n.º 4
0
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,   # 'ship', 'spacenet', 'pascal', 'coco'
                           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])

        with tf.name_scope('draw_gtboxes'):
            gtboxes_in_img = draw_box_with_color(
                img_batch,
                tf.reshape(gtboxes_and_label, [-1, 6])[:, :-1],
                text=tf.shape(gtboxes_and_label_batch)[1])

        # ***********************************************************************************************
        # *                                         shared CNN                                          *
        # ***********************************************************************************************
        _, 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                                              *
        # * Note: here the rpn is Feature Pyramid Networks                                              *
        # ***********************************************************************************************
        rpn = build_rpn.RPN(
            net_name=cfgs.NET_NAME,
            inputs=img_batch,
            gtboxes_and_label=gtboxes_and_label,
            is_training=True,
            share_head=cfgs.SHARED_HEADS,
            share_net=share_net,
            anchor_ratios=cfgs.ANCHOR_RATIOS,
            anchor_scales=cfgs.ANCHOR_SCALES,
            anchor_angles=cfgs.ANCHOR_ANGLES,
            scale_factors=cfgs.
            SCALE_FACTORS,  # this parameter will affect the performance
            base_anchor_size_list=cfgs.
            BASE_ANCHOR_SIZE_LIST,  # P2, P3, P4, P5, P6
            level=cfgs.LEVEL,
            anchor_stride=cfgs.ANCHOR_STRIDE,
            top_k_nms=cfgs.RPN_TOP_K_NMS,
            kernel_size=cfgs.KERNEL_SIZE,
            use_angles_condition=False,
            anchor_angle_threshold=cfgs.RPN_ANCHOR_ANGLES_THRESHOLD,
            nms_angle_threshold=cfgs.RPN_NMS_ANGLES_THRESHOLD,
            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=cfgs.
            IS_FILTER_OUTSIDE_BOXES,  # whether remove anchors outside
            rpn_weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME],
            scope='')

        rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals(
        )  # rpn_score shape: [300, ]

        rpn_location_loss, rpn_classification_loss, rpn_predict_boxes, rpn_predict_scores = rpn.rpn_losses(
        )
        rpn_total_loss = rpn_classification_loss + rpn_location_loss

        with tf.name_scope('draw_proposals'):
            # score > 0.6 is object
            rpn_object_boxes_indices = tf.reshape(
                tf.where(
                    tf.greater(rpn_proposals_scores,
                               cfgs.FINAL_SCORE_THRESHOLD)), [-1])
            rpn_object_boxes = tf.gather(rpn_proposals_boxes,
                                         rpn_object_boxes_indices)
            rpn_object_soxres = tf.gather(rpn_proposals_scores,
                                          rpn_object_boxes_indices)

            rpn_proposals_objcet_boxes_in_img = draw_boxes_with_scores(
                img_batch, rpn_object_boxes, scores=rpn_object_soxres)

            # 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_rcnn.FastRCNN(
            img_batch=img_batch,
            feature_pyramid=rpn.feature_pyramid,
            rpn_proposals_boxes=rpn_proposals_boxes,
            rpn_proposals_scores=rpn_proposals_scores,
            stop_gradient_for_proposals=False,
            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,
            fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD,
            top_k_nms=cfgs.FAST_RCNN_TOP_K_NMS,
            nms_angle_threshold=cfgs.FAST_RCNN_NMS_ANGLES_THRESHOLD,
            use_angle_condition=False,
            level=cfgs.LEVEL,
            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
            boxes_angle_threshold=cfgs.FAST_RCNN_BOXES_ANGLES_THRESHOLD,
            use_dropout=cfgs.USE_DROPOUT,
            weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME],
            is_training=True)

        fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category = \
            fast_rcnn.fast_rcnn_predict()
        fast_rcnn_location_loss, fast_rcnn_classification_loss = fast_rcnn.fast_rcnn_loss(
        )
        fast_rcnn_total_loss = fast_rcnn_location_loss + fast_rcnn_classification_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)

        # 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(70000), np.int64(120000)],
            values=[cfgs.LR, cfgs.LR / 10, cfgs.LR / 100])

        # optimizer = tf.train.MomentumOptimizer(lr, momentum=cfgs.MOMENTUM)
        optimizer = tf.train.MomentumOptimizer(lr, momentum=cfgs.MOMENTUM)

        if cfgs.RPN_TRAIN:
            train_op = slim.learning.create_train_op(rpn_total_loss, optimizer,
                                                     global_step)
        else:
            train_op = slim.learning.create_train_op(total_loss, optimizer,
                                                     global_step)
        # 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/fast_rcnn_predict',
                         fast_rcnn_predict_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_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)

        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)

            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_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_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_total_loss:{} |
                          total_loss:{} |\t per_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_total_loss, _total_loss,
                                  (end - start)))

                if (step % 50 == 0) and (step % 10000 != 0):  # 50
                    summary_str = sess.run(summary_op)
                    summary_writer.add_summary(summary_str, _global_step)
                    summary_writer.flush()

                if (step > 0
                        and step % 10000 == 0) or (step
                                                   == cfgs.MAX_ITERATION - 1):

                    summary_str = sess.run(summary_op)
                    summary_writer.add_summary(summary_str, _global_step)
                    summary_writer.flush()

                    save_dir = os.path.join(cfgs.TRAINED_CKPT, 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)
Exemplo n.º 5
0
def eval_with_plac(num_imgs, eval_dir, img_root, showbox, annotation_dir):

    # with open('/home/yjr/DataSet/VOC/VOC_test/VOC2007/ImageSets/Main/aeroplane_test.txt') as f:
    #     all_lines = f.readlines()
    # test_imgname_list = [a_line.split()[0].strip() for a_line in all_lines]

    test_imgname_list = [
        item for item in os.listdir(eval_dir)
        if item.endswith(('.jpg', 'jpeg', '.png', '.tif', '.tiff'))
    ]
    if num_imgs == np.inf:
        real_test_imgname_list = test_imgname_list
    else:
        real_test_imgname_list = test_imgname_list[:num_imgs]

    img_plac = tf.placeholder(dtype=tf.float32, shape=[None, None, 3])
    img = img_plac - tf.constant([103.939, 116.779, 123.68])
    img_batch = short_side_resize_for_inference_data(img, cfgs.SHORT_SIDE_LEN)
    h, w = img.shape[0], img.shape[1]
    gt_boxes_label = tf.placeholder(dtype=tf.float32, shape=[None, 5])
    gt_boxes_label_batch = tf.expand_dims(gt_boxes_label, axis=0)

    image_height, image_width = tf.shape(img_batch)[1], tf.shape(img_batch)[2]

    _, share_net = get_network_byname(net_name=cfgs.NET_NAME,
                                      inputs=img_batch,
                                      num_classes=None,
                                      is_training=False,
                                      output_stride=None,
                                      global_pool=False,
                                      spatial_squeeze=False)

    feature_pyramid = build_fpn.build_feature_pyramid(share_net)
    rpn = build_rpn.RPN(feature_pyramid=feature_pyramid,
                        image_height=image_height,
                        image_width=image_width,
                        gtboxes_and_label=gt_boxes_label_batch,
                        is_training=False)

    rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals(
        is_training=False)

    fast_rcnn = build_fast_rcnn.FAST_RCNN(
        feature_pyramid=feature_pyramid,
        rpn_proposals_boxes=rpn_proposals_boxes,
        gtboxes_and_label=gt_boxes_label_batch,
        origin_image=img_batch,
        is_training=True,
        image_height=image_height,
        image_width=image_width)

    detections = fast_rcnn.head_detection()
    detection_boxes, detection_category,detection_scores = tf.squeeze(detections[:,:,:4],axis=0),\
                                                            tf.squeeze(detections[:,:,4],axis=0),\
                                                            tf.squeeze(detections[:,:,5],axis=0)

    indices = tf.reshape(
        tf.where(tf.greater_equal(detection_scores,
                                  cfgs.FINAL_SCORE_THRESHOLD)), [-1])
    detection_boxes = tf.gather(detection_boxes, indices)
    detection_scores = tf.gather(detection_scores, indices)
    detection_category = tf.gather(detection_category, indices)

    init_op = tf.group(tf.global_variables_initializer(),
                       tf.local_variables_initializer())

    restorer, restore_ckpt = restore_model.get_restorer(
        test=True, checkpoint_path=cfgs.chekpoint_path)

    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')

        all_boxes = []
        for i, a_img_name in enumerate(real_test_imgname_list):
            raw_img = cv2.imread(os.path.join(img_root, a_img_name))
            raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]

            start = time.time()
            resized_img, detected_boxes, detected_scores, detected_categories = \
                sess.run(
                    [img, detection_boxes, detection_scores, detection_category],
                    feed_dict={img_plac: raw_img}  # cv is BGR. But need RGB
                )
            print(a_img_name, detected_boxes, detected_scores,
                  detected_categories)
            end = time.time()
            ymin, xmin, ymax, xmax = detected_boxes[:, 0], detected_boxes[:, 1], \
                                     detected_boxes[:, 2], 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

            boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
            dets = np.hstack((detected_categories.reshape(-1, 1),
                              detected_scores.reshape(-1, 1), boxes))
            all_boxes.append(dets)

            view_bar('{} image cost {}s'.format(a_img_name, (end - start)),
                     i + 1, len(real_test_imgname_list))

    voc_evaluate_detections(all_boxes=all_boxes,
                            test_annotation_path=annotation_dir,
                            test_imgid_list=real_test_imgname_list)
Exemplo n.º 6
0
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)

        # ***********************************************************************************************
        # *                                         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(
            feature_pyramid=rpn.feature_pyramid,
            rpn_proposals_boxes=rpn_proposals_boxes,
            rpn_proposals_scores=rpn_proposals_scores,
            img_shape=tf.shape(img_batch),
            img_batch=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,
            head_quadrant=None)

        fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category, \
        fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, fast_rcnn_head_quadrant, \
        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()
                # gdal.AllRegister()
                # ds = gdal.Open(img_path, gdalconst.GA_ReadOnly)
                # if ds is None:
                #     print("Image %s open failed!" % img_path)
                #     sys.exit()
                img = cv2.imread(img_path)

                box_res = []
                label_res = []
                score_res = []
                box_res_rotate = []
                label_res_rotate = []
                score_res_rotate = []
                head_rotate = []
                # imgH = ds.RasterYSize
                # imgW = ds.RasterXSize
                imgH = img.shape[0]
                imgW = img.shape[1]
                for hh in range(0, imgH, h_len):
                    h_size = min(h_len, imgH - hh - h_overlap)
                    if h_size < 10:
                        break
                    for ww in range(0, imgW, w_len):
                        w_size = min(w_len, imgW - ww - w_overlap)
                        if w_size < 10:
                            break

                        # src_img = ds.ReadAsArray(ww, hh, w_size, h_size)
                        src_img = img[hh:(hh + h_size), ww:(ww + w_size), :]
                        # if len(src_img.shape) == 2:
                        #     src_img = cv2.cvtColor(src_img, cv2.COLOR_GRAY2RGB)
                        # else:
                        #     src_img = chw2hwc(src_img)

                        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, _fast_rcnn_head_quadrant = \
                            sess.run([fast_rcnn_decode_boxes_rotate, detection_category_rotate,
                                      fast_rcnn_score_rotate,
                                      fast_rcnn_head_quadrant],
                                     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])
                        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])
                                head_rotate.append(
                                    _fast_rcnn_head_quadrant[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]
                head_rotate = np.array(head_rotate)[inx]

                # ds = None
                time_elapsed = timer() - start
                print("{} detection time : {:.4f} sec".format(
                    img_path.split('/')[-1].split('.')[0], time_elapsed))

                # if target_name == 'aircraft':
                # img = cv2.imread(img_path)
                # if len(img.shape) == 2:
                #     img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
                # elif len(img.shape) == 3:
                #     img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
                #     img[:, :, 0] = img[:, :, 1] = img[:, :, 2] = img_gray
                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),
                    head=np.argmax(head_rotate, axis=1))
                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)
                # clip_obj_imgs(src_img, box_res, label_res, score_res, des_folder)
                # print(img_path)
                # det_xml_path =img_path.replace(".tif", ".det.xml")
                # obj_to_det_xml(img_path, box_res, label_res, score_res, det_xml_path)

            coord.request_stop()
            coord.join(threads)
Exemplo n.º 7
0
def eval_ship(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_rcnn.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.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)

            gtboxes_horizontal_dict = {}
            predict_horizontal_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 \
                    = sess.run([img_name_batch, img_batch, gtboxes_and_label, gtboxes_and_label_minAreaRectangle,
                                fast_rcnn_decode_boxes, fast_rcnn_score, detection_category])
                end = time.time()

                # gtboxes convert dict
                gtboxes_horizontal_dict[str(_img_name_batch[0])] = []
                predict_horizontal_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)

                gtboxes_horizontal_dict[str(
                    _img_name_batch[0])].extend(gtbox_horizontal_list)
                predict_horizontal_dict[str(
                    _img_name_batch[0])].extend(predict_horizontal_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')
            pickle.dump(gtboxes_horizontal_dict, fw1)
            pickle.dump(predict_horizontal_dict, fw2)
            fw1.close()
            fw2.close()
            coord.request_stop()
            coord.join(threads)
Exemplo n.º 8
0
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=False,
            share_net=share_net,
            anchor_ratios=cfgs.ANCHOR_RATIOS,
            anchor_scales=cfgs.ANCHOR_SCALES,
            anchor_angles=cfgs.ANCHOR_ANGLES,
            scale_factors=cfgs.SCALE_FACTORS,
            base_anchor_size_list=cfgs.
            BASE_ANCHOR_SIZE_LIST,  # P2, P3, P4, P5, P6
            level=cfgs.LEVEL,
            anchor_stride=cfgs.ANCHOR_STRIDE,
            top_k_nms=cfgs.RPN_TOP_K_NMS,
            kernel_size=cfgs.KERNEL_SIZE,
            use_angles_condition=True,
            anchor_angle_threshold=cfgs.RPN_ANCHOR_ANGLES_THRESHOLD,
            nms_angle_threshold=cfgs.RPN_NMS_ANGLES_THRESHOLD,
            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],
            scope='')

        # 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,
            stop_gradient_for_proposals=False,
            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,
            fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD,
            top_k_nms=cfgs.FAST_RCNN_TOP_K_NMS,
            nms_angle_threshold=cfgs.FAST_RCNN_NMS_ANGLES_THRESHOLD,
            use_angle_condition=False,
            level=cfgs.LEVEL,
            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,
            boxes_angle_threshold=cfgs.FAST_RCNN_BOXES_ANGLES_THRESHOLD,
            use_dropout=cfgs.USE_DROPOUT,
            weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME],
            is_training=False)

        fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category = \
            fast_rcnn.fast_rcnn_predict()

        # test
        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_name, a_img in get_imgs():
                start = time.time()

                _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category \
                    = sess.run([fast_rcnn_decode_boxes, fast_rcnn_score, detection_category],
                               feed_dict={img_plac: a_img})
                end = time.time()
                print("{} cost time : {} ".format(img_name, end - start))
                final_detections = help_utils.draw_box_cv(
                    np.array(a_img, dtype=np.float32) -
                    np.array([103.939, 116.779, 123.68]),
                    boxes=_fast_rcnn_decode_boxes,
                    labels=_detection_category,
                    scores=_fast_rcnn_score)

                cv2.imwrite(cfgs.INFERENCE_SAVE_PATH + '/' + img_name,
                            final_detections)

            coord.request_stop()
            coord.join(threads)
Exemplo n.º 9
0
def eval_dict_convert(args):
    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)

        # ***********************************************************************************************
        # *                                         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=True,
            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=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,
            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()

        # train
        init_op = tf.group(tf.global_variables_initializer(),
                           tf.local_variables_initializer())

        restorer, restore_ckpt = restore_model.get_restorer(
            checkpoint_path=args.weights)

        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)

            gtbox_dict = {}
            predict_dict = {}

            for i in range(args.img_num):

                start = time.time()

                _img_name_batch, _img_batch, _gtboxes_and_label_batch, _fast_rcnn_decode_boxes, \
                    _fast_rcnn_score, _detection_category \
                    = sess.run([img_name_batch, img_batch, gtboxes_and_label_batch, fast_rcnn_decode_boxes,
                                fast_rcnn_score, detection_category])
                end = time.time()

                # gtboxes convert dict
                gtbox_dict[str(_img_name_batch[0])] = []
                predict_dict[str(_img_name_batch[0])] = []

                for j, box in enumerate(_gtboxes_and_label_batch[0]):
                    bbox_dict = {}
                    bbox_dict['bbox'] = np.array(
                        _gtboxes_and_label_batch[0][j, :-1], np.float64)
                    bbox_dict['name'] = LABEl_NAME_MAP[int(
                        _gtboxes_and_label_batch[0][j, -1])]
                    gtbox_dict[str(_img_name_batch[0])].append(bbox_dict)

                for label in NAME_LABEL_MAP.keys():
                    if label == 'back_ground':
                        continue
                    else:
                        temp_dict = {}
                        temp_dict['name'] = label

                        ind = np.where(
                            _detection_category == NAME_LABEL_MAP[label])[0]
                        temp_boxes = _fast_rcnn_decode_boxes[ind]
                        temp_score = np.reshape(_fast_rcnn_score[ind], [-1, 1])
                        temp_dict['bbox'] = np.array(
                            np.concatenate([temp_boxes, temp_score], axis=1),
                            np.float64)
                        predict_dict[str(_img_name_batch[0])].append(temp_dict)

                view_bar(
                    '{} image cost {}s'.format(str(_img_name_batch[0]),
                                               (end - start)), i + 1,
                    args.img_num)

            fw1 = open('gtboxes_dict.pkl', 'w')
            fw2 = open('predict_dict.pkl', 'w')
            pickle.dump(gtbox_dict, fw1)
            pickle.dump(predict_dict, fw2)
            fw1.close()
            fw2.close()
            coord.request_stop()
            coord.join(threads)