用 PyTorch 和 CNN 进行验证码识别

在本教程中，我们将使用 PyTorch 和卷积神经网络（CNN）来构建一个验证码识别系统。PyTorch 是一个广泛使用的深度学习框架，特别适合研究和原型设计。卷积神经网络（CNN）是处理图像数据的强大工具，它可以自动从图像中学习特征，并执行图像分类等任务。

环境准备

首先，确保你已安装以下依赖包：

pip install torch torchvision opencv-python numpy matplotlib pillow
更多内容访问ttocr.com或联系1436423940

torch：PyTorch 框架，支持深度学习模型的训练和推理。

torchvision：提供常用的图像处理功能，及数据集和模型。

opencv-python：用于图像预处理和处理。

numpy：处理数组和矩阵。

matplotlib：用于绘制图表。

数据集准备与图像预处理

验证码的图像通常包含多个字符，且图像噪声较大，可能还有旋转和扭曲。我们需要对图像进行一些基本的预处理，主要包括灰度化、二值化、去噪和字符分割。

(1) 图像预处理

我们将图像转换为灰度图，然后使用 Otsu 的二值化方法将图像转化为二值图，这将有助于提升字符的识别效果。

import cv2
import numpy as np

def preprocess_image(img_path):
# 读取图像
img = cv2.imread(img_path)

# 转换为灰度图
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# 使用 Otsu 的方法进行二值化
_, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

# 高斯模糊去噪
blurred = cv2.GaussianBlur(binary, (5, 5), 0)

return blurred

示例图像路径

img_path = 'captcha_images/test1.png'
processed_img = preprocess_image(img_path)

显示处理后的图像

cv2.imshow('Processed Image', processed_img)
cv2.waitKey(0)
cv2.destroyAllWindows()

(2) 提取字符区域

验证码中的每个字符通常由轮廓框住。我们使用 OpenCV 的轮廓检测方法来提取字符区域，并将其分割开。

def extract_characters(processed_img):
contours, _ = cv2.findContours(processed_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
char_images = []
for contour in contours:
x, y, w, h = cv2.boundingRect(contour)
if w > 10 and h > 10: # 忽略小的噪点
char_img = processed_img[y:y+h, x:x+w]
char_images.append(char_img)

# 按照字符从左到右的顺序排序
char_images.sort(key=lambda x: x[0][0])  # 排序依据是字符的左上角 x 坐标
return char_images

提取字符区域

char_images = extract_characters(processed_img)

显示提取的字符

for i, char_img in enumerate(char_images):
cv2.imshow(f'Character {i+1}', char_img)
cv2.waitKey(0)

cv2.destroyAllWindows()

构建卷积神经网络（CNN）

接下来，我们使用 PyTorch 来构建一个卷积神经网络（CNN）。CNN 由多个卷积层和池化层组成，能够自动从图像中提取特征，并对图像进行分类。

(1) 定义 CNN 模型

我们将构建一个简单的 CNN，包括两个卷积层、两个池化层，并在最后使用全连接层输出验证码的类别。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

class CNNModel(nn.Module):
def init(self):
super(CNNModel, self).init()
# 卷积层1
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
# 卷积层2
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
# 全连接层
self.fc1 = nn.Linear(64 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 36) # 36 类（假设验证码由数字和字母组成）

def forward(self, x):
    # 卷积层1 -> 池化层
    x = self.pool(torch.relu(self.conv1(x)))
    # 卷积层2 -> 池化层
    x = self.pool(torch.relu(self.conv2(x)))
    # 展平
    x = x.view(-1, 64 * 7 * 7)
    # 全连接层1
    x = torch.relu(self.fc1(x))
    # 输出层
    x = self.fc2(x)
    return x

初始化模型

model = CNNModel()
print(model)

数据加载与训练准备

为了训练模型，我们需要准备数据集。我们将加载图像并将其调整为 28x28 的大小，这样便于 CNN 进行处理。同时，我们需要将图像像素值规范化到 [0, 1] 之间。

(1) 创建数据集类

我们需要定义一个自定义的 PyTorch 数据集类，用于加载和处理验证码图像。

class CaptchaDataset(Dataset):
def init(self, image_paths, labels, transform=None):
self.image_paths = image_paths
self.labels = labels
self.transform = transform

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

def __getitem__(self, idx):
    img_path = self.image_paths[idx]
    label = self.labels[idx]
    
    # 读取图像
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    img = cv2.resize(img, (28, 28))  # 调整图像大小
    img = np.expand_dims(img, axis=0)  # 添加频道维度 (28, 28, 1)
    img = img.astype(np.float32) / 255.0  # 归一化

    # 转换为 tensor
    img = torch.tensor(img)
    label = torch.tensor(label)

    return img, label

示例数据集路径和标签

image_paths = ['captcha_images/train1.png', 'captcha_images/train2.png']
labels = [0, 1] # 假设标签是整数类别（0-9 或字母）

创建数据集

dataset = CaptchaDataset(image_paths, labels)

数据加载器

dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

训练模型

接下来，我们使用训练集训练模型。我们将使用交叉熵损失和 Adam 优化器来训练模型。

(1) 训练循环
def train_model(model, dataloader, num_epochs=5):
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    running_loss = 0.0
    for images, labels in dataloader:
        # 将图像和标签传入模型
        images = images.unsqueeze(1)  # 添加频道维度 (batch_size, 1, 28, 28)
        optimizer.zero_grad()  # 清除之前的梯度

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        loss.backward()
        optimizer.step()

        # 记录损失
        running_loss += loss.item()

    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(dataloader)}")

训练模型

train_model(model, dataloader)

测试与预测

训练完成后，我们可以使用测试数据对模型进行评估，并进行预测。

(1) 测试模型
def evaluate_model(model, dataloader):
model.eval() # 设置为评估模式
correct = 0
total = 0
with torch.no_grad(): # 不计算梯度
for images, labels in dataloader:
images = images.unsqueeze(1) # 添加频道维度
outputs = model(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()

accuracy = correct / total
print(f'Accuracy: {accuracy * 100:.2f}%')

测试模型

evaluate_model(model, dataloader)

(2) 预测新的验证码
def predict(model, img_path):
img = preprocess_image(img_path)
img = cv2.resize(img, (28, 28))
img = np.expand_dims(img, axis=0) # 添加批次维度
img = np.expand_dims(img, axis=0) # 添加频道维度
img = torch.tensor(img, dtype=torch.float32) / 255.0

model.eval()
with torch.no_grad():
    outputs = model(img)
    _, predicted = torch.max(outputs, 1)

return predicted.item()

预测新的验证码

predicted_label = predict(model, 'captcha_images/test1.png')
print(f"Predicted label: {predicted_label}")

posted @ 2025-09-17 21:49 ttocr、com 阅读(12) 评论(0) 收藏举报

刷新页面返回顶部