Tensorflow制作tfcerod数据集文件

将标记好的图片，制作成tfrecord图片

import tensorflow as tf
from PIL import Image
import numpy as np
import os

classes = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
data_path = "D:\\python学习\\神经网络动物分类\\"
tfRecord_train="traindata.tfrecords"

def write_tfRecord(tfRecordName, image_path, label_path):
    writer = tf.python_io.TFRecordWriter(tfRecordName)
    num_pic = 0
    f = open(label_path,'r')
    contents = f.readlines()     # 读取整个文件内容
    f.close()
    for content in contents:
        value = content.split()   # 按空格分开
        temp = int(value[1])
        img_path = image_path +"train\\"  +classes[temp] + "\\" + value[0]
        img = Image.open(img_path)
        img_raw = img.tobytes()     # 转化为二进制
        labels = [0] * 10
        labels [int(value[1])] = 1
        example = tf.train.Example(features=tf.train.Features(
            feature={
                'label': tf.train.Feature(int64_list=tf.train.Int64List(value=labels)),
                'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
            }))
        writer.write(example.SerializeToString())


        num_pic +=1
        if not num_pic % 1000:
            print('the number of picture:', num_pic)


    writer.close()
    print("write tfrecord successful")


def generate_tfRecord(data_path):
    isExists = os.path.exists(data_path)
    if not isExists:
        os.makedirs(data_path)
        print("The directory was created successfully")
    else:
        print("directory already exists")
    label_path = data_path + "dataset.txt"
    write_tfRecord(tfRecord_train, data_path, label_path)



def make_path():
    with open("dataset.txt","w") as f:
        for classe in classes:
            path = "D:\\python学习\\神经网络动物分类\\train" + classe
            for filename in os.listdir(path):
                f.write(filename+" "+str(classes.index(classe)))
                f.write("\n")



if __name__ == '__main__':
    generate_tfRecord(data_path)