TensorFlow 制作自己的TFRecord数据集

准备图片数据

网上下载了2类吉他和房子的图片，全部 resize成64*64大小
如下图，保存项目下：

现在利用这2 类共108张图片制作TFRecord文件

制作TFRECORD文件

1 先聊一下tfrecord, 这是一种将图像数据和标签放在一起的二进制文件，能更好的利用内存，在tensorflow中快速的复制，移动，读取，存储等等…
这里注意，tfrecord会根据你选择输入文件的类，自动给每一类打上同样的标签
如在本例中，只有0,1 两类

import os
import tensorflow as tf
from PIL import Image  #注意Image,后面会用到
import matplotlib.pyplot as plt
import numpy as np
# 制作TFRECORD文件===============================================
cwd='photo'
classes={'guitar','house'} #人为 设定 2 类
writer= tf.python_io.TFRecordWriter("train.tfrecords") #要生成的文件
for index,name in enumerate(classes):print('index:',index)# print('name:',name)# photo的子文件夹class_path=cwd+'/'+name+'/'# print('class_path:',class_path)for img_name in os.listdir(class_path):# print('img_name:',img_name)img_path=class_path+img_name #每一个图片的地址print('img_path:',img_path)img=Image.open(img_path)# print('img:',img)# 调整图片大小img= img.resize((64,64))img_raw=img.tobytes()#将图片转化为二进制字符格式# print('img_raw:',img_raw)# tf.train.Features：#           输入：键，值#                 int64_list：输入：整数#                 bytes_list：输入：字节（二进制数）#           输出：特征# tf.train.Example 协议内存块包含了Features字段# 通过feature将图片的二进制数据和label进行统一封装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()

tf.train.Example 协议内存块包含了Features字段，通过feature将图片的二进制数据和label进行统一封装，然后将example协议内存块转化为字符串， tf.python_io.TFRecordWriter 写入到TFRecords文件中。

读取TFRECORD文件

在制作完tfrecord文件后，将该文件读入到数据流中。
代码如下

# 读取TFRECORD文件=======================================================
def read_and_decode(filename): # 读入train.tfrecordsfilename_queue = tf.train.string_input_producer([filename])#生成一个queue队列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),})#将image数据和label取出来
# 注意，feature的属性“label”和“img_raw”名称要和制作时统一 ，返回的img数据和label数据一一对应。img = tf.decode_raw(features['img_raw'], tf.uint8)img = tf.reshape(img, [128, 128, 3])  #reshape为128*128的3通道图片img = tf.cast(img, tf.float32) * (1. / 255) - 0.5 #在流中抛出img张量label = tf.cast(features['label'], tf.int32) #在流中抛出label张量return img, label

注意，feature的属性“label”和“img_raw”名称要和制作时统一，返回的img数据和label数据一一对应。返回的img和label是2个 tf 张量

显示tfrecord格式的图片

有些时候我们希望检查分类是否有误，或者在之后的网络训练过程中可以监视，输出图片，来观察分类等操作的结果，那么我们就可以session回话中，将tfrecord的图片从流中读取出来，再保存。紧跟着一开始的代码写：

# 显示tfrecord格式的图片=========================================================
# 读入流中
filename_queue = tf.train.string_input_producer(["train.tfrecords"])
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),})  #取出包含image和label的feature对象
# 解码
image = tf.decode_raw(features['img_raw'], tf.uint8)
# 大小转化
image = tf.reshape(image, [64, 64, 3])
label = tf.cast(features['label'], tf.int32)
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(coord=coord)# 54副图片for i in range(108):example, l = sess.run([image,label])#在会话中取出image和labelimg=Image.fromarray(example, 'RGB')#这里Image是之前提到的# print('path:',cwd+str(i)+'_''Label_'+str(l)+'.jpg')img.save(cwd+str(i)+'_''Label_'+str(l)+'.jpg')#存下图片# print(example, l)coord.request_stop()coord.join(threads)

代码运行完后，从tfrecord中取出的文件被保存在photo文件夹了。

总代码如下：

import os
import tensorflow as tf
from PIL import Image  #注意Image,后面会用到
import matplotlib.pyplot as plt
import numpy as np
# 制作TFRECORD文件===============================================
cwd='photo'
classes={'guitar','house'} #人为 设定 2 类
writer= tf.python_io.TFRecordWriter("train.tfrecords") #要生成的文件
for index,name in enumerate(classes):print('index:',index)# print('name:',name)# photo的子文件夹class_path=cwd+'/'+name+'/'# print('class_path:',class_path)for img_name in os.listdir(class_path):# print('img_name:',img_name)img_path=class_path+img_name #每一个图片的地址print('img_path:',img_path)img=Image.open(img_path)# print('img:',img)# 调整图片大小img= img.resize((64,64))img_raw=img.tobytes()#将图片转化为二进制字符格式# print('img_raw:',img_raw)# tf.train.Features：#           输入：键，值#                 int64_list：输入：整数#                 bytes_list：输入：字节（二进制数）#           输出：特征# tf.train.Example 协议内存块包含了Features字段# 通过feature将图片的二进制数据和label进行统一封装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()
# 读取TFRECORD文件=======================================================
def read_and_decode(filename): # 读入train.tfrecordsfilename_queue = tf.train.string_input_producer([filename])#生成一个queue队列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),})#将image数据和label取出来img = tf.decode_raw(features['img_raw'], tf.uint8)img = tf.reshape(img, [128, 128, 3])  #reshape为128*128的3通道图片img = tf.cast(img, tf.float32) * (1. / 255) - 0.5 #在流中抛出img张量label = tf.cast(features['label'], tf.int32) #在流中抛出label张量return img, label
# 显示tfrecord格式的图片=========================================================
# 读入流中
filename_queue = tf.train.string_input_producer(["train.tfrecords"])
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),})  #取出包含image和label的feature对象
# 解码
image = tf.decode_raw(features['img_raw'], tf.uint8)
# 大小转化
image = tf.reshape(image, [64, 64, 3])
label = tf.cast(features['label'], tf.int32)
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(coord=coord)# 54副图片for i in range(108):example, l = sess.run([image,label])#在会话中取出image和labelimg=Image.fromarray(example, 'RGB')#这里Image是之前提到的# print('path:',cwd+str(i)+'_''Label_'+str(l)+'.jpg')img.save(cwd+str(i)+'_''Label_'+str(l)+'.jpg')#存下图片# print(example, l)coord.request_stop()coord.join(threads)