利用深度学习实现验证码识别-4-ResNet18+imagecaptcha

在当今的数字化世界中，验证码（CAPTCHA）是保护网站免受自动化攻击的重要工具。然而，对于用户来说，验证码有时可能会成为一种烦恼。为了解决这个问题，我们可以利用深度学习技术来自动识别验证码，从而提高用户体验。本文将介绍如何使用ResNet18模型来识别ImageCaptcha生成的验证码。

1. 环境设置与数据准备

首先，我们需要检查CUDA是否可用，以便利用GPU加速训练过程。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')

接下来，我们定义一个数据生成器CaptchaDataset，它使用imagecaptcha库生成验证码图像。

class CaptchaDataset(Dataset):
    def __init__(self, length=1000, charset=None, captcha_length=5, transform=None):
        self.length = length
        self.transform = transform
        self.charset = charset if charset is not None else string.ascii_letters + string.digits
        self.captcha_length = captcha_length
        self.num_classes = len(self.charset)
        self.image_generator = ImageCaptcha(width=160, height=60)

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        text = ''.join(random.choices(self.charset, k=self.captcha_length))
        image = self.image_generator.generate_image(text)
        if self.transform:
            image = self.transform(image)
        label = [self.charset.index(c) for c in text]
        return image, torch.tensor(label, dtype=torch.long)

2. 数据增强与预处理

为了提高模型的泛化能力，我们使用了一系列的数据增强和预处理步骤。

transform = transforms.Compose([
    transforms.Grayscale(),  # 将图像转换为灰度
    transforms.Resize((40, 100)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

3. 数据集划分与加载

我们将数据集划分为训练集和验证集，并使用DataLoader进行批量加载。

dataset = CaptchaDataset(length=2000, charset=charset, captcha_length=captcha_length, transform=transform)
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)

4. 模型定义与迁移学习

我们使用预训练的ResNet18模型，并对其进行微调以适应验证码识别任务。

class CaptchaModel(nn.Module):
    def __init__(self, num_classes, captcha_length):
        super(CaptchaModel, self).__init__()
        self.captcha_length = captcha_length
        self.resnet = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
        self.resnet.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        num_ftrs = self.resnet.fc.in_features
        self.resnet.fc = nn.Linear(num_ftrs, num_classes * self.captcha_length)

    def forward(self, x):
        x = self.resnet(x)
        return x.view(-1, self.captcha_length, num_classes)

5. 训练与评估

我们定义了训练函数train_model，并在每个epoch结束时保存模型检查点。

def train_model(epochs, resume=False):
    start_epoch = 0
    if resume and os.path.isfile("captcha_model_checkpoint.pth.tar"):
        checkpoint = load_checkpoint()
        model.load_state_dict(checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']

    scaler = torch.cuda.amp.GradScaler()

    for epoch in range(start_epoch, epochs):
        model.train()
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()

            with torch.cuda.amp.autocast():
                outputs = model(images)
                loss = sum(criterion(outputs[:, i, :], labels[:, i]) for i in range(captcha_length))

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            running_loss += loss.item()

        val_accuracy = evaluate_accuracy(val_loader)
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {running_loss / len(train_loader):.4f}, Val Accuracy: {val_accuracy:.4f}')

        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'optimizer': optimizer.state_dict(),
        })

6. 可视化预测结果

最后，我们定义了一个函数visualize_predictions来可视化模型的预测结果。

def visualize_predictions(num_samples=16):
    model.eval()
    samples, labels = next(iter(DataLoader(val_dataset, batch_size=num_samples, shuffle=True)))
    samples, labels = samples.to(device), labels.to(device)

    with torch.no_grad():
        outputs = model(samples)
        predicted = torch.argmax(outputs, dim=2)

    samples = samples.cpu()
    predicted = predicted.cpu()
    labels = labels.cpu()

    fig, axes = plt.subplots(4, 4, figsize=(10, 10))
    for i in range(16):
        ax = axes[i // 4, i % 4]
        ax.imshow(samples[i].squeeze(), cmap='gray')
        true_text = ''.join([dataset.charset[l] for l in labels[i]])
        pred_text = ''.join([dataset.charset[p] for p in predicted[i]])
        ax.set_title(f'True: {true_text}\nPred: {pred_text}')
        ax.axis('off')
    plt.show()

7. 训练与可视化

最后，我们调用train_model函数进行模型训练，并使用visualize_predictions函数来可视化模型的预测结果。

train_model(epochs=180, resume=True)
visualize_predictions()

通过上述步骤，我们成功地使用ResNet18模型来识别ImageCaptcha生成的验证码。这种方法不仅提高了验证码识别的准确性，还提升了用户体验。希望本文能为您在验证码识别领域的研究和应用提供有价值的参考。