import os

import sys

import numpy as np

import matplotlib.pyplot as plt

caffe_root = '/home/ubuntu/code/ssd/'

sys.path.insert(0, caffe_root + 'python')

import caffe

MODEL_FILE = '/home/ubuntu/code/ssd/examples/mnist/lenet.prototxt'

PRETRAINED = '/home/ubuntu/code/ssd/examples/mnist/lenet_iter_10000.caffemodel'

IMAGE_FILE = '/home/ubuntu/code/ssd/examples/images/test4.bmp'

input_image = caffe.io.load_image(IMAGE_FILE, color=False)

net = caffe.Classifier(MODEL_FILE, PRETRAINED)

prediction = net.predict([input_image], oversample = False)

caffe.set_mode_cpu()

print 'predicted class:', prediction[0].argmax()

返回