基于pytorch的验证码识别

我们在爬虫的时候，经常会遇到验证码的识别，尽管我在上一篇博客中介绍了使用超级鹰来进行验证码的识别（这里所指的验证码是指只有数字、字母所组成的照片），但是毕竟是需要花钱的，花钱谁不心疼。既然人工智能这么火的，我们也来蹭蹭热度，使用深度学习来进行验证码的识别。深度学习解决问题基本上都是四步走：获取数据集、定义模型、训练模型以及测试模型。下面跟着小墨一起来学习吧。

一 、获取和加载数据

1. 生成验证码照片

在生成验证码照片时，我们直接使用captcha包下的image文件中的ImageCaptcha类，具体代码如下所示：

from captcha.image import ImageCaptcha
image = ImageCaptcha()  #实例化一个Image对象，生成默认大小的验证码照片
#image = ImageCaptcha(width=width, height=height)  生成指定大小的验证码照片
image.write('2eRt', '2eRtjpg')   #（验证码上的信息，保存的路径）

生成的验证码照片如下图所示：

因此，我们可以通过循环生成几千、几万张照片。

2. one_hot编码

在本项目中，标签就是长度为4的数字和字母组合的文本信息，我们需要对每个文本信息进行编码。具体做法是初始化一个长度为（1，4*62）的向量（其中，4代表验证码上文本信息的长度，62代表数字和大小写字母的总个数），然后让相应字符的位置上为1，其余为0即可。

def one_hot(text):
    vector = np.zeros(4*62)  #(10+26+26)*4
    def char2pos(c):
        if c =='_':
            k = 62
            return k
        k = ord(c)-48
        if k > 9:
            k = ord(c) - 55
            if k > 35:
                k = ord(c) - 61
                if k > 61:
                    raise ValueError('No Map')
        return k
    for i, c in enumerate(text):
        idx = i * 62 + char2pos(c)
        vector[idx] = 1.0
    return vector

3. 数据加载

通过上面的两步，照片和标签都已经准备好了，下面我们直接使用Pytorch提供的DataLoader（）方法来加载数据，其步骤主要分为两步：
（1）定义一个继承Dataset的类，用来获取数据。其中函数__init__(self)主要用来初始化一些参数、函数__getitem__(self.index)主要是通过Index来加载相应的数据，并进行相应的转化、函数__len__(self)是返回数据集的大小。

class MyDataSet(Dataset):
    def __init__(self,dir):
        self.dir=dir
        self.img_name= next(os.walk(self.dir))[2]
    def __getitem__(self, index):
        img_path = os.path.join(self.dir,self.img_name[index])
        img = cv2.imread(img_path,0)
        img = img/255.
        img = torch.from_numpy(img).float()
        img = torch.unsqueeze(img,0)
        label = torch.from_numpy(one_hot(self.img_name[index][:-4])).float()
        return img,label
    def __len__(self):
        return len(self.img_name)

（2）使用 DataLoader（)以batch_size的批次加载数据。

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

至此，数据集已准备完备。

二 、定义模型

这里直接贴上模型的代码，有兴趣的自己也可以试着去搭建一些模型，调一调相关的参数。

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(inplace=True)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, stride=1, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(stride=2, kernel_size=2),  # 30 80
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, stride=1, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2),   # 15 40
        )

        self.fc = nn.Sequential(
            nn.Linear(128 * 15 * 40, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 248)
        )
        
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

三 、训练模型

 net = CNN_Network()
 optimizer = torch.optim.Adam(net.parameters(), lr=lr)
 loss = nn.MultiLabelSoftMarginLoss()    #多分类损失函数
def train(net, train_iter, test_iter, optimizer, loss,device, num_epochs):
    net = net.to(device)
    for epoch in range(num_epochs):
        for X, y in train_iter:
            X = X.to(device)
            y = y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
    torch.save(net.state_dict(), "yanzhengma.pkl")   #保存模型

四、测试模型

number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u',
            'v', 'w', 'x', 'y', 'z']
ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U',
            'V', 'W', 'X', 'Y', 'Z']
path="./images/test/0UE8.jpg"  #要测试照片的路径
net=CNN_Network()           
net.load_state_dict(torch.load("yanzhengma.pkl"))   #加载上步保存的模型       
img= cv2.imread(path,0)
img = img/255.
img = torch.from_numpy(img).float()
img = torch.unsqueeze(img,0)
img = torch.unsqueeze(img,0)
pred = net(img)
a1 = torch.argmax(pred[0,:62],dim=0)    
a2 = torch.argmax(pred[0,62:124],dim=0)	
a3 = torch.argmax(pred[0,124:186],dim=0)
a4 = torch.argmax(pred[0,186:],dim=0)
pred = [a1,a2,a3,a4]
labels=number+ALPHABET+alphabet
for i in pred:
    print(labels[i.item()],end='')
print('/n')

运行上面的代码，输出结果为：

五、数据可视化

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

#计算准确度
def get_acc(net,data_iter,device):
    acc_sum,n=0,0
    for X,y in data_iter:
        X=X.to(device)
        y=y.to(device)
        y_hat=net(X)
        pre1=torch.argmax(y_hat[:,:62],dim=1)
        real1=torch.argmax(y[:,:62],dim=1)
        pre2 = torch.argmax(y_hat[:, 62:124], dim=1)
        real2 = torch.argmax(y[:, 62:124], dim=1)
        pre3 = torch.argmax(y_hat[:, 124:186], dim=1)
        real3 = torch.argmax(y[:, 124:186], dim=1)
        pre4 = torch.argmax(y_hat[:, 186:], dim=1)
        real4 = torch.argmax(y[:, 186:], dim=1)
        pre_lable=torch.cat((pre1,pre2,pre3,pre4),0).view(4,-1)
        real_label=torch.cat((real1,real2,real3,real4),0).view(4,-1)
        bool_=(pre_lable==real_label).transpose(0,1)
        n+=y.shape[0]
        for i in range(0,y.shape[0]):
            if bool_[i].int().sum().item()==4:
                acc_sum+=1
    return acc_sum/n
    
 #可视化准确度   
def draw_acc(train_acc,test_acc,epoch):
    plt.clf()
    x=[i for i in range(epoch)]
    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')

以上我们完成所有的工作。这里小墨建议大家在生成验证码照片时多生成一些，否则可能会出现严重的过拟合现象，训练集的准确度接近100%，测试集准确度却很低，而我们更加看重的是测试集的准确率，测试集的泛化能力。如下图所示：

下面贴上完整代码：

train.py

import tensorflow as tf
import cv2
from captcha.image import ImageCaptcha
import numpy as np
from PIL import Image
import random
import os
from torch.utils.data import DataLoader,Dataset
import torchvision.transforms as transforms
import torch
import torch.nn as nn
import matplotlib.pyplot as plt


number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u',
            'v', 'w', 'x', 'y', 'z']
ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U',
            'V', 'W', 'X', 'Y', 'Z']


#width=180
#height=60

def random_captcha_text(char_set=number+alphabet+ALPHABET, captcha_size=4):
    captcha_text = []
    for i in range(captcha_size):
        c = random.choice(char_set)
        captcha_text.append(c)
    return captcha_text
    
def gen_captcha_text_and_image(total_number,flag):
    image = ImageCaptcha()
    #image = ImageCaptcha(width=width, height=height)
    for i in range(total_number):
        captcha_text = random_captcha_text()
        captcha_text = ''.join(captcha_text)
        captcha = image.generate(captcha_text)
        image.write(captcha_text,'./images/' + flag + '/' + captcha_text + '.jpg')

def one_hot(text):
    vector = np.zeros(4*62)  #(10+26+26)*4
    def char2pos(c):
        if c =='_':
            k = 62
            return k
        k = ord(c)-48
        if k > 9:
            k = ord(c) - 55
            if k > 35:
                k = ord(c) - 61
                if k > 61:
                    raise ValueError('No Map')
        return k
    for i, c in enumerate(text):
        idx = i * 62 + char2pos(c)
        vector[idx] = 1
    return vector

class MyDataSet(Dataset):
    def __init__(self,dir):
        self.dir=dir
        self.img_name= next(os.walk(self.dir))[2]
    def __getitem__(self, index):
        img_path = os.path.join(self.dir,self.img_name[index])
        img = cv2.imread(img_path,0)
        img = img/255.
        img = torch.from_numpy(img).float()
        img = torch.unsqueeze(img,0)
        label = torch.from_numpy(one_hot(self.img_name[index][:-4])).float()
        return img,label
    def __len__(self):
        return len(self.img_name)

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(inplace=True)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, stride=1, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(stride=2, kernel_size=2),  # 30 80
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, stride=1, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2),   # 15 40
        )

        self.fc = nn.Sequential(
            nn.Linear(128 * 15 * 40, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 248)
        )
        
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x
def train(net, train_iter, test_iter, optimizer, loss,device, num_epochs):
    net = net.to(device)
    train_acc_list,test_acc_list,train_loss_list,test_loss_list=[],[],[],[]
    flag=0.0
    for epoch in range(num_epochs):
        train_loss,n=0.0,0
        for X, y in train_iter:
            X = X.to(device)
            y = y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            train_loss+= l.cpu().item()
            n+=y.shape[0]
        print("train_loss=",train_loss)
        train_acc=get_acc(net,train_iter,device)
        test_acc=get_acc(net,test_iter,device)
        print("test_acc=",test_acc)
        if epoch>=80:
            if test_acc>flag:
                torch.save(net.state_dict(), "yanzhengma.pkl")
        flag=test_acc
        train_acc_list.append(train_acc)
        test_acc_list.append(test_acc)
        #train_loss_list.append(train_loss)
    #torch.save(net.state_dict(), "yanzhengma.pkl")

    #draw_losss(train_loss_list,test_loss_list,num_epochs)
    draw_acc(train_acc_list,test_acc_list,num_epochs)
def get_acc(net,data_iter,device):
    acc_sum,n=0,0
    for X,y in data_iter:
        X=X.to(device)
        y=y.to(device)
        y_hat=net(X)
        pre1=torch.argmax(y_hat[:,:62],dim=1)
        real1=torch.argmax(y[:,:62],dim=1)
        pre2 = torch.argmax(y_hat[:, 62:124], dim=1)
        real2 = torch.argmax(y[:, 62:124], dim=1)
        pre3 = torch.argmax(y_hat[:, 124:186], dim=1)
        real3 = torch.argmax(y[:, 124:186], dim=1)
        pre4 = torch.argmax(y_hat[:, 186:], dim=1)
        real4 = torch.argmax(y[:, 186:], dim=1)
        pre_lable=torch.cat((pre1,pre2,pre3,pre4),0).view(4,-1)
        real_label=torch.cat((real1,real2,real3,real4),0).view(4,-1)
        bool_=(pre_lable==real_label).transpose(0,1)
        n+=y.shape[0]
        for i in range(0,y.shape[0]):
            if bool_[i].int().sum().item()==4:
                acc_sum+=1
    return acc_sum/n

def draw_losss(train_loss,test_loss,epoch):
    plt.clf()
    x=[i for i in range(epoch)]
    plt.plot(x,train_loss,label='train_loss')
    plt.plot(x,test_loss,label='test_loss')
    plt.legend()
    plt.title("loss goes by epoch")
    plt.xlabel('eopch')
    plt.ylabel('loss_value')
    plt.savefig('loss.png')

def draw_acc(train_acc,test_acc,epoch):
    plt.clf()
    x=[i for i in range(epoch)]
    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')
if __name__ == "__main__":
   
    if not os.path.exists('images/train'):
        os.mkdir('images/train')
    if not os.path.exists('images/test'):
        os.mkdir('images/test')

    gen_captcha_text_and_image(15000, 'train')
    gen_captcha_text_and_image(3000, 'test')

    epoch=200
    batch_size=256
    lr=0.001

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

    net = CNN_Network()
    optimizer = torch.optim.Adam(net.parameters(), lr=lr)

    loss = nn.MultiLabelSoftMarginLoss()
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    train(net, train_loader, test_loader, optimizer, loss, device, epoch)

test.py

import tensorflow as tf
import cv2
from captcha.image import ImageCaptcha
import numpy as np
from PIL import Image
import random
import os
from torch.utils.data import DataLoader,Dataset
import torchvision.transforms as transforms
import torch
import torch.nn as nn
from train import CNN_Network

number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u',
            'v', 'w', 'x', 'y', 'z']
ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U',
            'V', 'W', 'X', 'Y', 'Z']
path="./images/test/0UE8.jpg"
net=CNN_Network()           
net.load_state_dict(torch.load("yanzhengma.pkl"))        
img= cv2.imread(path,0)
if img.shape!=(60,160):
    cropImage(path)
    img= cv2.imread(path,0)
    
img = img/255.
img = torch.from_numpy(img).float()
img = torch.unsqueeze(img,0)
img = torch.unsqueeze(img,0)
pred = net(img)
a1 = torch.argmax(pred[0,:62],dim=0)    
a2 = torch.argmax(pred[0,62:124],dim=0)	
a3 = torch.argmax(pred[0,124:186],dim=0)
a4 = torch.argmax(pred[0,186:],dim=0)
pred = [a1,a2,a3,a4]
labels=number+ALPHABET+alphabet
for i in pred:
    print(labels[i.item()],end='')
print('/n')