基于ResNet50的车辆识别模型: Python实现

本教程介绍如何使用Python和TensorFlow构建基于ResNet50的车辆识别模型。

代码示例

以下代码演示了如何使用ResNet50作为基础模型构建一个车辆识别模型：

# 导包
import os
import random
import numpy as np
import pandas as pd
import tensorflow as tf
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPooling2D, BatchNormalization, GlobalAveragePooling2D
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications import ResNet50
from sklearn.metrics import accuracy_score

# 设置参数
train_dir = 'C:/Users/27546/Desktop/Vehicle Identification/train' # 训练集目录
val_dir = 'C:/Users/27546/Desktop/Vehicle Identification/test' # 验证集目录
classes = os.listdir(train_dir)  # 获取类别列表
batch_size = 64  # 批大小
IMG_HEIGHT = 150  # 图像高度
IMG_WIDTH = 150  # 图像宽度
epochs = 10 # 训练轮数

# 创建图像生成器
train_image_generator = ImageDataGenerator(
    rescale=1./255,  # 归一化
    horizontal_flip=True  # 水平翻转
)

val_image_generator = ImageDataGenerator(
    rescale=1./255  # 归一化
)

train_data_gen = train_image_generator.flow_from_directory(
    batch_size=batch_size,  # 批大小
    directory=train_dir,  # 训练集目录
    shuffle=True,  # 是否打乱数据
    target_size=(IMG_HEIGHT, IMG_WIDTH),  # 图像大小
    class_mode='categorical'  # 分类方式
)

val_data_gen = val_image_generator.flow_from_directory(
    batch_size=batch_size,  # 批大小
    directory=val_dir,  # 验证集目录
    shuffle=True,  # 是否打乱数据
    target_size=(IMG_HEIGHT, IMG_WIDTH),  # 图像大小
    class_mode='categorical'  # 分类方式
)

total_train = train_data_gen.samples  # 训练集样本数
total_val = val_data_gen.samples  # 验证集样本数

# 加载 ResNet50 模型并进行微调
base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3))  # 加载 ResNet50 模型
x = base_model.output  # 获取 ResNet50 模型的输出
x = GlobalAveragePooling2D()(x)  # 添加全局平均池化层
x = Dense(1024, activation='relu')(x)  # 添加全连接层
x = Dropout(0.5)(x)  # 添加Dropout层
x = Dense(512, activation='relu')(x)  # 添加全连接层
x = Dropout(0.5)(x)  # 添加Dropout层
predictions = Dense(len(classes), activation='softmax')(x)  # 添加输出层
model = Model(inputs=base_model.input, outputs=predictions) 

# 冻结 ResNet50 的前几层
for layer in base_model.layers:
    layer.trainable = False

# 编译模型
opt = tf.keras.optimizers.Adam(learning_rate=0.0001)  # 定义优化器
model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])  # 编译模型

model.summary()  # 输出模型结构

# 训练模型
history = model.fit(
    train_data_gen,  # 训练集数据生成器
    steps_per_epoch=total_train // batch_size,  # 每轮迭代的步数
    epochs=epochs,  # 训练轮数
    validation_data=val_data_gen,  # 验证集数据生成器
    validation_steps=total_val // batch_size  # 验证集每轮迭代的步数
)

# 绘制训练过程中的准确率和损失曲线
acc = history.history['accuracy']  # 获取训练准确率列表
val_acc = history.history['val_accuracy']  # 获取验证准确率列表
loss = history.history['loss']  # 获取训练损失列表
val_loss = history.history['val_loss']  # 获取验证损失列表
epochs_range = range(epochs)  # 创建轮数范围

# 绘制准确率曲线
plt.figure(figsize=(16, 8))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label='Training Accuracy')  # 绘制训练准确率曲线
plt.plot(epochs_range, val_acc, label='Validation Accuracy')  # 绘制验证准确率曲线
plt.legend(loc='lower right')  # 添加图例，位置在右下角
plt.title('Training and Validation Accuracy')  # 设置标题

# 绘制损失曲线
plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label='Training Loss')  # 绘制训练损失曲线
plt.plot(epochs_range, val_loss, label='Validation Loss')  # 绘制验证损失曲线
plt.legend(loc='upper right')  # 添加图例，位置在右上角
plt.title('Training and Validation Loss')  # 设置标题

plt.show()  # 显示图像

# 模型测试
# 加载测试集数据生成器
test_image_generator = ImageDataGenerator(
    rescale=1. / 255  # 归一化
)

test_data_gen = test_image_generator.flow_from_directory(
    batch_size=1,  # 批大小为1，每次测试一个样本
    directory=val_dir,  # 测试集目录
    shuffle=False,  # 不打乱数据
    target_size=(IMG_HEIGHT, IMG_WIDTH),  # 图像大小
    class_mode='categorical'  # 分类方式
)

# 获取测试集样本数
total_test = len(test_data_gen)
print("Total testing data batches : ", total_test)

# 进行预测并计算准确率
y_pred = model.predict(test_data_gen)  # 对测试集进行预测
y_pred_class = np.argmax(y_pred, axis=1)  # 获取预测结果的类别
y_true_class = test_data_gen.classes  # 获取真实标签的类别
class_names = list(test_data_gen.class_indices.keys())  # 获取类别名称
accuracy = accuracy_score(y_true_class, y_pred_class)  # 计算准确率
print('Accuracy: {:.2f}%'.format(accuracy * 100))

# 随机选择一些测试样本进行展示
for i in range(5):
    rand_num = random.randint(0, total_test - 1)  # 随机选择一个样本
    img, label = test_data_gen[rand_num]  # 获取样本图像和标签
    img = img[0]  # 从batch中取出单个图像
    label_name = class_names[np.argmax(label)]  # 获取标签名称
    pred_name = class_names[y_pred_class[rand_num]]  # 获取预测结果名称
    plt.title('true label: {}, predicted label: {}'.format(label_name, pred_name))
    plt.imshow(img)
    plt.show()

模型说明

该模型使用ResNet50作为基础模型，并添加了自定义的全连接层和输出层。ResNet50是一个强大的卷积神经网络，已经在ImageNet数据集上进行了预训练，因此可以有效地提取图像特征。

总结

本文介绍了一个基于ResNet50的车辆识别模型，并提供了完整的Python代码实现。该模型可以有效地识别不同类型的车辆，并可以根据需要进行修改和扩展。