交通标志识别是计算机视觉领域的一个重要任务,尤其对于自动驾驶汽车和智能交通系统来说。在本文中,我们将介绍如何使用 LeNet 和 VGG 模型在德国交通标志识别数据集(GTSRB)上进行交通标志识别任务。我们将从数据预处理开始,然后分别构建和训练 LeNet 和 VGG 模型,最后对比两种模型的性能。
1. 数据集介绍
GTSRB 数据集包含 43 种不同类别的交通标志,共有 50,000 张图片。数据集中的图片大小不一,但为了训练模型,我们需要将它们调整为相同的大小。
2. 数据预处理
首先,我们需要加载数据集并进行预处理。数据预处理包括:加载图片,将图片大小调整为相同的尺寸,归一化,将标签转换为独热编码等。
2.1 加载数据集
我们可以使用以下代码从文件中加载 GTSRB 数据集。
import os
import numpy as np
import cv2
from skimage import io
from sklearn.model_selection import train_test_split
def load_data(data_dir, img_size):
images = []
labels = []
for label in os.listdir(data_dir):
label_dir = os.path.join(data_dir, label)
for img_file in os.listdir(label_dir):
img_path = os.path.join(label_dir, img_file)
img = io.imread(img_path)
img = cv2.resize(img, img_size)
images.append(img)
labels.append(int(label))
images = np.array(images)
labels = np.array(labels)
return images, labels
data_dir = "GTSRB/Final_Training/Images"
img_size = (32, 32)
images, labels = load_data(data_dir, img_size)
2.2 数据预处理
接下来,我们需要对图片进行归一化处理,并将标签转换为独热编码。
from tensorflow.keras.utils import to_categorical
# 归一化图片
images = images / 255.0
# 将标签转换为独热编码
num_classes = len(np.unique(labels))
labels = to_categorical(labels, num_classes)
# 划分训练集和测试集
train_X, test_X, train_y, test_y = train_test_split(images, labels, test_size=0.2, random_state=42)
3. 构建 LeNet 模型
LeNet 是一种经典的卷积神经网络,适用于处理小尺寸的图片。我们将使用 TensorFlow 构建一个 LeNet 模型。
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
def build_lenet_model(input_shape, num_classes):
model = Sequential()
model.add(Conv2D(filters=32, kernel_size=(5, 5), activation='relu', input_shape=input_shape))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(filters=64, kernel_size=(3, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(units=120, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(units=84, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(units=num_classes, activation='softmax'))
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
return model
input_shape = train_X.shape[1:]
lenet_model = build_lenet_model(input_shape, num_classes)
lenet_model.summary()
在这个模型中,我们首先添加两个卷积层和最大池化层,分别使用 32 个 5x5 和 64 个 3x3 的卷积核。接着,我们添加一个 Flatten 层将输出展平,并使用两个全连接层进行分类。为了防止过拟合,我们在全连接层之间添加了 Dropout 层。最后,我们使用一个全连接层将输出映射到 43 个交通标志类别,并使用 Softmax 激活函数。模型使用 Adam 优化器和分类交叉熵损失函数进行编译。
4. 训练 LeNet 模型
接下来,我们使用训练数据训练 LeNet 模型。
lenet_history = lenet_model.fit(train_X, train_y, epochs=20, batch_size=64, validation_split=0.2, verbose=1)
训练完成后,我们可以使用测试数据集对模型进行评估。为了量化模型性能,我们计算预测准确率。
lenet_loss, lenet_accuracy = lenet_model.evaluate(test_X, test_y, verbose=1)
print(f"LeNet model test accuracy: {lenet_accuracy:.4f}")
接下来,我们将使用 VGG 模型进行同样的任务,并与 LeNet 模型进行性能比较。
5. 构建 VGG 模型
VGG 是另一种流行的卷积神经网络,已被证明在大型数据集上具有很好的性能。在这里,我们将构建一个适用于 GTSRB 数据集的简化版 VGG 模型。
from tensorflow.keras.layers import BatchNormalization
def build_vgg_model(input_shape, num_classes):
model = Sequential()
model.add(Conv2D(filters=64, kernel_size=(3, 3), activation='relu', padding='same', input_shape=input_shape))
model.add(BatchNormalization())
model.add(Conv2D(filters=64, kernel_size=(3, 3), activation='relu', padding='same'))
model.add(BatchNormalization())
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))
model.add(Conv2D(filters=128, kernel_size=(3, 3), activation='relu', padding='same'))
model.add(BatchNormalization())
model.add(Conv2D(filters=128, kernel_size=(3, 3), activation='relu', padding='same'))
model.add(BatchNormalization())
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))
model.add(Flatten())
model.add(Dense(units=512, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(units=num_classes, activation='softmax'))
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
return model
vgg_model = build_vgg_model(input_shape, num_classes)
vgg_model.summary()
在简化版的 VGG 模型中,我们添加了两个卷积层堆叠的结构,分别使用 64 个 3x3 和 128 个 3x3 的卷积核。在每个卷积层后,我们添加 BatchNormalization 层以加速训练过程。与 LeNet 模型相似,我们在最后添加了一个全连接层进行分类,并在全连接层之前添加了 Dropout 层以防止过拟合。
6. 训练 VGG 模型
我们使用相同的训练数据集训练 VGG 模型。
vgg_history = vgg_model.fit(train_X, train_y, epochs=20, batch_size=64, validation_split=0.2, verbose=1)
训练完成后,我们使用测试数据集评估 VGG 模型的性能。
vgg_loss, vgg_accuracy = vgg_model.evaluate(test_X, test_y, verbose=1)
print(f"VGG model test accuracy: {vgg_accuracy:.4f}")
7. 比较 LeNet 和 VGG 模型
现在我们已经训练了两种不同的模型,我们可以比较它们在交通标志识别任务上的性能。
print(f"LeNet model test accuracy: {lenet_accuracy:.4f}")
print(f"VGG model test accuracy: {vgg_accuracy:.4f}")
根据测试准确率,我们可以评估两个模型在 GTSRB 数据集上的性能表现。在这个例子中,你可能会发现 VGG 模型的性能优于 LeNet 模型。这可能是因为 VGG 模型具有更多的卷积层和参数,可以捕捉到更多的特征。
8. 可视化模型性能
为了更直观地展示模型的训练过程和性能,我们可以绘制训练和验证准确率以及损失的曲线。
import matplotlib.pyplot as plt
def plot_history(histories, titles):
plt.figure(figsize=(10, 6))
for history, title in zip(histories, titles):
plt.plot(history.history['accuracy'])
plt.plot(history.history['val_accuracy'])
plt.title(f"{title} Model Accuracy")
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend(['Train', 'Validation'], loc='upper left')
plt.show()
plt.figure(figsize=(10, 6))
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title(f"{title} Model Loss")
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend(['Train', 'Validation'], loc='upper right')
plt.show()
plot_history([lenet_history, vgg_history], ['LeNet', 'VGG'])
通过观察训练和验证准确率以及损失曲线,我们可以分析模型在训练过程中的表现。这有助于我们检测过拟合现象以及验证模型泛化能力。
9. 总结
在本文中,我们介绍了如何在 GTSRB 数据集上使用 LeNet 和 VGG 模型进行交通标志识别任务。我们从数据预处理开始,然后分别构建和训练了 LeNet 和 VGG 模型。最后,我们比较了两种模型在识别任务上的性能。
虽然在这个例子中 VGG 模型的性能优于 LeNet,但请注意,实际应用中的结果可能会有所不同,具体取决于数据集和问题的复杂性。在实践中,你可以尝试使用不同的模型、参数和数据增强技术来优化模型性能。