将数据转为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