TensorFlow杂记 - 生成和读取TFRecord（一）

利用闲暇时间，通过SSD-TensorFlow和网络资源上总结两种生成和读取TFRecord的方法，代码测试通过。

软件平台：Pycarn 2019.2 + Tensorflow 1.13.1 + cuda 10.0 + cudnn 7.6.5

GPU: 1080Ti

 

不啰嗦了，直接上代码，有问题可发送邮件 <tecsai@163.com>， 或直接留言讨论。

 

TFRecord生成：

import os
import tensorflow as tf 
from PIL import Image  # 注意Image,后面会用到
import xml.etree.ElementTree as ET

 
JpgFilePath='H:\\11_DataSet\\QD\\JPEGImages\\'
XmlFilePath='H:\\11_DataSet\\QD\\Annotations\\'
#文件路径
filepath = 'H:\\11_DataSet\\QD'
writer= tf.python_io.TFRecordWriter("H:\\11_DataSet\\QD\\qd.tfrecord")

VOC_LABELS = {
    'B01': (0, 'B01'),
    'D01': (1, 'D01'),
    'G01': (2, 'G01'),
    'W01': (3, 'W01'),
    'W02': (4, 'W02'),
    'W03': (5, 'W03'),
    'W04': (6, 'W04'),
    'T01': (7, 'T01'),
    'R01': (8, 'R01'),
    'F01': (9, 'F01'),
    'S01': (10, 'S01'),
    'G02': (11, 'G02'),
    'G03': (12, 'G03'),
    'W05': (13, 'W05'),
    'I18': (14, 'I18'),
    'I01': (15, 'I01'),
    'I02': (16, 'I02'),
    'I03': (17, 'I03'),
    'I04': (18, 'I04'),
    'I99': (19, 'I99'),
}

for img_name in os.listdir(JpgFilePath):
    img_path=JpgFilePath+img_name  # 每一个图片的地址

    # Jpeg
    img=Image.open(img_path)
    print(img_path)
    # img= img.resize((300, 300))
    image_data=img.tobytes()  # 将图片转化为二进制格式

    # Xml
    xml_name = img_name[:-4]  # 文件名，不带后缀
    xml_path = XmlFilePath+xml_name+'.xml'
    print(xml_path)
    tree = ET.parse(xml_path)
    root = tree.getroot()

    # Image shape.
    size = root.find('size')
    shape = [int(size.find('height').text), # shape是个list，shape[0]: height, shape[1]: width, shape[2]: depth
             int(size.find('width').text),
             int(size.find('depth').text)]
    # Find annotations.
    bboxes = []
    labels = []
    labels_text = []
    difficult = []
    truncated = []
    for obj in root.findall('object'):
        label = obj.find('name').text
        labels.append(int(VOC_LABELS[label][0]))
        labels_text.append(label.encode('ascii'))

        if obj.find('difficult'):
            difficult.append(int(obj.find('difficult').text))
        else:
            difficult.append(0)
        if obj.find('truncated'):
            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]
                       ))
    print(labels)

    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': tf.train.Feature(int64_list=tf.train.Int64List(value=[shape[0]])),
            'image/width': tf.train.Feature(int64_list=tf.train.Int64List(value=[shape[1]])),
            'image/channels': tf.train.Feature(int64_list=tf.train.Int64List(value=[shape[2]])),
            'image/shape': tf.train.Feature(int64_list=tf.train.Int64List(value=shape)),
            'image/object/bbox/xmin': tf.train.Feature(float_list=tf.train.FloatList(value=xmin)),
            'image/object/bbox/xmax': tf.train.Feature(float_list=tf.train.FloatList(value=xmax)),
            'image/object/bbox/ymin': tf.train.Feature(float_list=tf.train.FloatList(value=ymin)),
            'image/object/bbox/ymax': tf.train.Feature(float_list=tf.train.FloatList(value=ymax)),
            'image/object/bbox/label': tf.train.Feature(int64_list=tf.train.Int64List(value=labels)),
            'image/object/bbox/label_text': tf.train.Feature(bytes_list=tf.train.BytesList(value=labels_text)),
            'image/object/bbox/difficult': tf.train.Feature(int64_list=tf.train.Int64List(value=difficult)),
            'image/object/bbox/truncated': tf.train.Feature(int64_list=tf.train.Int64List(value=truncated)),
            'image/format': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_format])),
            'image/encoded': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_data]))}))

    writer.write(example.SerializeToString())  # 序列化为字符串

writer.close()

TFRecord读取：

import os 
import tensorflow as tf 
from PIL import Image  

storepath='H:\\11_DataSet\\QD\\'

filename = 'H:\\11_DataSet\\QD\\qd.tfrecord'

filename_queue = tf.train.string_input_producer([filename])  # 生成一个queue队列
 
reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue)  # 返回文件名和文件
features = tf.parse_single_example(serialized_example,
                                   features={
                                            'image/shape': tf.FixedLenFeature([3], tf.int64),
                                            'image/height': tf.FixedLenFeature([1], tf.int64),
                                            'image/width': tf.FixedLenFeature([1], tf.int64),
                                            'image/encoded': tf.FixedLenFeature((), tf.string),
                                             })  # 将image数据和label取出来

image = tf.decode_raw(features['image/encoded'], tf.uint8)
image = tf.reshape(image, [1080, 1920, 3])
#image = tf.cast(img, tf.float32) * (1. / 255) - 0.5
shape = tf.cast(features['image/shape'], tf.int64)
height = tf.cast(features['image/height'], tf.int64)
width = tf.cast(features['image/width'], tf.int64)

with tf.Session() as sess:
    init_op = tf.global_variables_initializer()
    sess.run(init_op)
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(coord=coord)

    for i in range(10):
        print("loop %d\n" % i)
        example, s, h, w = sess.run([image, shape, height, width])  # 在会话中取出image和label
        # h, w = sess.run([height, width])  # 在会话中取出image和label
        print(s)
        print(h)
        print(w)

        img=Image.fromarray(example, 'RGB')  # 这里Image是之前提到的
        img.save(storepath+str(i)+'.jpg') #  存下图片

    coord.request_stop()
    coord.join(threads)