通过OCR实现验证码识别

作者: GT_老张

时间: 2021.05

摘要: 本篇将介绍如何通过飞桨实现简单的CRNN+CTC自定义数据集OCR识别模型，数据集采用CaptchaDataset中OCR部分的9453张图像，其中前8453张图像在本案例中作为训练集，后1000张则作为测试集。
在更复杂的场景中推荐使用PaddleOCR产出工业级模型，模型轻量且精度大幅提升。
同样也可以在PaddleHub中快速使用PaddleOCR。

一、环境配置

本教程基于Paddle 2.1 编写，如果你的环境不是本版本，请先参考官网安装 Paddle 2.1 。

  1. import paddle
  2. print(paddle.__version__)
  1. 2.1.0

二、自定义数据集读取器

常见的开发任务中，我们并不一定会拿到标准的数据格式，好在我们可以通过自定义Reader的形式来随心所欲读取自己想要数据。

设计合理的Reader往往可以带来更好的性能，我们可以将读取标签文件列表、制作图像文件列表等必要操作在__init__特殊方法中实现。这样就可以在实例化Reader时装入内存，避免使用时频繁读取导致增加额外开销。同样我们可以在__getitem__特殊方法中实现如图像增强、归一化等个性操作，完成数据读取后即可释放该部分内存。
需要我们注意的是，如果不能保证自己数据十分纯净，可以通过try和expect来捕获异常并指出该数据的位置。当然也可以制定一个策略，使其在发生数据读取异常后依旧可以正常进行训练。

2.1 数据展示

https://ai-studio-static-online.cdn.bcebos.com/57d6c77aa5194cdca5c7edc533cc57e4d5070de95f6a4454b3cd1ca1e0eebe98

点此快速获取本节数据集，待数据集下载完毕后可使用!unzip OCR_Dataset.zip -d data/命令或熟悉的解压软件进行解压，待数据准备工作完成后修改本文“训练准备”中的DATA_PATH = 解压后数据集路径。

  1. # 解压数据集
  2. !unzip OCR_Dataset.zip -d data/
  1. import os
  3. import PIL.Image as Image
  4. import numpy as np
  5. from paddle.io import Dataset
  7. # 图片信息配置 - 通道数、高度、宽度
  8. IMAGE_SHAPE_C = 3
  9. IMAGE_SHAPE_H = 30
  10. IMAGE_SHAPE_W = 70
  11. # 数据集图片中标签长度最大值设置 - 因图片中均为4个字符，故该处填写为4即可
  12. LABEL_MAX_LEN = 4
  15. class Reader(Dataset):
  16.     def __init__(self, data_path: str, is_val: bool = False):
  17.         """
  18.         数据读取Reader
  19.         :param data_path: Dataset路径
  20.         :param is_val: 是否为验证集
  21.         """
  22.         super().__init__()
  23.         self.data_path = data_path
  24.         # 读取Label字典
  25.         with open(os.path.join(self.data_path, "label_dict.txt"), "r", encoding="utf-8") as f:
  26.             self.info = eval(f.read())
  27.         # 获取文件名列表
  28.         self.img_paths = [img_name for img_name in self.info]
  29.         # 将数据集后1000张图片设置为验证集，当is_val为真时img_path切换为后1000张
  30.         self.img_paths = self.img_paths[-1000:] if is_val else self.img_paths[:-1000]
  32.     def __getitem__(self, index):
  33.         # 获取第index个文件的文件名以及其所在路径
  34.         file_name = self.img_paths[index]
  35.         file_path = os.path.join(self.data_path, file_name)
  36.         # 捕获异常 - 在发生异常时终止训练
  37.         try:
  38.             # 使用Pillow来读取图像数据
  39.             img = Image.open(file_path)
  40.             # 转为Numpy的array格式并整体除以255进行归一化
  41.             img = np.array(img, dtype="float32").reshape((IMAGE_SHAPE_C, IMAGE_SHAPE_H, IMAGE_SHAPE_W)) / 255
  42.         except Exception as e:
  43.             raise Exception(file_name + "t文件打开失败，请检查路径是否准确以及图像文件完整性，报错信息如下:n" + str(e))
  44.         # 读取该图像文件对应的Label字符串，并进行处理
  45.         label = self.info[file_name]
  46.         label = list(label)
  47.         # 将label转化为Numpy的array格式
  48.         label = np.array(label, dtype="int32").reshape(LABEL_MAX_LEN)
  50.         return img, label
  52.     def __len__(self):
  53.         # 返回每个Epoch中图片数量
  54.         return len(self.img_paths)

三、模型配置

3.1 定义模型结构以及模型输入

模型方面使用的简单的CRNN-CTC结构，输入形为CHW的图像在经过CNN->Flatten->Linear->RNN->Linear后输出图像中每个位置所对应的字符概率。考虑到CTC解码器在面对图像中元素数量不一、相邻元素重复时会存在无法正确对齐等情况，故额外添加一个类别代表“分隔符”进行改善。

CTC相关论文：Connectionist Temporal Classification: Labelling Unsegmented Sequence Data with Recurrent Neu

https://ai-studio-static-online.cdn.bcebos.com/8ce1fce57f9f47c5b19a4e61caae7e6330b1b42dde1c4f4593d2181fbf474b8b

网络部分，因本篇采用数据集较为简单且图像尺寸较小并不适合较深层次网络。若在对尺寸较大的图像进行模型构建，可以考虑使用更深层次网络/注意力机制来完成。当然也可以通过目标检测形式先检出文本位置，然后进行OCR部分模型构建。

https://ai-studio-static-online.cdn.bcebos.com/19ddf6107e7f47ee9b3b84ee0c12de1e15f7ab8b88f04eed95232440c92fe0d7

PaddleOCR效果图

  1. import paddle
  3. # 分类数量设置 - 因数据集中共包含0~9共10种数字+分隔符，所以是11分类任务
  4. CLASSIFY_NUM = 11
  6. # 定义输入层，shape中第0维使用-1则可以在预测时自由调节batch size
  7. input_define = paddle.static.InputSpec(shape=[-1, IMAGE_SHAPE_C, IMAGE_SHAPE_H, IMAGE_SHAPE_W],
  8.                                    dtype="float32",
  9.                                    name="img")
  11. # 定义网络结构
  12. class Net(paddle.nn.Layer):
  13.     def __init__(self, is_infer: bool = False):
  14.         super().__init__()
  15.         self.is_infer = is_infer
  17.         # 定义一层3x3卷积+BatchNorm
  18.         self.conv1 = paddle.nn.Conv2D(in_channels=IMAGE_SHAPE_C,
  19.                                   out_channels=32,
  20.                                   kernel_size=3)
  21.         self.bn1 = paddle.nn.BatchNorm2D(32)
  22.         # 定义一层步长为2的3x3卷积进行下采样+BatchNorm
  23.         self.conv2 = paddle.nn.Conv2D(in_channels=32,
  24.                                   out_channels=64,
  25.                                   kernel_size=3,
  26.                                   stride=2)
  27.         self.bn2 = paddle.nn.BatchNorm2D(64)
  28.         # 定义一层1x1卷积压缩通道数，输出通道数设置为比LABEL_MAX_LEN稍大的定值可获取更优效果，当然也可设置为LABEL_MAX_LEN
  29.         self.conv3 = paddle.nn.Conv2D(in_channels=64,
  30.                                   out_channels=LABEL_MAX_LEN + 4,
  31.                                   kernel_size=1)
  32.         # 定义全连接层，压缩并提取特征（可选）
  33.         self.linear = paddle.nn.Linear(in_features=429,
  34.                                    out_features=128)
  35.         # 定义RNN层来更好提取序列特征，此处为双向LSTM输出为2 x hidden_size，可尝试换成GRU等RNN结构
  36.         self.lstm = paddle.nn.LSTM(input_size=128,
  37.                                hidden_size=64,
  38.                                direction="bidirectional")
  39.         # 定义输出层，输出大小为分类数
  40.         self.linear2 = paddle.nn.Linear(in_features=64 * 2,
  41.                                     out_features=CLASSIFY_NUM)
  43.     def forward(self, ipt):
  44.         # 卷积 + ReLU + BN
  45.         x = self.conv1(ipt)
  46.         x = paddle.nn.functional.relu(x)
  47.         x = self.bn1(x)
  48.         # 卷积 + ReLU + BN
  49.         x = self.conv2(x)
  50.         x = paddle.nn.functional.relu(x)
  51.         x = self.bn2(x)
  52.         # 卷积 + ReLU
  53.         x = self.conv3(x)
  54.         x = paddle.nn.functional.relu(x)
  55.         # 将3维特征转换为2维特征 - 此处可以使用reshape代替
  56.         x = paddle.tensor.flatten(x, 2)
  57.         # 全连接 + ReLU
  58.         x = self.linear(x)
  59.         x = paddle.nn.functional.relu(x)
  60.         # 双向LSTM - [0]代表取双向结果，[1][0]代表forward结果,[1][1]代表backward结果，详细说明可在官方文档中搜索'LSTM'
  61.         x = self.lstm(x)[0]
  62.         # 输出层 - Shape = (Batch Size, Max label len, Signal) 
  63.         x = self.linear2(x)
  65.         # 在计算损失时ctc-loss会自动进行softmax，所以在预测模式中需额外做softmax获取标签概率
  66.         if self.is_infer:
  67.             # 输出层 - Shape = (Batch Size, Max label len, Prob) 
  68.             x = paddle.nn.functional.softmax(x)
  69.             # 转换为标签
  70.             x = paddle.argmax(x, axis=-1)
  71.         return x

四、训练准备

4.1 定义label输入以及超参数

监督训练需要定义label，预测则不需要该步骤。

  1. # 数据集路径设置
  2. DATA_PATH = "./data/OCR_Dataset"
  3. # 训练轮数
  4. EPOCH = 10
  5. # 每批次数据大小
  6. BATCH_SIZE = 16
  8. label_define = paddle.static.InputSpec(shape=[-1, LABEL_MAX_LEN],
  9.                                     dtype="int32",
  10.                                     name="label")

4.2 定义CTC Loss

了解CTC解码器效果后，我们需要在训练中让模型尽可能接近这种类型输出形式，那么我们需要定义一个CTC Loss来计算模型损失。不必担心，在飞桨框架中内置了多种Loss，无需手动复现即可完成损失计算。

使用文档：CTCLoss

  1. class CTCLoss(paddle.nn.Layer):
  2.     def __init__(self):
  3.         """
  4.         定义CTCLoss
  5.         """
  6.         super().__init__()
  8.     def forward(self, ipt, label):
  9.         input_lengths = paddle.full(shape=[BATCH_SIZE, 1],fill_value=LABEL_MAX_LEN + 4,dtype= "int64")
  10.         label_lengths = paddle.full(shape=[BATCH_SIZE, 1],fill_value=LABEL_MAX_LEN,dtype= "int64")
  11.         # 按文档要求进行转换dim顺序
  12.         ipt = paddle.tensor.transpose(ipt, [1, 0, 2])
  13.         # 计算loss
  14.         loss = paddle.nn.functional.ctc_loss(ipt, label, input_lengths, label_lengths, blank=10)
  15.         return loss

4.3 实例化模型并配置优化策略

  1. # 实例化模型
  2. model = paddle.Model(Net(), inputs=input_define, labels=label_define)
  1. # 定义优化器
  2. optimizer = paddle.optimizer.Adam(learning_rate=0.0001, parameters=model.parameters())
  4. # 为模型配置运行环境并设置该优化策略
  5. model.prepare(optimizer=optimizer,
  6.                 loss=CTCLoss())

五、开始训练

  1. # 执行训练
  2. model.fit(train_data=Reader(DATA_PATH),
  3.             eval_data=Reader(DATA_PATH, is_val=True),
  4.             batch_size=BATCH_SIZE,
  5.             epochs=EPOCH,
  6.             save_dir="output/",
  7.             save_freq=1,
  8.             verbose=1)
  1. The loss value printed in the log is the current step, and the metric is the average value of previous steps.
  2. Epoch 1/10
  3. step 529/529 [==============================] - loss: 0.4842 - 11ms/step        
  4. save checkpoint at /home/aistudio/output/0
  5. Eval begin...
  6. step 63/63 [==============================] - loss: 0.4257 - 7ms/step        
  7. Eval samples: 1000
  8. ...
  9. Epoch 10/10
  10. step 529/529 [==============================] - loss: 0.0092 - 11ms/step        
  11. save checkpoint at /home/aistudio/output/9
  12. Eval begin...
  13. step 63/63 [==============================] - loss: 0.0061 - 7ms/step        
  14. Eval samples: 1000
  15. save checkpoint at /home/aistudio/output/final

六、预测前准备

6.1 像定义训练Reader一样定义预测Reader

  1. # 与训练近似，但不包含Label
  2. class InferReader(Dataset):
  3.     def __init__(self, dir_path=None, img_path=None):
  4.         """
  5.         数据读取Reader(预测)
  6.         :param dir_path: 预测对应文件夹（二选一）
  7.         :param img_path: 预测单张图片（二选一）
  8.         """
  9.         super().__init__()
  10.         if dir_path:
  11.             # 获取文件夹中所有图片路径
  12.             self.img_names = [i for i in os.listdir(dir_path) if os.path.splitext(i)[1] == ".jpg"]
  13.             self.img_paths = [os.path.join(dir_path, i) for i in self.img_names]
  14.         elif img_path:
  15.             self.img_names = [os.path.split(img_path)[1]]
  16.             self.img_paths = [img_path]
  17.         else:
  18.             raise Exception("请指定需要预测的文件夹或对应图片路径")
  20.     def get_names(self):
  21.         """
  22.         获取预测文件名顺序 
  23.         """
  24.         return self.img_names
  26.     def __getitem__(self, index):
  27.         # 获取图像路径
  28.         file_path = self.img_paths[index]
  29.         # 使用Pillow来读取图像数据并转成Numpy格式
  30.         img = Image.open(file_path)
  31.         img = np.array(img, dtype="float32").reshape((IMAGE_SHAPE_C, IMAGE_SHAPE_H, IMAGE_SHAPE_W)) / 255
  32.         return img
  34.     def __len__(self):
  35.         return len(self.img_paths)

6.2 参数设置

  1. # 待预测目录 - 可在测试数据集中挑出b3张图像放在该目录中进行推理
  2. INFER_DATA_PATH = "./sample_img"
  3. # 训练后存档点路径 - final 代表最终训练所得模型
  4. CHECKPOINT_PATH = "./output/final.pdparams"
  5. # 每批次处理数量
  6. BATCH_SIZE = 32

6.3 展示待预测数据

  1. import matplotlib.pyplot as plt
  2. plt.figure(figsize=(10, 10))
  3. sample_idxs = np.random.choice(50000, size=25, replace=False)
  5. for img_id, img_name in enumerate(os.listdir(INFER_DATA_PATH)):
  6.     plt.subplot(1, 3, img_id + 1)
  7.     plt.xticks([])
  8.     plt.yticks([])
  9.     im = Image.open(os.path.join(INFER_DATA_PATH, img_name))
  10.     plt.imshow(im, cmap=plt.cm.binary)
  11.     plt.xlabel("Img name: " + img_name)
  12. plt.show()

png

七、开始预测

飞桨2.1 CTC Decoder 相关API正在迁移中，本节暂时使用简易版解码器。

  1. # 编写简易版解码器
  2. def ctc_decode(text, blank=10):
  3.     """
  4.     简易CTC解码器
  5.     :param text: 待解码数据
  6.     :param blank: 分隔符索引值
  7.     :return: 解码后数据
  8.     """
  9.     result = []
  10.     cache_idx = -1
  11.     for char in text:
  12.         if char != blank and char != cache_idx:
  13.             result.append(char)
  14.         cache_idx = char
  15.     return result
  18. # 实例化推理模型
  19. model = paddle.Model(Net(is_infer=True), inputs=input_define)
  20. # 加载训练好的参数模型
  21. model.load(CHECKPOINT_PATH)
  22. # 设置运行环境
  23. model.prepare()
  25. # 加载预测Reader
  26. infer_reader = InferReader(INFER_DATA_PATH)
  27. img_names = infer_reader.get_names()
  28. results = model.predict(infer_reader, batch_size=BATCH_SIZE)
  29. index = 0
  30. for text_batch in results[0]:
  31.     for prob in text_batch:
  32.         out = ctc_decode(prob, blank=10)
  33.         print(f"文件名：{img_names[index]}，推理结果为：{out}")
  34.         index += 1
  1. Predict begin...
  2. step 1/1 [==============================] - 6ms/step
  3. Predict samples: 3
  4. 文件名：9450.jpg，推理结果为：[8, 2, 0, 5]
  5. 文件名：9452.jpg，推理结果为：[0, 3, 0, 0]
  6. 文件名：9451.jpg，推理结果为：[3, 4, 6, 3]