Notes on tensorflow（七）将数据集转换为TFRecord

TFRecord是tensorflow使用的数据格式， 类似于caffe的imdb，mxnet的recordio。使用框架定义的数据格式好处是有强大的框架支持，例如封装了数据解析、多线程等操作， 使用起来方便。坏处主要是需要数据转换，要占用额外的空间。
本文描述将pascal voc数据转换成tfrecord文件的过程，得到的tfrecord可用于训练。 代码主要参考 https://github.com/balancap/SSD-Tensorflow/blob/master/datasets/pascalvoc_to_tfrecords.py

读取图片

import tensorflow as tf
# path
data_dir = '/home/dengdan/dataset_nfs/pascal-voc/voc2007trainval/VOCdevkit/VOC2007'
image_idx = '008541'
image_path = '%s/JPEGImages/%s.jpg'%(data_dir, image_idx)
annotation_path = '%s/Annotations/%s.xml'%(data_dir, image_idx)
# read file
image_data = tf.gfile.FastGFile(image_path, 'r').read()
print type(image_data)

tf.gfile.FastGFile(image_path, 'r').read()读取图片数据并以字符串形式返回. 至于如何编码暂时不用关心， 因为解码的过程也是tf完成。

读取annotation

pascal voc的标注以xml格式的文件存在， 可使用xml.etree.ElementTree解析

import xml.etree.ElementTree as ET
tree = ET.parse(annotation_path)
root = tree.getroot()
size = root.find('size')
shape = [int(size.find('height').text),
         int(size.find('width').text),
         int(size.find('depth').text)]
print shape

# Find annotations.
bboxes = []
labels = []
labels_text = []
difficult = []
truncated = []
for obj in root.findall('object'):
    label = obj.find('name').text
    print label
    labels.append(1)#int(VOC_LABELS[label][0]) label对应的类别编号， 此处直接使用1， 没什么特殊含义。
    labels_text.append(label.encode('ascii'))

    if obj.find('difficult') is not None:
        difficult.append(int(obj.find('difficult').text))
    else:
        difficult.append(0)

    if obj.find('truncated') is not None:
        truncated.append(int(obj.find('truncated').text))
    else:
        truncated.append(0)

    bbox = obj.find('bndbox')
    bboxes.append((float(bbox.find('ymin').text) / shape[0],
                   float(bbox.find('xmin').text) / shape[1],
                   float(bbox.find('ymax').text) / shape[0],
                   float(bbox.find('xmax').text) / shape[1]
                   ))

[336, 500, 3]
person
person
person
person
person

将读取出来的数据转换成Example

一个example就是一条record，它包含一张图片及其对应的标注，一个tfrecord文件中包含多个example。Example接口将输入的数据序列化，以定入tfrecord文件。
tf.train.Feature的作用是指定数据格式转换的协议。

def int64_feature(value):
    """Wrapper for inserting int64 features into Example proto.
    """
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


def float_feature(value):
    """Wrapper for inserting float features into Example proto.
    """
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))


def bytes_feature(value):
    """Wrapper for inserting bytes features into Example proto.
    """
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))

def _convert_to_example(image_data, labels, labels_text, bboxes, shape,
                        difficult, truncated):
    """Build an Example proto for an image example.

    Args:
      image_data: string, JPEG encoding of RGB image;
      labels: list of integers, identifier for the ground truth;
      labels_text: list of strings, human-readable labels;
      bboxes: list of bounding boxes; each box is a list of integers;
          specifying [xmin, ymin, xmax, ymax]. All boxes are assumed to belong
          to the same label as the image label.
      shape: 3 integers, image shapes in pixels.
    Returns:
      Example proto
    """
    xmin = []
    ymin = []
    xmax = []
    ymax = []
    for b in bboxes:
        assert len(b) == 4
        # pylint: disable=expression-not-assigned
        [l.append(point) for l, point in zip([ymin, xmin, ymax, xmax], b)]
        # pylint: enable=expression-not-assigned

    image_format = b'JPEG'
    example = tf.train.Example(features=tf.train.Features(feature={
            'image/height': int64_feature(shape[0]),
            'image/width': int64_feature(shape[1]),
            'image/channels': int64_feature(shape[2]),
            'image/shape': int64_feature(shape),
            'image/object/bbox/xmin': float_feature(xmin),
            'image/object/bbox/xmax': float_feature(xmax),
            'image/object/bbox/ymin': float_feature(ymin),
            'image/object/bbox/ymax': float_feature(ymax),
            'image/object/bbox/label': int64_feature(labels),
            'image/object/bbox/label_text': bytes_feature(labels_text),
            'image/object/bbox/difficult': int64_feature(difficult),
            'image/object/bbox/truncated': int64_feature(truncated),
            'image/format': bytes_feature(image_format),
            'image/encoded': bytes_feature(image_data)}))
    return example

将Example写入文件

将Example写入tfrecord文件是通过tf.python_io.TFRecordWriter完成的。

tf_filename = 'path_to_file'
tfrecord_writer = tf.python_io.TFRecordWriter(tf_filename)
example = _convert_to_example(image_data, labels, labels_text,  bboxes, shape, difficult, truncated)
tfrecord_writer.write(example.SerializeToString())