Tensorflow 中TFRecord格式转换与读取

把csv格式文件转化为TFRecord

Tensorflow提供了TFRecord格式来存储数据，以下是将csv格式转化为TFRecord格式的代码

import tensorflow as tf
import pandas as pd
import numpy as np

train = pd.read_csv('train.csv')

label = train['label'].values
y_train = train.iloc[:,:-1].values


writer = tf.python_io.TFRecordWriter('train_csv.tfrecords')
print(y_train[1].shape)

for i in range(y_train.shape[0]):
    image_raw = y_train[i].tostring()
    example = tf.train.Example(
        # 需要主要此处是tf.train.Features，下面的是tf.train.Feature，差别在于一个's'
        features=tf.train.Features(
            feature = {
                'image_raw':tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_raw])),
                'label':tf.train.Feature(int64_list=tf.train.Int64List(value=[label[i]])),
                }
        )
    )
    writer.write(record=example.SerializeToString())
writer.close()

将图片存为TFRecord格式文件

import tensorflow as tf
from scipy import misc

img = misc.imread('im.jpg')
img_raw = img.tostring()
# 此处假定图片标签为1,实际中标签可能在图片名，文件中
label = 1
writer = tf.python_io.TFRecordWriter('img_to_TFRcord.tfrecords')

example = tf.train.Example(
    features = tf.train.Features(
        feature = {
            'img_raw':tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw])),
            'label':tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))
        }
    )

)
writer.write(record=example.SerializeToString())
writer.close()

TFRecord格式文件的读取

import tensorflow as tf
import numpy as np


filename_tfrecord = tf.train.string_input_producer(['img_to_TFRcord.tfrecords'])

reader = tf.TFRecordReader()

_,serialized_record = reader.read(filename_tfrecord)

features = tf.parse_single_example(
    serialized=serialized_record,
    features={
        'img_raw':tf.FixedLenFeature([],tf.string),
        'label':tf.FixedLenFeature([],tf.int64),
    }
)
img = tf.decode_raw(features['img_raw'],tf.uint8)
label = tf.cast(features['label'],tf.int32)

with tf.Session() as sess:
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess,coord=coord)
    image, label = sess.run([img, label])

    print(image.shape)
    print(label.shape)