TFRecord 存入图像和标签

#-*- coding:utf-8 -*-
import os
import tensorflow as tf
import cv2

'''
文件目录为
chiwawa/
     xx.jpg
     xx.jpg
     .....
japandog/
     xx.jpg
     xx.jpg
     .....
'''
cwd = 'f:/py/tfrecord/'
classes={'chiwawa','japandog'} # 需要存入的标签，尽量与文件名一致，方便操作

sess = tf.Session()
writer = tf.python_io.TFRecordWriter("f:/py/tfrecord/train.tfrecords") # 建立一个writer
for index, name in enumerate(classes):
    class_path = cwd + name + "/"           # 构建文件路径
    for img_name in os.listdir(class_path): # 遍历目录下的文件
        img_path = class_path + img_name     # 构建具体每一张图片的路径
        image = cv2.imread(img_path)        # 读取图片

        # 获取图片的宽，高和通道数
        img_w = image.shape[0]
        img_h = image.shape[1]
        img_c = image.shape[2]

        # tf读取图片
        img = tf.read_file(img_path)
        img = tf.image.decode_jpeg(img)

        # img = tf.image.resize_images(img,(224, 224)) 改变大小
        img_raw = sess.run(tf.cast(img,tf.uint8)).tostring()              #将图片转化为原生bytes
        

        label = name.encode('utf-8')  #将标签转化为bytes
        '''
        以下是Example类的常用固定格式，但要注意第一个features有s，对应的是tf.train.Features
        tf.train.Features里的feature是没有s的，bytes_list对应的是tf.train.BytesList，
        int64_list对应的是tf.train.Int64List，输入的value的格式也要一致，可输入的格式有int,float,bytes
        label和img_raw的格式是bytes,宽、高、通道数的格式是int
        '''
        example = tf.train.Example(features=tf.train.Features(feature={
            "label": tf.train.Feature(bytes_list=tf.train.BytesList(value=[label])),
            'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw])),
            'img_w': tf.train.Feature(int64_list=tf.train.Int64List(value=[img_w])),
            'img_h': tf.train.Feature(int64_list=tf.train.Int64List(value=[img_h])),
            'img_c': tf.train.Feature(int64_list=tf.train.Int64List(value=[img_c]))
        }))
        writer.write(example.SerializeToString())  #序列化为字符串
        
writer.close()