- 在 Caffe学习:pycaffe利用caffemodel进行分类(2)中,我们手动将图片裁剪并导入pycaffe中,以提高识别率
- 这里,我们采用批处理,以加快识别速度
import caffe
MODEL_FILE = 'model/_iter_10000.caffemodel'
DEPLOY_FILE = 'deploy.prototxt'
TEST_ROOT = 'datas/'
caffe.set_mode_gpu()
net = caffe.Net(DEPLOY_FILE, MODEL_FILE, caffe.TEST)
batch_size = 10
crop_dims = (32, 96)
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_transpose('data', (2, 0, 1))
transformer.set_raw_scale('data', 255)
transformer.set_channel_swap('data', (2, 1, 0))
net.blobs['data'].reshape(batch_size, 3, crop_dims[0], crop_dims[1])
img = caffe.io.load_image('temp.jpg')
img_shape = np.array(img.shape)
crop_dims = np.array(crop_dims)
w_range = img_shape[1] - crop_dims[1]
k_range = range(0, w_range + 1, crop_dims[1] / 4) + range(w_range, 1, -crop_dims[1] / 4)
_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()