Caffe学习:pycaffe利用caffemodel进行分类=>批处理

  • 在 Caffe学习:pycaffe利用caffemodel进行分类(2)中,我们手动将图片裁剪并导入pycaffe中,以提高识别率
  • 这里,我们采用批处理,以加快识别速度

  • 导入相关库
import caffe
  • 配置
# caffemodel文件
MODEL_FILE = 'model/_iter_10000.caffemodel'
# deploy文件,参考/caffe/models/bvlc_alexnet/deploy.prototxt
DEPLOY_FILE = 'deploy.prototxt'
# 测试图片存放文件夹
TEST_ROOT = 'datas/'
  • GPU模式测试
caffe.set_mode_gpu()
net = caffe.Net(DEPLOY_FILE, MODEL_FILE, caffe.TEST)
  • 数据输入处理配置
# 批处理量
batch_size = 10
# 图片裁剪大小
crop_dims = (32, 96)

# 'data'对应于deploy文件:
# input: "data"
# input_dim: 1
# input_dim: 3
# input_dim: 32
# input_dim: 96
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
# python读取的图片文件格式为H×W×K,需转化为K×H×W
transformer.set_transpose('data', (2, 0, 1))
# python中将图片存储为[0, 1],而caffe中将图片存储为[0, 255],
# 所以需要一个转换
transformer.set_raw_scale('data', 255)
# caffe中图片是BGR格式,而原始格式是RGB,所以要转化
transformer.set_channel_swap('data', (2, 1, 0))
# 将输入图片格式转化为合适格式(与deploy文件相同)
net.blobs['data'].reshape(batch_size, 3, crop_dims[0], crop_dims[1])
  • 读取图片
# 详见/caffe/python/caffe/io.py
img = caffe.io.load_image('temp.jpg')
# 读取的图片文件格式为H×W×K,需注意
  • 数据处理,输入
# 图片维度(高、宽)
img_shape = np.array(img.shape)
# 裁剪的大小(高、宽)
crop_dims = np.array(crop_dims)

# 这里使用的图片高度全部固定为32,长度可变,最小为96
# 裁剪起点为0,终点为w_range
w_range = img_shape[1] - crop_dims[1]

# 从左往右剪一遍,再从右往左剪一遍,步长为96/4=24
k_range = range(0, w_range + 1, crop_dims[1] / 4) + range(w_range, 1, -crop_dims[1] / 4)
# 一张图片需要识别次数为_batch_size
_batch_size = max(len(k_range), 1)
net.blobs['data'].reshape(_batch_size, 3, crop_dims[0], crop_dims[1])

# 数据处理,输入
for ik, k in enumerate(k_range):
    # 裁剪图片
    crop_img = img[:, k:k + crop_dims[1], :]
    # 数据输入
    net.blobs['data'].data[ik] = transformer.preprocess('data', crop_img)
  • 进行分类
# 前向迭代,即分类
out = net.forward()
# 每一次分类的概率分布叠加
pridects = np.sum(out['prob'], axis=0)
# 取最大的概率分布为最终结果
pridect = pridects.argmax()

# 上述'prob'来源于deploy文件:
# layer {
#   name: "prob"
#   type: "Softmax"
#   bottom: "ip2"
#   top: "prob"
# }

你可能感兴趣的:(python,caffe)