二维码图片的生成
利用captcha可以生成二维码图片
# -*- coding: UTF-8 -*-
from captcha.image import ImageCaptcha # pip install captcha
from PIL import Image
import random
import time
import os
# 验证码中的字符
# string.digits + string.ascii_uppercase
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']
ALL_CHAR_SET = NUMBER + ALPHABET #验证码所有的字符和数字
ALL_CHAR_SET_LEN = len(ALL_CHAR_SET)
MAX_CAPTCHA = 12 #每个验证码字符个数
# 图像大小
IMAGE_HEIGHT = 60
IMAGE_WIDTH = 160
train_path = 'dataset' + os.path.sep + 'train'
test_path = 'dataset' + os.path.sep + 'test'
predict_path = 'dataset' + os.path.sep + 'predict'
def random_captcha():
#生成随机字符串
captcha_text = []
for i in range(MAX_CAPTCHA):
c = random.choice(ALL_CHAR_SET)
captcha_text.append(c)
return ''.join(captcha_text)
# 生成字符对应的验证码
def gen_captcha_text_and_image():
image = ImageCaptcha()
captcha_text = random_captcha()
#通过随机字符串生成二维码图片
captcha_image = Image.open(image.generate(captcha_text))
return captcha_text, captcha_image
if __name__ == '__main__':
count = 300 #生成二维码的个数
path = train_path #通过改变此处目录，以生成 训练、测试和预测用的验证码集
if not os.path.exists(path):
os.makedirs(path)
for i in range(count):
now = str(int(time.time()))
text, image = gen_captcha_text_and_image()
filename = text+'_'+str(i)+'.png'
#保存二维码图片
image.save(path + os.path.sep + filename)
print('saved %d : %s' % (i+1,filename))dataset
# -*- coding: UTF-8 -*-
import os
import random
from torch.utils.data import DataLoader,Dataset
import torchvision.transforms as transforms
from PIL import Image
import one_hot_encoding as ohe
import captcha_setting
class mydataset(Dataset):
def __init__(self, folder, transform=None):
self.train_image_file_paths = [os.path.join(folder, image_file) for image_file in os.listdir(folder)]
self.transform = transform
def __len__(self):
return len(self.train_image_file_paths)
def __getitem__(self, idx):
image_root = self.train_image_file_paths[idx]
image_name = image_root.split(os.path.sep)[-1]
image = Image.open(image_root)
if self.transform is not None:
image = self.transform(image)
label = ohe.encode(image_name.split('_')[0]) # 为了方便，在生成图片的时候，图片文件的命名格式 "4个数字或者数字_时间戳.PNG", 4个字母或者即是图片的验证码的值，字母大写,同时对该值做 one-hot 处理
return image, label
transform = transforms.Compose([
transforms.ColorJitter(),
# transforms.Grayscale(),
transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def get_train_data_loader():
dataset = mydataset(captcha_setting.TRAIN_DATASET_PATH, transform=transform)
return DataLoader(dataset, batch_size=64, shuffle=True)
def get_test_data_loader():
dataset = mydataset(captcha_setting.TEST_DATASET_PATH, transform=transform)
return DataLoader(dataset, batch_size=1, shuffle=True)
def get_predict_data_loader():
dataset = mydataset(captcha_setting.PREDICT_DATASET_PATH, transform=transform)
return DataLoader(dataset, batch_size=1, shuffle=True)
if __name__=="__main__":
from matplotlib import pyplot as plt
dataset=mydataset('./dataset/train',transform=transform)
indexes=random.sample(range(len(dataset)-1),16)
image,label=dataset[0]
for i,index in enumerate(indexes):
image,label=dataset[index]
image=transforms.ToPILImage()(image)
plt.subplot(4,4,i+1)
plt.title(ohe.decode(label))
plt.xticks([])
plt.yticks([])
plt.imshow(image)
plt.show()
one-hot编码
# -*- coding: UTF-8 -*-
import numpy as np
import captcha_setting
def encode(text):
vector = np.zeros(captcha_setting.ALL_CHAR_SET_LEN * captcha_setting.MAX_CAPTCHA, dtype=float)
#每个字符都需要进行编码.
#每个字符都需要在字典中查询得到，所有变量的维度是 *max_captcha
def char2pos(c):
if c =='_':
k = 62
return k
k = ord(c)-48
#将acii码转换为字符
# hh=chr(k+48)
if k > 9:
k = ord(c) - 65 + 10
if k > 35:
k = ord(c) - 97 + 26 + 10
if k > 61:
raise ValueError('error')
return k
for i, c in enumerate(text):
idx = i * captcha_setting.ALL_CHAR_SET_LEN + char2pos(c)
vector[idx] = 1.0
return vector
def decode(vec):
char_pos = vec.nonzero()[0]
text=[]
for i, c in enumerate(char_pos):
char_at_pos = i #c/63
char_idx = c % captcha_setting.ALL_CHAR_SET_LEN
if char_idx < 10:
char_code = char_idx + ord('0')
elif char_idx <36:
char_code = char_idx - 10 + ord('A')
elif char_idx < 62:
char_code = char_idx - 36 + ord('a')
elif char_idx == 62:
char_code = ord('_')
else:
raise ValueError('error')
text.append(chr(char_code))
return "".join(text)
if __name__ == '__main__':
e = encode("9L7H")
print(e)
print(decode(e))[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. 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. 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. 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.]
9L7H
模型
# -*- coding: UTF-8 -*-
import torch.nn as nn
import captcha_setting
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, padding=1),
nn.BatchNorm2d(32),
nn.Dropout(0.5), # drop 50% of the neuron
nn.ReLU(),
nn.MaxPool2d(2))
self.layer2 = nn.Sequential(
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.Dropout(0.5), # drop 50% of the neuron
nn.ReLU(),
nn.MaxPool2d(2))
self.layer3 = nn.Sequential(
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.Dropout(0.5), # drop 50% of the neuron
nn.ReLU(),
nn.MaxPool2d(2))
self.fc = nn.Sequential(
nn.Linear((captcha_setting.IMAGE_WIDTH//8)*(captcha_setting.IMAGE_HEIGHT//8)*64, 1024),
nn.Dropout(0.5), # drop 50% of the neuron
nn.ReLU())
self.rfc = nn.Sequential(
nn.Linear(1024, captcha_setting.MAX_CAPTCHA*captcha_setting.ALL_CHAR_SET_LEN),
)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = self.layer3(out)
out = out.view(out.size(0), -1)
out = self.fc(out)
out = self.rfc(out)
return out
if __name__=='__main__':
import torch
model=CNN()
input=torch.randn(1,1,60,160)
gt_output=torch.randn(1,144)
output=model(input)
print('输出的向量大小',output.shape)
criterion = nn.MultiLabelSoftMarginLoss()
loss=criterion(gt_output,output)
print('损失的大小',loss.data)
输出的向量大小 torch.Size([1, 144])
损失的大小 tensor(0.7603)
模型的训练
# -*- coding: UTF-8 -*-
import torch
import torch.nn as nn
from torch.autograd import Variable
import my_dataset
from captcha_cnn_model import CNN
# Hyper Parameters
num_epochs = 30
batch_size = 100
learning_rate = 0.001
def main():
cnn = CNN()
cnn.train()
print('init net')
criterion = nn.MultiLabelSoftMarginLoss()
optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
# Train the Model
train_dataloader = my_dataset.get_train_data_loader()
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_dataloader):
images = Variable(images)
labels = Variable(labels.float())
predict_labels = cnn(images)
loss = criterion(predict_labels, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i+1) % 10 == 0:
print("epoch:", epoch, "step:", i, "loss:", loss.item())
if (i+1) % 100 == 0:
torch.save(cnn.state_dict(), "./model.pkl") #current is model.pkl
print("save model")
print("epoch:", epoch, "step:", i, "loss:", loss.item())
torch.save(cnn.state_dict(), "./model.pth")
print("save last model")
if __name__ == '__main__':
main()
init net
epoch: 0 step: 4 loss: 0.38898029923439026
epoch: 1 step: 4 loss: 0.19714970886707306
epoch: 2 step: 4 loss: 0.15116369724273682
epoch: 3 step: 4 loss: 0.1429138034582138
epoch: 4 step: 4 loss: 0.1360236257314682
epoch: 5 step: 4 loss: 0.12835916876792908
epoch: 6 step: 4 loss: 0.1267365664243698
epoch: 7 step: 4 loss: 0.12457828223705292
epoch: 8 step: 4 loss: 0.12483084201812744
epoch: 9 step: 4 loss: 0.11893215030431747
epoch: 10 step: 4 loss: 0.11814623326063156
epoch: 11 step: 4 loss: 0.11591014266014099
epoch: 12 step: 4 loss: 0.11125991493463516
epoch: 13 step: 4 loss: 0.10649068653583527
epoch: 14 step: 4 loss: 0.10284445434808731
epoch: 15 step: 4 loss: 0.10144951194524765
epoch: 16 step: 4 loss: 0.0985511839389801
epoch: 17 step: 4 loss: 0.08964875340461731
epoch: 18 step: 4 loss: 0.08870525658130646
epoch: 19 step: 4 loss: 0.0839766412973404
epoch: 20 step: 4 loss: 0.0823589637875557
epoch: 21 step: 4 loss: 0.07506724447011948
epoch: 22 step: 4 loss: 0.06370603293180466
epoch: 23 step: 4 loss: 0.06234220042824745
epoch: 24 step: 4 loss: 0.06265763193368912
epoch: 25 step: 4 loss: 0.05445406585931778
epoch: 26 step: 4 loss: 0.05590423569083214
epoch: 27 step: 4 loss: 0.0482553206384182
epoch: 28 step: 4 loss: 0.04553262144327164
epoch: 29 step: 4 loss: 0.03754893317818642
save last model
模型的测试
# -*- coding: UTF-8 -*-
import numpy as np
import torch
from torch.autograd import Variable
import captcha_setting
import my_dataset
from captcha_cnn_model import CNN
import one_hot_encoding
def main():
cnn = CNN()
cnn.eval()
cnn.load_state_dict(torch.load('model.pth'))
print("load cnn net.")
test_dataloader = my_dataset.get_test_data_loader()
correct = 0
total = 0
for i, (images, labels) in enumerate(test_dataloader):
image = images
vimage = Variable(image)
predict_label = cnn(vimage)
c0 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, 0:captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
c1 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, captcha_setting.ALL_CHAR_SET_LEN:2 * captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
c2 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, 2 * captcha_setting.ALL_CHAR_SET_LEN:3 * captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
c3 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, 3 * captcha_setting.ALL_CHAR_SET_LEN:4 * captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
predict_label = '%s%s%s%s' % (c0, c1, c2, c3)
true_label = one_hot_encoding.decode(labels.numpy()[0,:])
total += labels.size(0)
if(predict_label == true_label):
correct += 1
if(total%200==0):
print('Test Accuracy of the model on the %d test images: %f %%' % (total, 100 * correct / total))
print('Test Accuracy of the model on the %d test images: %f %%' % (total, 100 * correct / total))
if __name__ == '__main__':
main()
