Caffe用python接口分类的例子

如何用训练好的caffemodel测试图片分类呢？caffe官方提供python的分类例子在examples/00-classification.ipynb下，但是只介绍了caffe.Net分类，这篇博客记录了另一中方式caffe.Classifier，并介绍了两者的区别。

1. caffe.Net

步骤大概分为：模型导入->数据处理->分类和输出

模型导入

需要deploy（网络定义）和caffemodel（权值）文件。

model_def = caffe_root + 'models/bvlc_reference_caffenet/deploy.prototxt'
model_weights = caffe_root + 'models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel'

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)

数据处理

首先确定均值，可以导入均值文件mu=np.load()，然后经过两次mean()后变为了一个(3,1,1)的ndarray，代表BGR三个通道的均值，也可以自己指定均值，如创建一个ndarray：mu=np.float32([104.0, 117.0, 123.0])。

接下来是的transformer用于使输入图像与网络input层能匹配，比如set_transpose()使输入的(h, w, c)变成(c, h, w)。

# load the mean ImageNet image (as distributed with Caffe) for subtraction
mu = np.load(caffe_root + 'python/caffe/imagenet/ilsvrc_2012_mean.npy')
mu = mu.mean(1).mean(1)  # average over pixels to obtain the mean (BGR) pixel values

# create transformer for the input called 'data'
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})

transformer.set_transpose('data', (2,0,1))  # move image channels to outermost dimension
transformer.set_mean('data', mu)            # subtract the dataset-mean value in each channel
transformer.set_raw_scale('data', 255)      # rescale from [0, 1] to [0, 255]
transformer.set_channel_swap('data', (2,1,0))  # swap channels from RGB to BGR

分类和输出

这一步用caffe的io接口导入图片，然后通过之前定义的transformer对数据执行处理，将处理好的输入数据给到'data'层：net.blobs['data'].data[..]。经过一次net.forward()前向传播后，在最后输出层'prob'就会有分类结果，是一个(n, 1)的ndarray，n是类别数目。

image = caffe.io.load_image(caffe_root + 'examples/images/cat.jpg')
transformed_image = transformer.preprocess('data', image)

# copy the image data into the memory allocated for the net
net.blobs['data'].data[...] = transformed_image

### perform classification
output = net.forward()

output_prob = output['prob'][0]  # the output probability vector for the first image in the batch

print 'predicted class is:', output_prob.argmax()

2. caffe.Classifier

pycaffe同样提供另一种用Classifier来分类的方式，不过没有例子，博主这里自己实现了一个。

import caffe, cv2
# 初始化网络
caffe.set_mode_gpu()
caffe.set_device(0)
deploy_file = '.../deploy.prototxt'
caffemodel = '.../xxxnet.caffemodel'
mean_value = np.float32([104., 117., 123.])
net = caffe.Classifier(deploy_file, caffemodel, mean=mean_value)
# 数据输入和分类
input_h, input_w = (224, 224)
input_data = cv2.resize(image, (input_w, input_h))
input_data = input_data.transpose((2, 0, 1))
input_data = input_data - net.transformer.mean['data']
net.blobs['data'].data[0] = input_data
net.forward()
outputs = net.blobs['prob'].data
print 'predicted class is: ', outputs.argmax()

上述代码中，net的均值直接作为参数传到Classifier中，并使用cv2导入图片。

分类的时候其实可以用net.predict()，这个函数对图像做了上下左右中心5个crop以及它们的镜像，一共10个输入，最后对这10个输入的分类结果平均后输出。

net.Classifier()相当于对之前分类例子的一个封装，它的init初始化函数中对transformer进行了初始化，即数据预处理在内部完成。