把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)