【TensorFlow】tfrecord

1 生成tfrecord文件

（1）数据文件格式

（2）代码

import tensorflow as tf
from PIL import Image
import os


def make_tfrecords(data_path, classes, tfrecord_path):
    """生成tfrecord文件"""
    writer = tf.python_io.TFRecordWriter(tfrecord_path)
    for index, name in enumerate(classes):
        class_path = data_path + name + '/'
        for img_name in os.listdir(class_path):
            img_path = class_path + img_name
            img = Image.open(img_path)
            img = img.resize((256, 256))
            # 将图片转化为二进制格式
            img_data = img.tobytes()
            # example对象对label和image数据进行封装
            example = tf.train.Example(features=tf.train.Features(feature={
                "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[index])),
                'img_data': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_data]))}))
            # 序列化为字符串
            writer.write(example.SerializeToString())
    writer.close()
    return


if __name__ == '__main__':
    data_path = 'data/'
    classes = {'archery', 'riding'}
    make_tfrecords(data_path, classes, "train.tfrecords")
    pass

（3） 结果

2 读取tfrecord文件并显示图片

（1）代码

import tensorflow as tf
from PIL import Image


def display_tfrecord_file(filename):
    filename_queue = tf.train.string_input_producer([filename])  # 读入流中
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)  # 返回文件名和文件
    features = tf.parse_single_example(serialized_example,
                                       features={
                                           'label': tf.FixedLenFeature([], tf.int64),
                                           'img_data': tf.FixedLenFeature([], tf.string),
                                       })  # 取出包含image和label的feature对象
    image = tf.decode_raw(features['img_data'], tf.uint8)
    image = tf.reshape(image, [256, 256, 3])
    label = tf.cast(features['label'], tf.int32)
    with tf.Session() as sess:
        init_op = tf.initialize_all_variables()
        sess.run(init_op)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        for i in range(9):
            image_data, label_info = sess.run([image, label])
            img = Image.fromarray(image_data, 'RGB')
            img.save("display/" + str(i) + '_Label_' + str(label_info) + '.jpg')
        coord.request_stop()
        coord.join(threads)
    pass


if __name__ == '__main__':
    display_tfrecord_file("train.tfrecords")
    pass

（2）结果