解决图像数据转为TFrecord格式后文件太大的问题

解决图像数据转为TFrecord格式后文件太大的问题

上次用tensorflow手册上的方法，将我的图片数据转为TFrecord格式后，发现数据大小是原来数据大小的十倍，严重占用了我的内存，于是就查了一些资料改良一下图片转TFrecord的方式。
具体代码如下，转化后的图像大小比原来的图像数据集还小。而且读取速度快了很多。

import os
import tensorflow as tf
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import sys
import urllib

path = '../../data/Imgnet2015_32/'
def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def img_to_feature_2(img_path):
    with tf.gfile.GFile(img_path, 'rb') as fid:
        encoded_jpg = fid.read()
        feature = tf.train.Features(
            feature={'img': _bytes_feature(encoded_jpg),
                    'format': _bytes_feature('jpeg'.encode('utf8'))
            })
        return feature

def creat_tfrecord(class_path):
    writer_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)
    writer = tf.python_io.TFRecordWriter(path="Imgnet2015_32.tfrecords", options=writer_options)
    for img_name in os.listdir(class_path):  # os.listdir 返回指定目录下所有的文件和目录
        img_path = class_path + img_name  # 每一个图片的地址
        features = img_to_feature_2(img_path)
        example = tf.train.Example(features=features)
        writer.write(example.SerializeToString())  # 序列化为字符串
    writer.close()


slim_example_decoder = tf.contrib.slim.tfexample_decoder

def get_read_data(filename, size):
    filename_queue = tf.train.string_input_producer([filename])
    tfrecord_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)
    reader = tf.TFRecordReader(options=tfrecord_options)
    _, serialized_example = reader.read(filename_queue)
    keys_to_features = {
        'img': tf.FixedLenFeature((), tf.string, default_value=''),
        'format': tf.FixedLenFeature((), tf.string, default_value='jpeg')
    }
    items_to_handlers = {
        'img': slim_example_decoder.Image(image_key='img', format_key='format', channels=3),
        'format': slim_example_decoder.Tensor('format')
    }
    serialized_example = tf.reshape(serialized_example, shape=[])
    decoder = slim_example_decoder.TFExampleDecoder(keys_to_features, items_to_handlers)
    keys = decoder.list_items()
    tensors = decoder.decode(serialized_example, items=keys)
    tensor_dict = dict(zip(keys, tensors))
    # format = tensor_dict['format']
    img_tensor = tensor_dict['img']
    img_tensor = tf.reshape(img_tensor, [size, size, 3])
    img_float = tf.cast(img_tensor, tf.float32)
    return img_tensor, img_float
# 用tf.cast转为tf.float格式后保存无法还原成原图像
def run(input_path, size):
    img_tensor, img_float = get_read_data(input_path, size)
    img_batch = tf.train.shuffle_batch([img_float], batch_size=8, capacity=64, min_after_dequeue=32)
    init_op = tf.initialize_all_variables()
    with tf.Session() as sess:
        sess.run(init_op)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        for i in range(10):
            img_b, img = sess.run([img_batch, img_tensor])
            image = Image.fromarray(img, 'RGB')  # 这里Image是之前提到的
            image.save('./output/' + str(i) + '_''Label_' + '.jpg')  # 存下图片
            # print ('img:', img)
            print('img shape:', img_b.shape)
            print('img shape:', img.shape)
            # print('img format:', format_name)

        coord.request_stop()
        coord.join(threads)

creat_tfrecord(path)
# run('test32.tfrecords',32)