Chainer with Caffenet example

chainer_caffenet_prediction.py

#!/usr/bin/env python

"""
Chainer sample code to predict labels for a given image based on pre-trained CaffeNet model

Before running this code, you must execure
$ ./download_model.py caffenet
$ ./download_mean_file.py
$ wget http://dl.caffe.berkeleyvision.org/caffe_ilsvrc12.tar.gz
$ tar zxvf caffe_ilsvrc12.tar.gz
"""

from __future__ import print_function
import argparse
import os
import sys

import numpy as np
from PIL import Image

import chainer
from chainer import cuda
import chainer.functions as F
from chainer.links import caffe

parser = argparse.ArgumentParser(
    description='Predict a label for a image based on CaffeNet')
parser.add_argument('image', help='Path to image for prediction')
parser.add_argument('model', default='bvlc_reference_caffenet.caffemodel',
                    help='Path to the pretrained Caffe model')
parser.add_argument('--mean', '-m', default='ilsvrc_2012_mean.npy',
                    help='Path to the mean file')
parser.add_argument('--gpu', '-g', type=int, default=-1,
                    help='Zero-origin GPU ID (nevative value indicates CPU)')
args = parser.parse_args()

if args.gpu >= 0:
    cuda.check_cuda_available()
xp = cuda.cupy if args.gpu >= 0 else np

print('Loading Caffe model file %s...' % args.model, file=sys.stderr)
func = caffe.CaffeFunction(args.model)
print('Loaded', file=sys.stderr)

if args.gpu >= 0:
    cuda.get_device(args.gpu).use()
    func.to_gpu()

in_size = 227
mean_image = np.load(args.mean)

def forward(x):
    y, = func(inputs={'data': x}, outputs=['fc8'], train=False)
    return y

image = np.asarray(Image.open(args.image).resize((in_size, in_size))).transpose(2, 0, 1)[::-1].astype(np.float32)
image -= mean_image

pred = forward(chainer.Variable(xp.asarray([image]), volatile=True))

if args.gpu >= 0:
        pred = cuda.to_cpu(pred.data)
else:
    pred = pred.data

words = open('synset_words.txt').readlines()
words = [(w[0], ' '.join(w[1:])) for w in [w.split() for w in words]]
words = np.asarray(words)

top5 = np.argsort(pred)[0][::-1][:5]
probs = np.sort(pred)[0][::-1][:5]

for w, p in zip(words[top5], probs):
    print('{}\tprobability:{}'.format(w, p))