Cascade R-CNN Python测试脚本

传送门：

1. 《Cascade R-CNN: Delving into High Quality Object Detection》论文笔记
2. Cascade R-CNN训练自己的数据

1. 前言

之前的博客中讲到了Cascade R-CNN的原理与训练数据的准备，这里贴出其Python测试的脚本。
下面这份代码是从该地址：代码地址，做了一些修改。

2. demo代码

# -*- coding=utf-8 -*-
import os
import sys
import argparse
import numpy as np
from PIL import Image, ImageDraw
import cv2
import time
caffe_root = '/home/xxxxx/codes/cascade-rcnn'
sys.path.insert(0, os.path.join(caffe_root, 'python'))
import caffe

# from google.protobuf import text_format
# from caffe.proto import caffe_pb2

class CaffeDetection:
    def __init__(self, gpu_id, model_def, model_weights, cascade=0, FPN=0):
        if gpu_id < 0:
            caffe.set_mode_cpu()
        else:
            caffe.set_device(gpu_id)
            caffe.set_mode_gpu()

        # Load the net in the test phase for inference, and configure input preprocessing.
        self.net = caffe.Net(model_def,  # defines the structure of the model
                             model_weights,  # contains the trained weights
                             caffe.TEST)  # use test mode (e.g., don't perform dropout)

        self.cascade = cascade > 0
        self.FPN = FPN > 0
        print(cascade, FPN)
        if not self.cascade:
            # baseline model
            if self.FPN:
                self.proposal_blob_names = ['proposals_to_all']
            else:
                self.proposal_blob_names = ['proposals']

            self.bbox_blob_names = ['output_bbox_1st']
            self.cls_prob_blob_names = ['cls_prob_1st']
            self.output_names = ['1st']
        else:
            # cascade-rcnn model
            if self.FPN:
                self.proposal_blob_names = ['proposals_to_all', 'proposals_to_all_2nd',
                                            'proposals_to_all_3rd', 'proposals_to_all_2nd', 'proposals_to_all_3rd']
            else:
                self.proposal_blob_names = ['proposals', 'proposals_2nd', 'proposals_3rd',
                                            'proposals_2nd', 'proposals_3rd']

            self.bbox_blob_names = ['output_bbox_1st', 'output_bbox_2nd', 'output_bbox_3rd',
                                    'output_bbox_2nd', 'output_bbox_3rd']
            self.cls_prob_blob_names = ['cls_prob_1st', 'cls_prob_2nd', 'cls_prob_3rd',
                                        'cls_prob_2nd_avg', 'cls_prob_3rd_avg']
            self.output_names = ['1st', '2nd', '3rd', '2nd_avg', '3rd_avg']

        self.num_outputs = len(self.proposal_blob_names)
        assert (self.num_outputs == len(self.bbox_blob_names))
        assert (self.num_outputs == len(self.cls_prob_blob_names))
        assert (self.num_outputs == len(self.output_names))
        # detection configuration
        # detect_final_boxes = np.zeros(nImg, num_outputs)
        # self.det_thr = 0.001 # threshold for testing
        self.det_thr = 0.001  # threshold for demo
        self.max_per_img = 100  # max number of detections
        self.nms_thresh = 0.5  # NMS
        if FPN:
            self.shortSize = 800
            self.longSize = 1312
        else:
            self.shortSize = 608
            self.