PyTorch 实现验证码识别系统-CFANZ编程社区

本教程介绍如何使用 Python 和 PyTorch 从头实现一个验证码识别模型，包括数据生成、模型搭建、训练和测试。

第一步：安装必要的依赖确保安装了以下库：

pip install torch torchvision captcha pillow numpy 第二步：生成验证码图像数据使用 captcha 库生成图片数据，字符集包含数字和大写字母，每张验证码为 4 位长度。

from captcha.image import ImageCaptcha import random import string import os

characters = string.digits + string.ascii_uppercase length = 4 width, height = 160, 60

def create_captcha_dataset(count=10000, output_dir='dataset'): os.makedirs(output_dir, exist_ok=True) generator = ImageCaptcha(width=width, height=height) for i in range(count): text = ''.join(random.choices(characters, k=length)) image = generator.generate_image(text) image.save(os.path.join(output_dir, f'{text}_{i}.png'))

create_captcha_dataset() 第三步：定义数据集加载类使用 PyTorch 的 Dataset 和 DataLoader 加载图像和标签。

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

class CaptchaDataset(Dataset): def init(self, root): self.root = root self.images = [f for f in os.listdir(root) if f.endswith('.png')] self.char_to_idx = {c: i for i, c in enumerate(characters)} self.transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ])

def __getitem__(self, index):
    file = self.images[index]
    label = file.split('_')[0]
    image = Image.open(os.path.join(self.root, file)).convert('RGB')
    image = self.transform(image)
    label_tensor = torch.tensor([self.char_to_idx[c] for c in label], dtype=torch.long)
    return image, label_tensor

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

train_loader = DataLoader(CaptchaDataset('dataset'), batch_size=64, shuffle=True) 第四步：构建卷积 + 循环神经网络模型模型结构包括 CNN 提取图像特征，LSTM 处理字符序列，最后输出每个字符位置的分类。

import torch.nn as nn

class CaptchaRecognizer(nn.Module): def init(self): super().init() self.cnn = nn.Sequential( nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(), nn.MaxPool2d((2, 1)) ) self.rnn = nn.LSTM(input_size=128 * 7, hidden_size=128, num_layers=2, bidirectional=True, batch_first=True) self.fc = nn.Linear(256, len(characters))

def forward(self, x):
    x = self.cnn(x)
    x = x.permute(0, 3, 1, 2)  # shape: [batch, width, channels, height]
    b, w, c, h = x.size()
    x = x.view(b, w, c * h)
    x, _ = self.rnn(x)
    x = self.fc(x)
    return x

第五步：训练模型使用交叉熵损失计算每个字符位置的预测误差。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = CaptchaRecognizer().to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) loss_fn = nn.CrossEntropyLoss()

for epoch in range(10): model.train() total_loss = 0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = sum(loss_fn(outputs[:, i, :], labels[:, i]) for i in range(length)) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() print(f'Epoch {epoch+1} Loss: {total_loss:.4f}') 第六步：测试验证码识别效果定义预测函数，给定图像路径返回识别结果。

def predict(model, image_path): image = Image.open(image_path).convert('RGB') image = transforms.ToTensor()(image) image = transforms.Normalize((0.5,), (0.5,))(image) image = image.unsqueeze(0).to(device)

model.eval（)
with torch.no_grad():
    output = model(image)
    pred = output.argmax(dim=2)[0]
    result = ''.join([characters[i] for i in pred])
return result

test_image = 'dataset/ABCD_123.png' print('预测:', predict(model, test_image))