#daiizメモ

Scrapboxに夢中

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

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

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

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

f:id:daiiz:20161005153440j:plain:w128 f:id:daiiz:20161005153445j:plain:w128
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:
    with open() 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())

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