Pytorch之验证码识别

本文主要实现了两个工作：1.验证码生成   2.Pytorch识别验证码

一. 验证码生成

方法1. 利用PIL库的ImageDraw实现绘图，此法参考博客实现：

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 27 15:45:04 2018

@author: lps
"""

from PIL import Image, ImageDraw, ImageFont, ImageFilter
import random
import cv2
import numpy as np
import matplotlib.pyplot as plt
path = '/media/lps/python-3.5.2.amd64/Lib/site-packages/matplotlib/mpl-data/fonts/ttf/'     # 选择字体
data_path = '/home/lps/yanzm/'

# random chr
def rndChar():
      return chr(random.randint(65, 90))     # 随机字母

def rndInt():
      return str(random.randint(0,9))        # 随机数字

def rndColor():
      return (random.randint(64, 255), random.randint(64, 255), random.randint(64, 255))   # 随机颜色

def rndColor2():
      return (random.randint(32, 127), random.randint(32, 127), random.randint(32, 127))   # 随机颜色

def gaussian_noise():   # 高斯噪声
      mu =  125
      sigma = 20
      return tuple((np.random.normal(mu, sigma, 3).astype(int)))

def rotate(x, angle):  # 旋转
    M_rotate = cv2.getRotationMatrix2D((x.shape[0]/2, x.shape[1]/2), angle, 1)
    x = cv2.warpAffine(x, M_rotate, (x.shape[0], x.shape[1]))
    return x

width = 180 * 4
height = 180

def gen_image(num):

      for l in range(num): 

           image = Image.new('RGB', (width, height), (255, 255, 255))   # 先生成一张大图

           font = ImageFont.truetype(path+'cmb10.ttf', 36)

           draw = ImageDraw.Draw(image)    # 新的画板

           for x in range(0,width):
                 for y in range(0,height):
                       draw.point((x, y), fill=rndColor())

           label = []

           for t in range(4):    # 每一张验证码4个数字
                 numb = rndInt()
                 draw.text((180 * t + 60+10, 60+10), numb, font=font, fill=rndColor2())
                 label.append(numb)

           with open(data_path+"label.txt","a") as f:
                 for s in label:
                       f.write(s + ' ')
                 f.writelines("\n")     # 写入label

           img = image.filter(ImageFilter.GaussianBlur(radius=0.5))
           img = np.array(img)

           img1 = np.array([])

           for i in range(0,4):
                 img0 = img[:, 180*i: 180*i+180]   # 提取含有验证码的小图
                 angle = random.randint(-45, 45)
                 img0 = rotate(img0, angle)    # 对小图随机旋转

                 if img1.any():
                      img1 = np.concatenate((img1, img0[60:120, 60:120, :]), axis=1)

                 else:
                      img1 = img0[60:120, 60:120, :]

           plt.imsave(data_path+'src/' + str(l)+'.jpg', img1)     # 保存结果

if __name__=='__main__':
      gen_image(100)

结果大致：

方法2. 利用更专业的库：captcha

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Mar 25 19:06:46 2018

@author: lps
"""
from captcha.image import ImageCaptcha
import numpy as np
#import matplotlib.pyplot as plt
from PIL import Image
import random
import cv2

number = ['','','','','','','','','','']
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']

data_path = '/home/lps/yanzm/'

def random_captcha_text(char_set=number,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_capthcha_text_and_image(m):
    image = ImageCaptcha()
    captcha_text = random_captcha_text()     # 生成数字
    captcha_text = ' '.join(captcha_text)  # 生成标签

    captcha = image.generate(captcha_text)

#    image.write(captcha_text,captcha_text+'.jpg')

    captcha_image = Image.open(captcha)
    captcha_image = np.array(captcha_image)

    with open(data_path+"label.txt","a") as f:     # 写入标签
            f.write(captcha_text)
            f.writelines("\n")
    cv2.imwrite(data_path + '/src/'+'%.4d.jpg'%m,  captcha_image)   # 保存

#    return captcha_text,captcha_image

if __name__ == '__main__':

    for m in range(0,5000):
#          text,image = gen_capthcha_text_and_image()
          gen_capthcha_text_and_image(m)

#    f = plt.figure()
#    ax = f.add_subplot(212)
#    ax.text(0.1,0.1,text,ha='center',va='center',transform=ax.transAxes)
#    plt.imshow(image)
#    plt.show()
#    

结果大致：

二. pytorch实现

对于一张验证码来说作为一张单一的图片，每输入一张图片，得到四个数字作为输出，只有4个数字同时预测正确才表示预测正确。所以在每一张图上是四个多二分类器：因为验证码上面的数字为0-9，类似于mnist，只不过此时一张图片对应于4个数字。所以思路很简单，实现如下：

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Mar 30 15:46:09 2018

@author: lps
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.optim as optim
import torchvision.models as models
import torchvision
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import matplotlib.pyplot as plt
from PIL import Image
#import pandas as pd
import numpy as np
import os
import copy, time

file_path = '/home/lps/yanzm'
BATCH_SIZE = 16
EPOCH = 10

# Load data
class dataset(Dataset):

      def __init__(self, root_dir, label_file, transform=None):

            self.root_dir = root_dir
            self.label = np.loadtxt(label_file)
            self.transform = transform

      def __getitem__(self, idx):

            img_name = os.path.join(self.root_dir,'%.4d.jpg'%idx)
            image = Image.open(img_name)
            labels = self.label[idx,:]

#            sample = image

            if self.transform:
                 image = self.transform(image)

            return image, labels

      def __len__(self):

            return (self.label.shape[0])

data = dataset(file_path+'/src', file_path+'/label.txt',transform=transforms.ToTensor())

dataloader = DataLoader(data, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, drop_last=True)

dataset_size = len(data)

# Conv network
class ConvNet(nn.Module):

      def __init__(self):
            super(ConvNet, self).__init__()
            self.conv =nn.Sequential(
                        nn.Conv2d(3, 32, kernel_size=4, stride=1, padding=2), # in:(bs,3,60,160)
                        nn.BatchNorm2d(32),
                        nn.LeakyReLU(0.2, inplace=True),
                        nn.MaxPool2d(kernel_size=2),        # out:(bs,32,30,80)

                        nn.Conv2d(32, 64, kernel_size=4, stride=1, padding=2),
                        nn.BatchNorm2d(64),
                        nn.LeakyReLU(0.2, inplace=True),
                        nn.MaxPool2d(kernel_size=2),        # out:(bs,64,15,40)

                        nn.Conv2d(64, 64, kernel_size=3 ,stride=1, padding=1),
                        nn.BatchNorm2d(64),
                        nn.LeakyReLU(0.2, inplace=True),
                        nn.MaxPool2d(kernel_size=2)         # out:(bs,64,7,20)
                    )

            self.fc1 = nn.Linear(64*7*20, 500)
            self.fc2 = nn.Linear(500,40)

      def forward(self, x):
            x = self.conv(x)
            x = x.view(x.size(0), -1)    # reshape to (batch_size, 64 * 7 * 30)
            output = self.fc1(x)
            output = self.fc2(output)

            return output

# Train the net
class nCrossEntropyLoss(torch.nn.Module):

      def __init__(self, n=4):
            super(nCrossEntropyLoss, self).__init__()
            self.n = n
            self.total_loss = 0
            self.loss = nn.CrossEntropyLoss()

      def forward(self, output, label):
            output_t = output[:,0:10]
            label = Variable(torch.LongTensor(label.data.cpu().numpy())).cuda()
            label_t = label[:,0]

            for i in range(1, self.n):
                  output_t = torch.cat((output_t, output[:,10*i:10*i+10]), 0)   # 损失的思路是将一张图平均剪切为4张小图即4个多分类，然后再用多分类交叉熵方损失
                  label_t = torch.cat((label_t, label[:,i]), 0)
                  self.total_loss = self.loss(output_t, label_t)

            return self.total_loss

def equal(np1,np2):

      n = 0
      for i in range(np1.shape[0]):
            if (np1[i,:]==np2[i,:]).all():
                  n += 1 

      return n

net = ConvNet().cuda()
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
#loss_func = nn.CrossEntropyLoss()
loss_func = nCrossEntropyLoss()

best_model_wts = copy.deepcopy(net.state_dict())
best_acc = 0.0   

since = time.time()
for epoch in range(EPOCH):

      running_loss=0.0
      running_corrects=0

      for step,(inputs,label) in enumerate(dataloader):

            pred = torch.LongTensor(BATCH_SIZE,1).zero_()
            inputs = Variable(inputs).cuda()   # (bs, 3, 60, 240)
            label = Variable(label).cuda()   # (bs, 4)

            optimizer.zero_grad() 

            output = net(inputs)   # (bs, 40)
            loss = loss_func(output, label) 

            for i in range(4):
                  pre = F.log_softmax(output[:,10*i:10*i+10], dim=1)  # (bs, 10)
                  pred = torch.cat((pred, pre.data.max(1, keepdim=True)[1].cpu()), dim=1)    # 

            loss.backward()
            optimizer.step()

            running_loss += loss.data[0] * inputs.size()[0]
            running_corrects += equal(pred.numpy()[:,1:], label.data.cpu().numpy().astype(int))

      epoch_loss = running_loss / dataset_size
      epoch_acc  = running_corrects / dataset_size

      if  epoch_acc > best_acc:
            best_acc = epoch_acc
            best_model_wts = copy.deepcopy(net.state_dict())  

      if epoch == EPOCH-1:
            torch.save(best_model_wts, file_path+'/best_model_wts.pkl')

      print()

      time_elapsed = time.time() - since
      print('Training complete in {:.0f}m {:.0f}s'.format(
                   time_elapsed // 60, time_elapsed % 60))
      print('Train Loss:{:.4f} Acc: {:.4f}'.format(epoch_loss, epoch_acc))

随机生成5000张图片拿来训练，准确率也会有97%左右。