一.TFRecord简介

TensorFlow提供了TFRecord的格式来统一存储数据，它是一种能够将图像数据和标签放在一起的二进制文件， 这种格式可以使TensorFlow的数据集更容易与网络应用架构相匹配，并可以在Tensorflow中快速的复制，移动，读取，存储 。

二.TFRecord生成

以kaggle比赛猫狗数据集（下载链接：https://pan.baidu.com/s/13hw4LK8ihR6-6-8mpjLKDA 密码：dmp4）为例生成TFRecord文件。
1.首先下载数据集，通过如下data_provider.py将数据分成train和validation两部分，并将划分后的图片路径和label信息分别保存到train_imagePath_label.json和val_imagePath_label.json文件中。

import glob
import json
import os
import random
"""
Get image path and label
"""


def split_train_val_datasets(images_path, tarin_ratio=0.90):
    """Split image to training and validation"""
    if not os.path.exists(images_path):
        raise ValueError('images_path does not exit')

    image_files = glob.glob(os.path.join(images_path, '*.jpg'))
    random.shuffle(image_files)
    num_train_samples = int(len(image_files)*tarin_ratio)

    train_files = image_files[:num_train_samples]
    val_files = image_files[num_train_samples:]

    train_dict = get_label_dict(train_files)
    val_dict = get_label_dict(val_files)

    return train_dict, val_dict


def get_label_dict(image_files=None):
    if image_files is None:
        return None
    label_dict = {}
    for image_file in image_files:
        image_name = image_file.split('\\')[-1]
        if image_name.startswith('cat'):
            label_dict[image_file] = 0
        elif image_name.startswith('dog'):
            label_dict[image_file] = 1

    return label_dict


def write_annotation_json(images_path, train_json_output_path,
                          val_jason_output_path):
    """Saving training and validation annotions."""
    train_files_dict, val_files_dict = split_train_val_datasets(images_path)

    train_json = json.dumps(train_files_dict)
    with open(train_json_output_path, 'w') as writer:
        json.dump(train_json, writer)

    val_json = json.dumps(val_files_dict)
    with open(val_jason_output_path, 'w') as writer:
        json.dump(val_json, writer)


def provide(annotion_path=None):
    """Return images_path and class labels.
    Args:
        annotation_path:Path to an antation's json file.


    Returns:
        annotation_dict: A dictionary containing the paths of images and the class labels of each image.
    """
    if not os.path.exists(annotion_path):
        raise ValueError('annotion path does not exist.')

    with open(annotion_path, 'r') as reader:
        annotion_str = json.load(reader)
        annotion_d = json.loads(annotion_str)

    annotion_dict = {}

    for image_name, label in annotion_d.items():
        annotion_dict[image_name] = label
    return annotion_dict


if __name__ == '__main__':

    images_path = './data'
    train_annotation_path = './datasets/train/train_imagePath_label.json'
    val_annotation_path = './datasets/val/val_imagePath_label.json'

    write_annotation_json(
        images_path, train_annotation_path, val_annotation_path)

    provide(train_annotation_path)

2.通过generate_tfrecord.py生成.record文件：

"""
Generate tfrecord from images and labels.

Example Usage:

python3 generate_tfrecord.py
       --image_path: Path to images directory.
       --train_annotation_path: Path to train annotation .json file.
       --train_output_path: Path to save train tfrecord
       --val_annotation_path: Path to val annotation .json file.
       --val_output_path: Path to save val tfrecord.
       --resize_size: Resize image to fixed size.

"""
import tensorflow as tf
import data_provider
from PIL import Image
import io

flags = tf.app.flags

flags.DEFINE_string('image_path',
                    './data-org',
                    'Path to images directory.')
flags.DEFINE_string('train_annotation_path',
                    './datasets/train/train_imagePath_label.json',
                    'Path to train annotation json file.')
flags.DEFINE_string('train_output_path',
                    './datasets/train/train.record',
                    'Path to save train tfrecord.')

flags.DEFINE_string('val_annotation_path',
                    './datasets/val/val_imagePath_label.json',
                    'Path to val annotation .json file.')
flags.DEFINE_string('val_output_path',
                    './datasets/val/val.record',
                    'Path to save val tfrecord.')
flags.DEFINE_integer('resize_size', 224, 'Resize image to fixed size.')

FLAGS = flags.FLAGS


def int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def int64_list_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


def bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def bytes_list_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))


def float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))


def float_list_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))


def create_tf_example(image_path, label, resize_size=None):
    with tf.gfile.GFile(image_path, 'rb') as fid:  # 若要读取二进制文件，比如图片、视频等，使用'rb'模式打开文件
        encoded_jpg = fid.read()
    encoded_jpg_io = io.BytesIO(encoded_jpg)  # 从内存读取二进制数据
    image = Image.open(encoded_jpg_io)

    width, height = image.size
    # Resize
    if resize_size is not None:
        width = resize_size
        height = resize_size

        image = image.resize((width, height), Image.ANTIALIAS)
        bytes_io = io.BytesIO()
        image.save(bytes_io, format='JPEG')  # 往内存中写二进制数据
        encoded_jpg = bytes_io.getvalue()

   #tf.train.Example:用于写入tfrecords文件
   #features ： tf.train.Features类型的特征实例
   #返回example协议格式块
    tf_example = tf.train.Example(
        features=tf.train.Features(feature={   
            'image/encoded': bytes_feature(encoded_jpg),
            'image/format': bytes_feature('jpg'.encode()),
            'image/path': bytes_feature(image_path.encode('utf8')),
            'image/class/label': int64_feature(label),
            'image/height': int64_feature(height),
            'image/width': int64_feature(width)}))
    return tf_example


def generate_tfrecord(annotation_dict, output_path, resize_size=None):
    num_valid_tf_examples = 0
    #建立TFRecord存储器
    writer = tf.python_io.TFRecordWriter(output_path)
    for image_path, label in annotation_dict.items():
        if not tf.gfile.GFile(image_path):
            print('{} not exits'.format(image_path))
            continue
        tf_example = create_tf_example(image_path, label, resize_size)
       #序列转换成字符串
       writer.write(tf_example.SerializeToString())
        num_valid_tf_examples += 1

        if num_valid_tf_examples % 100 == 0:
            print('{} tf examples created.'.format(num_valid_tf_examples))
    writer.close()
    print('Total {} tf examples created.'.format(num_valid_tf_examples))


def main(_):
    images_path = FLAGS.image_path
    train_annotation_path = FLAGS.train_annotation_path
    train_record_path = FLAGS.train_output_path
    val_annotation_path = FLAGS.val_annotation_path
    val_record_path = FLAGS.val_output_path
    resize_size = FLAGS.resize_size

    # Write json
    data_provider.write_annotation_json(
        images_path, train_annotation_path, val_annotation_path)

    train_annotation_dict = data_provider.provide(train_annotation_path)
    val_annotation_dict = data_provider.provide(val_annotation_path)

    generate_tfrecord(train_annotation_dict, train_record_path, resize_size)
    generate_tfrecord(val_annotation_dict, val_record_path, resize_size)


if __name__ == '__main__':
    tf.app.run()

三.查看生成TFRecord是否ok

运行view_tfrecord.py，实现对.record文件中图像和label的解析：

"""
Created on 04/23/2019

@author:sunjiankang

View and check created tfrecord.

Example Usage:

python3 view_tfrecord.py
       --tfrecord_path: Path to tfrecord.
       --num_samples: Num of samples.
       --num_classes: Num of classes.

"""
import tensorflow as tf
import numpy as np
from PIL import Image
import cv2
slim = tf.contrib.slim

flags = tf.app.flags

flags.DEFINE_string('tfrecord_path',
                    './datasets/val/val.record',
                    'Path to tfrecord.')
flags.DEFINE_integer('num_samples', 2000, 'Num of samples.')
flags.DEFINE_integer('num_classes', 2, 'Num of classes.')

FLAGS = flags.FLAGS


def get_record_dataset(record_path,
                       reader=None,
                       num_samples=2000,
                       num_classes=2):
    """Get a tensorflow record file."""
    if not reader:
        reader = tf.TFRecordReader
    #将example反序列化成存储之前的格式，由tf完成
    keys_to_features = {
        'image/encoded':
            tf.FixedLenFeature((), tf.string, default_value=''),
        'image/format':
            tf.FixedLenFeature((), tf.string, default_value='jpg'),
        'image/path':
            tf.FixedLenFeature((), tf.string, default_value=''),
        'image/class/label':
            tf.FixedLenFeature(
                [1], tf.int64, default_value=tf.zeros([1], dtype=tf.int64)),
        'image/height':
            tf.FixedLenFeature([], tf.int64),
        'image/width':
            tf.FixedLenFeature([], tf.int64)}
    #将反序列化的数据组装成更高级的格式。由slim完成
    items_to_handlers = {
        'image': slim.tfexample_decoder.Image(image_key='image/encoded',
                                              format_key='image/format',
                                              channels=3),
        'label': slim.tfexample_decoder.Tensor('image/class/label', shape=[]),
        'path': slim.tfexample_decoder.Tensor('image/path', shape=[]),
        'height': slim.tfexample_decoder.Tensor('image/height', shape=[]),
        'width': slim.tfexample_decoder.Tensor('image/width', shape=[])}

    #解码器，进行解码
    decoder = slim.tfexample_decoder.TFExampleDecoder(
        keys_to_features, items_to_handlers)

    labels_to_names = None
    items_to_descriptions = {
        'image': 'An image with shape image_shape.',
        'label': 'A single integer between 0 and 1.',
        'path': 'image path',
        'height': 'Image height.',
        'width': 'Image width.'}
    #dataset对象定义了数据集的文件位置，解码方式等元信息
    return slim.dataset.Dataset(
        data_sources=record_path,
        reader=reader,
        decoder=decoder,
        num_samples=num_samples,
        num_classes=num_classes,
        items_to_descriptions=items_to_descriptions,
        labels_to_names=labels_to_names)


def main(_):
    num_samples = FLAGS.num_samples
    num_classes = FLAGS.num_classes
    tfrecord_path = FLAGS.tfrecord_path

    dataset = get_record_dataset(tfrecord_path,
                                 num_samples=num_samples,
                                 num_classes=num_classes)
    #provider对象根据dataset信息读取数据
    data_provider = slim.dataset_data_provider.DatasetDataProvider(dataset)
    #获取数据
    [image, label, path, height, width] = data_provider.get(
        ['image', 'label', 'path', 'height', 'width'])

    with tf.Session() as sess:

        init_op = tf.global_variables_initializer()
        sess.run(init_op)
        
        #调用 tf.train.Coordinator() 来创建一个线程协调器，用来管理之后在Session中启动的所有线程;
        #调用tf.train.start_queue_runners, 启动入队线程，由多个或单个线程，按照设定规则，把文件读入Filename Queue中。
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        for i in range(num_samples):
            img, l, p, h, w = sess.run([image, label, path, height, width])
            img = np.reshape(img, [256, 256, 3])
            img = np.array(img, dtype=np.uint8)
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) #PIL读取图像是RGB格式，opencv显示需要BGR，在这里进行格式转换
            print('label: {}, image path: {}'.format(l, p.decode('utf8')))
            cv2.imshow('img', img)
            if cv2.waitKey(1) == 27:
                break
            if i % 100 == 0:
                print('{} example checked.'.format(i))

        coord.request_stop()
        coord.join(threads)


if __name__ == '__main__':
    tf.app.run()