PyTorch识别验证码

## 一、生成测试集数据

pip install captcha

common.py

import random
import time
captcha_array = list("0123456789abcdefghijklmnopqrstuvwxyz")
captcha_size = 4from captcha.image import ImageCaptchaif __name__ == '__main__':for i in range(10):image = ImageCaptcha()image_text = "".join(random.sample(captcha_array, captcha_size))image_path = "./datasets/train/{}_{}.png".format(image_text, int(time.time()))image.write(image_text, image_path)

生成验证码

二、one-hot编码将类别变量转换为机器学习算法易于利用的一种形式的过程。

one_hot.py

import common
import torch
import torch.nn.functional as Fdef text2vec(text):# 将文本转换为变量vectors = torch.zeros((common.captcha_size, common.captcha_array.__len__()))# vectors[0,0] = 1# vectors[1,3] = 1# vectors[2,4] = 1# vectors[3, 1] = 1for i in range(len(text)):vectors[i, common.captcha_array.index(text[i])] = 1return vectorsdef vectotext(vec):vec=torch.argmax(vec, dim=1)text_label=""for v in vec:text_label+=common.captcha_array[v]return  text_labelif __name__ == '__main__':vec=text2vec("aab1")print(vec, vec.shape)print(vectotext(vec))

三、 然后继续添加

my_datasets.py

import osfrom PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter
import one_hotclass mydatasets(Dataset):def __init__(self,root_dir):super(mydatasets, self).__init__()self.list_image_path=[ os.path.join(root_dir,image_name) for image_name in os.listdir(root_dir)]self.transforms=transforms.Compose([transforms.Resize((60,160)),transforms.ToTensor(),transforms.Grayscale()])def __getitem__(self, index):image_path = self.list_image_path[index]img_ = Image.open(image_path)image_name=image_path.split("\\")[-1]img_tesor=self.transforms(img_)img_lable=image_name.split("_")[0]img_lable=one_hot.text2vec(img_lable)img_lable=img_lable.view(1,-1)[0]return img_tesor,img_labledef __len__(self):return self.list_image_path.__len__()if __name__ == '__main__':d=mydatasets("datasets/train")img,label=d[0]writer=SummaryWriter("logs")writer.add_image("img",img,1)print(img.shape)writer.close()

dataLoader 加载dataset
就是数据加载器，结合了数据集和取样器，并且可以提供多个线程处理数据集。在训练模型时使用到此函数，用来把训练数据分成多个小组，此函数每次抛出一组数据。直至把所有的数据都抛出，就是做一个数据的初始化。

四、训练

五、CNN卷积神经网络

model.py

import torch
from torch import nn
import common
class mymodel(nn.Module):def __init__(self):super(mymodel, self).__init__()self.layer1 = nn.Sequential(# 卷积层nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, padding=1),# 激活层nn.ReLU(),# 池化层nn.MaxPool2d(kernel_size=2)   #[6, 64, 30, 80])self.layer2 = nn.Sequential(# 卷积层nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1),# 激活层nn.ReLU(),# 池化层nn.MaxPool2d(2)     #[6, 128, 15, 40])self.layer3 = nn.Sequential(# 卷积层nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1),# 激活层nn.ReLU(),# 池化层nn.MaxPool2d(2)  # [6, 256, 7, 20])self.layer4 = nn.Sequential(# 卷积层nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1),# 激活层nn.ReLU(),# 池化层nn.MaxPool2d(2)  # [6, 512, 3, 10])# self.layer5 = nn.Sequential(#     nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),#     nn.ReLU(),#     nn.MaxPool2d(2)  # [6, 512, 1, 5]# )self.layer6 = nn.Sequential(# 展平nn.Flatten(),    #[6, 2560] [64, 15360]# 线性层nn.Linear(in_features=15360, out_features=4096),# 防止过拟合nn.Dropout(0.2),  # drop 20% of the neuron# 激活曾nn.ReLU(),# 线性层nn.Linear(in_features=4096, out_features=common.captcha_size*common.captcha_array.__len__()))def forward(self,x):x = self.layer1(x)x = self.layer2(x)x = self.layer3(x)x = self.layer4(x)#x = x.view(1,-1)[0]#[983040]x = self.layer6(x)# x = x.view(x.size(0), -1)return x;if __name__ == '__main__':data = torch.ones(64, 1, 60, 160)model = mymodel()x = model(data)print(x.shape)

六、训练

train.py

import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from my_datasets import mydatasets
from model import mymodelif __name__ == '__main__':train_datas = mydatasets("datasets/train")test_data = mydatasets("datasets/test")train_dataloader = DataLoader(train_datas, batch_size=64, shuffle=True)test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)# m = mymodel().cuda()    没有GPUm = mymodel()# MultiLabelSoftMarginLoss 多标签交叉熵损失函数# 优化器 Adam 一般要求学习率比较小# 先将梯度归零 zero_grad# 反向传播计算 backward# loss_fn = nn.MultiLabelSoftMarginLoss().cuda()   没有GPUloss_fn = nn.MultiLabelSoftMarginLoss()optimizer = torch.optim.Adam(m.parameters(), lr=0.001)w = SummaryWriter("logs")total_step = 0for i in range(10):# print("外层训练次数{}".format(i))for i,(imgs, targets) in enumerate(train_dataloader):# imgs = imgs.cuda()   没有GPU# targets = targets.cuda()   没有GPUoutputs = m(imgs)loss = loss_fn(outputs, targets)optimizer.zero_grad()loss.backward()optimizer.step()if i%10 == 0:total_step += 1print("训练{}次,loss:{}".format(total_step*10, loss.item()))w.add_scalar("loss", loss, total_step)w.close()torch.save(m, "model.pth")

tensorboard --logdir=logs
使用tensorboard 查看损失率，接近零了。

七、图片预测

model.train() 和 model.eval()一般在模型训练和评价的时候会加上这两句，主要是针对由于model在训练时和评价时Batch Normalization 和Dropout方法模式不同，例如model指定t因此，在使用PyTorch进行训练和测试时一定注意要把rain/eval
predict.py

from PIL import Image
from torch.utils.data import DataLoader
import one_hot
import model
import torch
import common
import my_datasets
from torchvision import transformsdef test_pred():#  m = torch.load("model.pth").cuda()  没有GPUm = torch.load("model.pth")m.eval()test_data = my_datasets.mydatasets("datasets/test")test_dataloader = DataLoader(test_data, batch_size=1, shuffle=False)test_length = test_data.__len__()correct = 0for i, (imgs, lables) in enumerate(test_dataloader):# imgs = imgs.cuda()   没有GPU# lables = lables.cuda()  没有GPUlables = lables.view(-1, common.captcha_array.__len__())lables_text = one_hot.vectotext(lables)predict_outputs = m(imgs)predict_outputs = predict_outputs.view(-1, common.captcha_array.__len__())predict_labels = one_hot.vectotext(predict_outputs)if predict_labels == lables_text:correct += 1print("预测正确：正确值:{},预测值:{}".format(lables_text, predict_labels))else:print("预测失败:正确值:{},预测值:{}".format(lables_text, predict_labels))# m(imgs)print("正确率{}".format(correct / test_length * 100))def pred_pic(pic_path):img = Image.open(pic_path)tersor_img = transforms.Compose([transforms.Grayscale(),transforms.Resize((60, 160)),transforms.ToTensor()])# img = tersor_img(img).cuda() 没有GPUimg = tersor_img(img)print(img.shape)img = torch.reshape(img, (-1, 1, 60, 160))print(img.shape)# m = torch.load("model.pth").cuda() 没有GPUm = torch.load("model.pth")outputs = m(img)outputs = outputs.view(-1, len(common.captcha_array))outputs_lable = one_hot.vectotext(outputs)print(outputs_lable)if __name__ == '__main__':# test_pred()pred_pic("./datasets/test/5ogl_1705418909.png")

预测值是一样的，需要找一些真实的验证码图片