项目背景
随着智能交通系统的快速发展，自动驾驶技术逐渐成为研究的热点。在自动驾驶领域中，准确识别道路上的交通信号灯是确保车辆安全行驶的关键技术之一。近年来，深度学习技术的发展为交通信号灯的识别提供了强大的支持。YOLO（You Only Look Once）作为一种高效的物体检测算法，在实时场景下有着广泛的应用。本文将介绍如何使用YOLOv8模型进行红绿灯检测，并结合Streamlit实现一个简单的Web应用。
YOLOv8 & Streamlit 简介
YOLOv8 简介
YOLOv8是YOLO系列的最新版本，它在继承了前几代YOLO的优点基础上进行了改进，具有更快的速度和更高的精度。YOLOv8采用了统一的架构设计，可以轻松地在不同的任务之间切换，如目标检测、实例分割等。该模型支持多种后端框架，包括PyTorch等。
特性
- 统一的架构：YOLOv8提供了一个统一的训练脚本，可以快速地调整模型以适应不同的任务需求。
- 高性能：相较于之前的版本，YOLOv8在速度与准确性方面都有所提升。
- 易于部署：YOLOv8支持多种部署方式，包括ONNX格式，使得模型可以在不同平台上运行。
Streamlit 简介
Streamlit是一个用于创建和共享数据应用的开源Python库。它简化了构建交互式数据可视化界面的过程，允许开发者通过简单的Python代码快速搭建功能丰富的Web应用。
数据集
为了训练YOLOv8模型进行红绿灯检测，我们需要一个包含红绿灯图像的数据集。一个合适的公开数据集可以从以下链接下载：
- Link: [假设的数据集链接]
模型训练
复制配置文件
从YOLOv8的GitHub仓库中下载官方提供的配置文件，并根据我们的需求进行调整。例如，我们可以选择一个预训练的基础模型作为起点，并修改类别数以适应红绿灯检测任务。
配置YOLOv8参数
编辑配置文件中的路径设置，指向我们准备好的数据集目录。此外，还需要设置训练相关的参数，如batch size、学习率等。
总共4各类别：
0: Traffic Light-Red Light
1: Traffic Light-Yellow Light
2: Traffic Light-Green Light
3: Traffic Light-Off
train: data_train/images # train images (relative to 'path') 4 images
val: data_valid/images # val images (relative to 'path') 4 images
test: # test images (optional)
nc: 4
# Classes
names:
0: Traffic Light-Red Light
1: Traffic Light-Yellow Light
2: Traffic Light-Green Light
3: Traffic Light-Off
挂载数据集
确保数据集能够被模型正确读取。通常，这涉及到将数据集按照训练、验证和测试集分开，并生成相应的.txt文件指示每个集合中图片的位置。
模型验证
在训练过程中定期保存检查点，并在验证集上评估模型性能。使用诸如mAP（mean Average Precision）这样的指标来衡量模型的有效性。
# Use the model
model.train(data="yolo8.yaml", epochs=120,batch=4)
模型导出及量化
模型量化
为了提高模型在边缘设备上的推理速度，我们可以对其进行量化处理。量化可以减少模型大小并加快推理速度，但可能会牺牲一些精度。
ONNX 推理验证
将训练好的模型转换成ONNX格式，这是一种开放的交换格式，能够在多个平台和框架间进行模型互操作。然后，在ONNX环境中验证模型的准确性和性能。
Streamlit 部署ONNX模型-web
使用Streamlit创建一个用户界面，用户可以通过上传图片或者输入URL来让模型进行实时的红绿灯检测。通过调用ONNX模型来进行预测，并将结果展示给用户。
代码
mport numpy as np
import onnxruntime
import cv2
import matplotlib.pyplot as plt
target_size = 640.
inputs = {}
img_path = 'data/data153372/trafficlight/JPEGImages/00009.jpg'
session = onnxruntime.InferenceSession('ppyolov2_infer_quant_dynamic.onnx')
input_names = [input.name for input in session.get_inputs()]
output_names = [output.name for output in session.get_outputs()]
img = cv2.imread(img_path)
origin_shape = img.shape[:2]
# im_scale_y = target_size / float(origin_shape[0])
# im_scale_x = target_size / float(origin_shape[1])
im_scale_x = im_scale_y = 1.0
scale_factor = np.array([[im_scale_y, im_scale_x]]).astype('float32')
im = cv2.resize(img,(int(target_size), int(target_size)))
im = im / 255.0
mean = [0.485, 0.456, 0.406]
std =[0.229, 0.224, 0.225]
im = (im - mean) / std
# im = im[:, :, ::-1]
im = np.expand_dims(np.transpose(im, (2, 0, 1)), axis=0)
inputs['im_shape'] = np.array([origin_shape]).astype('float32')
inputs['scale_factor'] = scale_factor
inputs['image'] = im.astype('float32')
np_boxes = session.run(output_names, inputs)[0]
expect_boxes = (np_boxes[:, 1] > 0.2) & (np_boxes[:, 0] > -1)
np_boxes = np_boxes[expect_boxes, :]
def draw_results(results, img):
for result in results:
class_id, scores, x_min, y_min, x_max, y_max = result
print(class_id)
cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (255, 0, 255))
draw_results(np_boxes, img)
cv2.imwrite('save.jpg', img)
plt.imshow(img)
plt.show()
总结
项目总结
通过上述步骤，我们成功地训练了一个用于红绿灯检测的YOLOv8模型，并且将其部署到了一个由Streamlit构建的简单Web应用中。这不仅展示了深度学习在实际问题中的应用潜力，也为未来的开发提供了基础。
阿利同学的博客
计算机视觉、图像处理、毕业辅导、作业帮助、代码获取请私信
