python小工具--tfrecords文件的制作工具

最近一直在捣弄TensorFlow，其中有关数据的读取有很多种方式。因为嫌麻烦，自己就把读取的方式统一为用TensorFlow自己的数据格式也就是tfrecords的格式，所以我自己弄写了个将数据和标签转化为tfrecords格式的python程序。

"""tfrecord文件的生成程序，能实现将分好类的图像数据打乱写入tfrecords类型文件中，并自动将数据集分成测试数据和训练数据@author:liao"""import tensorflow as tffrom PIL import Imageimport osimport randomimport shutil#######################如果要改主要是改这里test_number_ratio=0.6#用于生成测试数据的比例cwd=os.getcwd()path=cwd+os.sep+'tfrecoder-test'#path是保存要转化图像数据的所在文件夹的绝对路径#保存生成的tfrecords数据的文件夹的绝对路径tf_records_dir=path+os.sep+'tfrecord_file'#生成的测试tfrecords文件的文件名tf_train_name='train_emotion_image_data'#生成的训练tfrecords文件的文件名tf_test_name='test_emotion_image_data'#########################生成的train_tfrecords文件所在路径tf_train_records_path=tf_records_dir+os.sep+tf_train_name+'.tfrecords'#生成的test_tfrecords文件所在路径tf_test_records_path=tf_records_dir+os.sep+tf_test_name+'.tfrecords'if os.path.exists(tf_records_dir):#如果tf_records_dir文件夹存在    shutil.rmtree(tf_records_dir)#则递归地删除tf_records_dir文件夹classname= os.listdir(path)#    print("root is\n %s"%(root))classname_path=[]#每个类的文件夹的绝对路径for x in classname:    x=path+os.sep+x    classname_path.append(x)    print(classname_path)train_file_path=[]#用来存放每个训练图片的绝对路径test_file_path=[]#用来存放每个测试图片的绝对路径for x in classname_path:    file_name_list=os.listdir(x)#得到每个类下面的文件的文件名list    file_name_number=len(file_name_list)#得到每个类下面的图像的数量    test_file_name_number=int(float(file_name_number)*test_number_ratio)#得到每个类作为测试图像的数量    test_file_name_path=file_name_list[0:test_file_name_number]#取出测试文件名list    train_file_name_path=file_name_list[test_file_name_number:]#取出训练文件名list    for y in train_file_name_path:        train_file_path.append(x+os.sep+y)#得到训练数据文件的绝对路径    for y in test_file_name_path:        test_file_path.append(x+os.sep+y)#得到测试数据文件的绝对路径    #print(file_path)random.shuffle(train_file_path)#打乱训练的文件路径的排序，达到tfrecords乱序的效果#print(file_path)#生成tfrecord_file文件夹，用来保存生成的tfrecords文件os.mkdir(tf_records_dir)#生成训练的tfrecords数据image_number=0TFwriter = tf.python_io.TFRecordWriter(tf_train_records_path)#生成一个train 的TFwriterfor imgPath in train_file_path:    print(imgPath)    imgPath_list=imgPath.split("\\")#将文件绝对路径imgPath按“\”分离，得到对应list(不知道这里为什么是双斜杠)    label=int(imgPath_list[-2][-1])#取出类别    print("class %s"%(label))#显示类别#    print(type(label))    img = Image.open(imgPath)#得到image对象    print (img.size,img.mode)    image_number+=1    print("pictures of train%s"%image_number)#显示处理到第几张#    os.system("pause")#暂停    imgRaw = img.tobytes()#转化为字节类型    example = tf.train.Example(features=tf.train.Features(feature={    "label":tf.train.Feature(int64_list = tf.train.Int64List(value=[label])),    "img":tf.train.Feature(bytes_list = tf.train.BytesList(value=[imgRaw]))    }) )#将标签和图像数据写入example这个实例中    TFwriter.write(example.SerializeToString())#写入文件中TFwriter.close()#生成测试的tfrecords数据image_number=0TFwriter = tf.python_io.TFRecordWriter(tf_test_records_path)#生成一个test 的TFwriterfor imgPath in test_file_path:    print(imgPath)    imgPath_list=imgPath.split("\\")#将文件绝对路径imgPath按“\”分离，得到对应list(不知道这里为什么是双斜杠)    label=int(imgPath_list[-2][-1])#取出类别    print("class %s"%(label))#显示类别#    print(type(label))    img = Image.open(imgPath)#得到image对象    print (img.size,img.mode)    image_number+=1    print("pictures of test:%s"%image_number)#显示处理到第几张#    os.system("pause")#暂停    imgRaw = img.tobytes()#转化为字节类型    example = tf.train.Example(features=tf.train.Features(feature={    "label":tf.train.Feature(int64_list = tf.train.Int64List(value=[label])),    "img":tf.train.Feature(bytes_list = tf.train.BytesList(value=[imgRaw]))    }) )#将标签和图像数据写入example这个实例中    TFwriter.write(example.SerializeToString())#写入文件中TFwriter.close()print("\n")print("test_number_ratio=%s"%(test_number_ratio))print("total_pictures=%s"%((len(train_file_path)+len(test_file_path))))print("pictures of train:%s have converted"%(len(train_file_path)))#显示所得到的训练文件的数量print("pictures of test: %s have converted"%(len(test_file_path)))#显示所得到的测试文件的数量

这个程序主要实现了将分好类别的图像数据转为为tfrecords的数据得，方便使用TensorFlow训练的时候的数据读取。得到的tfrecords的数据分为train_image_data.tfrecords和test_image_data.tfrecords，显然train_image_data.records用于训练，而test_image_data.records用于测试。

修改的话一般只需要修改程序里指定的某些部分就行了，具体看程序。

有关tfrecords格式的相关知识请参考这里http://wiki.jikexueyuan.com/project/tensorflow-zh/how_tos/reading_data.html#AUTOGENERATED-reading-from-files