【TensorFlow】画像データをTFRecordファイルに詰める

学習や評価用の画像データをTensorFlowに与えるとき，JPEG，PNG画像をそのまま与えてもよいけれど，これらの教師ラベルや画像サイズ情報などとともにTFRecord形式のファイルに格納して使う方法が便利．公式ドキュメントで言う以下の説明の前半について，コード書いてみる．

writerを生成する & 画像を読み込む

TFRecordに詰め込みたい画像ファイルのパスとラベルの組を用意する．今回使用する画像は以下のような，576×576pxな正方形のJPEGファイル．

img_data = [
    ['./img/chicken_curry.jpg', 0],
    ['./img/tandoori_chicken.jpg', 1]
]

TFRecordファイルの出力パスを指定して，TFRecordWriterを生成する．

record_file = './foods.tfrecords'
writer = tf.python_io.TFRecordWriter(record_file)

書き込む

今回は，画像のバイナリデータと教師ラベルの他に，画像のwidth, height, depthも一緒に記録する．これらの値はPIL (Pillow)を使うと簡単に取得できる．

from PIL import Image

for img_datum in img_data:
    withopen() as f:
        img_file = img_datum[0]
        img_label = img_datum[1]
        
        # PIL で画像のサイズデータを取得
        img = Image.open(img_file)
        width, height = img.size
        
        # TFRecordに書き込む
        record = tf.train.Example(features=tf.train.Features(feature={
            'label' : tf.train.Feature(int64_list=tf.train.Int64List(value=[img_label])),
            'image' : tf.train.Feature(bytes_list=tf.train.BytesList(value=[img])),
            'height': tf.train.Feature(int64_list=tf.train.Int64List(value=[height])),
            'width' : tf.train.Feature(int64_list=tf.train.Int64List(value=[width])),
            'depth' : tf.train.Feature(int64_list=tf.train.Int64List(value=[3])),
        }))
        writer.write(record.SerializeToString())

次回は画像の読み出しと加工を行う．