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