验证码端到端的识别，是对《我的PaddlePaddle学习之路》笔记六——验证码端到端的识别 的升级，这篇文章是我18年初写的，基于当时的V2版本编写，现在有点过时了，突然想升级一下。

在线运行

在线运行：https://aistudio.baidu.com/aistudio/projectdetail/1679868

创建数据列表和词汇表

数据列表是为了方便训练是读取数据的。

import os
import cv2def createDataList(data_path, list_path):# 读取所有的图像路径imgs = os.listdir(data_path)with open(list_path, 'w', encoding='utf-8') as f:for img in imgs:name = img.split('.')[0]image_path = os.path.join(data_path, img)# 写入图像路径和label，用Tab隔开f.write(image_path + '\t' + name + '\n')createDataList('dataset/train_data/', 'dataset/train.txt')
createDataList('dataset/test_data/', 'dataset/test.txt')

还缺词汇表，执行下面的代码。

with open('dataset/train.txt', 'r', encoding='utf-8') as f:lines = f.readlines()
v = set()
for line in lines:_, label = line.replace('\n', '').split('\t')for c in label:v.add(c)vocabulary_path = 'dataset/vocabulary.txt'
with open(vocabulary_path, 'w', encoding='utf-8') as f:for c in v:f.write(c + '\n')

解码器

这个好似贪心解码方法，用于解码预测的输出的结果，将PaddlePaddle输出的结果转换为字符串。这里还提供了数据标签转字符串的，和计算字错率的。

%%writefile decoder.pyimport Levenshtein as Lev
from itertools import groupby
import paddledef ctc_greedy_decoder(probs_seq, vocabulary):"""CTC贪婪（最佳路径）解码器。由最可能的令牌组成的路径被进一步后处理删除连续的重复和所有的空白。:param probs_seq: 每个词汇表上概率的二维列表字符。每个元素都是浮点概率列表为一个字符。:type probs_seq: list:param vocabulary: 词汇表:type vocabulary: list:return: 解码结果字符串:rtype: baseline"""# 尺寸验证for probs in probs_seq:if not len(probs) == len(vocabulary) + 1:raise ValueError("probs_seq 尺寸与词汇不匹配")# argmax以获得每个时间步长的最佳指标max_index_list = paddle.argmax(probs_seq, -1).numpy()# 删除连续的重复索引index_list = [index_group[0] for index_group in groupby(max_index_list)]# 删除空白索引blank_index = len(vocabulary)index_list = [index for index in index_list if index != blank_index]# 将索引列表转换为字符串return ''.join([vocabulary[index] for index in index_list])[:4]def label_to_string(label, vocabulary):"""标签转文字:param label: 结果的标签，或者数据集的标签:type label: list:param vocabulary: 词汇表:type vocabulary: list:return: 解码结果字符串:rtype: baseline"""return ''.join([vocabulary[index] for index in label])def cer(out_string, target_string):"""通过计算两个字符串的距离，得出字错率Arguments:out_string (string): 比较的字符串target_string (string): 比较的字符串"""s1, s2, = out_string.replace(" ", ""), target_string.replace(" ", "")return Lev.distance(s1, s2)

数据读取器

这个是用于训练时读取数据的，用数据列表中读取图片和标签，将图片进行预处理，字符串标签转换为整型的标签输入的网络模型中。

%%writefile data.pyimport cv2
import numpy as np
from paddle.io import Dataset# 图像预处理
def process(path):image = cv2.imread(path)# 转灰度图image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)# 统一缩放大小image = cv2.resize(image, (72, 27))# 转换成CHWimage = image[np.newaxis, :]# 归一化image = (image - 128) / 128return image# 数据加载器
class CustomDataset(Dataset):def __init__(self, data_list_path, voc_path):super(CustomDataset, self).__init__()with open(data_list_path, 'r', encoding='utf-8') as f:self.lines = f.readlines()with open(voc_path, 'r', encoding='utf-8') as f:labels = f.readlines()self.vocabulary = [labels[i].replace('\n', '') for i in range(len(labels))]self.vocabulary_dict = dict([(labels[i].replace('\n', ''), i) for i in range(len(labels))])def __getitem__(self, idx):path, label = self.lines[idx].replace('\n', '').split('\t')img = process(path)# 将字符标签转换为int数据transcript = [self.vocabulary_dict.get(x) for x in label]img = np.array(img, dtype='float32')transcript = np.array(transcript, dtype='int32')return img, transcriptdef __len__(self):return len(self.lines)

模型结构

这个模型类型CRNN，前面使用卷积层提前图像特征，后面用一个GRU，他是LSTM的变种，最后的全连接层，输出的大小为词汇表+1，因为还有一个空格字符，这个是CTC需要的。

%%writefile model.pyimport paddle
import paddle.nn as nnclass Model(nn.Layer):def __init__(self, vocabulary):super(Model, self).__init__()self.conv1 = nn.Conv2D(in_channels=1, out_channels=32, kernel_size=3)self.relu1 = nn.ReLU()self.bn1 = nn.BatchNorm2D(32)self.pool1 = nn.MaxPool2D(kernel_size=2, stride=1)self.conv2 = nn.Conv2D(in_channels=32, out_channels=64, kernel_size=3)self.relu2 = nn.ReLU()self.bn2 = nn.BatchNorm2D(64)self.pool2 = nn.MaxPool2D(kernel_size=2, stride=1)self.conv3 = nn.Conv2D(in_channels=64, out_channels=128, kernel_size=3)self.relu3 = nn.ReLU()self.bn3 = nn.BatchNorm2D(128)self.pool3 = nn.MaxPool2D(kernel_size=2, stride=1)self.conv4 = nn.Conv2D(in_channels=128, out_channels=256, kernel_size=3)self.relu4 = nn.ReLU()self.bn4 = nn.BatchNorm2D(256)self.pool4 = nn.MaxPool2D(kernel_size=2, stride=1)self.conv5 = nn.Conv2D(in_channels=256, out_channels=256, kernel_size=3)self.relu5 = nn.ReLU()self.bn5 = nn.BatchNorm2D(256)self.pool5 = nn.MaxPool2D(kernel_size=2, stride=1)self.conv6 = nn.Conv2D(in_channels=256, out_channels=256, kernel_size=3)self.relu6 = nn.ReLU()self.bn6 = nn.BatchNorm2D(256)self.pool6 = nn.MaxPool2D(kernel_size=2, stride=1)self.conv7 = nn.Conv2D(in_channels=256, out_channels=256, kernel_size=3)self.relu7 = nn.ReLU()self.bn7 = nn.BatchNorm2D(256)self.pool7 = nn.MaxPool2D(kernel_size=2, stride=1)self.fc = nn.Linear(in_features=306, out_features=128)self.gru = nn.GRU(input_size=256, hidden_size=128)self.output = nn.Linear(in_features=128, out_features=len(vocabulary) + 1)def forward(self, x):x = self.relu1(self.bn1(self.conv1(x)))x = self.pool1(x)x = self.relu2(self.bn2(self.conv2(x)))x = self.pool2(x)x = self.relu3(self.bn3(self.conv3(x)))x = self.pool3(x)x = self.relu4(self.bn4(self.conv4(x)))x = self.pool4(x)x = self.relu5(self.bn5(self.conv5(x)))x = self.pool5(x)x = self.relu6(self.bn6(self.conv6(x)))x = self.pool6(x)x = self.relu7(self.bn7(self.conv7(x)))x = self.pool7(x)x = paddle.reshape(x, shape=(x.shape[0], x.shape[1], -1))x = self.fc(x)x = paddle.transpose(x, perm=[0, 2, 1])y, h = self.gru(x)x = self.output(y)return x

训练

这就开始训练了。三两下就可以训练完，主要是数据量少。每十轮训练结束都执行一次评估，输出的是字错率。保存的模型是静态模型，方便预测。

import paddle
import numpy as np
import os
from datetime import datetime
from model import Model
from decoder import ctc_greedy_decoder, label_to_string, cer
from paddle.io import DataLoader
from data import CustomDataset
from visualdl import LogWriter
from paddle.static import InputSpectrain_data_list_path = 'dataset/train.txt'
test_data_list_path = 'dataset/test.txt'
voc_path = 'dataset/vocabulary.txt'
save_model = 'models/model'
batch_size = 32
pretrained_model = None
num_epoch = 100
learning_rate = 1e-3
writer = LogWriter(logdir='log')def train():# 获取训练数据train_dataset = CustomDataset(train_data_list_path, voc_path)train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)# 获取测试数据test_dataset = CustomDataset(test_data_list_path, voc_path)test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size)# 获取模型model = Model(train_dataset.vocabulary)paddle.summary(model, input_size=(batch_size, 1, 27, 72))# 设置优化方法boundaries = [10, 20, 50]lr = [0.1 ** l * learning_rate for l in range(len(boundaries) + 1)]scheduler = paddle.optimizer.lr.PiecewiseDecay(boundaries=boundaries, values=lr, verbose=False)optimizer = paddle.optimizer.Adam(parameters=model.parameters(), learning_rate=scheduler)# 获取损失函数ctc_loss = paddle.nn.CTCLoss(blank=len(train_dataset.vocabulary))# 加载预训练模型if pretrained_model is not None:model.set_state_dict(paddle.load(os.path.join(pretrained_model, 'model.pdparams')))optimizer.set_state_dict(paddle.load(os.path.join(pretrained_model, 'optimizer.pdopt')))train_step = 0test_step = 0# 开始训练for epoch in range(num_epoch):for batch_id, (inputs, labels) in enumerate(train_loader()):out = model(inputs)out = paddle.transpose(out, perm=[1, 0, 2])input_lengths = paddle.full(shape=[out.shape[1]], fill_value=out.shape[0], dtype="int64")label_lengths = paddle.full(shape=[out.shape[1]], fill_value=4, dtype="int64")# 计算损失loss = ctc_loss(out, labels, input_lengths, label_lengths)loss.backward()optimizer.step()optimizer.clear_grad()# 多卡训练只使用一个进程打印if batch_id % 100 == 0:print('[%s] Train epoch %d, batch %d, loss: %f' % (datetime.now(), epoch, batch_id, loss))writer.add_scalar('Train loss', loss, train_step)train_step += 1if (epoch % 10 == 0 and epoch != 0) or epoch == num_epoch:# 执行评估model.eval()cer = evaluate(model, test_loader, train_dataset.vocabulary)print('[%s] Test epoch %d, cer: %f' % (datetime.now(), epoch, cer))writer.add_scalar('Test cer', cer, test_step)test_step += 1model.train()# 记录学习率writer.add_scalar('Learning rate', scheduler.last_lr, epoch)scheduler.step()# 保存模型paddle.jit.save(layer=model, path=save_model, input_spec=[InputSpec(shape=[None, 1, 27, 72], dtype='float32')])# 评估模型
def evaluate(model, test_loader, vocabulary):cer_result = []for batch_id, (inputs, labels) in enumerate(test_loader()):# 执行识别outs = model(inputs)outs = paddle.nn.functional.softmax(outs)# 解码获取识别结果labelss = []out_strings = []for out in outs:out_string = ctc_greedy_decoder(out, vocabulary)out_strings.append(out_string)for label in labels:labels = label_to_string(label, vocabulary)labelss.append(labels)for out_string, label in zip(*(out_strings, labelss)):print(label, out_string)# 计算字错率c = cer(out_string, label) / float(len(label))cer_result.append(c)cer_result = float(np.mean(cer_result))return cer_resulttrain()

预测

使用训练好的模型识别验证码图片。

import numpy as np
import paddlefrom data import process
from decoder import ctc_greedy_decoderwith open('dataset/vocabulary.txt', 'r', encoding='utf-8') as f:vocabulary = f.readlines()vocabulary = [v.replace('\n', '') for v in vocabulary]save_model = 'models/model'
model = paddle.jit.load(save_model)
model.eval()def infer(path):data = process(path)data = data[np.newaxis, :]data = paddle.to_tensor(data, dtype='float32')# 执行识别out = model(data)out = paddle.nn.functional.softmax(out)[0]# 解码获取识别结果out_string = ctc_greedy_decoder(out, vocabulary)print('预测结果：%s' % out_string)if __name__ == '__main__':image_path = 'dataset/test.png'infer(image_path)