tf 制作tfrecord笔记

import os
import tensorflow as tf

cwd = 'C:\\Users\\lenovo\Desktop\\datasets\\flower_photos\\'
classes = {'roses', 'sunflowers'}  # 类别
tfRecordfilename = "C:\\Users\\lenovo\Desktop\\datasets\\flower_photos\\flowers_use_train.tfrecords"
writer = tf.python_io.TFRecordWriter(tfRecordfilename)  # 要生成的文件

for index, name in enumerate(classes):
    class_path = cwd + name + '\\'
    for img_name in os.listdir(class_path):
        img_path = class_path + img_name  # 每一个图片

        img = Image.open(img_path)
        img = img.resize((128, 128))
        img_raw = img.tobytes()  # 将图片转化为二进制格式 tostring()转字符串
        example = tf.train.Example(features=tf.train.Features(feature={
            "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[index])),
            'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
        }))  # example对象对label和image数据进行封装
        writer.write(example.SerializeToString())  # 序列化为字符串

writer.close()