基于卷积神经网络的验证码识别

模式识别（Pattern Recognition）是人工智能领域的一个重要分支，它涉及研究如何让计算机系统能够识别和处理数据中的模式。模式识别技术试图从数据中提取有用的信息，识别出数据中的规律性或模式，并据此做出决策或预测。

在写爬虫脚本的时候，经常会遇到要输入图形验证码的场景。从图片格式的验证码到文本格式其实就是一种模式识别。为了能够有效的识别验证码，本作业设计了一个基于卷积神经网络的验证码识别模型，并通过Pytorch进行代码实现。

配置参数说明

我们采用4位验证码，字符集包括数字和小写字母混合。

CAPTCHA_LEN = 4  # 验证码长度
CAPTCHA_WIDTH = 160  # 验证码图片宽度
CAPTCHA_HEIGHT = 60  # 验证码图片高度
CHAR_SET = 'abcdefghijklmnopqrstuvwxyz0123456789'  # 字符集
CHAR_LEN = len(CHAR_SET)

数据集准备

验证码生成

Python的captcha库中有一个ImageCaptcha模块，可以用于生成图像验证码。默认生成的验证码图片尺寸为160*60，如下图所示。这个验证码为491e。

我们首先随机生成4位的验证码文本作为标签，然后用ImageCaptcha生成对应的验证码图片，并把验证码文本作为图片的名称。例如上面那个图片的名称就是“491e.jpg”。利用这个库，我们随机生成15w张验证码用于训练，1w张验证码用于测试。

生成数据集的代码如下所示：

from captcha.image import ImageCaptcha
import random
import os
from config import *

pic_path = './pic'  # 验证码保存文件夹


def random_captcha_text(length: int):
    captcha_text = []
    for i in range(length):
        c = random.choice(CHAR_SET)
        captcha_text.append(c)
    return ''.join(captcha_text)


def gen_one_captcha(path):
    image = ImageCaptcha(width=CAPTCHA_WIDTH, height=CAPTCHA_HEIGHT)
    captcha_text = random_captcha_text(CAPTCHA_LEN)
    img_path = os.path.join(path, f'{captcha_text}.jpg')
    image.write(captcha_text, img_path)


train_path = os.path.join(pic_path, 'train')
test_path = os.path.join(pic_path, 'test')
if not os.path.exists(train_path):
    os.mkdir(train_path)
if not os.path.exists(test_path):
    os.mkdir(test_path)

# 训练集生成150000张验证码
train_list = next(os.walk(train_path))[2]
cnt = len(train_list)
while cnt < 150000:
    gen_one_captcha(train_path)
    cnt += 1

# 测试集生成10000张验证码
test_list = next(os.walk(test_path))[2]
cnt = len(test_list)
while cnt < 10000:
    gen_one_captcha(test_path)
    cnt += 1

标签表示

标签就是长度为4的数字和字母组合的文本信息，我们需要对每个文本信息进行编码。具体来说，对于每个文本，我们采用字符集中的索引来表示。例如，“491e”就可以表示为长度为4的向量[30, 35, 27, 4]。再进一步对每个数字展开，采用其one hot编码（长度为36）。因此，每个标签最后可以表示为长度为4*36的向量，也就是4个one hot拼接之后的结果。

上述过程的代码如下：

def text2onehot(text):
    vector = np.zeros(CAPTCHA_LEN * CHAR_LEN)

    for i, c in enumerate(text):
        idx = CHAR_SET.index(c)
        vector[i * CHAR_LEN + idx] = 1
    return vector

定义数据集类

在定义数据集类之前，先确定图片读取进来之后要进行的预处理。由于颜色是不必要的信息，所以我们读取图片之后转为灰度模式。这样也能降低数据的复杂度。然后对图片进行二值化，突出图片里的主要信息。

img = cv2.imread('./491e.jpg', 0)  # 灰度模式读取图片
_, img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)  # 二值化

例如，之前的验证码经过处理后就如下图所示。

然后，就可以通过继承Pytorch提供的Dataset类来定义MyDataset：

import os
import cv2
import torch
import numpy as np
from config import *
from torch.utils.data import Dataset


def text2onehot(text):
    vector = np.zeros(CAPTCHA_LEN * CHAR_LEN)

    for i, c in enumerate(text):
        idx = CHAR_SET.index(c)
        vector[i * CHAR_LEN + idx] = 1
    return vector


class MyDataset(Dataset):

    def __init__(self, dir):
        self.dir = dir
        self.img_name = next(os.walk(dir))[2]  # 获取图片目录下所有图片的名称

    def __getitem__(self, index):
        img_path = os.path.join(self.dir, self.img_name[index])
        img = cv2.imread(img_path, 0)  # 灰度模式读取图片
        _, img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)  # 二值化
        img = torch.from_numpy(img).float()  # 转换为张量
        img = img.unsqueeze(0)  # 增加一个通道维度
        label = text2onehot(self.img_name[index].split('.')[0])  # 获取onehot标签
        label = torch.FloatTensor(label)
        return img, label

    def __len__(self):
        return len(self.img_name)

其中，__getitem__方法返回处理好的图片张量和标签向量。

定义模型

模型结构如下图所示：

模型代码如下：

import torch.nn as nn
from config import *


class CNN_Network(nn.Module):

    def __init__(self):
        super(CNN_Network, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, stride=1, kernel_size=3, padding=1),
            nn.BatchNorm2d(16), nn.ReLU())
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, stride=1, kernel_size=3, padding=1),
            nn.MaxPool2d(stride=2, kernel_size=2),  # 30*80
            nn.BatchNorm2d(32),
            nn.ReLU())
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, stride=1, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, stride=1, kernel_size=3, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 15*40
            nn.BatchNorm2d(128),
            nn.ReLU())
        self.layer4 = nn.Sequential(
            nn.Conv2d(128, 256, stride=1, kernel_size=3, padding=1),
            nn.BatchNorm2d(256), nn.ReLU())

        self.fc = nn.Sequential(nn.Linear(256 * 15 * 40, 2048), nn.ReLU(),
                                nn.Linear(2048, 1024), nn.ReLU(),
                                nn.Linear(1024, CAPTCHA_LEN * CHAR_LEN))

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

训练模型

训练模型之前首先加载数据。利用之前定义的数据集类分别加载训练集和测试集：

train_dataset = MyDataset('./pic/train')
test_dataset = MyDataset('./pic/test')
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

我们采用多标签软边界损失作为损失函数。多标签软边界损失（MultiLabelSoftMarginLoss）是一种在PyTorch中用于多标签分类问题的损失函数。它基于最大熵原理，通过优化一个针对每个类别的一对多（one-vs-all）损失来进行计算。

我们采用Adam作为优化器，并设置batch大小为200，学习率为1e-4，训练15个epoch。训练代码如下：

def train(train_loader):
    net = CNN_Network().to(device)
    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
    criterion = nn.MultiLabelSoftMarginLoss()  #多分类损失函数
    best_loss = math.inf
    loss_per_epoch = []
    train_acc, test_acc = [], []

    net.train()
    for epoch in range(n_epochs):
        loss_record = []
        labels, preds = [], []
        train_pbar = tqdm(train_loader)

        for img, label in train_pbar:
            img, label = img.to(device), label.to(device)
            labels += tensor2captcha(label)
            optimizer.zero_grad()

            pred = net(img)
            preds += tensor2captcha(pred)
            loss = criterion(pred, label)
            loss.backward()
            optimizer.step()
            loss_record.append(loss.detach().item())

            train_pbar.set_description(f'Epoch [{epoch+1}/{n_epochs}]')
            train_pbar.set_postfix({'loss': loss.detach().item()})

        mean_loss = sum(loss_record) / len(loss_record)
        acc = get_acc(labels, preds)

        loss_per_epoch.append(mean_loss)
        train_acc.append(acc)

        if mean_loss < best_loss:
            best_loss = mean_loss
            torch.save(net.state_dict(), "captcha.pth")
            print(
                f'Saving model with loss {best_loss:.4f}...,  accuracy {acc}')

        acc = test(test_loader)
        test_acc.append(acc)

    draw_loss_curve(loss_per_epoch)
    draw_acc(train_acc, test_acc)

测试模型

测试模型就是看看模型在测试集上的表现如何。在测试模型之前已经将之前训练好的模型保存了下来。测试模型的时候读取模型并设置为评估模式，计算模型在整个测试集上的准确率。具体代码如下：

def test(test_loader):
    model = CNN_Network().to(device)
    model.load_state_dict(torch.load('captcha.pth'))
    model.eval()

    labels, preds = [], []
    for img, label in tqdm(test_loader):
        img, label = img.to(device), label.to(device)
        labels += tensor2captcha(label)

        pred = model(img)
        preds += tensor2captcha(pred)

    acc = get_acc(labels, preds)
    return acc

数据可视化

除此之外，我们需要观察模型的效果怎么样，所以需要可视化一些数据，从而方便我们调整网络结构以及相应的参数，比如损失值，准确率等。

首先需要将模型输出的张量转化为4个字符的验证码。由于标签是用4个one hot向量拼接在一起的，所以需要按照one hot向量的长度将张量分割开。然后在每个切片中找到数值最大的对应的索引。数值最大的我们就认为是one hot中的1。最后，根据索引去字符集中找到对应的字符即可。

def tensor2captcha(tensor):
    tot_len = tensor.shape[1]
    assert CAPTCHA_LEN * CHAR_LEN == tot_len

    chars = np.array(list(CHAR_SET))

    captcha = []
    for i in range(CAPTCHA_LEN):
        slice = tensor[:, i * CHAR_LEN:(i + 1) * CHAR_LEN]
        idx = torch.argmax(slice, dim=1).detach().cpu().numpy()
        col = chars[idx][:, np.newaxis]
        captcha.append(col)
    captcha = list(np.hstack(captcha))
    captcha = [''.join(row) for row in captcha]

    return captcha

计算准确率就很简单了，预测值和标签一样就认为是预测正确。

def get_acc(label, pred):
    assert len(label) == len(pred)
    label, pred = np.array(label), np.array(pred)
    return sum(label == pred) / len(label)

最后是绘制损失曲线和准确率曲线。

def draw_loss_curve(loss_per_epoch):
    plt.figure()
    plt.plot(list(range(1, len(loss_per_epoch) + 1)), loss_per_epoch)
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Loss Curve')
    plt.savefig('loss.png')


def draw_acc(train_acc, test_acc):
    assert len(train_acc) == len(test_acc)

    x = list(range(1, len(train_acc) + 1))
    plt.clf()
    plt.plot(x, train_acc, label='train_acc')
    plt.plot(x, test_acc, label='test_acc')
    plt.legend()
    plt.title("acc goes by epoch")
    plt.xlabel('eopch')
    plt.ylabel('acc_value')
    plt.savefig('acc.png')

训练/测试结果

经过15轮的训练，最终在测试集上取得了93.77%的准确率。

在每轮训练中，训练集和测试集上的准确率如下图所示。

可以看到训练集上的准确率接近100%，测试集上的准确率略逊一筹，为93.77%。整个训练过程的损失变化如下图所示。

使用模型识别验证码

训练好模型之后就可以拿来使用了。首先读取图片并预处理，然后加载模型，最后模型输出识别到的验证码。

import torch
import cv2
from model import *
from utils import *

# 读取图片并预处理
filename = '491e.jpg'
img = cv2.imread(filename, 0)  # 灰度模式读取图片
_, img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)  # 二值化
img = torch.from_numpy(img).float()  # 转换为张量
img = img.unsqueeze(0)  # 增加一个通道维度
img = img.unsqueeze(0)  # 增加一个batch维度

# 加载模型
device = 'cuda:1' if torch.cuda.is_available() else 'cpu'
model = CNN_Network().to(device)
model.load_state_dict(torch.load('captcha.pth'))
model.eval()

# 识别验证码
img = img.to(device)
pred = model(img)
captcha = tensor2captcha(pred)
print(f'识别到{filename}的验证码为{captcha}')

例如，上面的491e.jpg就可以被模型正确识别。