caffe.py 文件源码

def __init__(self, prototxt, caffemodel, mean, use_gpu=False):
        self.net = caffe.Net(prototxt, caffemodel, caffe.TEST)
        if use_gpu:
            caffe.set_mode_gpu()

        if type(mean) == str:
            if mean.endswith('.binaryproto'):
                self.mean = CaffeAdapter._load_binaryproto(mean)
            elif mean.endswith('.npy'):
                self.mean = np.load(mean)
            else:
                raise ValueError('Unknown mean file format. Known formats: .binaryproto, .npy')
        elif type(mean) == np.ndarray:
            self.mean = mean
        elif mean is not None:
            raise ValueError('Unknown mean format. Expected .binaryproto/.npy file or numpy array.')

        self.transformer = caffe.io.Transformer({'data': self.net.blobs['data'].data.shape})
        self.transformer.set_transpose('data', (2, 0, 1))
        if self.mean is not None:
            self.transformer.set_mean('data', self.mean.mean(1).mean(1))
        else:
            print('Warning. No mean specified.')
        self.transformer.set_raw_scale('data', 255)  # the reference model operates on images in [0,255] range instead of [0,1]
        self.transformer.set_channel_swap('data', (2, 1, 0))  # the reference model has channels in BGR order instead of RGB

        self.layer_types = self._load_layer_types(prototxt)

        self.ready = False

        self.use_gpu = use_gpu