# -*- 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)