import os
import tensorflow as tf
from PIL import Image #注意Image,后面会用到
import matplotlib.pyplot as plt
import numpy as np
writer= tf.python_io.TFRecordWriter("cifar10_train.tfrecords")
writer_2 = tf.python_io.TFRecordWriter("cifar10_test.tfrecords")
root = '101_ObjectCategories'
class_num = []
path_all = []
i = 0
num_test = 0
num_train = 0
for dirpath, dirnames, filenames in os.walk(root):
j = 0
for filepath in filenames:
img_path = os.path.join(dirpath, filepath)
path_all.append(img_path)
img=Image.open(img_path)
img= img.resize((224,224))
img =img.convert("RGB")
img_raw=img.tobytes()
example = tf.train.Example(features=tf.train.Features(feature={
"label": tf.train.Feature(int64_list=tf.train.Int64List(value=[i])),
'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
}))
if j%10==0 :
writer_2.write(example.SerializeToString())
num_test=num_test+1
else:
writer.write(example.SerializeToString())
num_train = num_train+1
j=j+1
for dir_path in dirnames:
class_num.append(dir_path)
i=i+1
writer.close()
print(str(num_test)+'\n')
print(str(num_train))
with open('label.txt','w') as f:
for i in range(len(class_num)):
f.write(str(class_num[i])+'\n')
将文件夹下分类好的子文件夹中的图片转换成tf_record的形式,然后再生成一个label的txt文件方便后面分类后对分类物体具体类别进行对应。