classify_caffe.py

from __future__ import division, print_function

from pprint import pprint
import glob
import caffe

MODEL_FILE = '../models/lenet.prototxt'
PRETRAINED_FILE = '../models/lenet_pretrained.caffemodel'
IMAGE_DIR = 'test'

caffe.set_mode_cpu()
classifier = caffe.Classifier(MODEL_FILE, PRETRAINED_FILE,
        image_dims=(28, 28), # this can be changed to enable autocropping
        input_scale=255.0 # images are loaded as float arrays with values in 0.0..1.0
)

print('\n' + '*' * 16 + ' end of Caffe log ' + '*' * 16 + '\n')

input_paths = glob.glob('%s/*.%s' % (IMAGE_DIR, 'png'))
input_paths.sort()
pprint(input_paths)

# Caffe loads images with 3 color channels by default
# actually, under the hood it calls a routine from skimage,
# so this can be safely replaced with another implementation
# ref: https://github.com/BVLC/caffe/blob/master/python/caffe/io.py#L274
inputs = [caffe.io.load_image(im_f, color=False) for im_f in input_paths]

# oversampling means generating many crops and mirrored crops
# and running the net on them, which is useful for real-life photos
scores = classifier.predict(inputs, oversample=False)
print(scores)