caffe框架学习一——Image Classification and Filter Visualization

英文原文地址： https://nbviewer.jupyter.org/github/BVLC/caffe/blob/master/examples/00-classification.ipynb
代码及部分注释如下：
import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = (10, 10)  # 设置画图尺寸
plt.rcParams['image.interpolation'] = 'nearest'  # 设置插值类型
plt.rcParams['image.cmap'] = 'gray'  # 设置颜色类型

import sys
caffe_root = '../'
sys.path.insert(0, caffe_root+'python')  # 插在sys.path的最前面
import caffe

import os
if(os.path.isfile(caffe_root+'models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel')):
    print('CaffeNet found')
else:
    print('Downloading pre-trained CaffeNet model')
    os.system(
        '../scripts/download_model_binary.py ../models/bvlc_reference_caffenet')

caffe.set_mode_gpu()  # 设置为gpu模式
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,  # 定义模型结构
                model_weights,  # 包含训练权重
                caffe.TEST)  # 使用测试模式

mu = np.load(caffe_root+'python/caffe/imagenet/ilsvrc_2012_mean.npy')
mu = mu.mean(1).mean(1)  # 对像素求均值
print('mean-subtracted values:', zip('BGR', mu))

# 设定图片的shape格式(1,3,28,28)
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
# 将图片channel移动到最外侧维度上，即由原始图片(28,28,3)变为(3,28,28)
transformer.set_transpose('data', (2, 0, 1))
transformer.set_mean('data', mu)  # 减去每个channel的均值
transformer.set_raw_scale('data', 255)  # 将变量幅值从[0,1]变为[0,255]
transformer.set_channel_swap('data', (2, 1, 0))  # 将RGB变为BGR

net.blobs['data'].reshape(50,  # batch尺寸
                          3,  # BGR图片的3 channel
                          227, 227)  # 图片尺寸

image = caffe.io.load_image(caffe_root+'examples/images/cat.jpg')  # 加载图片
transformed_image = transformer.preprocess('data', image)  # 预处理
plt.imshow(image)  # 展示

net.blobs['data'].data[...] = transformed_image  # 将待分类图片放入模型
output = net.forward()  # 网络通过前向传播进行预测

output_prob = output['prob'][0]  # 获取最大的可能概率

print('predicted class is:', output_prob.argmax())  # 获取最大的可能类型

labels_file = caffe_root+'data/ilsvrc12/synset_words.txt'  # 获取标签文本
if not os.path.exists(labels_file):
    os.system('../data/ilsvrc12/get_ilsvrc_aux.sh')

labels = np.loadtxt(labels_file, str, delimiter='\t')  # 加载标签文本
print('output label:', labels[output_prob.argmax()])

top_inds = output_prob.argsort()[::-1][:5]  # 最有可能的前五种类型
print('probabilities and labels:')
print(zip(output_prob[top_inds], labels[top_inds]))
caffe框架学习一——Image Classification and Filter Visualization

猜你喜欢
