制作tensorflow标准数据集即制作.tfrecords格式文件

#coding=utf-8
import  tensorflow as tf
import  os
from PIL import Image
cwd = os.getcwd()
#将图片制作成.tfrecords数据集
def encode_to_tfrecords():
cwd = os.getcwd()
classes = {"NGreasy","Greasy"}#classes = {"data/Greasy","data/NGreasy"}
writer = tf.python_io.TFRecordWriter("/home/system/Python/data/train.tfrecords")
for index, name in enumerate(classes):
#print ([index,name])
class_path = cwd + "/data/" + name + "/"
#print class_path
for img_name in os.listdir(class_path):
img_path = class_path + img_name
img = Image.open(img_path)
img = img.resize((45, 45))
img_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=[img_raw]))
}))
writer.write(example.SerializeToString())
writer.close

#读取.tfrecords数据集
def read_and_decode(filename,num_epoch=None):
filename_queue = tf.train.string_input_producer([filename],num_epochs=num_epoch)
reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue)
features = tf.parse_single_example(serialized_example,
features={
'label': tf.FixedLenFeature([], tf.int64),
'img_raw' : tf.FixedLenFeature([], tf.string),
})
img = tf.decode_raw(features['img_raw'], tf.uint8)
img = tf.reshape(img, [45, 45, 3])
#   img = tf.cast(img, tf.float32) * (1. / 255) - 0.5
label = tf.cast(features['label'], tf.int64)
return img, label
#验证制作的数据集是否正确
with tf.Session() as sess:
init_op = tf.initialize_all_variables()
sess.run(init_op)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
for i in range(20):
example, l = sess.run([image, label])
img = Image.fromarray(example, 'RGB')
img.save(cwd + '/pic/' + str(i) + '_Label_' + str(l) + '.jpg')
coord.request_stop()
coord.join(threads)