Tensorflow学习10-1：验证码识别——生成验证码和tfrecord文件

1 生成验证码图片

这些生成的图片位于同一个文件夹下，而且图片名就是 label 值。

生成代码如下：

# 验证码生成库
from captcha.image import ImageCaptcha # pip install captcha
import numpy as np
from PIL import Image
import random
import sys

number = ['0','1','2','3','4','5','6','7','8','9']
# letter = ['a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z']
# LETTER = ['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']
CAPTCHA_SAVE_DIR = "D:/Tensorflow/captcha/images/"

'''随机生成4个数字的字符串成。
char_set：用于生成的字符list
captcha_size：生成的验证码位数
'''
def random_captcha_text(char_set=number, captcha_size=4):
    # 验证码列表
    captcha_text = []
    for i in range(captcha_size):
        #随机选择
        c = random.choice(char_set)
        #加入验证码列表
        captcha_text.append(c)
    return captcha_text

'''生成字符对应的验证码'''
def gen_captcha_text_and_iamge():
    image = ImageCaptcha()
    #获得随机生成的验证码
    captcha_text = random_captcha_text()
    #把验证码列表转为字符串
    captcha_text = "".join(captcha_text)
    #生成验证码
    captcha = image.generate(captcha_text)
    image.write(captcha_text, CAPTCHA_SAVE_DIR + captcha_text + ".jpg") #写到文件
    
    
#循环生成10000次，但是重复的会被覆盖，所以<10000
num = 10000
if __name__ == "__main__":
    for i in range(num):
        gen_captcha_text_and_iamge()
        sys.stdout.write("\r>> Creating image %d/%d" % (i+1, num))
        sys.stdout.flush()
    sys.stdout.write("\n")
    sys.stdout.flush()
    
    print("Generate finished.")

Creating image 10000/10000
Generate finished.

2 将这些图片转为tfrecord文件

我们生成tfrecord文件用于验证码识别程序的训练和测试，生成好后会产生2个.tfrecord文件

生成tfrecord代码如下：

-# image_to_tfrecord_by_filename.py——把验证码转换成tfrecord文件
#tfrecord文件，底层就是protobuf格式

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

#验证集数量
tf.app.flags.DEFINE_integer("num_validation", 500, 
    "the divisiory number of validation data")

#随机种子
tf.app.flags.DEFINE_integer("random_seed", 7, 
    "random seed")

#图片目录
tf.app.flags.DEFINE_string("dataset_dir", "D:/Tensorflow/captcha/images/", 
    "dir of images and save position")
    
#保存tfrecord目录
tf.app.flags.DEFINE_string("tfrecord_dir", "D:/Tensorflow/captcha/", 
    "dir of tfrecord")


FLAGS = tf.app.flags.FLAGS

#判断tfrecord文件是否存在
def _dataset_exists(dataset_dir):
    for split_name in ["train", "validation"]:
        output_filename = os.path.join(dataset_dir, split_name + ".tfrecord")
        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 message
    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", "validation"]
    
    with tf.Session() as sess:
        #定义tfrecord文件路径
        output_filename = os.path.join(FLAGS.tfrecord_dir, split_name + ".tfrecord")
        with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
            for i, filename in enumerate(filenames):
                try:
                    sys.stdout.write("\r>> Converting image(%s) %d/%d" % (split_name, 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]))

                    #生成tfrecord文件
                    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:", filenames[i])
                    print("Error:",e)
                    print("Skip the pic.\n")
    sys.stdout.write("\n")
    sys.stdout.flush()

def main(_):
    #判断tfrecord文件是否存在
    if _dataset_exists(FLAGS.tfrecord_dir):
        print("tfrecord文件已存在")
    else:
        #获得所有图片以及分类
        photo_filenames = _get_filenames_and_classes(FLAGS.dataset_dir)
        
        #数据切分为训练集和测试集
        random.seed(FLAGS.random_seed)
        random.shuffle(photo_filenames)
        training_filenames = photo_filenames[FLAGS.num_validation:] #500之后的图片作为训练
        validation_filenames = photo_filenames[:FLAGS.num_validation] #0-500的图片作为训练

        #数据转换
        _convert_dataset("train", training_filenames, FLAGS.dataset_dir)
        _convert_dataset("validation", validation_filenames, FLAGS.dataset_dir)
        print("finished.")
                            
if __name__ == "__main__":
    tf.app.run()

Converting image(train) 5858/5858
Converting image(validation) 500/500
finished.