使用 PyTorch 构建验证码识别模型

1. 导入必要的库
python

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
import numpy as np
import random
import string
import matplotlib.pyplot as plt
from captcha.image import ImageCaptcha
2. 定义参数和字符集
python

# 定义参数
characters = string.digits + string.ascii_uppercase
width, height, n_len, n_class = 170, 80, 4, len(characters) + 1  # 加1是为了空白类别
3. 生成验证码图像
python

# 生成随机验证码字符串和对应的图像
generator = ImageCaptcha(width=width, height=height)

def generate_random_string():
    return ''.join(random.choice(characters) for _ in range(n_len))

random_str = generate_random_string()
img = generator.generate_image(random_str)

plt.imshow(img)
plt.title(random_str)
plt.show()
4. 自定义数据集
python

class CaptchaDataset(Dataset):
    def __init__(self, characters, width, height, n_len, transform=None):
        self.characters = characters
        self.width = width
        self.height = height
        self.n_len = n_len
        self.n_class = len(characters) + 1
        self.transform = transform
        self.generator = ImageCaptcha(width=self.width, height=self.height)
        
    def __len__(self):
        return 10000  # 可以根据需要调整

    def __getitem__(self, idx):
        random_str = generate_random_string()
        img = self.generator.generate_image(random_str)
        img = np.array(img)
        label = [self.characters.find(x) for x in random_str]
        
        if self.transform:
            img = self.transform(img)
        
        return img, label

# 定义转换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

dataset = CaptchaDataset(characters, width, height, n_len, transform=transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
5. 定义模型架构
python

class CaptchaModel(nn.Module):
    def __init__(self, n_class):
        super(CaptchaModel, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.rnn = nn.GRU(128 * 10, 128, bidirectional=True, batch_first=True)
        self.dense = nn.Linear(128 * 2, n_class)
    
    def forward(self, x):
        x = self.conv_layers(x)
        batch, channel, height, width = x.size()
        x = x.permute(0, 3, 1, 2).contiguous().view(batch, width, -1)
        x, _ = self.rnn(x)
        x = self.dense(x)
        return x

model = CaptchaModel(n_class)
6. 定义CTC损失函数和优化器
python

criterion = nn.CTCLoss(blank=len(characters))
optimizer = optim.Adam(model.parameters(), lr=0.001)
7. 训练和评估
python

def train(model, dataloader, criterion, optimizer, epochs=10):
    model.train()
    for epoch in range(epochs):
        epoch_loss = 0更多内容联系1436423940
        for imgs, labels in dataloader:
            imgs = imgs.permute(0, 3, 1, 2)
            labels = torch.tensor(labels, dtype=torch.long)
            input_lengths = torch.full((imgs.size(0),), imgs.size(3), dtype=torch.long)
            label_lengths = torch.full((imgs.size(0),), n_len, dtype=torch.long)
            
            optimizer.zero_grad()
            outputs = model(imgs)
            outputs = outputs.log_softmax(2)
            loss = criterion(outputs, labels, input_lengths, label_lengths)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        
        print(f'Epoch {epoch+1}/{epochs}, Loss: {epoch_loss/len(dataloader)}')

train(model, dataloader, criterion, optimizer)
8. 测试模型
python

def decode_predictions(preds, characters):
    preds = preds.permute(1, 0, 2).argmax(2)
    preds = preds.cpu().detach().numpy()
    decoded = []
    for pred in preds:
        decoded.append(''.join([characters[i] for i in pred if i < len(characters)]))
    return decoded

model.eval()
with torch.no_grad():
    for imgs, labels in dataloader:
        imgs = imgs.permute(0, 3, 1, 2)
        outputs = model(imgs)
        preds = decode_predictions(outputs, characters)
        break

print(preds)