基于yolov8的红绿灯目标检测训练与Streamlit部署（代码+教程）

项目背景

随着智能交通系统的快速发展，自动驾驶技术逐渐成为研究的热点。在自动驾驶领域中，准确识别道路上的交通信号灯是确保车辆安全行驶的关键技术之一。近年来，深度学习技术的发展为交通信号灯的识别提供了强大的支持。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应用中。这不仅展示了深度学习在实际问题中的应用潜力，也为未来的开发提供了基础。

阿利同学的博客
计算机视觉、图像处理、毕业辅导、作业帮助、代码获取请私信