使用Pytorch识别字符验证码

使用Pytorch识别字符验证码

之前已经学习过利用Keras搭建神经网络模型来识别字符验证码，相关的文章：
字符验证码识别之数据预处理
涉及图像预处理和标签处理等操作

字符验证码识别之模型构建
涉及模型构建以及训练过程。

近期又学习了pytorch实现卷积神经网络相关的技术，正好遇到一个验证码识别的需求，所以尝试使用pytorch来实现。

数据预处理

要训练的验证码如下所示:

其为中文汉字的简单运算，实际上仅包括零壹贰叁肆伍陆柒捌玖加减乘等于，这15个汉字，等于可以不识别（其实识别也完全没问题，只不过问题能简化就尽量简化嘛），那最后也就是总共要识别13个汉字，分类数就是13。

另外我们可以用’0123456789±x’来代替汉字，避免文件名称无法使用汉字（windows下open-cv不能读取带有中文路径或文件名称）的问题。

下载并标注了1000张验证码，观察其字体颜色和干扰线、点均多变，无法根据特定规则将其区分；另外，尝试中值模糊、均值模糊和高斯模糊，均得不到较好的效果（肉眼观察）。只有灰度化和二值化后，感觉稍微变得清晰了一些。

另外，针对数据集我还统计了一下各个类别的数量是否均衡（主要怕有的文字训练样本太少，训练效果差）。

{'捌': 206, '减': 346, '肆': 220, '柒': 205, '零': 200, '伍': 214, '加': 358, 
'玖': 189, '壹': 195, '叁': 191, '陆': 206, '乘': 297, '贰': 176}

数据集的分布情况如上数据，总体还算均衡，那就表示可以开始处理数据和进行训练了。

在进一步处理数据之前，先划分数据集，800个训练集、验证集和测试集分别100，划分完数据集后，我也是统计了下各个数据集下的分类数量是否均衡（毕竟以前犯过划分数据集有问题的错误）。

以上都是一些简单的操作，下面仅展示后续将图片和标签转换成numpy矩阵的代码：

import os
import cv2
import numpy as np
import random
from os import remove
import math


class ImageProcess:
    channel = 1
    height = 40
    width = 90
    num_classes = 13  # 共13个汉字
    labels_len = 3  # 每个标签包含3个汉字
    words = '0123456789+-x'  # 用字符来代替汉字

    images_path_train = 'D:/captcha/shanghai/train/'
    images_path_val = 'D:/captcha/shanghai/val/'
    images_path_test = 'D:/captcha/shanghai/test/'
    images_train = os.listdir(images_path_train)
    images_val = os.listdir(images_path_val)
    images_test = os.listdir(images_path_test)

    def __init__(self):
        self.x_data_train = None
        self.y_data_train = None
        self.x_data_val = None
        self.y_data_val = None
        self.x_data_test = None
        self.y_data_test = None

        print('预处理图像...')
        self.process_image("train")
        self.process_image("test")
        self.process_image("val")

        print('预处理标签')
        self.process_label("train")
        self.process_label("test")
        self.process_label("val")

        print('处理完成')

    def process_label(self, which):
        """
        处理标签
        如果每个样本是单类别，每个类别就一个值，处理成一个长度为batch的列表就可以
        如果每个样本是多类别(假设为n, n>=2)，处理成[batch, n]的二维数组
        :param: which 处理哪个数据集
        :return:
        """
        labels_list = []
        if which == "train":
            images = self.images_train
        elif which == "test":
            images = self.images_test
        else:
            images = self.images_val
        for image in images:
            labels = image.split("_")[1].replace('.jpg', '')
            """
            这部分是ont-hot编码的处理逻辑，在pytorch种实际不需要这样处理，
            这主要取决于 nn.CrossEntropyLoss()的输入参数格式
            参数只需要标签即可, 不需要传one-hot向量
            """
            # 初始化一个 3x13 的矩阵，初始值为0.0
            # result = np.zeros((self.labels_len, self.num_classes), dtype='float32')
            # for i, c in enumerate(labels):
            #     result[i][self.words.index(c)] = 1

            """
            直接处理为 [batch, n]的二维数组 即可
            """
            result = []
            for label in labels:
                result.append(self.words.index(label))
            labels_list.append(result)
        if which == "train":
            self.y_data_train = np.array(labels_list, dtype='int32')
        elif which == "test":
            self.y_data_test = np.array(labels_list, dtype='int32')
        else:
            self.y_data_val = np.array(labels_list, dtype='int32')

    def process_image(self, which):
        """
        处理图片 处理目标 (batch, channel, height, width)
        :return:
        """
        images_list = []
        if which == "train":
            images = self.images_train
            images_path = self.images_path_train
        elif which == "test":
            images = self.images_test
            images_path = self.images_path_test
        else:
            images = self.images_val
            images_path = self.images_path_val
        for image in images:
            path = f'{images_path}{image}'
            img = cv2.imread(path)

            # 中值模糊
            # img = cv2.medianBlur(img, 3)
            # 均值模糊
            # img = cv2.blur(img, (2, 2))
            # 高斯模糊
            # img = cv2.GaussianBlur(img, (5, 5), 1)

            # 灰度化
            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
            # 二值化
            ret, img = cv2.threshold(img, 0, 255, cv2.THRESH_OTSU)

            # cv2.namedWindow('captcha', cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)
            # cv2.resizeWindow('captcha', 180, 80)
            # cv2.imshow('captcha', img)
            # cv2.waitKey(0)

            img = np.array(img, dtype='float32')
            # 归一化
            img /= 255
            images_list.append(np.reshape(img, (1, self.height, self.width)))

        if which == "train":
            self.x_data_train = np.array(images_list, dtype='float32')
        elif which == "test":
            self.x_data_test = np.array(images_list, dtype='float32')
        else:
            self.x_data_val = np.array(images_list, dtype='float32')

    def train_loader(self, batch_size=16):
        """
        按批次，将训练数据和标签 迭代返回
        :param batch_size:
        :return:
        """
        batch_nums = math.ceil(len(self.x_data_train)/batch_size)
        for i in range(batch_nums):
            x_train = self.x_data_train[i*batch_size:(i+1)*batch_size]
            y_train = self.y_data_train[i*batch_size:(i+1)*batch_size]
            yield x_train, y_train

    def test_loader(self, batch_size=16):
        """
        按批次，将测试数据和标签 迭代返回
        :param batch_size:
        :return:
        """
        batch_nums = math.ceil(len(self.x_data_test)/batch_size)
        for i in range(batch_nums):
            x_test = self.x_data_test[i*batch_size:(i+1)*batch_size]
            y_test = self.y_data_test[i*batch_size:(i+1)*batch_size]
            yield x_test, y_test

    def val_loader(self, batch_size=16):
        """
        按批次，将验证数据和标签 迭代返回
        :param batch_size:
        :return:
        """
        batch_nums = math.ceil(len(self.x_data_val)/batch_size)
        for i in range(batch_nums):
            x_val = self.x_data_val[i*batch_size:(i+1)*batch_size]
            y_val = self.y_data_val[i*batch_size:(i+1)*batch_size]
            yield x_val, y_val

关于代码核心的地方，在代码中都有注释。

另外需要注意的一点是，如果输入到神经网络中的图片为三维，则

images_list.append(np.reshape(rr_img, (1, self.height, self.width)))

要替换为

images_list.append(np.transpose(img, (2, 0, 1)))

否则reshape会导致整个数据错乱。

搭建模型

之前使用Keras做字符验证码识别的时候，得到的经验就是针对这种比较简单的字符验证码，无需过于复杂的模型，几层CNN就够了。

import torch
from torch import nn
from torch import optim

import os


class NeuralNetWork(nn.Module):
    def __init__(self, channel, num_classes):
        """
        :param channel: 输入图片的channel
        :param num_classes: 分类数量
        """
        super(NeuralNetWork, self).__init__()
        self.convin = nn.Sequential(
            nn.Conv2d(channel, 64, kernel_size=(3, 3), padding=1, bias=False),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=(3, 3), padding=1, bias=False),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Dropout(0.25)
        )
        self.convall = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=(3, 3), padding=1, bias=False),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=(3, 3), padding=1, bias=False),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Dropout(0.25)
        )
        # 承接卷积层和fc层
        self.fc1 = nn.Sequential(
            nn.Linear(64*5*11, 1024),  # 这个输入值需要计算，根据输入图像的尺寸决定（本次输入图像尺寸为40*90）
            nn.ReLU(),
            nn.Dropout(0.5)
        )
        self.dense1 = nn.Sequential(
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, num_classes),
            # nn.LogSoftmax()
        )
        self.dense2 = nn.Sequential(
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, num_classes),
            # nn.LogSoftmax()
        )
        self.dense3 = nn.Sequential(
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, num_classes),
            # nn.LogSoftmax()
        )
       
    
   def forward(self, n_input):
        # 进行卷积、激活和池化操作
        feature = self.convin(n_input)
        feature = self.convall(feature)
        feature = self.convall(feature)

        # 对特征层(Tensor类型)进行维度变换，变成两维
        feature = feature.view(n_input.size(0), -1)  # size(0)是批次大小

        # 进行全连接操作
        feature = self.fc1(feature)
        out_put1 = self.dense1(feature)
        out_put2 = self.dense2(feature)
        out_put3 = self.dense3(feature)
        # 每个样本有三个输出值
        return [out_put1, out_put2, out_put3]

关于模型代码，有以下几点说明：

1. 使用几层卷积、卷积核的数量、池化操作和dropout等并不是固定的，这要根据你的训练情况逐步调整；

2. 全连接层的地方的输入值是需要计算的，是由输入到全连接层的输出通道数量x你的图片经过你的卷积和池化层后得到的尺寸，比如这里输出通道数量为64，原始输入图片尺寸为40x90，经过padding=1的卷积层尺寸不变，经过三次(2, 2)的池化层，变为5x11。

   40x90 --> 20x45 --> 10x22 --> 5x11

   并且在全连接层之前要把feature转换为(batch, )形状的二维tensor。

3. 如何控制每个样本有3个输出值，这里是我遇到的难题，因为之前学习都是每个样本一个类型。

   这里经过咨询有经验的同事得知，实际上就是利用相同的线性层计算得到三个值，同时返回。

   不过需要注意的是，即使这三个输出值是经过了相同的线性层，就像这里的

   nn.Sequential(
       nn.Linear(1024, 512),
       nn.ReLU(),
       nn.Linear(512, num_classes),
   )
   

   但是一定是三个独立定义的层（层名称无所谓），如果均使用同一个层，那么输出的这三个值永远都是一样的（亲身踩坑）

编写训练代码

from image_process import ImageProcess


if __name__ == "__main__":

    net = NeuralNetWork(1, 13)  # channel=1，classes=13

    epochs = 100  # 设置训练轮次
    batch_size = 16

    # 训练部分代码
    criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数
    # 随机梯度下降优化
    # optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
    optimizer = optim.Adam(net.parameters(), lr=0.001, weight_decay=1e-6)

    ip = ImageProcess()

    val_loss_min = 0  # 保存训练过程中的最小损失（验证）
    for epoch in range(epochs):

        net.train()  # 训练与测试，BN和Dropout有区别
        # 如果没有BN和Dropout,或者只训练不验证，可以不执行该方法
        train_loss = 0.0  # 实时打印当前损失变化情况
        for batch_idx, data in enumerate(ip.train_loader(batch_size=batch_size)):
            inputs, labels = data
            inputs = torch.from_numpy(inputs)  # 从numpy array转成tensor
            labels = torch.from_numpy(labels).long()  # 输入损失函数要求type为long
            optimizer.zero_grad()  # 先将梯度设置为0

            out_puts = net(inputs)  # 前向传播
            # out_puts的shape(n, batch_size, num_classes) 3x16x13 n表示每个样本包含的分类数量
            # 这里因为输出多个值，所以计算损失把多个损失加在一起
            # labels的shape(batch_size, n) 16x3
            loss = (
                    criterion(out_puts[0], labels[:, 0]) +
                    criterion(out_puts[1], labels[:, 1]) +
                    criterion(out_puts[2], labels[:, 2])
            )

            loss.backward()  # 反向传播
            optimizer.step()

            # 查看网络训练状态(损失是计算几批数据的平均损失)
            train_loss += loss.item()

            # 800个训练样本，batch_size=16, 800/16 = 50（一共50批次）
            # 每10批，打印一次损失
            if (batch_idx+1) % 10 == 0:
                print(f'epoch: {epoch+1}, batch_inx: {batch_idx+1} train loss: {train_loss/160}')
                train_loss = 0.0

        state = {
            'net': net.state_dict(),
            'epoch': epoch+1
        }
        if not os.path.isdir('checkpoint'):
            os.mkdir('checkpoint')
        if (epoch+1) % 10 == 0:   # 每10轮保存一次权重
            print(f'saving epoch {epoch+1} mode ...')
            torch.save(state, f'./checkpoint/shanghai_epoch_{epoch+1}.pth')  # pth 与 ckpt

        # 验证部分
        net.eval()
        val_loss = 0.0
        for batch_idx, val_data in enumerate(ip.val_loader(16)):
            inputs, labels = val_data
            inputs = torch.from_numpy(inputs)  # 从numpy array转成tensor
            labels = torch.from_numpy(labels).long()  # 输入损失函数要求type为long
            out_puts = net(inputs)

            loss = (
                    criterion(out_puts[0], labels[:, 0]) +
                    criterion(out_puts[1], labels[:, 1]) +
                    criterion(out_puts[2], labels[:, 2])
            )

            val_loss += loss.item()

            # 100个训练样本，batch_size=16, 100/16 = 6（一共7批次）
            # 一轮计算一次平均损失
            if (batch_idx+1) % 7 == 0:
                print(f'epoch: {epoch+1}, batch_inx: {batch_idx+1} val loss: {val_loss/100}')
                if not val_loss_min:
                    val_loss_min = val_loss
                # 正常是每10轮保存一次权重，当发现这一轮验证损失更小时，也会保存一次权重
                elif val_loss_min >= val_loss:
                    val_loss_min = val_loss
                    print(f'saving epoch {epoch+1} mode ...')
                    torch.save(state, f'./checkpoint/shanghai_epoch_{epoch+1}.pth')
            val_loss = 0.0

    print('training task finished')

关于训练代码，有以下几点说明：

1. 这里批量加载训练集和验证集是我在前面数据预处理部分特别开发好的，我觉得还是蛮巧妙地；

2. pytorch中都是使用tensor，所以需要将加载的数据(numpy矩阵)转换成tensor：

   torch.from_numpy(inputs)

3. 最重要的一点是三个输出的情况下，如何计算损失，这是我开发过程中遇到的另一个难题。

   经过咨询有经验的同事得知，实际上就是将三个输出的损失加在一起，但是你要根据神经网络的数据输出格式和你自己的标签格式，将正确的数据输入到损失函数中进行计算，且要注意CrossEntropyLoss的输入参数格式。

   关于CrossEntropyLoss使用方式的介绍：

   import torch
   from torch import nn
   
   x = torch.tensor([[0.2, 0.3, 0.5, 0.1], [0.3, 0.01, 0.02, 0.4]])
   y = torch.tensor([2, 3])
   criterion = nn.CrossEntropyLoss()
   loss = criterion(x, y)
   print(loss)
   

   其y参数只需要标签即可, 不需要传one-hot向量，这也就是前面数据预处理时没有采用one-hot编码来处理标签的原因。另外out_puts的输出shape我在代码中也有注释。

记录一些训练过程中遇到的情况

1. 训练到第30轮，训练损失才开始明显下降，一度让我以为程序哪里有问题，经过上网查资料发现：损失函数(loss)在最初的几个epochs时没有下降，可能的原因是学习率设置的太低、正则参数太高和陷入局部最小值。

   我当时设置的学习率lr=0.0001，确实比较小，我尝试调整为lr=0.001再训练，发现在第20轮时损失就开始下降了，果然是学习率设置的太低。

   另外我觉得dropout的太多有可能也是导致损失延迟降低的原因，所以我尝试将dropout的值缩小，也能提前几轮损失开始下降，但是最后的训练效果却不如dropout较大的时候。

2. 在较前面的轮次，val_loss远小于train_loss，一开始我总结的原因是在网络中添加了dropout层，而dropout仅在训练时生效，测试时是不会dropout的。所以val_loss会小于train_loss，因为我这里是远小于，后经排查是计算的时候写了bug。

   net.train()  # 训练与测试，BN和Dropout有区别
   
   net.eval() # 验证部分
   

   也就是这两行代码的作用，执行后告诉神经网络接下来将进入训练模式还是测试模型；另外BN层也是仅在训练时生效，在测试时不使用。

3. 关于如何设计出较好的模型，目前是我能力欠缺的一个地方，上面代码使用的模型架构（经过80轮的训练准确率能达到80%，经过250轮的训练准确率能达到90%），是参考大佬的模型，而我自己设计的模型准确率最高仅能达到70%，并且我也经过多次调整和训练，效果也并没有显著提升。

测试

import torch

from train import NeuralNetWork
from image_process import ImageProcess


if __name__ == "__main__":
    # 测试
    net = NeuralNetWork(1, 13)
    # 如果有dropout和BN操作，这里一定执行该方法，表示网络接下来进行测试操作
    net.eval() 
    check_point = torch.load('./checkpoint/shanghai_epoch_27.pth')
    # check_point = torch.load('shanghai_epoch_80.pth')
    net.load_state_dict(check_point['net'])
    batch_size = 16

    ip = ImageProcess()
    total_image = 0  # 总的图片数量
    correct_image = 0

    total_label = 0  # 总的标签数量
    correct_label = 0
    for data in ip.test_loader(batch_size):
        images, labels = data
        images = torch.from_numpy(images)
        out_puts = net(images)
        # batch_result = []
        _, predicted1 = torch.max(out_puts[0], 1)
        _, predicted2 = torch.max(out_puts[1], 1)
        _, predicted3 = torch.max(out_puts[2], 1)
        # batch_result.append(temp_result)
        for i in range(labels.shape[0]):
            total_image += 1
            total_label += 3
            print(f'true label: {labels[i]}')
            true_label = labels[i]
            print(f'predicted label: {predicted1[i]}  {predicted2[i]}  {predicted3[i]}')
            predicted_label = [int(predicted1[i]), int(predicted2[i]), int(predicted3[i])]
            if list(true_label) == predicted_label:
                correct_image += 1
            if true_label[0] == predicted_label[0]:
                correct_label += 1
            if true_label[1] == predicted_label[1]:
                correct_label += 1
            if true_label[2] == predicted_label[2]:
                correct_label += 1

    print(f'correct_image / total_image: {correct_image}/{total_image}')
    print(f'correct_label / total_label: {correct_label}/{total_label}')

这里的测试是批量测试，与实际的预测方法还有区别，但是大同小异，只不过在预测的方法中要注意针对单张图片再增加一个维度表示批次，否则传入神经网络的数据格式会出问题。

尝试进一步优化

1. 添加BN层

   self.convin = nn.Sequential(
       nn.Conv2d(channel, 64, kernel_size=(3, 3), padding=1, bias=False),
       nn.BatchNorm2d(64),
       nn.ReLU(),
       nn.Conv2d(64, 64, kernel_size=(3, 3), padding=1, bias=False),
       nn.BatchNorm2d(64),
       nn.ReLU(),
       nn.MaxPool2d(2, 2),
       nn.Dropout(0.25)
   )
   
   

   添加批规范化层后，训练得到的模型效果并没有不添加之前好，不过训练损失则在第2~3轮就开始明显下降了，不使用BN层的话，要训练15~20轮，训练损失才开始明显下降。

   考虑到BN层的目标就是防止梯度消失或爆炸、加快训练速度，所以损失下降比较快就体现了BN层的用处，但是针对我这个项目，整体效果却并没有提升。

2. 旋转图片

   因为观察验证码会稍微有些倾斜，倾斜幅度很小，所以想着能不能利用数据增强（旋转一个很小的角度）来进一步提升准确率。

​ 为了进行数据增强，我是直接在process_image方法中，对每一张图片进行旋转，然后生成一张新的直接添加 到训练集中，另外标签也要添加两遍，这样我的训练集就变成了1600张，这种方法有个缺点就是一张图片和他的旋转图是挨着的两个样本，在训练时如果可以彻底打乱比较好，而且不知道是不是这个原因，再次训练时，损失吃吃降不下来了（到60轮没有下降，我就停了），但是添加BN层后快速下降（但最终效果没有提升）。

​ 这里记录下旋转用到的技术：

  from torchvision.transforms import transforms
      
  # 随机旋转图像
  def random_rotation(image):
      image = Image.fromarray(image)  # 传入的image为CV2对象，转换为PIL.Image格式
      # image.show()
      rr = transforms.RandomRotation(degrees=(5, 10))
      rr_image = rr(image)
      # rr_image.show()
      return rr_image  # 返回的依然是PIL.Image格式，但是同样可以直接转为np.array

后来我觉得进行数据增强实际上可以对训练集操作然后直接生成相应的图片保存下来，然后在读取的时候打乱数据集比较方便，感兴趣的可以自己尝试。

我的gtihub博客地址：https://forchenxi.github.io/

另外，如果对投资理财感兴趣的同学，可以关注我的微信公众号：运气与实力。