[TensorFlow] 学習済みモデルInception-v3から特徴量を抽出する

TensorFlowで、Inception-v3モデルを使って特徴量抽出をします。

1. 以下から学習済みモデルをダウンロードする
   http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
2. 上記ファイルを解凍して、classify_image_graph_def.pbを抽出
3. ソースコード
   feature_extraction.py
   

IMG_PATHとMODEL_PATHを環境に合わせて書き換えてください。

feature_extraction.py

import tensorflow as tf
import numpy as np

IMG_PATH = 'path/to/input/image.jpg'
MODEL_PATH = 'path/to/classify_image_graph_def.pb'

# 学習済みモデルの読み込み
inception_v3 = tf.gfile.FastGFile(MODEL_PATH, 'rb')
graph_def =tf.GraphDef()
graph_def.ParseFromString(inception_v3.read())
tf.import_graph_def(graph_def, name='')

with tf.Session() as sess:
    # 抽出層の指定
    pool3 = sess.graph.get_tensor_by_name('pool_3:0')

    # 入力画像の読み込み
    image_data = tf.gfile.FastGFile(IMG_PATH, 'rb').read()

    # 特徴量の抽出
    features = sess.run(pool3, {'DecodeJpeg/contents:0': image_data})
    print(np.squeeze(features))

PNGファイルを入力する場合

上記はJPEGファイルの入力のみを受け付けます。
PNGファイルを入力したい場合は、特徴量の抽出時に以下ように、'DecodeJpeg:0'ノードを指定します。

features = sess.run(pool3, {'DecodeJpeg:0':image_data})