【深度学习】四位定长图形验证码识别(基于pytorch)

题记

验证码在我们网络中已经无处不在了,再爬虫中遇到也非常头疼,大部分都采用第三方平台如"超级鹰",进行识别,有没有什么方法可以不调第三方API呢。当然有那就是深度学习,本文介绍一下从数据集采集
标注,构建网络,训练模型,测试结果,完整体系。找几张生活常见的验证码如登陆场景

【深度学习】四位定长图形验证码识别(基于pytorch)_第1张图片

【深度学习】四位定长图形验证码识别(基于pytorch)_第2张图片
在登陆模块最常见的就是这种四位数字+英文验证码接下来实现完整方案。

数据采集

深度学习的第一步就是采集数据,一个好的数据集就成功了一半了。数据采集也就是爬虫部分,大部分python学者因该都不陌生。代码如下:

header = headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/75.0.3770.142 Safari/537.36"
             }
def save_dataset(path, mode, img):
    with open(path, mode=mode) as f:
        f.write(img)
    
    
def spider(url, num):
    """
    4位验证码
    :url 网站登录地址
    :num 需要爬取图片个数
    :return:
    """
    for i in range(num):
        try:
            img = requests.get(url, headers=header, timeout=2).content
            #  img_name = ddddocrbz(img)
            img_name = str(i)
            # 保存路径
            save_path = os.path.join("./train_all", img_name)
            # 保存数据集
            save_dataset(save_path, "wb", img)
            print(f"采集成功4位验证码  {i} {img_name}")
        except:
            print("请求失败,重试")
            
 spider("你要爬取的网站", 20000) 

采集脚本很简单,就是将二进制文件,保存到本地文件夹当中。

数据标注

我们有了图片之后要干什么呢,当时是数据标注了,我们需要把图片重命名如:
请添加图片描述
这张图片我们需要让他的文件名是1a5z。
那我们因该怎么做呢,如果一张一张标注当然可以而且准确,但是非常耗费人工和时间,这里推荐一个开源库 ddddocr,里面有api可以识别验证码,如用如下:

pip install ddddocr -i https://pypi.tuna.tsinghua.edu.cn/simple

调用代码如下

import ddddocr
ocr = ddddocr.DdddOcr()
with open('1.png', 'rb') as f:
    img_bytes = f.read()
res = ocr.classification(img_bytes)
print(res)

实例化ocr对象将一张二进制图片传入classification api即可得到结果,是不是一级棒。

https://github.com/sml2h3/ddddocr

这是ddddocr的github 如果想详细了的可以查看连接,让我们致敬大佬。

建议将我们爬取到的图片直接进行标注存储,这样就省去时间。代码如下:

ocr = ddddocr.DdddOcr(show_ad=False)

def save_dataset(path, mode, img):
    with open(path, mode=mode) as f:
        f.write(img)

def check_text(text: str, len_num: int):
    if len(text) != len_num:
        return False
    text = text.lower()
    captcha_array = "0123456789abcdefghijklmnopqrstuvwxyz"
    for t in text:
        if t not in captcha_array:
            print("校验失败 {}".format(text))
            return False
    return text

def ddddocrbz(img_bytes):
    res = ocr.classification(img_bytes)
    return res

def spider(url, num):
    """
    4位验证码
    :url 网站登录地址
    :num 需要爬取图片个数
    :return:
    """
    for i in range(num):
        try:
            img = requests.get(url, headers=header, timeout=2).content
            # ddddocr识别结果
            img_name = ddddocrbz(img)
            # 校验识别结果是否在预期
            img_name = check_text(img_name, 4)
            if img_name:
	            # 保存路径
	            save_path = os.path.join("./qufushifan", img_name +"_" +str(time.time())+".png")
	            # 保存数据集
	            save_dataset(save_path, "wb", img)
	            print(f"采集成功4位验证码   {i} {img_name}")
        except:
            print("请求失败,重试")
            
 spider("你要爬取的网站", 20000) 

我们先实例化一次ocr对象写在第一行,免得每次都要实例化,可以看到 check_text用来校验ddddocr识别的内容,是否在预期之内,我这里将全部转为小写,并判断是否在其中和长度。然后直接将请求的二进制图片,传入ddddocr进行,将结果作为图片名称,为了保证名字不重复,加了时间戳。
采集+标注,结果如下
【深度学习】四位定长图形验证码识别(基于pytorch)_第3张图片
这样我们就有了训练集了,采集的时候记得再猜一部分作为测试集我得训练集是2w,测试集是2000。
加粗样式

深度学习开始

采用pytorch框架

Dataset

我们训练自己的数据要自己定义dataset

import os
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
import one_hot

class my_dataset(Dataset):
    def __init__(self, root_dir):
        super(my_dataset, self).__init__()
        self.image_path = [os.path.join(root_dir, image_name) for image_name in os.listdir(root_dir)]
        self.transforms = transforms.Compose(
            [
                transforms.ToTensor(), # 转为tensor类型
                transforms.Resize([100, 300]),  # 修改尺寸
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), # 归一化
            ]

        )

    def __len__(self):
        return self.image_path.__len__()

    def __getitem__(self, index):
        image_path = self.image_path[index]
        image = Image.open(image_path).convert("RGB")
        image = self.transforms(image)
        label = image_path.split("/")[-1]
        label = label.split("_")[0].lower()
        # print(label)
        try:
            label_tensor = one_hot.text2Vec(label)  # [4,36]
        except:
            label_tensor = one_hot.text2Vec("labe")
        return image, label_tensor


if __name__ == '__main__':
    train_data = my_dataset("数据集路径")
    img, label = train_data[0]

继承torch Dataset实现__getitem__ 和__init__
init主要的功能是将我们的图片用
转换为tensor类型**,对图片修改尺寸 ,归一化
getitem将图片取出
PIL库读取进来转化为"RGB"
图片标签作为结果同时返回(我们标注时候是用下划线分开,取的时候通过下划线分割)
将标签进行onehot编码方便训练,可以看到统一将标签转为小写
我们可以拿出一张看一下:
在这里插入图片描述

onehot编码

import torch
captcha_array = list("0123456789abcdefghijklmnopqrstuvwxyz")
captcha_size = 4

def text2Vec(text):
	"""
	将标注文本转为onehot编码
	:return vec
	"""
    vec = torch.zeros(captcha_size, len(captcha_array))
    for i in range(len(text)):
        vec[i, captcha_array.index(text[i])] = 1
    return vec


def vec2Text(vec):
	"""
	将onehot编码转为标注文本
	:return text
	"""
    vec = torch.argmax(vec, dim=1)
    text = ""
    for i in vec:
        text += captcha_array[i]
    return text


if __name__ == '__main__':
    vec = text2Vec("aabcc1")
    print(vec, vec.shape)
    print(vec2Text(vec))

这是我们onehot转码的部分,里面用到了torch里的api如torch.zeros、torch.argmax
torch.zeros返回一个形状为为size,类型为torch.dtype,里面的每一个值都是0的tensor
torch.argmax返回指定维度最大值的索引

测试一下:

vec = text2Vec("abcd")

# 结果为
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.,
         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.,
         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.,
         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.,
         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]) torch.Size([4, 36])

output = vec2Text(vec)
# 结果为
abcd

就是一个tensor onehot的编码和解码

模型引用和训练

这里我们使用resnet34进行训练

import torch
import time
from tqdm import tqdm
from mydataset import my_dataset
from torch.utils.data import DataLoader
from torch import nn
from torch.optim import Adam
from Verification import Verification
import torchvision.models as models


if __name__ == '__main__':
    train_dataset = my_dataset("/train_all")
    train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True, drop_last=True, num_workers=4)
    verification = models.resnet34(num_classes=4 * 36).cuda(0)
    loss_fn = nn.MSELoss().cuda(0)
    optim = Adam(verification.parameters(), lr=0.001)
    verification.train()
    for epoch in range(20):
        now = time.time()
        print("训练epoch次数{}".format(epoch+1))
        bar = tqdm(enumerate(train_dataloader), total=len(train_dataloader))
        for i, (images, labels) in bar:
            optim.zero_grad()
            images = images.cuda(0)
            labels = labels.cuda(0)
            outputs = verification(images)
            outputs = outputs.reshape(64, 4, 36)
            loss = loss_fn(outputs, labels)
            loss.backward()
            optim.step()
            if i % 1000 == 0:
                print("训练次数{},损失率{}".format(i, loss.item()))
        torch.save(verification, "Verification_resnet34_4w.pth")
  1. 实例化dataset路径为刚刚爬取标注的目录
  2. 创建训练集加载器train_dataloader将batch_size设置为64 shuffle=True随机打乱
    drop_last舍弃不够一个batch_size的数据
  3. 创建模型 models.resnet34(num_classes=4 *
    36).cuda(0)我们直接用resnet34为我们的训练模型num_classes为类型4位 x (0到9+a到z)
  4. nn.MSELoss().cuda(0)损失函数为MSE
  5. Adam(verification.parameters(), lr=0.001),优化器采用Adam,学习率为0.001
  6. 开始训练20轮
  7. bar = tqdm(enumerate(train_dataloader),
    total=len(train_dataloader))为进度条。具体可看tqdm库
  8. optim.zero_grad()梯度清零
  9. verification(images) 将图片传入模型
  10. outputs.reshape(64, 4, 36) 将结果reshape变为64batch_size,4_位数,36_类型
  11. loss = loss_fn(outputs, labels) 传入损失函数求loss值
  12. loss.backward()反向传播
  13. optim.step()更新参数

由于我是在GPU训练的所以要加上cuda(0)
如果没有GPU不加就好了

开始训练:
【深度学习】四位定长图形验证码识别(基于pytorch)_第4张图片
可以看出来loss一直在下降
在这里插入图片描述
训练结束

测试模型

from tqdm import tqdm

import common
from mydataset import my_dataset
from torch.utils.data import DataLoader
import one_hot
import torch
captcha_array = list("0123456789abcdefghijklmnopqrstuvwxyz")


if __name__ == '__main__':
    test_dataset = my_dataset("/home/yons/Verification_Code/yzm46/yzm4_fixed/qufushifan_test")
    test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=True)
    model = torch.load("Verification_resnet34_4w.pth").cuda(0)
    total = len(test_dataset)
    yuce = 0
    bar = tqdm(enumerate(test_dataloader), total=len(test_dataloader))
    for i, (images, labels) in bar:
        images = images.cuda(0)
        labels = labels.cuda(0)
        labels = labels.view(-1, len(captcha_array))
        label_text = one_hot.vec2Text(labels)
        output = model(images)
        output_text = output.view(-1, len(captcha_array))
        output_text = one_hot.vec2Text(output_text)
        if output_text == label_text:
            yuce += 1
    print("正确率: {}%".format(yuce/total*100))

qufushifan_test目录为:
【深度学习】四位定长图形验证码识别(基于pytorch)_第5张图片
两千张测试集
将图片模型测试,得到onhot的tensor类型,再通过vec2Text转为字符串,与标注文本进行对比,最终得到准确率。
【深度学习】四位定长图形验证码识别(基于pytorch)_第6张图片
针对单一类型resnet34有很好的效果,虽然训练集只有2w张,准确率可以到达非常理想。

结束

resnet34和resnet18都不错,我这里采用的是resnet34因为比18多了16层所以导出的模型要大一些大约90多m,如果是resnet18大约40多m,针对单一类型的图形验证码还是很好解决的,如果要是别五位或者六位的也只需换一下参数,对于想解决此类验证码的朋友是一个不错的方法。

https://blog.csdn.net/qq_36551453?type=blog
这里@一下我的同事,也给予我很大帮助。

希望可以给朋友们一点帮助。
【深度学习】四位定长图形验证码识别(基于pytorch)_第7张图片

你可能感兴趣的:(验证码识别,pytorch,深度学习,python)