ExtractImageFeat_caffe.py 文件源码

def get_caffe_model(caffe_dir, caffe_model, gpu=True,
                    image_dims=(256, 256),
                    mean_file='default',
                    raw_scale=255.0,
                    channel_swap=(2,1,0),
                    input_scale=None):
    if mean_file == 'default':
        mean_file = os.path.join(caffe_dir, 'python', 'caffe', 'imagenet', 'ilsvrc_2012_mean.npy')
    model_path = os.path.join(caffe_dir, 'models', caffe_model, '%s.caffemodel'%caffe_model)
    model_def = os.path.join(caffe_dir, 'models', caffe_model, 'deploy.prototxt')

    print('Loading mean file %s' % mean_file)
    mean = np.load(mean_file).mean(1).mean(1)
    if gpu:
        caffe.set_mode_gpu()
    else:
        caffe.set_mode_cpu()
    net = caffe.Classifier(model_def, model_path,
                           image_dims=image_dims, mean=mean,
                           input_scale=input_scale, raw_scale=raw_scale,
                           channel_swap=channel_swap)
    return net