基于YOLOv10算法的交通信号灯检测与识别

为了实现基于 YOLOv10 算法的交通信号灯检测与识别，我们可以将整个过程分为以下几个步骤，并在每个步骤中提供相应的代码示例。假设我们使用 PyTorch 框架来训练 YOLOv10 模型，下面是一个完整的实现流程。

1. 数据准备与标注

首先，收集一个交通信号灯数据集并进行标注。假设我们有一个名为 traffic_light_dataset 的文件夹，其中包含了多个交通信号灯的图片。你可以使用工具如 LabelImg 来标注图像，并生成 .xml 或 .txt 格式的标注文件。

对于 YOLO 格式的标注文件，文件内容大致如下：

xml 复制代码

<class_id> <x_center> <y_center> <width> <height>
php
49 Bytes
© 菜鸟-创作你的创作

其中：

<class_id>：类别 ID，红灯、绿灯、黄灯。
<x_center>、<y_center>：标注框中心点的相对坐标（归一化到 [0, 1] 区间）。
<width>、<height>：标注框的宽度和高度，归一化到 [0, 1] 区间。

2. YOLOv10 模型训练

我们使用 YOLOv5 的 PyTorch 实现训练 YOLOv10 模型（YOLOv10 通常为 YOLOv5 的升级版），主要步骤如下：

2.1. 安装必要的依赖

bash 复制代码

# 克隆 YOLOv5 仓库并安装依赖
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -U -r requirements.txt
php
112 Bytes
© 菜鸟-创作你的创作

2.2. 数据集结构

确保你的数据集符合 YOLO 的要求，数据集目录结构应如下：

kotlin 复制代码

traffic_light_dataset/
├── images/
│   ├── train/
│   │   ├── img1.jpg
│   │   ├── img2.jpg
│   ├── val/
│   │   ├── img1.jpg
│   │   ├── img2.jpg
├── labels/
│   ├── train/
│   │   ├── img1.txt
│   │   ├── img2.txt
│   ├── val/
│   │   ├── img1.txt
│   │   ├── img2.txt
php
270 Bytes
© 菜鸟-创作你的创作

2.3. 配置 `data.yaml` 文件

在 YOLOv5 中，需要一个 data.yaml 配置文件，指定训练数据集的位置和类别信息。

bash 复制代码

# traffic_light.yaml
train: ./traffic_light_dataset/images/train
val: ./traffic_light_dataset/images/val
nc: 3  # 类别数目：红灯、绿灯、黄灯
names: ['red', 'green', 'yellow']  # 类别标签
php
170 Bytes
© 菜鸟-创作你的创作

2.4. 训练 YOLOv10 模型

使用 YOLOv5 提供的训练脚本进行模型训练。

css 复制代码

# 开始训练模型
python train.py --img 640 --batch 16 --epochs 50 --data traffic_light.yaml --weights yolov5s.pt --cache
php
112 Bytes
© 菜鸟-创作你的创作

--img：图像大小，YOLOv10 默认 640×640。
--batch：批处理大小。
--epochs：训练轮数。
--data：指向数据配置文件。
--weights：预训练模型权重，yolov5s.pt 是一个小型的 YOLOv5 模型，适用于快速训练。

2.5. 模型评估

训练完成后，你可以使用 test.py 脚本评估模型的表现，检查模型在验证集上的 mAP（mean Average Precision） 等指标。

css 复制代码

python val.py --data traffic_light.yaml --weights runs/train/exp/weights/best.pt --img 640
php
90 Bytes
© 菜鸟-创作你的创作

3. 交通信号灯状态分类

虽然 YOLOv10 会检测出交通信号灯的位置和类别（如：红灯、绿灯、黄灯），但是如果需要对交通信号灯的状态进行进一步分类（如，判断红灯是否亮起），可以通过以下方式：

使用色彩分析：基于图像中的颜色，判断信号灯的状态。
训练额外的分类模型：使用 YOLOv10 检测出的交通信号灯区域作为输入，训练一个分类模型来进一步确认信号灯的状态。

这里我们展示一种使用颜色识别的方法：

ini 复制代码

import cv2
import numpy as np
# 颜色范围（H, S, V）
COLOR_RANGES = {
    "red": [(0, 150, 50), (10, 255, 255)],    # 红色
    "green": [(35, 150, 50), (85, 255, 255)], # 绿色
    "yellow": [(20, 150, 50), (40, 255, 255)] # 黄色
}
def detect_light_color(roi):
    # 转换到 HSV 色彩空间
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
    # 按照颜色范围过滤
    for color, (lower, upper) in COLOR_RANGES.items():
        mask = cv2.inRange(hsv, np.array(lower), np.array(upper))
        if np.any(mask):  # 如果有颜色区域
            return color
    return "unknown"
# 示例：对 YOLOv10 检测到的信号灯区域进行颜色分类
image = cv2.imread('traffic_light.jpg')
roi = image[100:200, 150:250]  # 假设这是 YOLOv10 检测到的红灯区域
light_color = detect_light_color(roi)
print(f"Detected Light Color: {light_color}")
php
744 Bytes
© 菜鸟-创作你的创作

4. 实时推理与部署

对于实时推理，我们将 YOLOv10 模型部署到边缘设备（如 Raspberry Pi、Jetson Nano）上，或者将其部署到云端服务中。这里演示如何在本地进行推理：

4.1. 加载 YOLOv10 模型进行推理