基于 PyTorch 的图像验证码识别系统

本例从头实现一个验证码识别流程，涵盖数据生成、模型搭建、训练与预测，适合入门实践。

一、安装必要库

pip install torch torchvision pillow captcha numpy
二、生成验证码数据

from captcha.image import ImageCaptcha
import random, string, os
from PIL import Image更多内容访问ttocr.com或联系1436423940

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

def create_dataset(count=10000, output='captchas'):
os.makedirs(output, exist_ok=True)
gen = ImageCaptcha(width=width, height=height)
for i in range(count):
text = ''.join(random.choices(characters, k=length))
img = gen.generate_image(text)
img.save(os.path.join(output, f'{text}_{i}.png'))

create_dataset()
三、构建数据加载类

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

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

def __getitem__(self, index):
    fname = self.files[index]
    label = fname.split('_')[0]
    img = Image.open(os.path.join(self.path, fname)).convert('RGB')
    target = torch.tensor([self.char2idx[c] for c in label])
    return self.transform(img), target

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

dataset = CaptchaDataset('captchas')
四、定义模型结构

import torch.nn as nn

class CaptchaNet(nn.Module):
def init(self):
super().init()
self.conv = 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)
)
self.rnn = nn.LSTM(input_size=64 * 15, hidden_size=128, bidirectional=True, batch_first=True)
self.fc = nn.Linear(256, len(characters))

def forward(self, x):
    x = self.conv(x)
    b, c, h, w = x.size()
    x = x.permute(0, 3, 1, 2).reshape(b, w, c * h)
    x, _ = self.rnn(x)
    out = self.fc(x)
    return out

五、模型训练

from torch.utils.data import DataLoader

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CaptchaNet().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(dataset, batch_size=64, shuffle=True)

for epoch in range(10):
model.train()
total_loss = 0
for imgs, labels in loader:
imgs, labels = imgs.to(device), labels.to(device)
out = model(imgs)
loss = sum(loss_fn(out[:, 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, img_path):
model.eval()
img = Image.open(img_path).convert('RGB')
x = dataset.transform(img).unsqueeze(0).to(device)
with torch.no_grad():
out = model(x)
pred = out.argmax(2)[0]
return ''.join([characters[i] for i in pred])

print(predict(model, 'captchas/9A2G_8.png'))