针对分割问题（图片）的TFRecord使用方法

转载自https://zhuanlan.zhihu.com/p/32078191

出于想把程序放在美团云的dls上跑的目的，学习了一下如何使用TF自带的数据格式TFRecord（吐槽一下美团云一次只能上传20个文件的设定，想把一个数据集上传上去真的太麻烦了）。说是针对分割问题的使用方法是因为在分割问题里label也是图片，所以就只有图片，其实TFRecord的使用方法非常灵活，只要能组织起自己的数据就可以随便用起来。

TFRecord的格式

TFRecord文件中的数据都是通过tf.train.Example Protocol Buffer的格式存储的：

message Example{
    Features features = 1;
}; 

message Features{
    map<string, Feature> feature = 1;
};

message Feature{
    oneof kind {
        BytesList bytes_list = 1;
        FloatList float_list = 2;
        Int64List int64_list = 3;
    }
}; 

这个格式也是需要在程序中使用的，仅做了解。

将已有的数据集转换成TFRecord格式

1. 定义一个将已有的数据转换成Feature数据结构的函数（官方教程中的函数）
2. 将数据集中的文件名写入一个txt文件中（为了在程序中方便读取），可以这样来组织：
   文件名1.jpg 文件名1.png
   文件名2.jpg 文件名2.png
   ......
3. 创建一个句柄来读这个txt文件
4. 定义一个TFRecord的writer
5. 逐个文件来写入TFRecoder文件

import tensorflow as tf

TXT_PATH = './dataset.txt'
TFRECORD_PATH = './dataset.tfrecord'

# 1.定义一个将已有的数据转换成Feature数据结构的函数（官方教程中的函数）
def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

# 3.创建一个句柄来读这个txt文件
f = open(TXT_PATH)

# 4.定义一个TFRecord的writer
writer = tf.python_io.TFRecordWriter(TFRECORD_PATH)

# 5.逐个文件来写入TFRecoder文件
for i in f.readlines():
    # 每一行的文件名是用空格隔开的，所以需要使用split方法把string映射成list
    item = i.split()
    
    # 这里可能有些教程里会使用skimage或者opencv来读取文件，但是我试了一下opencv的方法
    # 一个400MB左右的训练集先用cv2.imread转换成numpy数组再转成string，最后写入TFRecord
    # 得到的文件有17GB，所以这里推荐使用FastGFile方法，生成的tfrecord文件会变小，  
    # 唯一不好的一点就是这种方法需要在读取之后再进行一次解码。
    img = tf.gfile.FastGFile(item[0], 'rb').read()
    label = tf.gfile.FastGFile(item[1], 'rb').read()

    # 按照第一部分中Example Protocol Buffer的格式来定义要存储的数据格式
    example = tf.train.Example(features=tf.train.Features(feature={
        'raw_image': _bytes_feature(img),
        'label': _bytes_feature(label)
    }))
    # 最后将example写入TFRecord文件
    writer.write(example.SerializeToString())

writer.close()
f.close()

#自己修改代码

# 1.定义一个将已有的数据转换成Feature数据结构的函数（官方教程中的函数）
def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

# 获取文件列表
files_img = tf.gfile.Glob(os.path.join(images_dir,'*.jpg')) # 如我想列出images_dir下所有jpg文件路径
files_anno = tf.gfile.Glob(os.path.join(annotations_dir,'*.png')) # 如我想列出annotations_dir下所有png文件路径


# 4.定义一个TFRecord的writer
writer = tf.python_io.TFRecordWriter(TFRECORD_PATH)

# 5.逐个文件来写入TFRecoder文件
for i in range(len(files_img)):

    # 这里可能有些教程里会使用skimage或者opencv来读取文件，但是我试了一下opencv的方法
    # 一个400MB左右的训练集先用cv2.imread转换成numpy数组再转成string，最后写入TFRecord
    # 得到的文件有17GB，所以这里推荐使用FastGFile方法，生成的tfrecord文件会变小，
    # 唯一不好的一点就是这种方法需要在读取之后再进行一次解码。
    img = tf.gfile.FastGFile(files_img[i], 'rb').read()
    label = tf.gfile.FastGFile(files_anno[i], 'rb').read()

    # 按照第一部分中Example Protocol Buffer的格式来定义要存储的数据格式
    example = tf.train.Example(features=tf.train.Features(feature={
        'raw_image': _bytes_feature(img),
        'label': _bytes_feature(label)
    }))
    # 最后将example写入TFRecord文件
    writer.write(example.SerializeToString())

writer.close()

因为FastGFile读取的是图片没有解码过的的原始数据，所以在使用存在tfrecord中的这些原始数据时，需要对读取出来的图片原始数据进行解码。

读取TFRecord格式的数据

1. 把所有的TFRecord文件名列表写入队列中（只有一个就写一个文件名在列表中，多个就写多个）
2. 创建一个读取器
3. 将队列中的tfrecord文件读取为example格式
4. 根据定义数据的方式对应说明读取的方式
5. 对图片进行解码

# 1. 把所有的TFRecord文件名列表写入队列中（只有一个就写一个文件名在列表中，多个就写多个）
queue = tf.train.string_input_producer([TFRECORD_PATH], shuffle=True)

# 2. 创建一个读取器
reader = tf.TFRecordReader()

# 3. 将队列中的tfrecord文件读取为example格式
_, serialized_example = reader.read(queue)

# 4. 根据定义数据的方式对应说明读取的方式
features = tf.parse_single_example(serialized_example, 
                                   features={
                                       'raw_image': tf.FixedLenFeature([], tf.string)
                                       'label': tf.FixedLenFeature([], tf.string)
                                   })
img = features['raw_image']
label = features['label']

# 5. 对图片进行解码
img = tf.image_decode_jpeg(img, channels=3)
label = tf.image_decode_png(label, channels=1)

# 然后就可以用tf.train.batch方法来生成一个batch的image和label啦
img_batch, label_batch = tf.train.batch([img, label], BATCH_SIZE)

关于速度

因为在没有使用TFRecord时数据集是放在SSD上并且也是有解码的步骤的，所以在使用了TFRecord之后数据的读取没有明显变快，只是可以将上千张图片的数据集压缩到一个文件中这一点是比较方便的（拷一个巨大的数据集真的太慢了）。