使用tensorflow搭建自己的验证码识别系统

学习tensorflow有一段时间了，想做点东西来练一下手。为了更有意思点，下面将搭建一个简单的验证码识别系统。

准备验证码数据

下面将生成一万张四位英文字母的验证码，验证码的大小是100 * 30的图片，只包含大写的英文字母，并将目标值保存到csv文件。

import random
import pandas as pd
from PIL import Image, ImageDraw, ImageFont


def generate_captcha(filename, format):
    """
    生成四位验证码
    :param filename: 要保存的文件名
    :param format: 保存图片格式
    :return: 验证码的值
    """
    # 定义使用Image类实例化一个长为100px,宽为30px,基于RGB的(255,255,255)颜色的图片
    img = Image.new(mode="RGB", size=(100, 30), color=(255, 255, 255))
    # 实例化一支画笔
    draw = ImageDraw.Draw(img, mode="RGB")
    # 定义要使用的字体
    font = ImageFont.truetype("arial", 28)

    result = ""

    for i in range(4):
        # 每循环一次,从a到z中随机生成一个字母
        # 65到90为字母的ASCII码,使用chr把生成的ASCII码转换成字符
        # str把生成的数字转换成字符串
        char = random.choice([chr(random.randint(65, 90))])
        result += char

        # 每循环一次重新生成随机颜色
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

        # 把生成的字母或数字添加到图片上
        # 图片长度为100px,要生成4个数字或字母则每添加一个,其位置就要向后移动24px
        draw.text([i * 24 + 3, 0], char, color, font=font)

    # 保存生成的文件
    with open(filename, "wb") as f:
        img.save(f, format=format)

    return result


if __name__ == "__main__":

    data = []

    # 生成10000张验证码图片，并将目标值存入csv文件
    for j in range(10000):
        val = generate_captcha("./pics/{}.png".format(j), "png")
        data.append([val])

    # 将验证码的值保存到csv文件
    df = pd.DataFrame(data, columns=['label'])
    df.to_csv('./pics/data.csv', header=False)

生成的验证码图片是这样子的，如下：

csv文件内容：

0,EFGQ
1,ZDKO
2,UWLD
3,CPDH
....

保存为tfrecords文件

上面生成的图片和其目标值是分开的，在进行训练时不太方便（训练时每次都要单独的读取图片和特征值）。保存为tfrecords文件，在训练时会方便很多，读取出来的每条记录既有图片特征值又有目标值。

import tensorflow as tf
import os
import numpy as np


class CaptchaInput(object):

    def __init__(self, captcha_dir, letter, tfrecords_dir):
        """
        :param captcha_dir: 验证码路径
        :param letter: 验证码字符种类
        :param tfrecords_dir: tfrecords文件保存的目录
        """
        self.captcha_dir = captcha_dir
        self.letter = letter
        self.tfrecords_dir = tfrecords_dir

        # 列出图片文件，并进行排序
        self.file_list = os.listdir(self.captcha_dir)
        self.file_list = [i for i in self.file_list if i.endswith(".png")]
        self.file_list.sort(key=lambda x: int(x[0:-4]))
        self.file_list = [os.path.join(self.captcha_dir, i) for i in self.file_list]

        # 标签文件路径
        self.labels_path = os.path.join(self.captcha_dir, "data.csv")

    def read_captcha_image(self):
        """读取验证码图片数据"""
        # 构造文件队列
        file_queue = tf.train.string_input_producer(self.file_list, shuffle=False)

        # 构建阅读器
        reader = tf.WholeFileReader()

        # 读取图片内容
        key, value = reader.read(file_queue)
        # 解码图片
        image = tf.image.decode_png(value)
        image.set_shape([30, 100, 3])

        # 批量读取
        image_batch = tf.train.batch([image], batch_size=len(self.file_list),
                                     num_threads=1, capacity=len(self.file_list))
        return image_batch

    def read_captcha_label(self):
        """读取 验证码标签数据"""
        # 构造文件队列
        file_queue = tf.train.string_input_producer([self.labels_path], shuffle=False)

        # 构建文件阅读器
        reader = tf.TextLineReader()

        # 读取标签内容
        key, value = reader.read(file_queue)

        records = [[0], [""]]
        index, label = tf.decode_csv(value, record_defaults=records)

        # 批量读取
        label_batch = tf.train.batch([label], batch_size=len(self.file_list),
                                     num_threads=1, capacity=len(self.file_list))

        return label_batch

    def process_labels(self, labels):
        """将标签字符转换成数字张量"""
        # 构建字符索引
        num_letter_dict = dict(enumerate(list(self.letter)))
        letter_num_dict = dict(zip(num_letter_dict.values(), num_letter_dict.keys()))

        ret = []

        for label in labels:
            arr = [letter_num_dict[i] for i in label.decode("utf-8")]
            ret.append(arr)

        return np.array(ret)

    def write_to_tfrecords(self, images, labels):
        """
        将图片和标签写入到tfrecords文件中
        :param images: 特征值
        :param labels: 目标值
        :return:
        """
        # labels = tf.cast(labels, tf.uint8)
        # images = tf.cast(images, tf.uint8)

        # 建立存储文件
        fw = tf.python_io.TFRecordWriter(self.tfrecords_dir)
        for i in range(len(self.file_list)):
            # images[i]为numpy.ndarray
            image_bytes = images[i].tobytes()
            # labels[i]为numpy.ndarray
            label_bytes = labels[i].tobytes()

            example = tf.train.Example(features=tf.train.Features(feature={
                "image": tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_bytes])),
                "label": tf.train.Feature(bytes_list=tf.train.BytesList(value=[label_bytes]))
            }))

            print("保存第%d张图片" % (i, ))

            fw.write(example.SerializeToString())

        # 关闭
        fw.close()

    def execute(self):
        image_batch = self.read_captcha_image()
        label_batch = self.read_captcha_label()

        with tf.Session() as sess:

            coord = tf.train.Coordinator()

            threads = tf.train.start_queue_runners(sess, coord=coord)

            # [b'EFGQ' b'ZDKO' b'UWLD' ... b'TKPD' b'ZZEU' b'ATYA']
            labels = sess.run(label_batch)

            # labels为numpy.ndarray
            labels = self.process_labels(labels)
            # images为numpy.ndarray
            images = sess.run(image_batch)

            self.write_to_tfrecords(images, labels)

            coord.request_stop()
            coord.join(threads)


FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_string("captcha_dir", "./pics", "验证码图片路径")
tf.app.flags.DEFINE_string("letter", "ABCDEFGHIJKLMNOPQRSTUVWXYZ", "验证码字符种类")
tf.app.flags.DEFINE_string("tfrecords_dir", "./tfrecords/captcha.tfrecords", "验证码tfrecords文件")


if __name__ == "__main__":
    c = CaptchaInput(FLAGS.captcha_dir, FLAGS.letter, FLAGS.tfrecords_dir)
    c.execute()

需要注意：

• os.listdir返回的文件名称的顺序是按照ascii表的顺序（1.png, 10.png...）需要对其进行排序
• 使用tensorflow读取图片和标签文件时，需要加上shuffle=False，避免文件乱序了，图片和目标值对应不上。

验证码训练

import tensorflow as tf


FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_string("captcha_dir", "./tfrecords/captcha.tfrecords", "验证码数据文件")
tf.app.flags.DEFINE_integer("batch_size", 100, "每批次训练样本数")


def read_and_decode():
    """读取验证码数据
    :return image_batch, label_batch
    """
    # 文件队列
    file_queue = tf.train.string_input_producer([FLAGS.captcha_dir])

    # 文件读取器
    reader = tf.TFRecordReader()

    # 读取内容
    key, value = reader.read(file_queue)
    # 解析tfrecords
    features = tf.parse_single_example(value, features={
        "image": tf.FixedLenFeature([], tf.string),
        "label": tf.FixedLenFeature([], tf.string)
    })
    # 解码
    image = tf.decode_raw(features["image"], tf.uint8)
    label = tf.decode_raw(features["label"], tf.uint8)
    # print(image, label)

    # 改变形状
    image_reshape = tf.reshape(image, [30, 100, 3])
    label_reshape = tf.reshape(label, [4])
    # print(image_reshape, label_reshape)

    # 批处理
    image_batch, label_batch = tf.train.batch([image_reshape, label_reshape],
                                              batch_size=FLAGS.batch_size, num_threads=1, capacity=FLAGS.batch_size)
    return image_batch, label_batch


def weight_variables(shape):
    """权重初始化函数"""
    w = tf.Variable(tf.random_normal(shape=shape, mean=0.0, stddev=1.0))
    return w


def bias_variables(shape):
    """偏置初始化函数"""
    b = tf.Variable(tf.constant(0.0, shape=shape))
    return b


def fc_model(image):
    """全连接模型"""
    with tf.variable_scope("fc_model"):
        image_reshape = tf.reshape(image, [-1, 30 * 100 * 3])

        # 随机初始化权重和偏重
        weights = weight_variables([30 * 100 * 3, 4 * 26])
        bias = bias_variables([4 * 26])

        # 全连接计算
        y_predict = tf.matmul(tf.cast(image_reshape, tf.float32), weights) + bias

    return y_predict


def label_to_onehot(label):
    """目标值转换成one-hot编码"""
    label_onehot = tf.one_hot(label, depth=26, on_value=1.0, axis=2)
    return label_onehot


def captcharec():
    """验证码识别"""
    image_batch, label_batch = read_and_decode()
    # [100, 104]
    y_predict = fc_model(image_batch)

    y_true = label_to_onehot(label_batch)

    # softmax计算，交叉熵损失计算
    with tf.variable_scope("soft_cross"):
        loss = tf.nn.softmax_cross_entropy_with_logits(
            labels=tf.reshape(y_true, [-1, 4 * 26]),
            logits=y_predict
        )

    # 梯度下降损失优化
    with tf.variable_scope("optimizer"):
        train_op = tf.train.GradientDescentOptimizer(0.01).minimize(loss)

    # 准确率
    with tf.variable_scope("acc"):
        equal_list = tf.equal(tf.argmax(y_true, 2), tf.argmax(tf.reshape(y_predict, [-1, 4, 26]), 2))
        accuracy = tf.reduce_mean(tf.cast(equal_list, tf.float32))

    init_op = tf.global_variables_initializer()

    with tf.Session() as sess:

        sess.run(init_op)

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess, coord=coord)

        for i in range(3000):
            sess.run(train_op)

            print("第%d次训练的准确率为：%f" % (i, accuracy.eval()))

        coord.request_stop()
        coord.join(threads)


if __name__ == '__main__':
    captcharec()