猫狗大战的TFrecord数据集制作

Posted 2020-10-29 lenkong

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