将图片数据集制作成TFRecord文件与读取

一、制作TFRecord文件

import numpy
import tensorflow as tf
import os
import cv2
# 图片文件夹根目录
root_dir_img = '图片文件夹根目录'
# 获得该目录下所有文件的文件名
list_img = os.listdir(root_dir_img)
# 创建TFRecord写入对象
writer = tf.python_io.TFRecordWriter("train.tfrecords")
for index in range(len(list_img)):
    # 将文件名与根路径组合为全路径
    img_path = os.path.join(root_dir_img, name)
    # 使用OpenCV读取图片，注意：路径中不能有任何中文
    img = cv2.imread(img_path)
    # 将图片转化为二进制，注意：记住原始数据类型，OpenCV读取的为int8
    img_raw = img.tobytes()
    # 创建一条记录，label可以是一个数也可以是一个ndarray数组，但注意detype
    example = tf.train.Example(features=tf.train.Features(feature={
        'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[index])),
        'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
    }))
    writer.write(example.SerializeToString())
writer.close()

二、TFRecord文件读取

import tensorflow as tf


def read_TFRecord(filename)
    # filename是TFRecord文件路径，如果TFRecord和py文件在同一目录下可以只写文件名
    filename_queue = tf.train.string_input_producer([filename])
    # 创建一个读文件的对象
    reader = tf.TFRecordReader()
    # 返回文件名和文件
    _, serialized_example = reader.read(filename_queue) 
    # 得到一个记录，一定注意写的时候是int64解析的时候也要是int64
    features = tf.parse_single_example(serialized_example, 
                        features={
                           'label': tf.FixedLenFeature([], tf.int32)
                           'img_raw': tf.FixedLenFeature([], tf.string)
                                       })
    # 解码
    img = tf.decode_raw(features['img_raw'], tf.uint8)
    # 必须重新指定shape，否则会报错，如果label也是一个数组，也必须指定shape
    img = tf.reshape(img, [224, 224, 3])
    # 获得标签
    label = tf.cast(features["label"], tf.int32)
    return img, label


# 得到一个标签
img, label = read_TFRecord("train.tfrecords")
# 使用shuffle_batch可以随机打乱输入，一批次8张图片
img_batch, label_batch = tf.train.shuffle_batch([img, label],
                                                batch_size=8,       
                                                capacity=8*64,
                                                min_after_dequeue=8*32)
# 定义一个初始化op
init = tf.initialize_all_variables()
with tf.Session() as sess:
    sess.run(init)
    threads = tf.train.start_queue_runners(sess=sess)
    # 输出五个批次的内容
    for i in range(5):
        batch, label = sess.run([img_batch, label_batch])
        # 输出一个batch的图片和标签
        print(batch.shape, label.shape)

三、可能出现的问题

1、发现解码出来的出现错误

制作TFRecord文件时标签是20×1，解码出来却不是，可能是160×1或者80×1，这是因为数据类型不一致，例如使用np.zeros((20, 1))，默认detype=float64，解码时必须是tf.float64，否则就会出现不一致的问题

2、OpenCV读不出来图片

OpenCV读取的图片路径不能带有任何中文，必须全英文

3、异常

All shapes must be fully defined: [TensorShape([Dimension(224), Dimension(224), Dimension(3)]), TensorShape([Dimension(None)])]
或者
All shapes must be fully defined: [TensorShape([Dimension(None), Dimension(None), Dimension(None)]), TensorShape([Dimension(None)])]
必须显式制定读取出来的数据形状，图片要tf.reshape(img, [224, 224, 3])，如果label是一个数组也必须tf.reshape(label, [10, 1])。