import tensorflow as tf
import numpy as np
import os
from PIL import Image
#沒有下面兩句德華會出現 image file is truncate的問題
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

file_dir = r‘D:\old computer\contest\貓狗大戰\貓狗大戰‘
file_path = r‘D:\old computer\contest\貓狗大戰\貓狗大戰\cats_and_dogs.tfrecords‘
classes = [‘狗狗-訓練集‘, ‘貓-測試集‘]
writer = tf.python_io.TFRecordWriter(file_path)
#使用枚舉更好的將標簽與圖片聯系
for index, file in enumerate(classes):
    for f in os.listdir(file):
        #print(f)
        file_name = file_dir+r‘\\‘+file+r‘\\‘+f
        image = Image.open(file_name)
        image = image.resize((128, 128))
        image = image.tobytes()
        examples = tf.train.Example(features=tf.train.Features(feature={
            ‘labels‘: tf.train.Feature(int64_list=tf.train.Int64List(value=[index])),
            ‘image‘: tf.train.Feature(bytes_list=tf.train.BytesList(value=[image]))}))
        writer.write(examples.SerializeToString())
writer.close()

貓狗大戰的TFrecord數據集制作