TensorFlow的学习之路--TfTFrecords的制作

先给出一个图像压缩成数据集的程序

import numpy as np
import pandas as pd
import tensorflow as tf
from tqdm import tqdm
from PIL import Image
import os
import io
import argparse


parser = argparse.ArgumentParser()
parser.add_argument(
    '--train_dir', type=str,
    default="E:/Python/tensorflow/猫狗识别/trainandtest/train",
    help='A path to a folder with training data.'
)
#parser.add_argument(
#    '--val_dir', type=str,
#    default="E:/liangbx/python/flower_class/flower_test",
#    help='A path to a folder with validation data.'
#)
parser.add_argument(
    '--save_dir', type=str,
    default='E:/Python/tensorflow/猫狗识别/trainandtest',
    help='A path to a folder where to save results.'
)
args = parser.parse_args()


"""The purpose of this script is
to convert image dataset that looks like:
    class1/image1.jpg
    class1/image44.jpg
    class1/image546.jpg
    ...
    class6/image55.jpg
    class6/image12.jpg
    class6/image76.jpg
    ...
to tfrecords format.

1. It assumes that each folder is separate class and
that the number of classes equals to the number of folders.

2. Also it assumes that validation and training folders
have the same subfolders (the same classes).

3. Additionally it outputs 'class_encoder.npy' file
that contains dictionary: folder_name -> class_index (integer).
"""


def main():
    encoder = create_encoder(args.train_dir)
    # now you can get a folder's name from a class index

    np.save(os.path.join(args.save_dir, 'class_encoder.npy'), encoder)
    convert(args.train_dir, encoder, os.path.join(args.save_dir, 'train.tfrecords'))
    #convert(args.val_dir, encoder, os.path.join(args.save_dir, 'val.tfrecords'))

    print('\nCreated two tfrecords files:')
    print(os.path.join(args.save_dir, 'train.tfrecords'))
    #print(os.path.join(args.save_dir, 'val.tfrecords'))


def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


# here you can also just use `return array.tostring()`
# but it will make tfrecords files a lot larger and
# you will need to change the input pipeline
def to_bytes(array):
    image = Image.fromarray(array)
    tmp = io.BytesIO()
    image.save(tmp, format='jpeg')
    return tmp.getvalue()


def convert(folder, encoder, tfrecords_filename):
    """Convert a folder with directories of images to tfrecords format.

    Arguments:
        folder: A path to a folder where directories with images are.
        encoder: A dict, folder_name -> integer.
        tfrecords_filename: A path where to save tfrecords file.
    """

    images_metadata = collect_metadata(folder, encoder)
    writer = tf.python_io.TFRecordWriter(tfrecords_filename)

    for _, row in tqdm(images_metadata.iterrows()):

        file_path = os.path.join(folder, row.img_path)

        # read an image
        image = Image.open(file_path)

        # convert to an array
        array = np.asarray(image, dtype='uint8')

        # some images are grayscale
        if array.shape[-1] != 3:
            array = np.stack([array, array, array], axis=2)

        # get class of the image
        target = int(row.class_number)

        feature = {
            'image': _bytes_feature(to_bytes(array)),
            'target': _int64_feature(target),
        }

        example = tf.train.Example(features=tf.train.Features(feature=feature))
        writer.write(example.SerializeToString())

    writer.close()


def create_encoder(folder):
    """Encode directories in the folder with integer values.
    Values are in the range 0..(n_directories - 1).

    Arguments:
        folder: A path to a folder where directories with images are.
            Each directory - separate class.
    Returns:
        A dict.
    """
    classes = os.listdir(folder)
    encoder = {n: i for i, n in enumerate(classes)}
    return encoder


def collect_metadata(folder, encoder):
    """Collect paths to images. Collect their classes.
    All paths must be with respect to 'folder'.

    Arguments:
        folder: A path to a folder where directories with images are.
            Each directory - separate class.
        encoder: A dict, folder_name -> integer.
    Returns:
        A pandas dataframe.
    """

    subdirs = list(os.walk(folder))[1:]
    metadata = []

    for dir_path, _, files in subdirs:
        dir_name = dir_path.split('\\')[-1]  #
        for file_name in files:
            image_metadata = [dir_name, os.path.join(dir_name, file_name)]
            metadata.append(image_metadata)

    M = pd.DataFrame(metadata)
    M.columns = ['class_name', 'img_path']

    # encode folder names by integers 
    M['class_number'] = M.class_name.apply(lambda x: encoder[x])


    # shuffle the dataframe
    M = M.sample(frac=1).reset_index(drop=True)

    return M


if __name__ == '__main__':
    main()

下面这个是解压的程序

from PIL import Image
import tensorflow as tf

SHUFFLE_BUFFER_SIZE = 10000
PREFETCH_BUFFER_SIZE = 1000
NUM_THREADS = 4
BATCH_SIZE = 32
N_EPOCH = 40
cwd='E:\\Python\\tensorflow\\猫狗识别\\trainandtest\\train\\'
def _parse_and_preprocess(example_proto, image_size, augmentation=False):

    features = {
        'image': tf.FixedLenFeature([], tf.string),
        'target': tf.FixedLenFeature([], tf.int64)
    }
    parsed_features = tf.parse_single_example(example_proto, features)

    image = tf.image.decode_jpeg(parsed_features['image'], channels=3)
    image = tf.image.convert_image_dtype(image, tf.float32)
    target = parsed_features['target']

    if augmentation:
        image = tf.image.random_flip_left_right(image)
        image = tf.image.random_brightness(image, 0.15)
        image = tf.image.random_contrast(image, 0.8, 1.25)
        image = tf.image.random_hue(image, 0.1)
        image = tf.image.random_saturation(image, 0.8, 1.25)
#        image = tf.random_crop(image, [image_size, image_size, 3])
        image = tf.image.resize_images(images=image, size=[image_size, image_size])
    else:
#        image = tf.image.resize_image_with_crop_or_pad(image, image_size, image_size)
        image = tf.image.resize_images(images=image, size=[image_size, image_size])
        

    image = tf.clip_by_value(image, 0.0, 1.0)
    return image, target

def get_data(file_path, image_size, is_augmentation):
    
    dataset = tf.data.TFRecordDataset(file_path).repeat() #
    dataset = dataset.map(
        lambda x: _parse_and_preprocess(x, image_size, augmentation=is_augmentation),
        num_parallel_calls=NUM_THREADS
    ).prefetch(PREFETCH_BUFFER_SIZE)
    
    dataset = dataset.shuffle(buffer_size=SHUFFLE_BUFFER_SIZE)
    dataset = dataset.batch(BATCH_SIZE)
    
    iterator = dataset.make_one_shot_iterator()
    image_batch, label_batch = iterator.get_next()
    
    return image_batch, label_batch    

得到的image_batch以及label_batch就可以用在训练上了，此程序再session中，不需要再增加

        coord=tf.train.Coordinator()

        threads= tf.train.start_queue_runners(coord=coord)  ，在猫狗识别中，此训练集能产生较为随机的数字集，方便训练。

第二种程序较为简单：网上很多例子都一样：

import os 
import tensorflow as tf 
from PIL import Image  #注意Image,后面会用到
import matplotlib.pyplot as plt 
import numpy as np
cwd='E:\\Python\\tensorflow\\猫狗识别\\trainandtest\\train\\'
classes={'dog','cat'} #人为 设定 2 类
writer= tf.python_io.TFRecordWriter("dog_train.tfrecords") #要生成的文件 
for index,name in enumerate(classes):
    class_path=cwd+name+'\\'
    for img_name in os.listdir(class_path): 
        img_path=class_path+img_name #每一个图片的地址
      
        img=Image.open(img_path)
        img= img.resize((128,128))
        img_raw=img.tobytes()#将图片转化为二进制格式
        example = tf.train.Example(features=tf.train.Features(feature={
             "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[index])),
             'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
             })) #example对象对label和image数据进行封装
        writer.write(example.SerializeToString())  #序列化为字符串 
writer.close() 
def read_and_decode(filename): # 读入dog_train.tfrecords
    filename_queue = tf.train.string_input_producer([filename])#生成一个queue队列
  
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)#返回文件名和文件
    features = tf.parse_single_example(serialized_example,
                                        features={
                                            'label': tf.FixedLenFeature([], tf.int64),
                                            'img_raw' : tf.FixedLenFeature([], tf.string),
                                        })#将image数据和label取出来
  
    img = tf.decode_raw(features['img_raw'], tf.uint8)
    img = tf.reshape(img, [128, 128, 3])  #reshape为128*128的3通道图片
    img = tf.cast(img, tf.float32) * (1. / 255) - 0.5 #在流中抛出img张量
    label = tf.cast(features['label'], tf.int32) #在流中抛出label张量
    return img, label

filename_queue = tf.train.string_input_producer(["dog_train.tfrecords"]) #读入流中
reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue)   #返回文件名和文件
features = tf.parse_single_example(serialized_example,
                                   features={
                                       'label': tf.FixedLenFeature([], tf.int64),
                                       'img_raw' : tf.FixedLenFeature([], tf.string),
                                   })  #取出包含image和label的feature对象
image = tf.decode_raw(features['img_raw'], tf.uint8)
image = tf.reshape(image, [128, 128, 3])
label = tf.cast(features['label'], tf.float32)


img_batch, label_batch = tf.train.shuffle_batch([image, label],
                                                    batch_size=32, capacity=8000,
                                                    min_after_dequeue=2000)
with tf.Session() as sess: #开始一个会话
    init_op = tf.initialize_all_variables()
    sess.run(init_op)
    coord=tf.train.Coordinator()
    threads= tf.train.start_queue_runners(coord=coord)
    for i in range(20000):
        example, l = sess.run([image,label])#在会话中取出image和label
        img=Image.fromarray(example, 'RGB')#这里Image是之前提到的
        img.save(cwd+str(i)+'_''Label_'+str(l)+'.jpg')#存下图片
        print(l)
    coord.request_stop()
    coord.join(threads)

这个程序在运行过程中，不管我怎么改变tf.train.shuffle_batch的参数，所得到的数据，不是很随机，在一个batch中，很多均为1111或者000，具体原因待查找，但得到的图像和label都是正常