TensorFlowによるVGG16を用いた画像分類

最終更新: 2017-03-16 14:05

TensorFlowによるVGG16を用いた画像分類

インストール

TensorFlowに標準で用意されているSlimというライブラリを使いました。執筆時点でのTensorFlowではモデルは標準で入っていないので別途ダウンロードする必要があります。ダウンロード先は好きなところでいいです。

sudo pip install tensorflow
cd /home/kivantium/workspace/
git clone https://github.com/tensorflow/models/

モデルのダウンロード

Pre-trained Modelsからvgg_16.tar.gzをダウンロードして解凍しました。

コード

公式ドキュメントのTF-Slim Walkthroughが参考になりますが、バージョン違いで動かない部分があったので修正しました。

#!/usr/bin/python
# -*- coding: utf-8 -*-

import tensorflow as tf
import numpy as np
import argparse
import os
import sys
import tensorflow.contrib.slim as slim
import tensorflow.contrib.slim.python.slim.nets.vgg as vgg

sys.path.append("/home/kivantium/workspace/models/slim") # you need to change
from preprocessing import vgg_preprocessing
from datasets import imagenet

parser = argparse.ArgumentParser()
parser.add_argument('jpeg_image')
args = parser.parse_args()

image_size = vgg.vgg_16.default_image_size
with tf.Graph().as_default():
    image = tf.image.decode_jpeg(tf.read_file(args.jpeg_image), channels=3)
    processed_image = vgg_preprocessing.preprocess_image(image, image_size, image_size, is_training=False)
    processed_images  = tf.expand_dims(processed_image, 0)
    
    with slim.arg_scope(vgg.vgg_arg_scope()):
        logits, _ = vgg.vgg_16(processed_images, num_classes=1000, is_training=False)
    probabilities = tf.nn.softmax(logits)
    
    checkpoints_dir = '/home/kivantium/workspace/models/' # you need to change
    init_fn = slim.assign_from_checkpoint_fn(
        os.path.join(checkpoints_dir, 'vgg_16.ckpt'),
        slim.get_model_variables('vgg_16'))
    
    with tf.Session() as sess:
        init_fn(sess)
        np_image, probabilities = sess.run([image, probabilities])
        probabilities = probabilities[0, 0:]
        sorted_inds = [i[0] for i in sorted(enumerate(-probabilities), key=lambda x:x[1])]
    
    names = imagenet.create_readable_names_for_imagenet_labels()
    for i in range(5):
        index = sorted_inds[i]
        print('Probability {:.2f}% => [{}]'.format(probabilities[index]*100., names[index+1]))

vgg16.pyという名前で保存し、Wikipediaの象の画像で実行したところ以下のような結果になりました。

$ python vgg16.py elephants.jpg 
Probability 75.12% => [tusker]
Probability 15.56% => [Indian elephant, Elephas maximus]
Probability 9.31% => [African elephant, Loxodonta africana]
Probability 0.00% => [triceratops]
Probability 0.00% => [chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour]

関連ページ