PyTorch 和 CNN-LSTM 的验证码识别系统实现

在本文中，我们将介绍如何使用 PyTorch 框架结合卷积神经网络（CNN）和长短时记忆网络（LSTM）构建一个验证码识别系统。通过卷积神经网络来提取图像特征，再利用 LSTM 处理字符序列的依赖关系，最终实现对验证码的准确识别。

环境准备
首先，确保已安装必要的库。可以使用以下命令来安装：

bash

pip install torch torchvision numpy opencv-python pillow
PyTorch：用于构建和训练深度学习模型。
NumPy：用于数据处理。
OpenCV：用于图像处理。
Pillow：用于图像处理。
2. 数据准备与预处理
验证码图像通常包含多个字符，我们需要将图像转换为神经网络可以接受的格式。以下是我们所需的步骤：灰度化、归一化、缩放以及字符编码。

(1) 数据集加载与预处理
假设验证码图像目录下的文件名即为标签（例如：AB12.png 对应标签为 AB12）。我们将定义一个数据集类来加载图像，并对图像进行预处理。

python

import os
import cv2
import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader

class CaptchaDataset(Dataset):
def init(self, image_dir, char_set, image_size=(128, 64), sequence_length=4):
self.image_dir = image_dir
self.image_paths = [os.path.join(image_dir, fname) for fname in os.listdir(image_dir)]
self.char_set = char_set
self.image_size = image_size
self.sequence_length = sequence_length

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

def __getitem__(self, idx):
    # 读取图像
    img_path = self.image_paths[idx]
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    img = cv2.resize(img, self.image_size)  # 缩放图像
    img = img.astype('float32') / 255.0  # 归一化处理
    img = np.expand_dims(img, axis=0)  # 扩展维度 [H, W] -> [1, H, W]
    
    # 获取标签：文件名作为标签
    label = os.path.basename(img_path).split('.')[0]
    label_encoded = [self.char_set.index(c) for c in label]  # 将字符编码为数字

    # 填充标签长度（确保所有标签长度一致）
    label_encoded = label_encoded + [0] * (self.sequence_length - len(label_encoded))

    return torch.tensor(img), torch.tensor(label_encoded)

字符集定义：包含大写字母和数字

char_set = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" # 可以根据需要修改字符集

创建数据集和数据加载器

train_dataset = CaptchaDataset(image_dir="captcha_images", char_set=char_set)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
(2) 标签编码
验证码中的每个字符将通过索引转换为数字。例如，“A”对应0，“B”对应1，以此类推。我们还通过填充来确保所有标签长度一致，方便输入到模型中。

构建 CNN-LSTM 模型
在本模型中，我们将使用卷积神经网络（CNN）提取图像特征，再通过长短时记忆网络（LSTM）来捕捉字符的时序关系。

(1) 定义模型架构
python

import torch.nn as nn

class CNN_LSTM_Model(nn.Module):
def init(self, input_channels=1, output_size=36, sequence_length=4):
super(CNN_LSTM_Model, self).init()

    # 卷积层：提取图像特征
    self.conv1 = nn.Conv2d(input_channels, 32, kernel_size=3, padding=1)
    self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
    self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
    self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
    self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
    self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)

    # LSTM层：处理字符序列
    self.lstm = nn.LSTM(128, 128, batch_first=True)
    
    # 全连接层：输出字符类别
    self.fc = nn.Linear(128, output_size)

def forward(self, x):
    # 经过卷积层提取特征
    x = self.pool1(torch.relu(self.conv1(x)))
    x = self.pool2(torch.relu(self.conv2(x)))
    x = self.pool3(torch.relu(self.conv3(x)))
    
    # 展平后送入LSTM层
    x = x.view(x.size(0), -1, 128)  # [batch_size, seq_len, feature_size]
    
    # LSTM处理
    x, _ = self.lstm(x)
    
    # 输出预测字符
    x = self.fc(x)
    
    return x

模型实例化

model = CNN_LSTM_Model(input_channels=1, output_size=len(char_set), sequence_length=4)
(2) 模型架构解析
卷积层（Conv2D）：用于提取图像中的空间特征。
池化层（MaxPool2D）：减少图像尺寸，降低计算量。
LSTM层：处理图像中各字符的时序关系（例如：字符1到字符4的顺序依赖）。
全连接层（Linear）：将LSTM输出的特征映射为字符类别。
4. 模型训练
我们将使用交叉熵损失函数和 Adam优化器来训练模型。

(1) 定义损失函数与优化器
python

import torch.optim as optim

损失函数

criterion = nn.CrossEntropyLoss()

优化器

optimizer = optim.Adam(model.parameters(), lr=0.001)
(2) 训练模型
python

num_epochs = 10
for epoch in range(num_epochs):
model.train()
total_loss = 0

for i, (images, labels) in enumerate(train_loader):
    images, labels = images.squeeze(1).float(), labels
    
    # 清除之前的梯度
    optimizer.zero_grad()
    
    # 前向传播
    outputs = model(images)
    
    # 计算损失
    loss = 0
    for t in range(outputs.size(1)):  # 针对每个时间步的预测
        loss += criterion(outputs[:, t, :], labels[:, t])
    
    # 反向传播
    loss.backward()
    optimizer.step()
    
    total_loss += loss.item()

print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(train_loader):.4f}")

模型评估与预测
训练完成后，我们可以使用该模型对新的验证码图像进行预测。

(1) 评估模型
python

评估模型

model.eval()
test_dataset = CaptchaDataset(image_dir="test_captcha_images", char_set=char_set)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

correct = 0
total = 0

with torch.no_grad():
for images, labels in test_loader:
images, labels = images.squeeze(1).float(), labels
outputs = model(images)

    _, predicted = torch.max(outputs, 2)
    predicted = predicted.squeeze(1)  # 去除批次维度
    
    correct += (predicted == labels).sum().item()
    total += labels.size(0)

print(f"Test Accuracy: {100 * correct / total:.2f}%")

(2) 进行预测
python
更多内容访问ttocr.com或联系1436423940
def predict_captcha(model, img_path, char_set, sequence_length=4):
img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
img = cv2.resize(img, (128, 64))
img = img.astype('float32') / 255.0
img = np.expand_dims(img, axis=0) # 增加批次维度
img = np.expand_dims(img, axis=0) # 增加通道维度 [1, 128, 64] -> [1, 1, 128, 64]

# 转换为PyTorch张量
img_tensor = torch.tensor(img).float()

# 模型预测
output = model(img_tensor)
predicted_label = []
for t in range(output.size(1)):  # 针对每个时间步
    _, pred = torch.max(output[0, t, :], 0)
    predicted_label.append(char_set[pred.item()])

return ''.join(predicted_label)

测试预测

test_image_path = "captcha_images/test1.png"
predicted_label = predict_captcha(model, test_image_path, char_set)
print(f"Predicted CAPTCHA label: {predicted_label}")