使用 PyTorch 实现验证码识别系统

在本文中，我们将使用 PyTorch 框架实现一个验证码识别系统。通过卷积神经网络（CNN），我们可以有效地从验证码图像中提取特征，并对其进行字符分类。PyTorch 作为一种灵活且强大的深度学习框架，能够帮助我们高效地构建、训练和评估神经网络模型。

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

bash

pip install torch torchvision opencv-python numpy pillow
PyTorch：深度学习框架，提供了强大的张量计算和自动求导功能。
TorchVision：用于计算机视觉任务的 PyTorch 扩展。
OpenCV：用于图像预处理和数据增强。
NumPy 和 Pillow：用于处理数据和图像。
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)

# 二值化处理
_, 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) 提取字符区域
通过轮廓检测来分割验证码中的每个字符：

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）。该模型将处理每个字符的图像，并对每个字符进行分类。

(1) 构建 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, num_classes=36):
super(CNNModel, self).init()

    # 卷积层1
    self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
    self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
    
    # 卷积层2
    self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
    self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
    
    # 卷积层3
    self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
    self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
    
    # 全连接层
    self.fc1 = nn.Linear(128 * 3 * 3, 512)
    self.fc2 = nn.Linear(512, num_classes)
    
def forward(self, x):
    x = self.pool1(torch.relu(self.conv1(x)))
    x = self.pool2(torch.relu(self.conv2(x)))
    x = self.pool3(torch.relu(self.conv3(x)))
    x = x.view(-1, 128 * 3 * 3)  # 展平
    x = torch.relu(self.fc1(x))
    x = self.fc2(x)
    return x

创建模型

model = CNNModel(num_classes=36) # 假设字符集包含 0-9 和 A-Z 总共 36 个字符
(2) 数据预处理与训练
对于输入数据，我们需要将其调整为适合神经网络处理的形式。通过 transforms 进行标准化处理，并将图像转换为 PyTorch 的 Tensor 格式。

python

from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

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):
    img = cv2.imread(self.image_paths[idx], cv2.IMREAD_GRAYSCALE)
    img = cv2.resize(img, (28, 28))  # 将图像大小调整为 28x28
    img = img.astype('float32') / 255.0  # 归一化

    if self.transform:
        img = self.transform(img)
    
    label = self.labels[idx]
    return img, label

数据转换：归一化和转为Tensor

transform = transforms.Compose([transforms.ToTensor()])

假设你已经准备好了图像路径和标签

train_dataset = CaptchaDataset(image_paths=train_image_paths, labels=train_labels, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
(3) 训练模型
python

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

定义损失函数和优化器

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

训练过程

num_epochs = 10
for epoch in range(num_epochs):
model.train() # 设置模型为训练模式
running_loss = 0.0
correct = 0
total = 0
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)

    # 前向传播
    outputs = model(images)
    loss = criterion(outputs, labels)

    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 统计结果
    running_loss += loss.item()
    _, predicted = torch.max(outputs, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum().item()

print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(train_loader)}, Accuracy: {100 * correct / total}%')

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

(1) 评估模型
python
更多内容访问ttocr.com或联系1436423940

使用测试集评估模型

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

model.eval() # 设置模型为评估模式
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()

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

def predict_captcha(model, img_path, char_set="0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"):
model.eval() # 设置模型为评估模式
img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
img = cv2.resize(img, (28, 28))
img = img.astype('float32') / 255.0 # 归一化
img = transforms.ToTensor()(img).unsqueeze(0).to(device) # 转换为Tensor并增加批次维度

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

return char_set[predicted.item()]

对图像进行预测

captcha_image = 'captcha_images/test1.png'
predicted_label = predict_captcha(model, captcha_image)
print(f"Predicted CAPTCHA label: {predicted_label}")