tfrecord构建自己的数据集

根据tfrecord构建本地自己的数据集

套话：TFRecord文件中的数据都是通过tf.train.Example Protocol Buffer的格式存储的。其中包含一个从属性名到取值的字典，属性的取值可以为字符串（BytesList）,实数列表（FloatList）或整数列表（Int64List）。

比如将，将一个图片存为字符串，其label值存为整数。

message Feature{
    oneof kind{
    BytesList    bytes_list = 1;
    BytesList    int64_list = 1;
    }
}

给出代码：

我定义的 Feature包含三个字段：

message Feature{
    oneof kind{
    BytesList    bytes_list = 1; // 图片值
    BytesList    int64_list = 1;  // label
    BytesList    bytes_list = 1; // 图片的名称
    }
}

代码：

#encoding：utf-8
import tensorflow as tf 
import numpy as np
import os
def _int64_feature(label):
	return tf.train.Feature(int64_list=tf.train.Int64List(value=[label])) 

def _byte_feature(value):
	return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) 

#指定使用显卡0
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
#图像文件的路径
image_path = '/home/ccf/data/v2/256/0/train/image/'
#设置tfrecord输出目录
out_path = '/home/ccf/'
name ='test2'

'''
将文件路径和标签放在list中
'''
def get_image_path_label(image_path,label_path):
	image=[]
	label=[]
	for one_path in  os.listdir(image_path):		
		image.append(image_path+one_path)
		label.append(one_path.split('_')[0]) 
	return image,label

'''
image_path:图片路径,list
labels： 图片标签,list 
save_dir：tfrecord保存路径
name: 保存名称
'''
def convert_to_tfrecord(images,labels,save_dir,name):
	#创建一个writer来写tfrecord
	writer = tf.python_io.TFRecordWriter(save_dir+name+'.tfrecords')
	with tf.Session() as sess:

		for i in  range(len(images)):			  
			#读取图片并编码
			image_raw_data = tf.gfile.FastGFile(images[i],'rb').read()
			#若是jpeg格式的图片，换为tf.decode_jpeg(image_raw_data)
			image_data = tf.image.decode_png(image_raw_data)
			label = int(labels[i])
			print(images[i].encode())
			example = tf.train.Example(features = tf.train.Features(feature= {
				'label':_int64_feature(label),
				'name':_byte_feature(images[i].encode()),
				'image_raw':_byte_feature(image_data.eval().tostring())
				}
				))
			#写入
			writer.write(example.SerializeToString())
			if(i>10):
				break
			print(images[i],i)
		
	writer.close()
	print('writer done')

def read_from_tfrecord(tfrecord_path):
	#创建一个队列从tfrecord中读取数据
	file_queue = tf.train.string_input_producer([tfrecord_path])
	reader = tf.TFRecordReader()

	_,serilazed_example = reader.read(file_queue)

	image_features = tf.parse_single_example(serilazed_example,features ={
		'label':tf.FixedLenFeature([],tf.int64),
		'name':tf.FixedLenFeature([],tf.string),
		'image_raw':tf.FixedLenFeature([],tf.string)
		})

	image = image_features['image_raw']
	label = image_features['label']
	image_name = image_features['name']

	decode_image = tf.decode_raw(image,tf.uint8)
	reshape_image = tf.reshape(decode_image,[256,256,3])

	init = tf.initialize_all_variables()
	with tf.Session() as sess:
		sess.run(init)
		#启动多线程处理输入数据
		coord = tf.train.Coordinator()
		threads = tf.train.start_queue_runners(sess=sess,coord=coord)
		for i in range(10):
			images,labels,image_names = sess.run(image,label,image_name)
			print(image_names)
		coord.request_stop()
		coord.join(threads)
if __name__ == '__main__':
	images,labels = get_image_path_label(image_path,image_path)  
	convert_to_tfrecord(images,labels,out_path,name)   //写入tfrecord
	read_from_tfrecord('/home/ccf/test2.tfrecords')  //读取tfrecord