使用PyTorch识别简单验证码

前言

在这篇文章中，我们将演示如何使用PyTorch来识别简单的数字图形CAPTCHA。示例比较简单，主要演示图片预处理及简单的CNN网络。

环境准备

安装依赖包

conda install pytorch torchvision torchaudio cpuonly -c pytorch

sudo apt-get install libgl1  # for opencv
pip install requests matplotlib opencv-python

下载验证码图片（我们将验证码的值放在HTTP头中返回，方便的对原始数据集进行标注，更一般的情况需要对图片进行人工标注）：

CAPTCHA_URL = 'https://captcha.tomo.wang'
r = requests.get(CAPTCHA_URL)
captcha = r.headers['X-Captcha']
with open('{}.png'.format(captcha), 'wb') as f:
    f.write(r.content)

图像处理及训练

准备

首先我们导入需要用到的包

import os
import re
import sys
import argparse
import glob
from io import BytesIO

import requests
import numpy as np
import cv2
import matplotlib.pyplot as plt

import torch
from torch.autograd import Variable
import torch.nn.functional as F

定义程序运行的相关常量

• 字符个数
• 验证码宽高
• 裁剪后字符的宽高
• …

NUM_CHARS = 4
CAPTCHA_WIDTH = 200
CAPTCHA_HEIGHT = 62
CH_WIDTH = 20
CH_HEIGHT = 28

CAPTCHA_DIR = './images'
TORCH_NET_PATH = 'captcha.torch'
BG_COLOR = (243, 251, 254)  # captcha backgroud color
BG_THRESHOLD = 245

BLANK_THRESHHOLD = 1
DOTS_THRESHOLD = 3
CH_MIN_WIDTH = 8

获取验证码图片并展示

def get_captcha():
    CAPTCHA_URL = 'https://captcha.tomo.wang'
    r = requests.get(CAPTCHA_URL)
    return r.content

img = get_captcha()
plt.imshow(plt.imread(BytesIO(img)))

<matplotlib.image.AxesImage at 0x7f1fcc4d7c40>

一个验证码图片包含背景和不同颜色的字符，共四个字符，在训练前需要对其进行灰化处理，并进行切割

img_array = np.asarray(bytearray(img), dtype=np.uint8)
img = cv2.imdecode(img_array, cv2.IMREAD_GRAYSCALE)
assert img.shape == (CAPTCHA_HEIGHT, CAPTCHA_WIDTH)
img = cv2.threshold(img, BG_THRESHOLD, 255, cv2.THRESH_BINARY_INV)[1]
plt.imshow(cv2.cvtColor(img, cv2.COLOR_GRAY2BGR))

<matplotlib.image.AxesImage at 0x7f1fcc3e7e80>

切割

将图片切割成四个独立的字符

def _denoise(img):
    img = cv2.threshold(img, BG_THRESHOLD, 255, cv2.THRESH_BINARY_INV)[1]
    return img


def _preprocess(img