将数据转为tfrecord格式

假设emo文件夹下，有1,2,3,4等文件夹，每个文件夹代表一个类别

 import tensorflow as tf
 from PIL import Image
 from glob import glob
 import os
 import progressbar
 import time
 
 class TFRecord():
     def __init__(self, path=None, tfrecord_file=None):
         self.path = path
         self.tfrecord_file = tfrecord_file
 
     def _convert_image(self, idx, img_path, is_training=True):
         label = idx
 
         with tf.gfile.FastGFile(img_path, 'rb') as fid:
             img_str = fid.read()
 
         # img_data = Image.open(img_path)
         # img_data = img_data.resize((224, 224))
         # img_str = img_data.tobytes()
 
         file_name = img_path
 
         if is_training:
             feature_key_value_pair = {
                 'file_name': tf.train.Feature(bytes_list=tf.train.BytesList(
                     value=[file_name.encode()])),
                 'img': tf.train.Feature(bytes_list=tf.train.BytesList(
                     value=[img_str])),
                 'label': tf.train.Feature(int64_list=tf.train.Int64List(
                     value=[label]))
             }
         else:
             feature_key_value_pair = {
                 'file_name': tf.train.Feature(bytes_list=tf.train.BytesList(
                     value=[file_name.encode()])),
                 'img': tf.train.Feature(bytes_list=tf.train.BytesList(
                     value=[img_str])),
                 'label': tf.train.Feature(int64_list=tf.train.Int64List(
                     value=[-1]))
             }
 
         feature = tf.train.Features(feature=feature_key_value_pair)
         example = tf.train.Example(features=feature)
         return example
 
     def convert_img_folder(self):
 
         folder_path = self.path
         tfrecord_path = self.tfrecord_file
         img_paths = []
         for file in os.listdir(folder_path):
             for img_path in os.listdir(os.path.join(folder_path, file)):
                 img_paths.append(os.path.join(folder_path, file, img_path))
 
         with tf.python_io.TFRecordWriter(tfrecord_path) as tfwrite:
             widgets = ['[INFO] write image to tfrecord: ', progressbar.Percentage(), " ",
                        progressbar.Bar(), " ", progressbar.ETA()]
             pbar = progressbar.ProgressBar(maxval=len(img_paths), widgets=widgets).start()
 
             cate = [folder_path + '/' + x for x in os.listdir(folder_path) if
                     os.path.isdir(folder_path + '/' + x)]
 
             i = 0
             for idx, folder in enumerate(cate):
                 for img_path in glob(folder + '/*.jpg'):
                     example = self._convert_image(idx, img_path)
                     tfwrite.write(example.SerializeToString())
                     pbar.update(i)
                     i += 1
 
             pbar.finish()
 
     def _extract_fn(self, tfrecord):
         feautres = {
             'file_name': tf.FixedLenFeature([], tf.string),
             'img': tf.FixedLenFeature([], tf.string),
             'label': tf.FixedLenFeature([], tf.int64)
         }
         sample = tf.parse_single_example(tfrecord, feautres)
         img = tf.image.decode_jpeg(sample['img'])
         img = tf.image.resize_images(img, (224, 224), method=1)
         label = sample['label']
         file_name = sample['file_name']
         return [img, label, file_name]
 
     def extract_image(self, shuffle_size, batch_size):
         dataset = tf.data.TFRecordDataset([self.tfrecord_file])
         dataset = dataset.map(self._extract_fn)
         dataset = dataset.shuffle(shuffle_size).batch(batch_size)
         print("---------", type(dataset))
         return dataset
 
 if __name__=='__main__':
 
     # start = time.time()
     # t = GenerateTFRecord('/')
     # t.convert_img_folder('/media/xia/Data/emo', '/media/xia/Data/emo.tfrecord')
     # print("Took %f seconds." % (time.time() - start))
 
     t =TFRecord('/media/xia/Data/emo', '/media/xia/Data/emo.tfrecord')
     t.convert_img_folder()
     dataset = t.extract_image(100, 64)
     for(batch, batch_data) in enumerate(dataset):
         data, label, _ = batch_data
         print(label)
         print(data.shape)

ps: tf.enable_eager_execution()

tf.__version__==1.8.0

参考：https://zhuanlan.zhihu.com/p/30751039

https://lonepatient.top/2018/06/01/tensorflow_tfrecord.html

https://zhuanlan.zhihu.com/p/51186668