制作tfrecords文件

# -*- coding: utf-8 -*-
import os.path
import tensorflow as tf
from pathlib import Path
import time
from PIL import Image

image_size = 32

INPUT_PATH = Path(r'C:\Users\admin\Desktop\cifar10')
OUTPUT_PATH = r'.\tfrecords\cifar10\train.tfrecords'

classes = []
for i in os.listdir(INPUT_PATH):
    classes.append(i)


def create_tfrecords():
    writer = tf.python_io.TFRecordWriter(OUTPUT_PATH)
    for index, name in enumerate(classes):
        class_path = INPUT_PATH / name
        for img_name in os.listdir(class_path):
            print("正在处理" + img_name)
            img = Image.open(str(class_path / img_name), 'r')
            # img = img.resize((image_size, image_size), Image.ANTIALIAS)
            image_raw = img.tobytes()  # 将图片转化为二进制格式

            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=[image_raw]))
                }))
            writer.write(example.SerializeToString())
    writer.close()


if __name__ == '__main__':
    with tf.Session() as sess:
        time_start = time.time()
        create_tfrecords()
        time_end = time.time()
        print("总时长：" + str(time_end-time_start))
        print(classes)

虽然不推荐，但如果想随机从图像中划分测试集和验证集，可以略加修改，代码如下：

# -*- coding: utf-8 -*-
import os.path
import numpy as np
import tensorflow as tf
from pathlib import Path
import time
from PIL import Image

image_size = 32

INPUT_PATH = Path(r'C:\Users\admin\Desktop\train')
TRAIN_OUTPUT_PATH = './tfrecords/cifar10/train.tfrecords'
TEST_OUTPUT_PATH = './tfrecords/cifar10/test.tfrecords'

# 测试数据和验证数据的比例7:3
TEST_PERCENTAGE = 30
classes = []
for i in os.listdir(INPUT_PATH):
    classes.append(i)


def create_tfrecords():
    writer = tf.python_io.TFRecordWriter(TRAIN_OUTPUT_PATH)
    writer_test = = tf.python_io.TFRecordWriter(TEST_OUTPUT_PATH)
    for index, name in enumerate(classes):
        class_path = INPUT_PATH / name
        for img_name in os.listdir(class_path):
            print("正在处理" + img_name)
            img = Image.open(str(class_path / img_name), 'r')
            # img = img.resize((image_size, image_size), Image.ANTIALIAS)
            image_raw = img.tobytes()  # 将图片转化为二进制格式

            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=[image_raw]))
                }))
            chance = np.random.randint(100)
            # 随机划分数据集
            if chance < TEST_PERCENTAGE:
                writer_test.write(example.SerializeToString())
            else:
                writer.write(example.SerializeToString())
    writer.close()


if __name__ == '__main__':
    with tf.Session() as sess:
        time_start = time.time()
        create_tfrecords()
        time_end = time.time()
        print("总时长：" + str(time_end-time_start))
        print(classes)