tensorflow学习笔记——验证码识别（二）生成tfrecord文件

通常情况下，我们在训练模型之前会生成数据集中的train，test，val这些含有训练信息的文件。这些文件被散列的存着，它们并不是计算机能直接读取的信息，在大小网络的训练中，文件读取速度的问题尤为明显。因此在数据集制作完成后，需要将数据集转换成tfrecord文件，tfrecord文件使用于tensorflow网络模型训练的二进制数据编码文件，训练时只需一次性加载一个tfrecord文件的即可，编码简单，速度较快，非常适合用于大型训练数据。

import tensorflow as tf
import os
import random
import math
import sys
from PIL import Image
import numpy as np

#验证集数量
_NUM_TEST = 500

#随机种子
_RANDOM_SEED = 0

#数据集路径（路径里'\\'和'/'等效）
DATASET_DIR = "C:\\Users\\45374\\tensorflow\\captcha\\images"

#tfrecord文件存放路径
TFRECORD_DIR = "C:\\Users\\45374\\tensorflow\\captcha"


#判断tfrecord文件是否存在
def _dataset_exists(dataset_dir):
    for split_name in ['train', 'test']:
        output_filename = os.path.join(dataset_dir,split_name + '.tfrecords')
        if not tf.gfile.Exists(output_filename):
            return False
    return True

#获取所有验证码图片
def _get_filenames_and_classes(dataset_dir):
    photo_filenames = []
    for filename in os.listdir(dataset_dir):
        #获取文件路径
        path = os.path.join(dataset_dir, filename)
        photo_filenames.append(path)
    return photo_filenames

def int64_feature(values):
    if not isinstance(values, (tuple, list)):
        values = [values]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=values))

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

def image_to_tfexample(image_data, label0, label1, label2, label3):
    #Abstract base class for protocol messages.
    return tf.train.Example(features=tf.train.Features(feature={
      'image': bytes_feature(image_data),
      'label0': int64_feature(label0),
      'label1': int64_feature(label1),
      'label2': int64_feature(label2),
      'label3': int64_feature(label3),
    }))

#把数据转为TFRecord格式
def _convert_dataset(split_name, filenames, dataset_dir):
    assert split_name in ['train', 'test']

    with tf.Session() as sess:
        #定义tfrecord文件的路径+名字
        output_filename = os.path.join(TFRECORD_DIR,split_name + '.tfrecords')
        with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
            for i,filename in enumerate(filenames):
                try:
                    sys.stdout.write('\r>> Converting image %d/%d' % (i+1, len(filenames)))
                    sys.stdout.flush()

                    #读取图片
                    image_data = Image.open(filename)  
                    #根据模型的结构resize(使之适应于所用模型的输入大小)
                    image_data = image_data.resize((224, 224))
                    #灰度化（验证码识别中色彩对于识别意义不大，因此转化位灰度图像）
                    image_data = np.array(image_data.convert('L'))
                    #将图片转化为bytes
                    image_data = image_data.tobytes()              

                    #获取label（因为我前面的路径中用的'\\'，所以这里只能以'\\'分割）
                    labels = filename.split('\\')[-1][0:4]
                    num_labels = []
                    for j in range(4):
                        num_labels.append(int(labels[j]))
                                            
                    #生成protocol数据类型
                    example = image_to_tfexample(image_data, num_labels[0], num_labels[1], num_labels[2], num_labels[3])
                    tfrecord_writer.write(example.SerializeToString())
                    
                except IOError as e:
                    print('Could not read:',filename)
                    print('Error:',e)
                    print('Skip it\n')
    sys.stdout.write('\n')
    sys.stdout.flush()

#判断tfrecord文件是否存在
if _dataset_exists(TFRECORD_DIR):
    print('tfcecord文件已存在')
else:
    #获得所有图片
    photo_filenames = _get_filenames_and_classes(DATASET_DIR)

    #把数据切分为训练集和测试集,并打乱
    random.seed(_RANDOM_SEED)
    random.shuffle(photo_filenames)
    training_filenames = photo_filenames[_NUM_TEST:]
    testing_filenames = photo_filenames[:_NUM_TEST]

    #数据转换
    _convert_dataset('train', training_filenames, DATASET_DIR)
    _convert_dataset('test', testing_filenames, DATASET_DIR)
    print('生成tfcecord文件')

 最后在对应路径下生成tfrecord文件：