Tensorflow中创建自己的TFRecord格式数据集

参考文献《TensorFlow实战Google深度学习框架》

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;
    }
};

它实际上存储了一个从属性名到取值的字典。其中属性名为一个字符串，属性取值可以为字符串（ByteList），实数列表（FloatList）和整数列表（Int64List）。比如对于一幅图像而言，可以将图像的像素信息保存成一个字符串，将图像对应的标签保存成整数列表。

创建TFRecord文件

先导入一些必要的库：（jupyter-notebook中实现）

import numpy as np
import tensorflow as tf
from PIL import Image
import os
import matplotlib.pyplot as plt
%matplotlib inline 

数据预处理

我自己从网上下载了10张图片（3张猫，4张狗，3张马），分别存放在cat, dog和horse文件夹下，因为从网上下载的图片大小格式不统一，先将这些图片做预处理，函数如下（这里只是附上函数部分代码，文末会附上完整测试代码）：

def preprocess(imageRawDir, imageDir):
    """
    images preprocess

    Arguments:
    imageRawDir -- directory of primary images.
    imageDir -- directory of processed images.

    Return: none.
    """
    imageNames = os.listdir(imageRawDir)
    label = imageDir.split("/")[-2] # directory format："./data/cat/"
    for index, imageName in enumerate(imageNames):
        image = Image.open(os.path.join(imageRawDir,imageName))
        image = image.resize((256, 256))
        savePath = os.path.join(imageDir, str(label+"_"+str(index))+".jpg")
        image.save(savePath)  

预处理后的图片会保存在另一个指定的文件夹下。

写入到TFRecord文件

下面两个函数会在创建TFRecord文件的时候用到。因为如果不写成函数的形式，代码会很长，看起来也很头疼。

def _int64_feature(value):
    """
    generate int64 feature.
    """
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def _bytes_feature(value):
    """
    generate byte feature.
    """
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

下面这个函数完成的功能是读取我们之前预处理后的所有图片，并依次将每张图片写入到TFRecord文件中,这里因为数据很少，只写入到了一个TFRecord文件中，当数据量很大时，也可以写入多个文件中。

def createRecord(imageDir):
    """
    create TFRecord data.

    Arguments:
    imageDir -- image directory.
    Return: none.
    """
    # create a writer to write TFRecord file
    writer = tf.python_io.TFRecordWriter(os.path.join(imageDir, "train.tfrecords"))
    classNames = ["cat", "dog", "horse"]

    for classIndex, className in enumerate(classNames):
        print "class name = ",className
        currentClassDir = os.path.join(imageDir,className)
        print "current dir = ",currentClassDir
        for index, imageName in enumerate(os.listdir(currentClassDir)):
            image = Image.open(os.path.join(currentClassDir,imageName))
            image_raw = image.tobytes() # convert image to binary format
            print index, imageName

            # write image data(pixel values and label) to Example Protocol Buffer
            example = tf.train.Example(features = tf.train.Features(feature = {
            "label": _int64_feature(classIndex),
            "image_raw": _bytes_feature(image_raw),
            }))

            # write an example to TFRecord file
            writer.write(example.SerializeToString())
    writer.close()

读取TFRecord文件

def readRecord(recordName):
    """
    read TFRecord data (images).

    Arguments:
    recordName -- the TFRecord file to be read.
    return: data saved in recordName (image and label).
    """
    filenameQueue = tf.train.string_input_producer([recordName])
    reader = tf.TFRecordReader()
    _, serializedExample = reader.read(filenameQueue)
    features = tf.parse_single_example(serializedExample, features={
        "label": tf.FixedLenFeature([], tf.int64),
        "image_raw": tf.FixedLenFeature([], tf.string)
    })

    label = features["label"]
    image = features["image_raw"]
    image = tf.decode_raw(image, tf.uint8)
    image = tf.reshape(image,[256,256,3])
    label = tf.cast(label, tf.int32)
    return image, label

注意，这里我们得到的返回值都是张量，需要在tensorflow中创建session后才能得到实际的数据。如下：

##test code
image, label =  readRecord("./data/train.tfrecords")
print image, label
imageBatch, labelBatch = tf.train.shuffle_batch([image, label], batch_size=4, capacity=10, min_after_dequeue=5)


init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)
thread = tf.train.start_queue_runners(sess=sess)
for i in range(10):
    #print image_batch.shape, label.shape
    images, labels = sess.run([imageBatch, labelBatch])
    print "batch shape = ", images.shape,"labels = ", labels
print "label = ", labels    
for i in range(4):
    plt.subplot(1,4,i+1)
    plt.axis("off")
    plt.imshow(images[i])

输出结果：

batch shape =  (4, 256, 256, 3) labels =  [0 1 1 0]
batch shape =  (4, 256, 256, 3) labels =  [1 0 0 0]
batch shape =  (4, 256, 256, 3) labels =  [1 2 2 1]
batch shape =  (4, 256, 256, 3) labels =  [1 2 1 2]
batch shape =  (4, 256, 256, 3) labels =  [0 0 1 0]
batch shape =  (4, 256, 256, 3) labels =  [2 1 2 1]
batch shape =  (4, 256, 256, 3) labels =  [1 1 0 0]
batch shape =  (4, 256, 256, 3) labels =  [0 2 1 1]
batch shape =  (4, 256, 256, 3) labels =  [2 2 1 0]
batch shape =  (4, 256, 256, 3) labels =  [1 2 0 1]
label =  [1 2 0 1]

可以发现，最后一个的batch中的图像和标签是一一对应的（0: cat; 1: dog; 2: horse）,说明我们已经成功从TFRecord文件中读出了数据。

完整样例代码

import numpy as np
import tensorflow as tf
from PIL import Image
import os
import matplotlib.pyplot as plt

currentDir = os.getcwd()
os.chdir(currentDir)
print currentDir

def preprocess(imageRawDir, imageDir):
    """
    images preprocess

    Arguments:
    imageRawDir -- directory of primary images.
    imageDir -- directory of processed images.

    Return: none.
    """
    imageNames = os.listdir(imageRawDir)
    label = imageDir.split("/")[-2] # directory format："./data/cat/"
    for index, imageName in enumerate(imageNames):
        image = Image.open(os.path.join(imageRawDir,imageName))
        image = image.resize((256, 256))
        savePath = os.path.join(imageDir, str(label+"_"+str(index))+".jpg")
        image.save(savePath)    

##test code
catRawDir = "./data_raw/cat/"
catDir = "./data/cat/"
preprocess(catRawDir, catDir)

dogRawDir = "./data_raw/dog/"
dogDir = "./data/dog/"
preprocess(dogRawDir, dogDir)

horseRawDir = "./data_raw/horse/"
horseDir = "./data/horse/"
preprocess(horseRawDir, horseDir)


def _int64_feature(value):
    """
    generate int64 feature.
    """
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def _bytes_feature(value):
    """
    generate byte feature.
    """
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def createRecord(imageDir):
    """
    create TFRecord data.

    Arguments:
    imageDir -- image directory.
    Return: none.
    """
    writer = tf.python_io.TFRecordWriter(os.path.join(imageDir, "train.tfrecords"))
    classNames = ["cat", "dog", "horse"]

    for classIndex, className in enumerate(classNames):
        print "class name = ",className
        currentClassDir = os.path.join(imageDir,className)
        print "current dir = ",currentClassDir
        for index, imageName in enumerate(os.listdir(currentClassDir)):
            image = Image.open(os.path.join(currentClassDir,imageName))
            image_raw = image.tobytes() # convert image to binary format
            print index, imageName

            example = tf.train.Example(features = tf.train.Features(feature = {
            "label": _int64_feature(classIndex),
            "image_raw": _bytes_feature(image_raw),
            }))
            writer.write(example.SerializeToString())
    writer.close()


##test code
createRecord(os.path.join(currentDir, "data/"))


def readRecord(recordName):
    """
    read TFRecord data (images).

    Arguments:
    recordName -- the TFRecord file to be read.
    return: data saved in recordName (image and label).
    """
    filenameQueue = tf.train.string_input_producer([recordName])
    reader = tf.TFRecordReader()
    _, serializedExample = reader.read(filenameQueue)
    features = tf.parse_single_example(serializedExample, features={
        "label": tf.FixedLenFeature([], tf.int64),
        "image_raw": tf.FixedLenFeature([], tf.string)
    })

    label = features["label"]
    image = features["image_raw"]
    image = tf.decode_raw(image, tf.uint8)
    image = tf.reshape(image,[256,256,3])
    label = tf.cast(label, tf.int32)
    return image, label


##test code
image, label =  readRecord("./data/train.tfrecords")
print image, label
imageBatch, labelBatch = tf.train.shuffle_batch([image, label], batch_size=4, capacity=10, min_after_dequeue=5)


##test code
init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)
thread = tf.train.start_queue_runners(sess=sess)
for i in range(10):
    #print image_batch.shape, label.shape
    images, labels = sess.run([imageBatch, labelBatch])
    print "batch shape = ", images.shape,"labels = ", labels
print "label = ", labels    
for i in range(4):
    plt.subplot(1,4,i+1)
    plt.axis("off")
    plt.imshow(images[i])