EeBlog(テクニカルブログ)

AI（人工知能）実践 第11回　TFRecordの作成

前回は、ある画像をベースに複数の画像を作成し、
学習データを嵩増しする方法を紹介しました。

今回も、学習データ絡みの話を続けて紹介していきます。
 

TensorFlowでは、学習データとして複数のファイル形式が選択出来ますが、
その中でも、TFRecordというバイナリフォーマット形式が推奨されています。

TFRecordを利用することで、複数のファイルを1つにまとめることもでき、
また、Dataset API(https://www.tensorflow.org/api_docs/python/tf/data/TFRecordDataset?hl=ja)で簡単に操作出来るようにもなります。

tf.train.Featuresを使用してデータを詰め込む形となりますが、
中身としてはキーと値のセットを複数持っているデータとなります。
今回のサンプルでは、ラベル名のディレクトリに入った画像を詰め込む形を想定しますので、
キーは「height」「width」「label」「value」とします。
尚、ラベル名については数値としますので、ディレクトリ名も数値のみで構成されているものとします。
 

ソースコード

import os
import sys
import glob
import tensorflow as tf
from PIL import Image

TFRECORD_FILE = 'c:\\work\\test\\example.tfrecord'
IMG_DIR = 'c:\\work\\test\\src\\'
IMG_EXTENTION = '.jpg'

#features作成メソッド
def make_features(image_file, label, height=300, width=300):
    image_content = Image.open(image_file).resize((width, height))
    return tf.train.Example(features=tf.train.Features(feature={
        'height': tf.train.Feature(int64_list=tf.train.Int64List(value=[height])),
        'width' : tf.train.Feature(int64_list=tf.train.Int64List(value=[width])),
        'label' : tf.train.Feature(int64_list=tf.train.Int64List(value=[int(label)])),
        'image' : tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_content.tobytes()]))
    }))

#tfrecord作成メソッド
def write_tfrecord(filename, src_dir):
    writer = tf.python_io.TFRecordWriter(filename)
    for image_dir in os.listdir(src_dir):
        if os.path.isdir(src_dir + image_dir):
            for file in glob.glob(src_dir + image_dir + '\\\\*' + IMG_EXTENTION):
                features = make_features(file, image_dir)
                writer.write(features.SerializeToString())
    writer.close()


write_tfrecord(TFRECORD_FILE, IMG_DIR)

data_count = len(list(tf.python_io.tf_record_iterator(TFRECORD_FILE)))
print("データ件数：{}".format(data_count))

 

解説

TFRecordの詰め込み用のソースを順に説明していきます。

ライブラリインポート

import os
import sys
import glob
import tensorflow as tf
from PIL import Image

今回初めてPILというものが出てきました。
PIL(http://www.pythonware.com/products/pil/)とは、Python Image Libraryで画像データを扱うため機能をまとめたライブラリです。
尚、開発は止まっており、PILに関しては、Python2までしか使用できません・・が、
Pillowというライブラリがフォークしているため、同様の機能が使用できます。
ですので、Python3 を使用する場合、事前準備として
pip install Pillow
を実行してください。
※但し、PILとPillowは共存できない為、PILが入っている場合は先にアンインストールしてください。
 

データパス等の定義

TFRECORD_FILE = 'c:\\work\\test\\example.tfrecord'
IMG_DIR = 'c:\\work\\test\\src\\'
IMG_EXTENTION = '.jpg'

後で使うTFRecordの保存先、TFRecordに詰め込むデータのベースとなるフォルダ、
詰め込むデータの拡張子を定義しています。
今回は読み込みデータをjpgのみに絞るため、読み込むデータの拡張子は.jpgとしています。
 

TFRecordのデータ(Features)作成メソッドの定義

#features作成メソッド
def make_features(image_file, label, height=300, width=300):
    image_content = Image.open(image_file).resize((width, height))
    return tf.train.Example(features=tf.train.Features(feature={
        'height': tf.train.Feature(int64_list=tf.train.Int64List(value=[height])),
        'width' : tf.train.Feature(int64_list=tf.train.Int64List(value=[width])),
        'label' : tf.train.Feature(int64_list=tf.train.Int64List(value=[int(label)])),
        'image' : tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_content.tobytes()]))
    }))

TFRecordに詰め込むデータの作成を行っています。
今回は、指定されたサイズに画像をリサイズ処理を行った後に
高さ・幅・ラベル・リサイズ後の画像をバイナリ形式で保存します。

TFRecordを作成する時には、tf.train.Exampleを使用します。
tf.train.Exampleにはfeaturesを渡しますが、ここは特段別の使い方は無いかと思いますので、
下記をセットで覚えれば良いかと思います。

tf.train.Example(features=tf.train.Features(feature={ 【格納したいキーと値のペアを複数個】 }))

【格納したいキーと値のペアを複数個】の部分については、
キー(文字列) : 値(tf.train.Feature)
の形式で記述します。
tf.train.Featureは詳細や解説が何故か公式ページでは見つかりませんが、
公式ページに貼ってあるソースのリンク(https://github.com/tensorflow/tensorflow/blob/r1.10/tensorflow/core/example/feature.proto)
を見る限り、使用出来る型は
・int64
・float
・bytes
の3つとなり、実際に指定する場合、リストで持つ事となりますので、それぞれの型の格納する名前は
・Int64List
・FloatList
・BytesList
となります。
格納したい値に合わせた型を上記から選びましょう。
 

TFRecord作成メソッドの定義

#tfrecord作成メソッド
def write_tfrecord(filename, src_dir):
    writer = tf.python_io.TFRecordWriter(filename)
    for image_dir in os.listdir(src_dir):
        if os.path.isdir(src_dir + image_dir):
            for file in glob.glob(src_dir + image_dir + '\\\\*' + IMG_EXTENTION):
                features = make_features(file, image_dir)
                writer.write(features.SerializeToString())
    writer.close()

今回作成するwrite_tfrecordメソッドでは、
第一引数で渡されたファイルパスに、
第二引数で渡されたディレクトリ配下のラベル名が付けられたディレクトリの画像を全て読み込んで
TFRecordに変換したものを保存します。

まずは、TFRecordを書き出すためのWriter(https://www.tensorflow.org/api_docs/python/tf/python_io/TFRecordWriter)を開きます。

    writer = tf.python_io.TFRecordWriter(filename)

TFRecordWriterの第一引数がファイルパス、第二引数(defaul=None)がオプションとなります。
オプションとして使用できるものは
TFRecordCompressionType.ZLIP
TFRecordCompressionType.GZIP
TFRecordCompressionType.NONE
の3つです。
こちらも公式で使用ガイド等が見当たりませんでしたが、ソース(https://github.com/tensorflow/tensorflow/blob/r1.10/tensorflow/python/lib/io/tf_record.py)
を見ると分かりやすいかもしれません。

続いて、第二引数で渡されたディレクトリの配下にあるディレクトリを処理するためのループを回します。

    for image_dir in os.listdir(src_dir):
        if os.path.isdir(src_dir + image_dir):

os.listdirメソッドで、引数に渡したパスのディレクトリにあるファイル・ディレクトリの一覧を
取得します。
今回はディレクトリのみを対象としたいため、
読み込んだファイル・ディレクトリがディレクトリかどうかをos.path.isdirメソッドで判定しています。
ディレクトリだった場合のみ次の処理へと移ります。

ディレクトリと判別できた場合、ラベル名のディレクトリであり、
その中には学習用の画像(jpg)が格納されている想定となりますので、
該当のディレクトリの中にあるjpgを全てfor文で処理します。
for文の中では、
先ほど作成したwrite_tfrecordメソッドを使用してTFRecordに格納するfeaturesを作成し、
writerで書き込んでいます。

            for file in glob.glob(src_dir + image_dir + '\\\\*' + IMG_EXTENTION):
                features = make_features(file, image_dir)
                writer.write(features.SerializeToString())

最後は勿論、後始末としてwriterをcloseしましょう。

writer.close()

次回は、作成したTFRecordを読み込み、画像に再度変換する処理を紹介します。

2018年8月29日 | カテゴリー：AI（人工知能）

│トップ│