如何使用 Python 和 PyTorch 构建一个完整的图像验证码识别系统

本文介绍如何使用 Python 和 PyTorch 构建一个完整的图像验证码识别系统，包括数据生成、模型设计、训练和预测。

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

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

def generate_captcha_images(num=10000, save_dir='captcha_images'):
os.makedirs(save_dir, exist_ok=True)
generator = ImageCaptcha(width=width, height=height)
for i in range(num):
label = ''.join(random.choices(characters, k=length))
image = generator.generate_image(label)
image.save(os.path.join(save_dir, f'{label}_{i}.png'))

generate_captcha_images()
3. 数据集加载与处理
定义一个自定义的 PyTorch 数据集类来加载图像和标签。

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

class CaptchaDataset(Dataset):
def init(self, root, transform=None):
self.root = root
self.files = [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 = transform

def __getitem__(self, idx):
    file_name = self.files[idx]
    label = file_name.split('_')[0]
    image = Image.open(os.path.join(self.root, file_name)).convert('RGB')
    if self.transform:
        image = self.transform(image)
    target = torch.tensor([self.char_to_idx[c] for c in label], dtype=torch.long)
    return image, target

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

transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])

dataset = CaptchaDataset('captcha_images', transform=transform)
loader = DataLoader(dataset, batch_size=64, shuffle=True)
4. 构建模型
我们使用卷积网络提取图像特征，然后通过 LSTM 进行字符序列建模。

import torch.nn as nn

class CaptchaModel(nn.Module):
def init(self, num_classes):
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, num_classes)

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

5. 模型训练
   设置训练设备、优化器和损失函数，开始训练模型。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CaptchaModel(num_classes=len(characters)).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

for epoch in range(10):
model.train()
total_loss = 0
for images, labels in loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images) # shape: [batch, width, num_classes]
loss = sum(criterion(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}')
6. 验证码预测
加载模型，读取图片进行预测，输出识别的验证码字符串。

def predict(model, image_path):
image = Image.open(image_path).convert('RGB')
image = transform(image).unsqueeze(0).to(device)
model.eval()
with torch.no_grad():
output = model(image)
pred = output.argmax(dim=2)[0]
return ''.join([characters[i] for i in pred])

示例

image_file = 'captcha_images/A9B2_123.png'
print('预测结果:', predict(model, image_file))
7. 后续优化建议