在本文中,我们将使用 PyTorch 框架和 卷积递归神经网络(CRNN) 架构来构建一个验证码识别系统。验证码(CAPTCHA)是防止自动化攻击的常用方式,通常包含数字、字母或两者的组合。为了提高系统的识别能力,我们将结合 卷积神经网络(CNN) 和 递归神经网络(RNN) 来处理验证码图像,特别是 CTC损失函数 用于序列标注任务。
- 环境准备
首先,确保已安装 PyTorch 和其他必要的依赖项。可以使用以下命令安装:
bash
更多内容访问ttocr.com或联系1436423940
pip install torch torchvision numpy opencv-python pillow
PyTorch:用于构建和训练深度学习模型。
NumPy:用于数据处理。
OpenCV:用于图像处理。
Pillow:用于图像加载和处理。
2. 数据集准备与预处理
验证码通常由多个字符组成,因此我们需要将图像进行预处理,以便输入到模型中。这包括图像的灰度化、调整大小和归一化处理。同时,需要将标签进行编码,以便于训练。
(1) 数据加载与预处理
我们首先定义一个数据集类,用于加载验证码图像并进行预处理。
python
import os
import cv2
import numpy as np
from torch.utils.data import Dataset
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, index):path = self.image_paths[index]img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)img = cv2.resize(img, self.image_size)img = img.astype('float32') / 255.0 # 归一化处理img = np.expand_dims(img, axis=0) # 增加一个维度以适应CNN输入# 获取标签:文件名作为标签label = os.path.basename(path).split('.')[0]label_encoded = [self.char_set.index(c) for c in label]# 转换为torch张量return np.array(img), np.array(label_encoded)
字符集定义
char_set = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
初始化数据集
train_dataset = CaptchaDataset(image_dir="captcha_images", char_set=char_set)
(2) 标签编码
验证码中的每个字符通过索引转换为数字。例如,“A” 映射为 0,“B” 映射为 1,以此类推。标签的长度通常是固定的,假设每个验证码由4个字符组成。
- 构建 CRNN 模型
为了识别验证码,我们将使用 CRNN(卷积递归神经网络)模型。该模型结合了 卷积神经网络(CNN) 和 递归神经网络(RNN),并使用 CTC(Connectionist Temporal Classification) 损失函数进行训练。CNN 用于提取图像特征,而 RNN 用于处理字符序列。
(1) 定义模型架构
python
import torch
import torch.nn as nn
class CRNNModel(nn.Module):
def init(self, num_classes, sequence_length=4):
super(CRNNModel, self).init()
# 卷积层self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)self.pool = nn.MaxPool2d(kernel_size=2, stride=2)# 循环层(LSTM)self.lstm = nn.LSTM(128, 128, bidirectional=True, batch_first=True)# 全连接层self.fc = nn.Linear(128 * 2, num_classes) # 128*2因为是双向LSTMdef forward(self, x):# 卷积层x = self.pool(torch.relu(self.conv1(x)))x = self.pool(torch.relu(self.conv2(x)))x = self.pool(torch.relu(self.conv3(x)))# 改变形状以适应LSTM层x = x.permute(0, 2, 3, 1) # [batch_size, height, width, channels] 转 [batch_size, width, height * channels]x = x.view(x.size(0), x.size(2), -1) # [batch_size, width, height * channels]# LSTM层x, _ = self.lstm(x)# 输出层x = self.fc(x)return x
定义模型
model = CRNNModel(num_classes=len(char_set))
(2) 模型架构解释
卷积层(Conv2D):用于提取图像中的局部特征。
池化层(MaxPool2D):用于减小图像尺寸,减少计算量。
LSTM层:用于处理图像序列中的字符依赖关系。
全连接层:用于将LSTM的输出映射到字符类别空间。
4. CTC损失函数
CTC损失函数(Connectionist Temporal Classification)是一种用于处理序列到序列问题的损失函数。它适用于输出序列的长度不等于输入序列的情况,如在OCR(光学字符识别)中,字符数目和图像宽度可能不一致。
python
import torch.optim as optim
定义损失函数和优化器
criterion = nn.CTCLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
5. 模型训练
在训练模型时,我们需要提供训练数据和标签,计算损失并更新权重。这里我们使用 PyTorch 中的 DataLoader 进行批量处理。
(1) 训练代码
python
from torch.utils.data import DataLoader
初始化DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
训练过程
epochs = 10
for epoch in range(epochs):
model.train()
running_loss = 0.0
for imgs, labels in train_loader:
imgs = torch.tensor(imgs).float().unsqueeze(1) # [batch_size, 1, height, width]
labels = torch.tensor(labels).long()
# 前向传播optimizer.zero_grad()output = model(imgs)# CTC损失计算output = output.permute(1, 0, 2) # [T, N, C] 转置为 [seq_len, batch_size, num_classes]loss = criterion(output, labels)# 反向传播和优化loss.backward()optimizer.step()running_loss += loss.item()print(f"Epoch [{epoch + 1}/{epochs}], Loss: {running_loss / len(train_loader):.4f}")
(2) 训练过程解释
DataLoader:负责按批次加载数据。
CTC损失函数:计算每个字符的损失并进行反向传播。
优化器(Adam):更新网络权重。
6. 模型评估与预测
训练完成后,我们可以在测试集上评估模型的性能,并对验证码图像进行预测。
(1) 评估模型
python
测试模型
def evaluate_model(model, test_dataset):
model.eval()
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
correct = 0
total = 0
with torch.no_grad():for imgs, labels in test_loader:imgs = torch.tensor(imgs).float().unsqueeze(1)labels = torch.tensor(labels).long()output = model(imgs)output = output.permute(1, 0, 2)predicted = torch.argmax(output, 2)for i in range(len(labels)):if torch.equal(predicted[i], labels[i]):correct += 1total += 1accuracy = correct / total * 100
print(f"Test Accuracy: {accuracy:.2f}%")
假设我们有测试数据集 test_dataset
evaluate_model(model, test_dataset)
(2) 单个验证码预测
python
def predict_captcha(model, img_path, char_set):
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]
# 模型预测
output = model(torch.tensor(img).float())
output = output.permute(1, 0, 2) # 转置为 [seq_len, batch_size, num_classes]pred_label = []
for i in range(output.size(0)):pred_label.append(char_set[torch.argmax(output[i])])return ''.join(pred_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}")