当前位置：首页 > ds >正文

# 交通标志识别：使用卷积神经网络的完整实现

ds 2025/5/3 13:35:51

交通标志识别：使用卷积神经网络的完整实现

在当今数字化时代，计算机视觉技术在各个领域都发挥着重要作用，交通标志识别就是其中的一个典型应用。通过让机器自动识别交通标志，不仅可以辅助自动驾驶系统，还能为交通管理提供更高效的支持。本文将详细介绍如何使用卷积神经网络（CNN）来实现交通标志的识别，并提供完整的代码实现。

1. 数据准备

数据是机器学习项目的基础，对于交通标志识别任务，我们需要准备包含交通标志图像及其对应标签的数据集。在本项目中，我们假设数据集已经准备好，并且图像和标签分别存储在指定的文件夹和文本文件中。

图像文件夹：所有交通标志图像存储在同一个文件夹中，文件名与标签文件中的记录相对应。
标签文件：标签文件是一个文本文件，每一行包含一个图像文件名和对应的标签。标签以空格分隔的形式表示，例如 "image1.jpg 0 1 0 0 0 0"，表示该图像属于第二个类别。

我们通过自定义的 TrafficDataset 类来加载数据集，并应用数据预处理操作。数据预处理包括将图像缩放到统一的大小（64x64）、转换为张量，并进行归一化处理。归一化使用了 ImageNet 数据集的均值和标准差，这有助于提高模型的训练效果。

1.1 数据集类

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import numpy as npclass TrafficDataset(Dataset):def __init__(self, data_dir, label_file, transform=None):self.data_dir = data_dirself.label_file = label_fileself.transform = transformself.images = []self.labels = []# 读取标签文件with open(label_file, 'r') as f:for line in f.readlines():parts = line.strip().split()image_path = os.path.join(data_dir, parts[0])label = list(map(int, parts[1:]))self.images.append(image_path)self.labels.append(label)def __len__(self):return len(self.images)def __getitem__(self, idx):image_path = self.images[idx]label = self.labels[idx]image = Image.open(image_path).convert('RGB')if self.transform:image = self.transform(image)return image, torch.tensor(label, dtype=torch.float32)

1.2 数据预处理

transform = transforms.Compose([transforms.Resize((64, 64)),  # 缩放到64x64transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

1.3 数据集路径

data_dir = r'D:\Users\妄生\PycharmProjects\人工智能\作业（2023）\task4'
train_label_file = r'D:\Users\妄生\PycharmProjects\人工智能\作业（2023）\task4\trainlabels.txt'
test_label_file = r'D:\Users\妄生\PycharmProjects\人工智能\作业（2023）\task4\testlabels.txt'

1.4 创建数据集和数据加载器

train_dataset = TrafficDataset(data_dir, train_label_file, transform=transform)
test_dataset = TrafficDataset(data_dir, test_label_file, transform=transform)train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)

2. 模型构建

交通标志识别的核心是卷积神经网络（CNN）。CNN 在图像识别任务中表现出色，因为它能够自动提取图像的特征。我们设计了一个简单的 CNN 模型，包含三个卷积层和两个全连接层。

卷积层：每个卷积层后接一个最大池化层，用于降低特征图的空间维度，同时保留重要特征。
全连接层：在卷积层之后，我们将特征图展平，然后通过两个全连接层进行分类。输出层有 7 个神经元，对应 7 个类别。

2.1 定义卷积神经网络模型

class TrafficCNN(nn.Module):def __init__(self):super(TrafficCNN, self).__init__()self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)self.conv3 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)self.fc1 = nn.Linear(64 * 8 * 8, 128)self.fc2 = nn.Linear(128, 7)  # 输出层对应7个类别self.pool = nn.MaxPool2d(kernel_size=2, stride=2)self.relu = nn.ReLU()self.sigmoid = nn.Sigmoid()def forward(self, x):x = self.pool(self.relu(self.conv1(x)))x = self.pool(self.relu(self.conv2(x)))x = self.pool(self.relu(self.conv3(x)))x = x.view(-1, 64 * 8 * 8)  # 展平x = self.relu(self.fc1(x))x = self.sigmoid(self.fc2(x))return x

3. 训练过程

训练是让模型学习数据特征的过程。我们使用随机梯度下降（SGD）优化器和二元交叉熵损失函数来训练模型。训练过程中，模型会不断调整参数，以最小化损失函数的值。

3.1 初始化模型、优化器和损失函数

model = TrafficCNN()
optimizer = optim.SGD(model.parameters(), lr=0.001, weight_decay=0.0001)
criterion = nn.BCELoss()  # 使用二元交叉熵损失函数

3.2 训练模型

num_epochs = 10
for epoch in range(num_epochs):model.train()running_loss = 0.0for images, labels in train_loader: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(train_loader):.4f}')

在每个 epoch 结束时，我们打印出平均损失值，以便观察训练过程的进展。

4. 测试与评估

训练完成后，我们需要对模型进行测试，以评估其性能。测试过程中，模型不会进行反向传播，我们只计算损失值。

4.1 测试模型

model.eval()
with torch.no_grad():test_loss = 0.0for images, labels in test_loader:outputs = model(images)loss = criterion(outputs, labels)test_loss += loss.item()print(f'Test Loss: {test_loss/len(test_loader):.4f}')