keras 视觉 captcha_ocr 案例分析

captcha_ocr 是一个Keras实现的验证码识别示例，目的是识别包含数字和字母的验证码图像，将其转换为文本。该示例主要包括以下步骤：

1. 将原始的验证码图像进行预处理，包括二值化、去噪、裁剪、缩放等操作，以便于后续的模型训练和预测。

2. 使用卷积神经网络（CNN）对预处理后的图像进行特征提取，提取出相应的特征向量。

3. 使用循环神经网络（RNN）对特征向量序列进行处理，以便于将其转化为文本序列。

4. 使用类似于CTC（Connectionist Temporal Classification）的方法对文本序列进行解码，得到识别结果。

模型的评价准则采用准确率和检测率来评估。整个模型的训练过程使用了类似于迁移学习的方法，从已训练过的ResNet网络上预训练模型参数，加快模型训练的速度和效果。

导入依赖

import os
import numpy as np
import matplotlib.pyplot as plt

from pathlib import Path
from collections import Counter

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

加载数据：Captcha Images

参考下面的代码，下载并解压演示用的验证码数据集。

该数据集包含1040个作为 png 图像的 captcha 文件。每个样本的标签是一个字符串，文件的名称（减去文件扩展名）。
我们将把字符串中的每个字符映射到一个整数来训练模型。类似地，我们需要将模型的预测映射回字符串。为此目的我们将分别地维护两个字典，分别用于将字符映射到整数，分别用于将整数映射到字符。

# 数据目录的路径
data_dir = Path("./captcha_images_v2/")

# 获取所有图像的列表
images = sorted(list(map(str, list(data_dir.glob("*.png")))))
labels = [img.split(os.path.sep)[-1].split(".png")[0] for img in images]
characters = set(char for label in labels for char in label)  # 建立用于验证码中字符的字符集合
characters = sorted(list(characters))

print("发现的图像数量：", len(images))
print("发现的标签数量：", len(labels))
print("独特字符的数量：", len(characters))
print("存在的字符：", characters)

# 训练和验证使用的批量大小
batch_size = 16

# 所需图像大小
img_width = 200
img_height = 50

# 用于卷积块降采样的因子。
# 我们将使用两个卷积块，每个块都有一个池化层，
# 池化层将特征图降采样 2 倍。
# 因此总的降采样因子为 4。
downsample_factor = 4

# 数据集中任何验证码的最大长度
max_length = max([len(label) for label in labels])

预处理

• 将图像转换为灰度图像：将 RGB 彩色图像转换为灰度图像，以便更好地处理和降低计算成本。
• 将灰度图像二值化 (binaryzation): 使用 adaptive_threshold 策略将灰度图像转换为黑白二值图像，使得字符和背景更加明显。
• 前景框选 (foreground extraction): 在二值图像中找到字符的前景区域，同时去除掉无关的背景。
• 图像缩放 (image resizing): 将前景区域的图像缩放到预定的大小，以便后续训练和预测过程中输入图像的尺寸保持一致。

# 将字符映射到整数
char_to_num = layers.StringLookup(vocabulary=list(characters), mask_token=None)

# 将整数映射回原始字符
num_to_char = layers.StringLookup(
    vocabulary=char_to_num.get_vocabulary(), mask_token=None, invert=True
)


def split_data(images, labels, train_size=0.9, shuffle=True):
    # 1. 获取数据集的总大小
    size = len(images)
    # 2. 创建索引数组并进行洗牌（如果需要）
    indices = np.arange(size)
    if shuffle:
        np.random.shuffle(indices)
    # 3. 获取训练样本的大小
    train_samples = int(size * train_size)
    # 4. 将数据分成训练集和验证集
    x_train, y_train = images[indices[:train_samples]], labels[indices[:train_samples]]
    x_valid, y_valid = images[indices[train_samples:]], labels[indices[train_samples:]]
    return x_train, x_valid, y_train, y_valid


# 将数据集拆成训练集和验证集
x_train, x_valid, y_train, y_valid = split_data(np.array(images), np.array(labels))


def encode_single_sample(img_path, label):
    # 1. 读取图片
    img = tf.io.read_file(img_path)
    # 2. 解码并转换为灰度图
    img = tf.io.decode_png(img, channels=1)
    # 3. 转为 [0, 1] 范围内的 float32 格式
    img = tf.image.convert_image_dtype(img, tf.float32)
    # 4. 调整大小到所需大小
    img = tf.image.resize(img, [img_height, img_width])
    # 5. 转置图片，因为我们希望时间维度对应于图片的宽度。
    img = tf.transpose(img, perm=[1, 0, 2])
    # 6. 将标签中的字符映射为数字
    label = char_to_num(tf.strings.unicode_split(label, input_encoding="UTF-8"))
    # 7. 因为我们的模型预期两个输入，所以返回一个字典
    return {"image": img, "label": label}

创建 Dataset（数据集）对象

• 定义字符集：定义将要识别的字符集，即所有合法字符的列表，包括数字和字母等。
• 生成字符序列：使用字符集生成随机字符序列，即模拟验证码生成的过程，生成指定数量的字符序列。
• 生成图像及标签：使用字符序列生成对应的图像和标签，其中图像包括字符和背景，标签则为字符序列的文本字符串。
• 数据划分：将生成的样本数据集按照训练集、验证集、测试集的比例进行划分。
• 创建 Dataset 对象：对划分后的样本数据集，创建 Dataset 对象，其中包含图像和标签，在训练模型时可以直接使用该对象进行训练。

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))  # 将训练集数据和标签转换为tensor slices并创建Dataset对象
train_dataset = (
    train_dataset.map(encode_single_sample,
                      num_parallel_calls=tf.data.AUTOTUNE)  # 对每个样本进行编码，使用tf.data.AUTOTUNE参数自动调整并行度
    .batch(batch_size)  # 将数据集分成batch_size大小的批次
    .prefetch(buffer_size=tf.data.AUTOTUNE)  # 在训练期间异步预取下一个批次，使用tf.data.AUTOTUNE参数自动设置缓冲区大小
)

validation_dataset = tf.data.Dataset.from_tensor_slices((x_valid, y_valid))  # 将验证集数据和标签转换为tensor slices并创建Dataset对象
validation_dataset = (
    validation_dataset.map(encode_single_sample,
                           num_parallel_calls=tf.data.AUTOTUNE)  # 对每个样本进行编码，使用tf.data.AUTOTUNE参数自动调整并行度
    .batch(batch_size)  # 将数据集分成batch_size大小的批次
    .prefetch(buffer_size=tf.data.AUTOTUNE)  # 在训练期间异步预取下一个批次，使用tf.data.AUTOTUNE参数自动设置缓冲区大小
)

将数据可视化

在这段代码中，我们使用 plt.subplots() 创建了一个包含 4 行 4 列的图像网格。然后，我们使用 train_dataset.take(1) 从训练数据集中获取一个 batch，遍历其中的 16 张图像，并展示它们。

在展示每张图像之前，我们首先将图像转换为 numpy 数组，并使用 astype() 方法将浮点数数据类型转换为整数类型。这一步是因为 matplotlib要求图像数据为整数。然后，我们将标签转换为字符串，并在每个字符之间插入空格。

最后，我们使用 imshow() 方法展示图像，并使用 set_title() 方法为其设置标签。我们使用 axis(“off”) 方法取消坐标轴的显示，以使图像更加美观。最后，我们使用 plt.show() 方法展示生成的图像网格。

# 生成4x4的图像网格
_, ax = plt.subplots(4, 4, figsize=(10, 5))

# 从训练数据集中取出一个batch
for batch in train_dataset.take(1):
    # 获取该batch的图像和标签
    images = batch["image"]
    labels = batch["label"]

    # 对16张图像进行遍历并展示
    for i in range(16):
        # 将图像转换为numpy数组，范围从0到1的浮点数，并乘以255来转换为0到255的整数
        img = (images[i] * 255).numpy().astype("uint8")

        # 将数字标签转换为字符串，并在每个字符之间插入空格
        label = tf.strings.reduce_join(num_to_char(labels[i])).numpy().decode("utf-8")

        # 展示图像并设置标题
        # 使用第一维度（高度）与第二维度（宽度）的灰度值来绘制图像
        ax[i // 4, i % 4].imshow(img[:, :, 0].T, cmap="gray")
        ax[i // 4, i % 4].set_title(label)
        ax[i // 4, i % 4].axis("off")

# 展示图像网格
plt.show()

构建模型

• 定义 input 和 output：使用 Keras Input 类定义神经网络的输入，即验证码图片；使用 Keras Output 类定义神经网络的输出，即识别的字符序列。
• 构建卷积神经网络结构：考虑到验证码识别的特殊性，建议使用卷积神经网络进行训练。通过堆叠卷积层、最大池化层、全连接层等，构建卷积神经网络的结构。
• 定义损失函数：在本例中，采用 Categorical Crossentropy 损失函数进行训练，经验表明 Categorical Crossentropy 作为多分类任务中比较好的一种损失函数，可以有效地提高分类准确率。
• 编译模型：在完成以上操作之后，使用 Keras 的 compile 方法编译模型，指定优化器和损失函数。

class CTCLayer(layers.Layer):
    """
    这是一个自定义的损失函数层 CTCLayer 类。该类定义了 call() 方法，其中传递了两个参数：y_true 是真实标签，y_pred 是预测值。该损失函数的作
    用是计算 CTC 损失函数。在训练时，计算损失值，并将其添加到层中。在测试时，直接返回预测值。
    """

    def __init__(self, name=None):
        super().__init__(name=name)
        # 设置损失函数
        self.loss_fn = keras.backend.ctc_batch_cost

    def call(self, y_true, y_pred):
        # 计算训练时的损失值并将其添加到层中 using `self.add_loss()`.
        batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64")
        input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64")  # 预测序列的长度
        label_length = tf.cast(tf.shape(y_true)[1], dtype="int64")  # 正确标签序列的长度

        input_length = input_length * tf.ones(shape=(batch_len, 1), dtype="int64")  # 区分每个序列的预测序列长度
        label_length = label_length * tf.ones(shape=(batch_len, 1), dtype="int64")  # 区分每个序列的正确标签序列长度

        # 使用 CTC 损失函数计算损失值
        loss = self.loss_fn(y_true, y_pred, input_length, label_length)
        self.add_loss(loss)  # 将损失值添加到层中

        # 在测试时，直接返回预测值
        return y_pred


def build_model():
    """
    其中，CNN 部分提取图像特征，RNN 部分将特征序列化。这个神经网络模型用于验证码识别任务。
    """
    # 定义输入层
    input_img = layers.Input(
        shape=(img_width, img_height, 1), name="image", dtype="float32"
    )
    # 定义 label 层
    labels = layers.Input(name="label", shape=(None,), dtype="float32")

    # 第一卷积块，32 个 3x3 卷积核，采用 relu 激活函数，同样大小的 padding，初始化方式为 He
    x = layers.Conv2D(
        32,
        (3, 3),
        activation="relu",
        kernel_initializer="he_normal",
        padding="same",
        name="Conv1",
    )(input_img)
    # 对卷积层的输出进行 2x2 最大池化
    x = layers.MaxPooling2D((2, 2), name="pool1")(x)

    # 第二卷积块，64 个 3x3 卷积核，采用 relu 激活函数，同样大小的 padding，初始化方式为 He
    x = layers.Conv2D(
        64,
        (3, 3),
        activation="relu",
        kernel_initializer="he_normal",
        padding="same",
        name="Conv2",
    )(x)
    # 对卷积层的输出进行 2x2 最大池化
    x = layers.MaxPooling2D((2, 2), name="pool2")(x)

    # 我们使用了两个 2x2 大小的最大池化和 stride，因此缩小了特征映射的大小 4 倍。最后一层的
    # 过滤器数量为 64。在将输出传递到模型的 RNN 部分之前进行相应的重塑
    new_shape = ((img_width // 4), (img_height // 4) * 64)
    x = layers.Reshape(target_shape=new_shape, name="reshape")(x)
    # 添加密集层, 有 64 个节点，使用 relu 激活函数
    x = layers.Dense(64, activation="relu", name="dense1")(x)
    # DropOut
    x = layers.Dropout(0.2)(x)

    # RNNs，使用双向 LSTM，第一个 LSTM 层返回完整的输出序列
    x = layers.Bidirectional(layers.LSTM(128, return_sequences=True, dropout=0.25))(x)
    # 第二个 LSTM 层也返回完整的输出序列
    x = layers.Bidirectional(layers.LSTM(64, return_sequences=True, dropout=0.25))(x)

    # 输出层，对应字符表的数量加一，使用 softmax 激活函数
    x = layers.Dense(
        len(char_to_num.get_vocabulary()) + 1, activation="softmax", name="dense2"
    )(x)

    # 添加 CTC 层来在每个时间步计算 CTC 损失
    output = CTCLayer(name="ctc_loss")(labels, x)

    # 定义模型
    model = keras.models.Model(
        inputs=[input_img, labels], outputs=output, name="ocr_model_v1"
    )
    # Adam 优化器
    opt = keras.optimizers.Adam()
    # 编译模型并返回
    model.compile(optimizer=opt)
    return model


# 获取模型
model = build_model()
model.summary()

训练

在下述代码中，我们添加了 EarlyStopping 的早停回调，用于监测验证集的损失值 val_loss，当损失值连续 early_stopping_patience 个周期
没有下降时，训练就会提前终止，并在最好的训练结果处恢复权重，这可以避免过拟合问题。

在模型的训练中，我们用 fit 方法进行训练，并设置了训练 epochs 数，即训练周期数，同时将 early_stopping 回调函数添加到训练集中，用于约
束训练进度并在需要时提前终止训练。

epochs = 100
early_stopping_patience = 10
# 添加 EarlyStopping 早停回调
early_stopping = keras.callbacks.EarlyStopping(
    monitor="val_loss", patience=early_stopping_patience, restore_best_weights=True
)

# 训练模型
history = model.fit(
    train_dataset,
    validation_data=validation_dataset,
    epochs=epochs,  # 训练周期数
    callbacks=[early_stopping],  # 添加回调函数，用于监控并约束训练进度
)

使用训练好的模型

介绍如何使用已经训练好的模型进行验证码识别的具体过程。

首先，通过加载已训练好的模型，得到模型的预测函数，即将输入的验证码图像转换为字符序列的函数。然后，读取待识别验证码图像文件，并进行预处理，包括将图像转换为灰度图像、二值化、前景框选和缩放等，与训练数据预处理过程相似。

接着，通过预测函数对处理后的图像进行预测，得到预测结果。最后，将预测结果与真实标签进行比对，计算识别准确率，同时输出预测结果和真实标签等信息。

# 从模型中提取出输出层以前的所有层，用于预测
prediction_model = keras.models.Model(
    model.get_layer(name="image").input, model.get_layer(name="dense2").output
)
prediction_model.summary()


# 解码网络的输出的辅助函数
def decode_batch_predictions(pred):
    # 获取预测的输出序列的长度
    input_len = np.ones(pred.shape[0]) * pred.shape[1]
    # 使用贪心搜索来获取最佳匹配结果（对于复杂任务，可以使用 Beam Search 算法）
    results = keras.backend.ctc_decode(pred, input_length=input_len, greedy=True)[0][0][
              :, :max_length
              ]
    # 将预测结果转成文本
    output_text = []
    for res in results:
        # num_to_char 是数字到字符的转化函数
        res = tf.strings.reduce_join(num_to_char(res)).numpy().decode("utf-8")
        output_text.append(res)
    return output_text


# 用于在一些验证集样本上测试准确性
for batch in validation_dataset.take(1):
    batch_images = batch["image"]
    batch_labels = batch["label"]
    # 预测输出
    preds = prediction_model.predict(batch_images)
    pred_texts = decode_batch_predictions(preds)

    orig_texts = []
    for label in batch_labels:
        label = tf.strings.reduce_join(num_to_char(label)).numpy().decode("utf-8")
        orig_texts.append(label)

    # 用matplotlib在一个子图中画出真实的标签和预测的标签
    _, ax = plt.subplots(4, 4, figsize=(15, 5))
    for i in range(len(pred_texts)):
        img = (batch_images[i, :, :, 0] * 255).numpy().astype(np.uint8)
        img = img.T
        title = f"Prediction: {pred_texts[i]}"
        ax[i // 4, i % 4].imshow(img, cmap="gray")
        ax[i // 4, i % 4].set_title(title)
        ax[i // 4, i % 4].axis("off")
plt.show()