Pytorch 验证码识别

阿里云双11来了!从本博客参与阿里云,服务器最低只要86元/年!

思路

网上有很多验证码的生成方式,而在python中最常用的验证码生成模块就是:captchagvcode,前者的示例可以参考:链接,这里用后者生成验证码并识别(其实原理都一样就是了…)

导入模块

如果在colab环境下,需要先安装验证码生成模块,代码如下:

# !pip install graphic-verification-code

然后导入所有需要的模块

import torch
from torch import nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import gvcode
import random

初始化定义

这里首先配置一下GPU,然后还有所有的验证码字符集,以及训练的batch等,代码如下:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# torch.cuda.is_available()
all_char = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
    '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', 
    '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']
features = len(all_char)
batch_size = 32

数据预处理

为了节约内存,这里验证码使用生成器生成,生成器定义如下:

def gen(batch_size=32):
    data_img = []
    data_x = torch.zeros(batch_size, 80, 170, 3)
    data_y = torch.zeros(batch_size, 4, features)
    for i in range(batch_size):
        img, chars = gvcode.generate(size=(170, 80), chars=all_char, fg_color=(random.randint(0, 256), random.randint(0, 256), random.randint(0, 256)))
        data_img.append((img, chars))
        data_x[i] = torch.tensor(np.array(img))
        for j in range(4):
            data_y[i][j][all_char.index(chars[j])] = 1
            # 验证码4个字符的one-hot
        # break
    data_x.transpose_(1, 3)
    # 转成格式:batch_size, 3, 120, 30,因为输入格式里好像通道要在宽高前面
    data_y.transpose_(0, 1)
    # 转成格式:4, batch_size, features,因为多输出模型,4个(batch_size, features)分别对应验证码4个位置的字符one-hot
    return data_x, data_y, data_img
# 测试生成器
# x, y, _ = gen(32)
# x.shape, y.shape
# plt.imshow(x[0].transpose_(0, 2))
# plt.show()

定义网络模型

模型的前面先用几次卷积提取特征,最后由于验证码有4个值,所以用4个全连接分别训练,最后用softmax进行判断,代码如下:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1 = nn.Conv2d(3, 32, 3, 1).to(device)
        self.layer2 = nn.Conv2d(32, 32, 3, 1)
        self.layer3 = nn.MaxPool2d(2, 2)
        self.layer4 = nn.Conv2d(32, 64, 3, 1)
        self.layer5 = nn.Conv2d(64, 64, 3, 1)
        self.layer6 = nn.MaxPool2d(2, 2)
        self.layer7 = nn.Conv2d(64, 128, 3, 1)
        self.layer8 = nn.Conv2d(128, 128, 3, 1)
        self.layer9 = nn.MaxPool2d(2, 2)
        self.layer10 = nn.Conv2d(128, 256, 3, 1)
        self.layer11 = nn.Conv2d(256, 256, 3, 1)
        self.layer12 = nn.MaxPool2d(2, 2)
        # 多输出模型,输出4个,分别对应4个字母,格式都是:(batch_size, features)
        self.layer13_1 = nn.Linear(batch_size*48, features)
        self.layer13_2 = nn.Linear(batch_size*48, features)
        self.layer13_3 = nn.Linear(batch_size*48, features)
        self.layer13_4 = nn.Linear(batch_size*48, features)
        self.layer14 = nn.Softmax()
    
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = self.layer3(x)
        x = F.relu(self.layer4(x))
        x = F.relu(self.layer5(x))
        x = self.layer6(x)
        x = F.relu(self.layer7(x))
        x = F.relu(self.layer8(x))
        x = self.layer9(x)
        x = F.relu(self.layer10(x))
        x = F.relu(self.layer11(x))
        x = self.layer12(x)
        
        x = torch.flatten(x).reshape(batch_size, -1)
        # x = F.dropout(x, 0.3)
        x1 = self.layer14(self.layer13_1(x))
        x2 = self.layer14(self.layer13_2(x))
        x3 = self.layer14(self.layer13_3(x))
        x4 = self.layer14(self.layer13_4(x))
        return x1, x2, x3, x4
model = Net().to(device)

定义损失函数和优化器

loss_fun = nn.MSELoss()
# loss_fun = nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.parameters(), lr=0.0001)

训练模型

model.train()
li_loss = []
for epoch in range(1, 5000):
    x, y, _ = gen(batch_size)
    output1, output2, output3, output4 = model(x.to(device))
    # print(output1.shape, y[0].shape)
    loss1 = loss_fun(output1, y[0].to(device))
    # 32,10
    loss2 = loss_fun(output2, y[1].to(device))
    loss3 = loss_fun(output3, y[2].to(device))
    loss4 = loss_fun(output4, y[3].to(device))
    loss = loss1 + loss2 + loss3 + loss4
    optim.zero_grad()
    loss.backward()
    optim.step()
    print("epoch:{}, loss1:{}, loss2:{}, loss3:{}, loss4:{}".format(epoch, loss1, loss2, loss3, loss4))
    torch.save(model.state_dict(), "gvcode_{}.mdl".format(batch_size))
    li_loss.append(loss)
    if epoch % 30 == 0:
        # 每30趟看一下损失的变化
        plt.plot(li_loss)
        plt.show()

测试模型

产生batch的数据测试看看结果如何,代码如下:

x, y, _ = gen(batch_size)
model.eval()
output1, output2, output3, output4 = model(x.to(device))
for i in range(32):
    print(_[i][1], "{}{}{}{}".format(all_char[output1[i].argmax()], all_char[output2[i].argmax()], all_char[output3[i].argmax()], all_char[output4[i].argmax()]))

完整代码示例-gvcode识别

# 安装模块
# !pip install graphic-verification-code
# -----------------------------
# 导入模块
import torch
from torch import nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import gvcode
import random
# -----------------------------
# 设置gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.cuda.is_available()
# -----------------------------
# 验证码字符集和基本参数设置
all_char = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
    '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', 
    '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']
features = len(all_char)
batch_size = 32
# -----------------------------
# 定义验证码生成器
def gen(batch_size=32):
    data_img = []
    data_x = torch.zeros(batch_size, 80, 170, 3)
    data_y = torch.zeros(batch_size, 4, features)
    for i in range(batch_size):
        img, chars = gvcode.generate(size=(170, 80), chars=all_char, fg_color=(random.randint(0, 256), random.randint(0, 256), random.randint(0, 256)))
        data_img.append((img, chars))
        data_x[i] = torch.tensor(np.array(img))
        for j in range(4):
            data_y[i][j][all_char.index(chars[j])] = 1
            # 验证码4个字符的one-hot
        # break
    data_x.transpose_(1, 3)
    # 转成格式:batch_size, 3, 120, 30,因为输入格式里好像通道要在宽高前面
    data_y.transpose_(0, 1)
    # 转成格式:4, batch_size, features,因为多输出模型,4个(batch_size, features)分别对应验证码4个位置的字符one-hot
    return data_x, data_y, data_img
# 测试生成器
# x, y, _ = gen(32)
# x.shape, y.shape
# plt.imshow(x[0].transpose_(0, 2))
# plt.show()
# -----------------------------
# 定义网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1 = nn.Conv2d(3, 32, 3, 1).to(device)
        self.layer2 = nn.Conv2d(32, 32, 3, 1)
        self.layer3 = nn.MaxPool2d(2, 2)
        self.layer4 = nn.Conv2d(32, 64, 3, 1)
        self.layer5 = nn.Conv2d(64, 64, 3, 1)
        self.layer6 = nn.MaxPool2d(2, 2)
        self.layer7 = nn.Conv2d(64, 128, 3, 1)
        self.layer8 = nn.Conv2d(128, 128, 3, 1)
        self.layer9 = nn.MaxPool2d(2, 2)
        self.layer10 = nn.Conv2d(128, 256, 3, 1)
        self.layer11 = nn.Conv2d(256, 256, 3, 1)
        self.layer12 = nn.MaxPool2d(2, 2)
        # 多输出模型,输出4个,分别对应4个字母,格式都是:(batch_size, features)
        self.layer13_1 = nn.Linear(batch_size*48, features)
        self.layer13_2 = nn.Linear(batch_size*48, features)
        self.layer13_3 = nn.Linear(batch_size*48, features)
        self.layer13_4 = nn.Linear(batch_size*48, features)
        self.layer14 = nn.Softmax()
    
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = self.layer3(x)
        x = F.relu(self.layer4(x))
        x = F.relu(self.layer5(x))
        x = self.layer6(x)
        x = F.relu(self.layer7(x))
        x = F.relu(self.layer8(x))
        x = self.layer9(x)
        x = F.relu(self.layer10(x))
        x = F.relu(self.layer11(x))
        x = self.layer12(x)
        
        x = torch.flatten(x).reshape(batch_size, -1)
        # x = F.dropout(x, 0.3)
        # print(x.shape)
        x1 = self.layer14(self.layer13_1(x))
        x2 = self.layer14(self.layer13_2(x))
        x3 = self.layer14(self.layer13_3(x))
        x4 = self.layer14(self.layer13_4(x))
        return x1, x2, x3, x4
model = Net().to(device)
# model
# -----------------------------
# 损失函数和优化器
loss_fun = nn.MSELoss()
# loss_fun = nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.parameters(), lr=0.0001)
# -----------------------------
# 训练数据
model.train()
li_loss = []
for epoch in range(1, 5000):
    x, y, _ = gen(batch_size)
    output1, output2, output3, output4 = model(x.to(device))
    # print(output1.shape, y[0].shape)
    loss1 = loss_fun(output1, y[0].to(device))
    # 32,10
    loss2 = loss_fun(output2, y[1].to(device))
    loss3 = loss_fun(output3, y[2].to(device))
    loss4 = loss_fun(output4, y[3].to(device))
    loss = loss1 + loss2 + loss3 + loss4
    optim.zero_grad()
    loss.backward()
    optim.step()
    print("epoch:{}, loss1:{}, loss2:{}, loss3:{}, loss4:{}".format(epoch, loss1, loss2, loss3, loss4))
    torch.save(model.state_dict(), "gvcode_{}.mdl".format(batch_size))
    li_loss.append(loss)
    if epoch % 30 == 0:
        plt.plot(li_loss)
        plt.show()
# -----------------------------
# 测试数据
x, y, _ = gen(batch_size)
model.eval()
output1, output2, output3, output4 = model(x.to(device))
for i in range(32):
    print(_[i][1], "{}{}{}{}".format(all_char[output1[i].argmax()], all_char[output2[i].argmax()], all_char[output3[i].argmax()], all_char[output4[i].argmax()]))

https://www.jianshu.com/p/359fd77d17f2

Python量化投资网携手4326手游为资深游戏玩家推荐:《天国旅立下载

「点点赞赏,手留余香」

    还没有人赞赏,快来当第一个赞赏的人吧!
0 条回复 A 作者 M 管理员
    所有的伟大,都源于一个勇敢的开始!
欢迎您,新朋友,感谢参与互动!欢迎您 {{author}},您在本站有{{commentsCount}}条评论