@hanxiaoyang
2016-10-10T13:08:45.000000Z
字数 1005
阅读 2739
未分类
example code by @寒小阳
#example code by @Hanxiaoyang
import numpy as np
import sys
import matplotlib.pyplot as plt
# 这一步一定要保证,从caffe路径import
caffe_root = '../' # 你的caffe编译文件夹路径
sys.path.insert(0, caffe_root + 'python')
import caffe
# 我们如果已经训练好模型了,只是想做个图像识别的demo
# 模型网络结构
MODEL_FILE = '../models/bvlc_reference_caffenet/deploy.prototxt'
# 预训练好的模型
PRETRAINED = '../models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel'
# 要识别的图像文件
IMAGE_FILE = 'images/dog.jpg'
# 选择使用cpu还是gpu
caffe.set_mode_cpu()
#使用gpu
#caffe.set_device(0)
#caffe.set_mode_gpu()
# 用网络结构定义文件,预训练模型,训练集均值文件初始化分类器
net = caffe.Classifier(MODEL_FILE, PRETRAINED,
mean=np.load(caffe_root + 'python/caffe/imagenet/ilsvrc_2012_mean.npy').mean(1).mean(1),
channel_swap=(2,1,0),
raw_scale=255,
image_dims=(256, 256))
# 载入待预测图片
input_image = caffe.io.load_image(IMAGE_FILE)
# 展示一下图片
plt.imshow(input_image)
# predict阶段,其实就是做前向运算,softmax输出概率
prediction = net.predict([input_image])
print 'prediction shape:', prediction[0].shape
plt.plot(prediction[0])
print 'predicted class:', prediction[0].argmax()
plt.show()