classifier.py 文件源码

def __init__(self, deploy, pretrained, mean, labels, gpu = False):
        if gpu:
            caffe.set_mode_gpu()
        else:
            caffe.set_mode_cpu()    # in windows, only CPU mode supported

        self.__labels = self.load_labels(labels);
        mean_ar = self.convert(mean)

        if True:
            self.__net = caffe.Classifier(deploy, pretrained,
                    mean = mean_ar.mean(1).mean(1),
                    channel_swap = (2, 1, 0), 
                    raw_scale = 255,
                    image_dims = (256, 256))
        else: 
            self.__net = caffe.Net(deploy, pretrained, caffe.TEST)
            print self.__net.blobs['data'].data.shape    

            self.__transformer = caffe.io.Transformer({'data': self.__net.blobs['data'].data.shape})
            self.__transformer.set_transpose('data', (2,0,1)) # height*width*channel -> channel*height*width
            self.__transformer.set_mean('data', mean_ar)
            self.__transformer.set_raw_scale('data', 255)
            self.__transformer.set_channel_swap('data', (2,1,0)) # RGB -> BGR