tensorfolw_训练验证码

使用captcha.image.Image 生成随机验证码，随机生成的验证码为0到9的数字，验证码有4位数字组成,这是一个自己生成验证码，自己不断训练的模型

使用三层卷积层，三层池化层，二层全连接层来进行组合
导入库

import numpy as np  
import tensorflow.compat.v1 as tf
from captcha.image import ImageCaptcha
import numpy as np  
import matplotlib.pyplot as plt  
from PIL import Image  
import random   

1、定义生成随机验证码图片

number = ['0','1','2','3','4','5','6','7','8','9']
# alphabet = ['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']
# ALPHABET = ['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']

def random_captcha_text(char_set=number, captha_size=4):
    captha_texts = []
    for i in range(captha_size):
        # 随机抽取数字，添加到列表中
        captha_texts.append(random.choice(char_set))
    return captha_texts

def gen_captcha_text_and_image():
    image = ImageCaptcha()
    captcha_texts = random_captcha_text()
    # 列表转换为字符串
    captcha_texts = ''.join(captcha_texts)
    # 产生图片
    captcha = image.generate(captcha_texts)
    captcha_image = Image.open(captcha)
    captcha_image = np.array(captcha_image)
    # 返回字符串和图片
    return captcha_texts, captcha_image

2、 生成训练样本

# 把彩图转换为灰度图
def convert2gray(image):
    if len(image.shape)> 2:
        grey = np.mean(image, -1)
        return grey
    else:
        return image
# 把文本转换为可用的标签维度是40
def text2vec(text):
    text_len = len(text)
    int(text[0])
    if text_len > MAX_CAPTCHA:
        raise ValueError('验证码最长4个字符')
    vec = np.zeros(MAX_CAPTCHA*CHAR_SET_LEN)
    for index, c in enumerate(text):

        now_index = index * CHAR_SET_LEN + int(c.strip())
        vec[now_index] = 1
    return vec
# 生成训练样本
def get_next_batch(batch_size=128):
    batch_x = np.zeros([batch_size, IMAGE_HEIGHT*IMAGE_WEIGHT])
    batch_y = np.zeros([batch_size, MAX_CAPTCHA*CHAR_SET_LEN])
    # 有时候生成的图像大小不是(60, 160, 3), 重新生成
    def wrap_gen_captcha_text_and_image():
        text, image = gen_captcha_text_and_image()
        while True:
            if image.shape == (60, 160, 3):
                return text, image

    for i in range(batch_size):
        text, image = wrap_gen_captcha_text_and_image()
        image = convert2gray(image)
        # 转换成的一维的灰度图，使得其范围为(0, 1)
        batch_x[i, :] = image.flatten() / 255  # (image.flatten()-128)/128  mean为0
       # 把输入的文本转换为标签类型
        batch_y[i, :] = text2vec(text)

    return batch_x, batch_y

3、 定义CNN，这里的CNN为3层卷积，3层池化， 2层全连接

# 定义CNN
def crack_captcha_cnn(w_alpha=0.01, b_alpha=0.1):
    # [-1, IMAGE_HEIGHT, IMAGE_WEIGHT, 1] -1表示batch_size,1表示样本深度，也就是RGB通道的个数 
    x = tf.reshape(X, [-1, IMAGE_HEIGHT, IMAGE_WEIGHT, 1])

    # 创建w_c1和b_c1的初始化变量
    w_c1 = tf.Variable(w_alpha*tf.random_normal([3, 3, 1, 32]))
    b_c1 = tf.Variable(b_alpha*tf.random_normal([32]))
    # 进行卷积操作
    conv1 = tf.nn.relu(tf.nn.bias_add(tf.nn.conv2d(x, w_c1, strides=(1, 1, 1, 1), padding='SAME'), b_c1))
    # 进行池化操作
    conv1 = tf.nn.max_pool(conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    conv1 = tf.nn.dropout(conv1, keep_prob)

    w_c2 = tf.Variable(w_alpha * tf.random_normal([3, 3, 32, 64]))
    b_c2 = tf.Variable(b_alpha * tf.random_normal([64]))
    conv2 = tf.nn.relu(tf.nn.bias_add(tf.nn.conv2d(conv1, w_c2, strides=(1, 1, 1, 1), padding='SAME'), b_c2))
    conv2 = tf.nn.max_pool(conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    conv2 = tf.nn.dropout(conv2, keep_prob)

    w_c3 = tf.Variable(w_alpha * tf.random_normal([3, 3, 64, 64]))
    b_c3 = tf.Variable(b_alpha * tf.random_normal([64]))
    conv3 = tf.nn.relu(tf.nn.bias_add(tf.nn.conv2d(conv2, w_c3, strides=(1, 1, 1, 1), padding='SAME'), b_c3))
    conv3 = tf.nn.max_pool(conv3, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    conv3 = tf.nn.dropout(conv3, keep_prob)

    # 第一个全连接层
    #8*20*64表示conv3的维度， 60/2/2/2 = 8 160/2/2/2=20 
    w_d = tf.Variable(w_alpha * tf.random_normal([8 * 20 * 64, 1024]))
    b_d = tf.Variable(b_alpha * tf.random_normal([1024]))
    dense = tf.reshape(conv3, [-1, w_d.get_shape().as_list()[0]])
    print(tf.matmul(dense, w_d).shape, b_d.shape)
    dense = tf.nn.relu(tf.add(tf.matmul(dense, w_d), b_d))
    dense = tf.nn.dropout(dense, keep_prob)

    # 第二个全连接层， 不需要激活层
    w_out = tf.Variable(w_alpha * tf.random_normal([1024, MAX_CAPTCHA*CHAR_SET_LEN]))
    b_out = tf.Variable(b_alpha * tf.random_normal([MAX_CAPTCHA*CHAR_SET_LEN]))

    out = tf.add(tf.matmul(dense, w_out), b_out)
    return out

4、 定义训练CNN函数

def train_crack_captcha_cnn():
    output = crack_captcha_cnn()
    loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=output, labels=Y))
    optimizer = tf.train.AdamOptimizer(learning_rate=0.001).minimize(loss)
    predict = tf.reshape(output, [-1, MAX_CAPTCHA, CHAR_SET_LEN])
    max_idx_p = tf.argmax(predict, 2)
    max_idx_l = tf.argmax(tf.reshape(Y, [-1, MAX_CAPTCHA, CHAR_SET_LEN]), 2)
    correct_pred = tf.equal(max_idx_p, max_idx_l)
    accr = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

    saver = tf.train.Saver()
    with tf.Session() as sess:
        #变量初始化
        sess.run(tf.global_variables_initializer())
        step = 0
        # 让它一直都训练直到精度大于0.5
        while True:
            # 生成64个样本
            batch_x, batch_y = get_next_batch(batch_size=64)
            __, _loss = sess.run([optimizer, loss], feed_dict={X:batch_x, Y:batch_y, keep_prob: 0.75})
            print(step, _loss)

            # 每一百次循环计算一次返回值
            if step%100 == 0 :
                batch_text_x, batch_text_y = get_next_batch(batch_size=128)
                acc = sess.run(accr, feed_dict={X:batch_text_x, Y:batch_text_y, keep_prob:1.})
                print(acc)
                # 如果准确率大于0.5就保存模型
                if acc > 0.5:
                    saver.save(sess, '.model/crack_captcha/model')
                    break
            step += 1

5、定义训练好后的预测模型

# 用于训练好后的模型进行预测
def crack_captcha(captcha_image):
    
    output = crack_captcha_cnn()
    # 初始化保存数据
    saver = tf.train.Saver()
    with tf.Session() as sess:
        # 重新加载sess
        saver.restore(sess, '.model/crack_captcha/model')
        
        predict = tf.reshape(output, [-1, MAX_CAPTCHA, CHAR_SET_LEN], 2)
        # 获得CNN之后的结果
        text_list = sess.run(predict, feed_dict={X:[captcha_image], keep_prob:1})
        # 让输出结果变成一个列表
        text = text_list[0].tolist()
        return text

6、主要函数用来进行训练，或者测试

if __name__ == '__main__':
    #获得文本和图片
    train = 0
    # 当train=0时进行训练
    if train==0:
        number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
        text, image = gen_captcha_text_and_image()

        IMAGE_HEIGHT = 60
        IMAGE_WEIGHT = 160
        MAX_CAPTCHA = len(text)
        print('验证码文本最长字符数', MAX_CAPTCHA)
        char_set = number
        CHAR_SET_LEN = len(char_set)

        X = tf.placeholder(tf.float32, [None, IMAGE_HEIGHT*IMAGE_WEIGHT])
        Y = tf.placeholder(tf.float32, [None, MAX_CAPTCHA*CHAR_SET_LEN])
        keep_prob = tf.placeholder(tf.float32)

        train_crack_captcha_cnn()
    # 当trian=1时进行测试    
    elif train == 1:
        text, image = gen_captcha_text_and_image()
        # 将模型转换为灰度图以后再进行测试
        image = convert2gray(image)
        image = image.flatten() / 255
        IMAGE_HEIGHT = 60
        IMAGE_WEIGHT = 160
        MAX_CAPTCHA = len(text)
        print('验证码文本最长字符数', MAX_CAPTCHA)
        char_set = number
        CHAR_SET_LEN = len(char_set)
        X = tf.placeholder(tf.float32, [None, IMAGE_HEIGHT * IMAGE_WEIGHT])
        Y = tf.placeholder(tf.float32, [None, MAX_CAPTCHA * CHAR_SET_LEN])
        keep_prob = tf.placeholder(tf.float32)
        pred_text = crack_captcha(image)
        print('真实值', text, '测试值', pred_text)