使用 Python 构建验证码识别模型

1. 导入必要的库
python

import numpy as np
import random
import string
import matplotlib.pyplot as plt
from captcha.image import ImageCaptcha
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Reshape, Dense, Dropout, LSTM, BatchNormalization, Lambda
from tensorflow.keras.optimizers import Adam
import tensorflow.keras.backend as K
2. 定义参数和字符集
python

# 定义参数
width, height, n_len, n_class = 170, 80, 4, len(characters) + 1  # 加1是为了空白类别

# 字符集
characters = string.digits + string.ascii_uppercase
3. 生成验证码图像
python

# 生成随机验证码字符串和对应的图像
generator = ImageCaptcha(width=width, height=height)

def generate_random_string():
    return ''.join(random.choice(characters) for _ in range(n_len))

random_str = generate_random_string()
img = generator.generate_image(random_str)

plt.imshow(img)
plt.title(random_str)
plt.show()
4. 定义CTC损失函数
python

def ctc_lambda_func(args):
    y_pred, labels, input_length, label_length = args
    y_pred = y_pred[:, 2:, :]  # 去掉前两个输出，通常无意义
    return K.ctc_batch_cost(labels, y_pred, input_length, label_length)

# 模型输入
input_tensor = Input((width, height, 3))
labels = Input(name='the_labels', shape=[n_len], dtype='float32')
input_length = Input(name='input_length', shape=[1], dtype='int64')
label_length = Input(name='label_length', shape=[1], dtype='int64')

# CTC损失层
loss_out = Lambda(ctc_lambda_func, output_shape=(1,), name='ctc')([x, labels, input_length, label_length])
5. 构建模型架构
python
复制代码
# CNN特征提取
x = input_tensor
x = Lambda(lambda x: (x - 127.5) / 127.5)(x)
for i in range(3):
    for j in range(2):
        x = Conv2D(32 * 2 ** i, 3, kernel_initializer='he_uniform')(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
    x = MaxPooling2D((2, 2))(x)

# Reshape为RNN输入
conv_shape = x.get_shape().as_list()
rnn_length = conv_shape[1]
rnn_dimen = conv_shape[2] * conv_shape[3]
x = Reshape(target_shape=(rnn_length, rnn_dimen))(x)
x = Dense(rnn_size, kernel_initializer='he_uniform')(x)
x = BatchNormalization()(x)
x = Activation('relu')(x)
x = Dropout(0.2)(x)

# 双向LSTM层
gru_1 = LSTM(rnn_size, return_sequences=True, kernel_initializer='he_uniform', name='gru1')(x)
gru_1b = LSTM(rnn_size, return_sequences=True, kernel_initializer='he_uniform', go_backwards=True, name='gru1_b')(x)
x = tf.keras.layers.Concatenate()([gru_1, gru_1b])

gru_2 = LSTM(rnn_size, return_sequences=True, kernel_initializer='he_uniform', name='gru2')(x)
gru_2b = LSTM(rnn_size, return_sequences=True, kernel_initializer='he_uniform', go_backwards=True, name='gru2_b')(x)
x = tf.keras.layers.Concatenate()([gru_2, gru_2b])

x = Dropout(0.2)(x)
x = Dense(n_class, activation='softmax')(x)

# 定义基础模型和编译
base_model = Model(inputs=input_tensor, outputs=x)
model = Model(inputs=[input_tensor, labels, input_length, label_length], outputs=[loss_out])
model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer=Adam())
6. 数据生成器用于训练
python
复制代码
def gen(batch_size=128):
    X = np.zeros((batch_size, width, height, 3), dtype=np.uint8)
    y = np.zeros((batch_size, n_len), dtype=np.uint8)
    generator = ImageCaptcha(width=width, height=height)
    while True:
        for i in range(batch_size):
            random_str = ''.join([random.choice(characters) for j in range(n_len)])
            X[i] = np.array(generator.generate_image(random_str)).transpose(1, 0, 2)
            y[i] = [characters.find(x) for x in random_str]
        yield [X, y, np.ones(batch_size) * rnn_length, np.ones(batch_size) * n_len], np.ones(batch_size)
7. 训练和评估
python
复制代码
h = model.fit_generator(gen(128), steps_per_epoch=400, epochs=20,
                        callbacks=[evaluator],
                        validation_data=gen(128), validation_steps=20)

# 评估模型
def evaluate(batch_size=128, steps=10):
    batch_acc = 0
    generator = gen(batch_size)
    for i in range(steps):
        [X_test, y_test, _, _], _ = next(generator)
        y_pred = base_model.predict(X_test)
        shape = y_pred[:, 2:, :].shape
        ctc_decode = K.ctc_decode(y_pred[:, 2:, :], input_length=np.ones(shape[0]) * shape[1])[0][0]
        out = K.get_value(ctc_decode)[:, :n_len]
        if out.shape[1] == n_len:
            batch_acc += (y_test == out).all(axis=1).mean()
    return batch_acc / steps