基于多模型的车辆检测与识别技术的开发（车型检测）

88conch

于 2024-06-25 14:30:00 发布

阅读量1.1k

点赞数 17

文章标签：人工智能机器学习神经网络 python 计算机视觉 cnn opencv

本文链接：https://blog.youkuaiyun.com/2301_78972534/article/details/139948211

版权

车辆识别

准备
导入工具包
创建EfficientNet 的模型
加载训练好的模型参数
定义图像转换
类别标签
初始化全局变量
上传图片
车型检测

准备

本篇将展示车型检测功能。

导入工具包

import torch.nn as nn
import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk,ImageDraw,ImageFont
import torch
from torchvision import transforms, models
from efficientnet_pytorch import EfficientNet
import numpy as np
import cv2
import tkinter.ttk as ttk
import time
from tkinter import filedialog
import subprocess
import os
import re
import threading

要点:

torch.nn：这是PyTorch的一个模块，用于构建神经网络
tkinter：Python的标准GUI库，用于创建图形用户界面
PIL（Pillow）：一个图像处理库，用于打开、处理和保存多种不同格式的图像
torch：PyTorch的核心库，用于构建和训练神经网络
torchvision：包含图像处理、数据集和预训练模型等工具的库
efficientnet_pytorch：一个提供EfficientNet模型的PyTorch实现
numpy：一个强大的数学库，用于进行科学计算
cv2：OpenCV库的Python接口，用于计算机视觉任务
tkinter.ttk：tkinter的一个模块，提供了改进的Tk控件
threading：Python的标准线程库，用于创建和管理线程

创建EfficientNet 的模型

加载预训练模型

class EfficientNetModel(nn.Module):
    def __init__(self, num_classes=10, pretrained=True):
        super(EfficientNetModel, self).__init__()
        # 加载预训练的EfficientNet模型
        self.efficientnet = EfficientNet.from_name('efficientnet-b3')
        num_ftrs = self.efficientnet._fc.in_features
        # 将EfficientNet模型的最后一层全连接层替换为一个新的全连接层，输出特征数量设置为num_classes
        self.efficientnet._fc = nn.Linear(num_ftrs, num_classes)

要点:

类继承自 torch.nn.Module，用于创建一个基于 EfficientNet 的模型，可以用于图像分类任务。
这个类接受两个参数：num_classes 表示输出分类的数量，pretrained 表示是否使用预训练的模型权重。

向前传播

def forward(self, x):
    return self.efficientnet(x)

定义EfficientNet类

model = EfficientNetModel(num_classes=12)
# num_classes = 12  # 根据实际类别数修改

加载训练好的模型参数

model_path = 'best_EfficientNet_b3_updata6.pth'
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()

要点:
model_path 变量指定了模型权重文件的路径，这个文件通常是一个 PyTorch 模型的状态字典，它包含了模型的参数
torch.load() 函数用于加载模型权重文件。map_location=torch.device(‘cpu’) 参数指定了加载模型权重时使用的设备。在这里，它被设置为 CPU。这意味着即使模型是在 GPU 上训练的，加载时也会将其参数移动到 CPU 上
model.load_state_dict() 方法用于将加载的状态字典加载到模型中，这样模型就有了训练好的参数
model.eval() 方法将模型设置为评估模式。这个方法对模型没有什么影响，但是如果模型中有定义为仅在训练模式下使用的层（如 Dropout 或 BatchNorm），则这个调用会使其在评估模式下正确地行为

定义图像转换

transform = transforms.Compose([
    transforms.Resize((300, 300)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # ResNet-50 的标准化
])

要点:

transforms.Resize((300, 300))：将输入图像的大小调整为 300x300 像素。这个步骤确保所有输入到模型的图像都具有相同的尺寸，这是许多深度学习模型的要求
transforms.ToTensor()：将 PIL 图像或 NumPy 数组转换为浮点数张量，并对其进行归一化，使得其数值范围在 [0, 1] 内。此外，它还会将图像的维度从 (H, W, C) 转换为 (C, H, W)，这与 PyTorch 的默认图像维度格式相匹配
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])：`对图像进行标准化处理。这个步骤根据提供的均值和标准差对图像的每个通道进行标准化。这些均值和标准差通常来自预训练模型的训练数据集，对于 ResNet-50 模型，这些值是常用的。标准化可以使得模型在训练和推理时对输入数据的分布有更稳定的预期，这有助于模型的性能和收敛

类别标签

classes =  ['皮卡', '敞篷车', '跑车', '掀背两箱车', '小型面包车', 'SUV', '轿车', '厢式货车', '旅行车', '公共汽车', '消防车', '出租车']

初始化全局变量

selected_image_path = None
label_text = None
right_canvas_image = None

要点:

selected_image_path：这个变量可能用于存储用户选择的图像文件的路径。在 GUI 应用程序中，通常会提供一个文件对话框让用户选择图像文件，然后应用程序将选择的文件的路径存储在这个变量中
label_text：这个变量可能用于存储与图像相关联的标签文本。例如，如果应用程序是一个图像分类器，label_text 可能用于显示图像的最可能的类别
right_canvas_image：这个变量可能用于存储一个图像对象，该对象将在 GUI 的画布上显示。在 Tkinter（Python 的 GUI 库）中，画布（Canvas）是一个可以用来绘制图形和图像的组件。right_canvas_image 可能是一个 ImageTk.PhotoImage 对象，它是一个可以在 Tkinter 中显示的图像

上传图片

def upload_image():
    global selected_image_path, label_text
    # 获取'.\output'目录下的所有文件
    files = os.listdir(r'.\output')

    # 遍历文件
    for file in files:
        # 拼接文件的完整路径
        file_path = os.path.join(r'.\output', file)

        # 检查文件是否为图片文件（这里假设只需要删除jpg和png文件）
        if file.endswith('.jpg') or file.endswith('.png'):
            # 删除文件
            os.remove(file_path)
    file_path = filedialog.askopenfilename()
    if file_path:
        selected_image_path = file_path
        image = Image.open(file_path)

        # 调整图片大小为500x400
        image = image.resize((500, 400), Image.Resampling.LANCZOS)

        # 居中图片
        photo = ImageTk.PhotoImage(image)
        canvas_left.create_image(0, 0, anchor='nw', image=photo)
        canvas_left.image = photo  # Keep a reference!

        # 创建图片的标签
        if label_text is None:
            label_text = tk.Label(window, text="", font=("Arial", 16))
            label_text.place(x=1185,y=155)

要点：

函数首先获取 .\output 目录下的所有文件，并遍历这些文件。如果文件是图片文件（假设只有 .jpg 和 .png
格式），则将其删除。这可能是为了清理之前的输出，为新的图像处理做准备
使用tkinter.filedialog.askopenfilename()
弹出一个文件对话框，让用户选择一个图像文件。如果用户选择了文件，则将文件的路径存储在 selected_image_path 变量中
使用PIL.Image.open() 打开用户选择的图像文件
调整图像的大小为 500x400 像素，使用Image.Resampling.LANCZOS 重采样方法，这是一种高质量的重采样滤波器
将调整大小后的图像居中显示在一个名为canvas_left 的画布上。这里使用了 ImageTk.PhotoImage 来创建一个可以在 Tkinter中显示的图像对象，并将其附加到画布上
如果 label_text 是 None，则创建一个新的标签label_text，用于显示图像的标签文本。标签被放置在窗口的特定位置（坐标为 x=1185, y=155），并设置了字体样式和大小

车型检测

def start_detection():
    global right_canvas_image
    if selected_image_path is not None:
        image = Image.open(selected_image_path)
        input_image = transform(image).unsqueeze(0)
        with torch.no_grad():
            outputs = model(input_image)
            _, predicted = torch.max(outputs, 1)
            label = classes[predicted.item()]
            probabilities = torch.nn.functional.softmax(outputs, dim=1)
            max_probability = probabilities[0][predicted].item() * 100  # 将概率值乘以100
            label_text.config(text=f"{label} - {max_probability:.2f}%")  # 显示为百分比格式

        # 调整图片大小为500x400
        image = image.resize((500, 400), Image.Resampling.LANCZOS)

        # 显示图片在右侧画布
        photo = ImageTk.PhotoImage(image)

        # 检查是否已经创建了右侧画布的图片
        if right_canvas_image is None:
            right_canvas_image = canvas_right.create_image(0, 0, anchor='nw', image=photo)
        else:
            canvas_right.itemconfig(right_canvas_image, image=photo)
        canvas_right.image = photo  # 保持引用
    else:
        messagebox.showwarning("警告", "请先选择一张图像")

    # 将标签放置在图片上
    label_text.place(x=1185,y=155)