【tensorflow】tfrecord文件生成保存读取

tfreocrds数据将原始图像数据和标签数据以二进制格式存储。存储内容以如下形式存储：

example = tf.train.Example(features=tf.train.Features(feature={
"label": tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
"height": tf.train.Feature(int64_list=tf.train.Int64List(value=[shape[0]])),
"width": tf.train.Feature(int64_list=tf.train.Int64List(value=[shape[1]])),
"channel": tf.train.Feature(int64_list=tf.train.Int64List(value=[shape[2]])),
'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
}))

例子（生成与保存）：

from random import shuffle
import numpy as np
import glob
import tensorflow as tf
import cv2
import sys
import os


os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
shuffle_data = True
image_path = 'IOU.jpg'
# 取得该路径下所有图片的路径，type（addrs）= list
addrs = glob.glob(image_path)
# 标签数据的获得具体情况具体分析，type（labels）= list
labels = [1]
# 这里是打乱数据的顺序
if shuffle_data:
    c = list(zip(addrs, labels))
    shuffle(c)
    addrs, labels = zip(*c)
# 按需分割数据集
train_addrs = addrs[0:int(len(addrs))]
train_labels = labels[0:int(len(labels))]
# 上面不是获得了image的地址么，下面这个函数就是根据地址获取图片
def load_image(addr):  # A function to Load image
    img = cv2.imread(addr)
    img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_CUBIC)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    # 这里/255是为了将像素值归一化到[0，1]
    img = img / 255.
    img = img.astype(np.float32)
    return img
# 将数据转化成对应的属性
def _int64_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 _float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
# 下面这段就开始把数据写入TFRecods文件
train_filename = 'train.tfrecords'  # 输出文件地址
# 创建一个writer来写 TFRecords 文件
writer = tf.python_io.TFRecordWriter(train_filename)
print(len(train_addrs))
for i in range(len(train_addrs)):
    # 这是写入操作可视化处理
    if not i % 1:
        print('Train data: {}/{}'.format(i, len(train_addrs)))
        sys.stdout.flush()
    # 加载图片
    img = load_image(train_addrs[i])
    label = train_labels[i]
    # 创建一个属性（feature）
    feature = {'train/label': _int64_feature(label),
               'train/image': _bytes_feature(tf.compat.as_bytes(img.tostring()))}
    # 创建一个 example protocol buffer
    example = tf.train.Example(features=tf.train.Features(feature=feature))
    # 将上面的example protocol buffer写入文件
    writer.write(example.SerializeToString())

例子（读取与显示）：

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

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
data_path = 'train.tfrecords'  # tfrecords 文件的地址

with tf.Session() as sess:
    # 先定义feature，这里要和之前创建的时候保持一致
    feature = {
        'train/image': tf.FixedLenFeature([], tf.string),
        'train/label': tf.FixedLenFeature([], tf.int64)
    }
    # 创建一个队列来维护输入文件列表
    filename_queue = tf.train.string_input_producer([data_path])
    # 定义一个 reader ，读取下一个 record
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)
    # 解析读入的一个record
    features = tf.parse_single_example(serialized_example, features=feature)
    # 将字符串解析成图像对应的像素组
    image = tf.decode_raw(features['train/image'], tf.float32)

    # 将标签转化成int32
    label = tf.cast(features['train/label'], tf.int32)

    # 这里将图片还原成原来的维度
    #image = tf.reshape(image, [1152, 648, 3])
    with tf.Session() as sess:  # 开始一个会话
        init_op = tf.initialize_all_variables()
        sess.run(init_op)
        # 启动多线程
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        single, l = sess.run([image, label])  # 在会话中取出image和label
        print(single.shape)
        image1 = tf.reshape(single, [224, 224, 3])
        data_numpy = image1.eval()
        img = Image.fromarray(data_numpy, 'RGB')  # 这里Image是之前提到的
        img.save(str(0) + '_''Label_' + str(l) + '.jpg')

    # 关闭线程
    coord.request_stop()
    coord.join(threads)
    sess.close()