PyTorch 构建验证码识别系统

在本教程中，我们将使用 PyTorch 构建一个验证码识别系统。通过使用卷积神经网络（CNN），我们能够从验证码图像中提取特征，并进行字符分类。PyTorch 是一个强大的深度学习框架，提供了灵活的模型构建和训练机制，非常适合进行验证码识别任务。

1. 环境准备
   首先，确保你已经安装了所需的库：

bash

pip install torch torchvision opencv-python numpy matplotlib pillow
torch 和 torchvision：PyTorch 和其计算机视觉工具包，支持深度学习任务。
opencv-python：用于图像加载、处理和增强。
numpy：用于数据处理和数组运算。
matplotlib：用于可视化图像和训练过程中的损失和准确率。
2. 数据集准备与图像预处理
验证码图像通常包含噪声、干扰线以及不规则的字符形状，因此我们需要进行一系列预处理。主要的预处理步骤包括：灰度化、二值化、去噪、字符切割等。

(1) 图像加载与预处理
首先，我们加载验证码图像并进行灰度化和二值化处理。灰度化可以减少颜色干扰，而二值化可以将图像转为黑白图像，从而提高识别的准确性。

python

import cv2
import numpy as np

def preprocess_image(img_path):
# 读取图像
img = cv2.imread(img_path)

# 转换为灰度图
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# 二值化处理，使用 Otsu 的方法自动选择阈值
_, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

# 高斯模糊去噪
blurred = cv2.GaussianBlur(binary, (5, 5), 0)

return blurred

示例图像路径

img_path = 'captcha_images/test1.png'
processed_img = preprocess_image(img_path)

显示处理后的图像

cv2.imshow('Processed Image', processed_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
(2) 提取字符区域
为了识别每个字符，我们需要提取图像中的字符区域。我们可以使用 轮廓检测 来提取字符区域。OpenCV 的 findContours 函数可以帮助我们找到每个字符的边界框。

python

def extract_characters(processed_img):
contours, _ = cv2.findContours(processed_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
char_images = []
for contour in contours:
x, y, w, h = cv2.boundingRect(contour)
if w > 10 and h > 10: # 忽略小的噪点
char_img = processed_img[y:y+h, x:x+w]
char_images.append(char_img)

# 按照字符的从左到右顺序排序
char_images.sort(key=lambda x: x[0][0])  # 排序依据是字符的左上角 x 坐标
return char_images

提取字符区域

char_images = extract_characters(processed_img)

显示提取的字符

for i, char_img in enumerate(char_images):
cv2.imshow(f'Character {i+1}', char_img)
cv2.waitKey(0)

cv2.destroyAllWindows()
3. 构建卷积神经网络（CNN）
接下来，我们将使用 PyTorch 构建一个卷积神经网络（CNN）模型。CNN 是处理图像数据的经典模型，它能够通过多个卷积层和池化层提取图像特征，并通过全连接层进行分类。

(1) 构建 CNN 模型
我们将构建一个简单的 CNN 模型，包括多个卷积层和池化层。最后通过一个全连接层输出字符的分类结果。

python

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms

class CNNModel(nn.Module):
def init(self):
super(CNNModel, self).init()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1) # 卷积层1
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) # 卷积层2
self.pool = nn.MaxPool2d(2, 2) # 池化层
self.fc1 = nn.Linear(64 * 7 * 7, 128) # 全连接层1
self.fc2 = nn.Linear(128, 36) # 输出层，假设字符集为数字0-9和字母A-Z，总共36个字符

def forward(self, x):
    x = self.pool(torch.relu(self.conv1(x)))  # 第一层卷积
    x = self.pool(torch.relu(self.conv2(x)))  # 第二层卷积
    x = x.view(-1, 64 * 7 * 7)  # 展平层
    x = torch.relu(self.fc1(x))  # 全连接层1
    x = self.fc2(x)  # 输出层
    return x

实例化模型

model = CNNModel()

打印模型结构

print(model)
(2) 数据加载与训练
在训练模型之前，我们需要将图像数据转换为张量，并对图像进行归一化处理。假设我们使用的是28x28的灰度图像，并且字符集包括数字0-9和字母A-Z，总共36个字符。

python

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

class CaptchaDataset(Dataset):
def init(self, image_paths, labels, transform=None):
self.image_paths = image_paths
self.labels = labels
self.transform = transform

def __len__(self):
    return len(self.image_paths)

def __getitem__(self, idx):
    image = Image.open(self.image_paths[idx]).convert('L')
    label = self.labels[idx]

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

    return image, label

数据预处理

transform = transforms.Compose([
transforms.Resize((28, 28)),
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)) # 归一化
])

假设你有训练图像路径和对应的标签

train_image_paths = ['captcha_images/train1.png', 'captcha_images/train2.png'] # 示例路径
train_labels = [0, 1] # 示例标签

创建数据集和数据加载器

train_dataset = CaptchaDataset(train_image_paths, train_labels, transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

选择优化器和损失函数

optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

训练模型

num_epochs = 10
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for inputs, labels in train_loader:
optimizer.zero_grad()

    # 前向传播
    outputs = model(inputs)
    
    # 计算损失
    loss = criterion(outputs, labels)
    loss.backward()
    
    # 更新参数
    optimizer.step()

    running_loss += loss.item()

print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(train_loader)}")

4. 模型评估与预测
   训练完成后，我们可以评估模型的准确性，并对新的验证码图像进行预测。

(1) 评估模型
python

假设你有测试图像和标签

test_image_paths = ['captcha_images/test1.png']
test_labels = [0]

创建数据集和数据加载器

test_dataset = CaptchaDataset(test_image_paths, test_labels, transform)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

评估模型

model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
outputs = model(inputs)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()

print(f"Accuracy: {100 * correct / total:.2f}%")
(2) 对验证码进行预测
python

def predict_captcha(model, img_path, char_set="0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"):
img = Image.open(img_path).convert('L')
img = transform(img).unsqueeze(0) # 增加批次维度

model.eval()
with torch.no_grad():
    outputs = model(img)
    _, predicted = torch.max(outputs, 1)

return char_set[predicted.item()]