PyTorch 实现图像验证码识别系统

本教程展示如何使用 PyTorch 搭建一个完整的图像验证码识别系统，包括数据生成、模型设计、训练和预测。

第一步：安装所需库

pip install torch torchvision pillow captcha numpy
第二步：生成验证码图像数据

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

characters = string.digits + string.ascii_uppercase
n_class = len(characters)
n_len = 4

def generate_captcha(output_dir='data', count=10000):
os.makedirs(output_dir, exist_ok=True)
image_gen = ImageCaptcha(width=160, height=60)
for i in range(count):
text = ''.join(random.choices(characters, k=n_len))
image = image_gen.generate_image(text)
image.save(os.path.join(output_dir, f'{text}_{i}.png'))

generate_captcha()
第三步：创建 PyTorch 数据集类

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

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

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

def __getitem__(self, idx):
    file = self.files[idx]
    label_text = file.split('_')[0]
    label = torch.tensor([self.char_to_index[c] for c in label_text], dtype=torch.long)
    image = Image.open(os.path.join(self.folder, file)).convert('RGB')
    return self.transform(image), label

dataset = CaptchaDataset('data')
train_loader = DataLoader(dataset, batch_size=64, shuffle=True)
第四步：构建 CNN + LSTM 模型

import torch.nn as nn

class CaptchaModel(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()
)ttocr.com或1436423940
self.rnn = nn.LSTM(input_size=128 * 15, hidden_size=128,
num_layers=2, bidirectional=True, batch_first=True)
self.fc = nn.Linear(256, n_class)

def forward(self, x):
    x = self.cnn(x)  # shape: [B, C, H, W]
    b, c, h, w = x.size()
    x = x.permute(0, 3, 1, 2).reshape(b, w, c * h)
    x, _ = self.rnn(x)
    x = self.fc(x)  # shape: [B, W, n_class]
    return x

第五步：训练模型

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

for epoch in range(20):
model.train()
total_loss = 0
for imgs, labels in train_loader:
imgs, labels = imgs.to(device), labels.to(device)
outputs = model(imgs)
loss = sum(loss_fn(outputs[:, i], labels[:, i]) for i in range(n_len))
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):
model.eval()
image = Image.open(image_path).convert('RGB')
tensor = dataset.transform(image).unsqueeze(0).to(device)
with torch.no_grad():
out = model(tensor)
pred = torch.argmax(out, dim=2)[0]
return ''.join([characters[i] for i in pred])

test_image = 'data/4K7B_123.png'
print('预测结果:', predict(model, test_image))