feature.py 文件源码

def init_detection_net(self, gpu_id=0, prototxt=None, caffemodel=None):
        """init extraction network"""
        cfg.TEST.HAS_RPN = True  # Use RPN for proposals
        if prototxt is None:
            prototxt = os.path.join(cfg.ROOT_DIR, 'models', NETS['zf'][0],
                            'faster_rcnn_alt_opt', 'faster_rcnn_test.pt')
        if caffemodel is None:
            caffemodel = os.path.join(cfg.ROOT_DIR, 'output/default/train',
                              NETS['zf'][1])

        if not os.path.isfile(caffemodel):
            raise IOError(('{:s} not found.\nDid you run ./data/script/'
                       'fetch_faster_rcnn_models.sh?').format(caffemodel))

        #np.random.seed(cfg.RNG_SEED)
        caffe.set_random_seed(cfg.RNG_SEED)       
        caffe.set_mode_gpu()
        caffe.set_device(gpu_id)        
        self.net_d = caffe.Net(prototxt, caffemodel, caffe.TEST)