Python + YOLOv8 疲劳驾驶 AI 视觉检测入门:从模型训练到 ONNX 实时摄像头检测完整实战

文章来源说明 :本文由 ZeroOne AI 整理发布,官网与 CSDN 双端同步首发,原文见 www.zeroone-ai.com

下载资源地址:https://download.csdn.net/download/guimotion/93398843?spm=1001.2014.3001.5503https://download.csdn.net/download/guimotion/93398843?spm=1001.2014.3001.5503

一、前言

随着人工智能视觉技术的发展,摄像头已经不仅仅用于视频监控,还可以通过深度学习模型自动分析驾驶员状态。

例如:

  • 驾驶员闭眼
  • 驾驶员打瞌睡
  • 驾驶员疲劳状态
  • 驾驶员异常状态
  • 驾驶员注意力下降

这些任务都可以通过目标检测模型进行识别。

对于刚开始学习人工智能视觉检测的开发者来说,一个完整的 AI 项目通常会涉及很多内容:

复制代码
数据集 → 标注 → 模型训练 → 模型验证 → 模型导出 → 推理 → 摄像头实时检测

如果一开始就研究复杂的网络结构,很容易把整个流程搞复杂。

因此本文采用一个更加适合入门学习的方案:

Python + YOLOv8 + ONNX Runtime + OpenCV

通过三个 Python 程序,把整个疲劳驾驶 AI 检测流程串起来:

复制代码
fatigueTrain.py
     ↓
训练 YOLOv8 模型
     ↓
fatiguebest.pt
     ↓
exportonnx.py
     ↓
fatiguebest.onnx
     ↓
fatigueRun.py
     ↓
OpenCV 摄像头
     ↓
实时疲劳驾驶检测

整个项目非常适合 AI 视觉、Python、YOLO 目标检测初学者作为第一个完整实战项目。

二、项目最终实现什么?

本文最终实现一个简单的疲劳驾驶 AI 视觉检测程序。

程序打开电脑摄像头:

复制代码
摄像头
   ↓
OpenCV读取视频
   ↓
图像预处理
   ↓
YOLOv8 ONNX模型
   ↓
目标检测
   ↓
NMS
   ↓
绘制检测框
   ↓
实时显示

例如:

复制代码
┌──────────────────────────────┐
│                              │
│        摄像头实时画面         │
│                              │
│   ┌────────────┐             │
│   │    人脸     │             │
│   │             │             │
│   └────────────┘             │
│                              │
│        疲劳状态检测           │
│                              │
└──────────────────────────────┘

需要特别说明:本文程序属于 AI 视觉检测入门案例,并不能直接等同于经过安全认证的汽车驾驶员监测系统。

真正用于车辆安全控制时,还需要进一步考虑:

  • 连续时间判断
  • 闭眼持续时间
  • 打哈欠持续时间
  • 驾驶员身份
  • 光照变化
  • 夜间环境
  • 遮挡
  • 眼镜
  • 红外摄像头
  • 多帧融合
  • 误报与漏报
  • 报警策略

本文主要解决的是:如何用 Python 从零跑通一个 YOLO 疲劳驾驶目标检测 AI 项目。

三、为什么选择 YOLOv8?

对于初学者来说,YOLO 是非常适合学习目标检测的算法框架。

它可以直接完成:

复制代码
图片
 ↓
YOLO
 ↓
目标位置 + 类别 + 置信度

例如训练一个疲劳驾驶数据集之后,模型可以学习:

复制代码
类别:
fatigue
normal

也可以根据自己的数据集设计:

复制代码
normal
closed_eye
yawn
sleep

具体类别并不是程序固定的,而是由数据集决定。

所以学习 YOLO 的关键不是记住某一个类别,而是理解整个:

复制代码
数据集 → 训练 → 模型 → 推理

流程。

四、项目环境

建议使用 Python 3.x 环境。

安装主要依赖:

复制代码
pip install ultralytics
pip install onnx
pip install onnxruntime
pip install opencv-python
pip install numpy
pip install matplotlib

其中:

  • Ultralytics:负责 YOLO 模型训练和导出
  • ONNX:用于读取和检查导出的 ONNX 模型
  • ONNX Runtime:负责运行 ONNX 模型
  • OpenCV:负责摄像头采集、图像处理、图像缩放、绘制检测框、显示检测结果

五、项目目录

建议建立下面的目录:

复制代码
fatigue_ai/
│
├── fatigueTrain.py
├── exportonnx.py
├── fatigueRun.py
│
├── fatigue/
│   ├── fatiguedata.yaml
│   ├── images/
│   │   ├── train/
│   │   └── val/
│   │
│   └── labels/
│       ├── train/
│       └── val/
│
└── fatiguebest.pt

其中:

  • fatigueTrain.py:负责训练
  • exportonnx.py:负责模型导出
  • fatigueRun.py:负责实时推理

六、准备疲劳驾驶数据集

YOLO 训练首先需要数据集。

数据集一般由图片 + 标签组成,例如:

复制代码
images/train/001.jpg
labels/train/001.txt

标签文件和图片名称对应,例如:

复制代码
001.jpg
001.txt

YOLO 标签通常采用:

复制代码
class x_center y_center width height

例如:

复制代码
0 0.512 0.432 0.245 0.381

这些数字都是经过归一化的坐标。

初学者不需要一开始研究 YOLO 内部网络结构。首先理解一个核心概念:

YOLO 训练的本质就是让模型学习"图片中什么位置出现了什么目标"。

七、编写 fatiguedata.yaml

训练脚本中指定了:

复制代码
data_config = 'fatigue/fatiguedata.yaml'

因此程序会从这里读取数据集配置。

例如:

复制代码
path: fatigue

train: images/train
val: images/val

names:
  0: normal
  1: fatigue

如果你的数据集只有一个类别,也可以:

复制代码
path: fatigue

train: images/train
val: images/val

names:
  0: fatigue

这里的类别名称必须和实际数据集保持一致。

八、第一个程序:fatigueTrain.py

本文提供的第一个程序就是模型训练程序。

核心代码:

复制代码
from ultralytics import YOLO

model = YOLO('yolov8n.pt')

data_config = 'fatigue/fatiguedata.yaml'

epochs = 20
imgsz = 640
batch_size = 16
workers = 8
device = 'cpu'

results = model.train(
    data=data_config,
    epochs=epochs,
    imgsz=imgsz,
    batch=batch_size,
    workers=workers,
    device=device,
    name='yolov8_fatigue_model',
    save=True,
    exist_ok=True,
    cache=False
)

print("训练完成:", results.save_dir)

这个程序的核心只有几个步骤。

九、加载 YOLOv8n 模型

程序首先:

复制代码
model = YOLO('yolov8n.pt')

这里使用 yolov8n,也就是 YOLOv8 Nano。

对于入门学习来说,选择轻量模型比较合适。

原因很简单:

复制代码
模型越小
   ↓
训练和推理压力越低
   ↓
越容易快速验证程序

如果后面需要提高检测精度,可以进一步尝试:

复制代码
YOLOv8n
YOLOv8s
YOLOv8m
YOLOv8l
YOLOv8x

初学者建议先把整个流程跑通。

十、设置训练参数

程序设置:

复制代码
epochs = 20
imgsz = 640
batch_size = 16
workers = 8
device = 'cpu'

分别表示:

epochs:训练轮数。这里设置 20,也就是说完整遍历训练数据 20 次。实际项目中可以根据数据量和训练效果调整。

imgsz:模型训练输入尺寸 640 × 640,这也是本文后续 ONNX 模型的输入尺寸。

batchbatch_size = 16 表示一次训练使用多少张图片。如果显存不足,可以降低:

复制代码
16 → 8 → 4 → 2

device :本文程序 device = 'cpu' 表示使用 CPU。如果电脑有可用 NVIDIA GPU,可以根据自己的环境配置 GPU。但是作为入门教程,CPU 也可以先把完整流程跑通。

十一、开始训练

真正执行训练的是:

复制代码
results = model.train(
    data=data_config,
    epochs=epochs,
    imgsz=imgsz,
    batch=batch_size,
    workers=workers,
    device=device,
    name='yolov8_fatigue_model',
    save=True,
    exist_ok=True,
    cache=False
)

训练完成后,程序输出:

复制代码
训练完成: ...

训练结果通常会保存到训练目录。最终我们需要找到训练出来的 best.pt,例如:

复制代码
runs/detect/yolov8_fatigue_model/weights/best.pt

然后复制或者重命名为 fatiguebest.pt,因为本文后面的 ONNX 导出程序使用:

复制代码
model = YOLO("fatiguebest.pt")

十二、为什么需要导出 ONNX?

训练出来的 fatiguebest.pt 是 PyTorch / Ultralytics 体系下的模型文件。

但是实际部署时,我们可能希望 Python、C++、C#、OpenCV、ONNX Runtime、边缘计算设备、工业电脑都能够使用。

因此可以把 fatiguebest.pt 转换成 fatiguebest.onnx。

ONNX 可以作为训练框架和部署框架之间的一个通用模型格式。

十三、第二个程序:exportonnx.py

本文第二个程序负责:YOLOv8 PT → ONNX

核心代码:

复制代码
from ultralytics import YOLO
import onnx

model = YOLO("fatiguebest.pt")

export_results = model.export(
    format="onnx",
    simplify=False,
    opset=18,
    imgsz=640,
    batch=1,
    device="CPU",
    verbose=False
)

这里有几个关键参数。

十四、format="onnx"

复制代码
format="onnx"

表示导出 ONNX。

最终目标:

复制代码
fatiguebest.pt
     ↓
fatiguebest.onnx

十五、为什么 simplify=False?

代码:

复制代码
simplify=False

这里特意关闭模型简化。

对于入门环境来说,可以减少对 onnxsim 额外依赖的要求。也就是说:先把模型正常导出来,再考虑模型优化。这是非常适合初学者的思路。

不要一开始就同时处理模型训练、模型量化、模型剪枝、模型简化、TensorRT、OpenVINO、CUDA、NPU,否则非常容易把问题搞复杂。

十六、opset=18

程序:

复制代码
opset=18

用于指定 ONNX 算子集版本。本文程序选择 ONNX Opset 18,并以 ONNX Runtime、TensorRT、OpenCV 等常见部署方向作为兼容性考虑。

十七、输入尺寸保持 640

程序:

复制代码
imgsz=640

也就是说训练和导出统一使用 640 × 640。后面的实时推理程序同样设置:

复制代码
height, width = 640, 640

这样整个流程更加容易理解:

复制代码
训练:640 × 640
导出:640 × 640
推理:640 × 640

对于第一个 YOLO 项目而言,这种统一配置非常重要。

十八、验证 ONNX 模型

导出之后程序还会:

复制代码
onnx_model = onnx.load("fatiguebest.onnx")

onnx.checker.check_model(onnx_model)

这一步非常有用。它不是简单地判断文件有没有生成,而是检查 ONNX 模型结构是否合法。

程序还会输出:

复制代码
print(f"模型输入形状:...")
print(f"模型输出形状:...")

因此第二个程序实际上完成了:

复制代码
加载 PT
  ↓
导出 ONNX
  ↓
读取 ONNX
  ↓
检查模型
  ↓
输出输入/输出结构

本文提供的导出脚本正是按照这个流程实现的。

十九、第三个程序:fatigueRun.py

模型训练和导出完成之后,终于进入实时检测。

本文第三个程序使用:

复制代码
import onnxruntime as ort
import numpy as np
import cv2

其中:

  • ONNX Runtime:负责模型推理
  • OpenCV:负责摄像头

二十、打开电脑摄像头

程序:

复制代码
cap = cv2.VideoCapture(0)

这里 0 代表默认摄像头。如果电脑存在多个摄像头,可以尝试:

复制代码
cv2.VideoCapture(1)

或者:

复制代码
cv2.VideoCapture(2)

程序首先判断:

复制代码
if not cap.isOpened():
    print("无法打开摄像头")
    exit()

这样可以避免摄像头没有打开时程序继续运行。

二十一、加载 ONNX 模型

核心代码:

复制代码
session = ort.InferenceSession(
    "fatiguebest.onnx",
    providers=["CPUExecutionProvider"]
)

这里指定 CPUExecutionProvider,因此当前版本使用 CPU 执行 ONNX 推理。这对于入门测试非常方便。

整个推理过程变成:

复制代码
摄像头
  ↓
OpenCV
  ↓
numpy
  ↓
ONNX Runtime
  ↓
YOLO模型
  ↓
检测结果

二十二、读取 ONNX 输入输出节点

程序:

复制代码
input_name = session.get_inputs()[0].name
label_name = session.get_outputs()[0].name

这里没有把输入节点名称硬编码,而是直接从 ONNX 模型读取。这种写法对于学习 ONNX Runtime 非常重要。

可以理解为:

复制代码
ONNX模型
  ↓
查询输入节点
  ↓
查询输出节点
  ↓
执行推理

二十三、实时读取视频

程序进入循环:

复制代码
while cap.isOpened():

    ret, frame = cap.read()

    if not ret:
        break

每次 cap.read() 获取一张视频帧。

所以摄像头视频实际上就是:

复制代码
第1帧
 ↓
第2帧
 ↓
第3帧
 ↓
第4帧
 ↓
......

YOLO 对每一帧进行检测。

二十四、图像预处理

程序:

复制代码
img = frame / 255.

把像素值从 0~255 转换到 0~1。

然后:

复制代码
img = cv2.resize(img, (width, height))

调整到 640 × 640。

接下来:

复制代码
img = np.transpose(img, (2, 0, 1))

将 OpenCV 常见的 H × W × C 转换为 C × H × W。

然后:

复制代码
data = np.expand_dims(img, axis=0)

增加 Batch 维度。

最终输入变成:

复制代码
1 × 3 × 640 × 640

这就是 YOLO 模型常见的输入形式。

二十五、执行 ONNX 推理

核心代码:

复制代码
pred = session.run(
    [label_name],
    {input_name: data.astype(np.float32)}
)[0]

这一步就是真正调用 AI 模型。

输入:640 × 640 RGB/BGR 图像数据

输出:YOLO 预测结果

之后程序:

复制代码
pred = np.squeeze(pred)
pred = np.transpose(pred, (1, 0))

对输出数据进行整理。

二十六、什么是置信度?

程序:

复制代码
pred_class = pred[..., 4:]
pred_conf = np.max(pred_class, axis=-1)

可以理解为:模型针对一个候选目标给出多个类别预测。

例如:

复制代码
normal   0.12
fatigue  0.86

那么最大类别置信度 = 0.86。

程序后面设置:

复制代码
nms(pred, 0.3, 0.45)

其中 0.3 是置信度阈值,也就是:置信度太低的目标直接过滤。

二十七、什么是 NMS?

YOLO 检测一个目标时,有可能产生多个重叠框。例如:

复制代码
┌───────────────┐
│     人脸      │
└───────────────┘

┌──────────────┐
│     人脸     │
└──────────────┘

其实两个框可能都是在检测同一个人脸。

所以需要 Non-Maximum Suppression,也就是非极大值抑制 NMS。

本文 fatigueRun.py 中自己实现了 NMS。

核心逻辑:

复制代码
iou = getIou(
    max_conf_box,
    current_box,
    interArea
)

if iou > iou_thres:
    del_index.append(j)

如果两个框的 IoU 太高,就删除重复框。

当前程序设置:

复制代码
conf_thres = 0.3
iou_thres = 0.45

NMS 的完整处理逻辑就在提供的 fatigueRun.py 中。

二十八、什么是 IoU?

IoU:Intersection over Union,即交并比。

简单理解:两个框重叠程度。

计算:

复制代码
IoU = 交集面积 / 并集面积

例如 IoU = 0 说明两个框没有重叠,而 IoU ≈ 1 说明两个框高度重合。

所以 NMS 可以利用 IoU 判断:两个框是不是在重复检测同一个目标。

二十九、把检测结果画回原图

检测模型使用 640 × 640,但是摄像头实际分辨率可能不是 640 × 640。

因此程序计算:

复制代码
x_scale = frame.shape[1] / width
y_scale = frame.shape[0] / height

用于将模型坐标转换回摄像头画面的坐标。

然后:

复制代码
cv2.rectangle(...)

绘制检测框。

对应代码使用了预测框的中心点和宽高,将其转换为左上角、右下角坐标。

三十、完整实时检测流程

现在整个 fatigueRun.py 就可以理解成:

复制代码
摄像头
 │
 ▼
OpenCV读取帧
 │
 ▼
Resize 640×640
 │
 ▼
/255 图像归一化
 │
 ▼
HWC → CHW
 │
 ▼
增加 Batch 维度
 │
 ▼
ONNX Runtime
 │
 ▼
YOLOv8 ONNX
 │
 ▼
获取预测结果
 │
 ▼
置信度筛选
 │
 ▼
NMS
 │
 ▼
坐标缩放
 │
 ▼
OpenCV画框
 │
 ▼
实时显示

这就是一个最基础的 YOLO AI 视觉推理系统。

三十一、运行整个项目

按照顺序运行。

第一步:训练

复制代码
python fatigueTrain.py

训练完成后获得 best.pt,重命名为 fatiguebest.pt

第二步:导出 ONNX

复制代码
python exportonnx.py

最终得到 fatiguebest.onnx,同时程序会检查 ONNX 模型是否合法。

第三步:启动实时检测

复制代码
python fatigueRun.py

程序打开电脑摄像头,然后进行实时疲劳驾驶目标检测。

退出按 q,程序中的退出逻辑也是通过 cv2.waitKey(1) 监听 q 键实现的。

三十二、为什么把项目拆成三个 Python 程序?

这是本文特别适合初学者的地方。

没有把训练、导出、推理全部写到一个 Python 文件中,而是:

这种结构非常容易理解。

以后如果要做工业视觉项目,也可以采用类似结构:

复制代码
train.py
export.py
run.py

甚至进一步拆分:

复制代码
dataset.py
train.py
validate.py
export.py
inference.py
camera.py

项目规模扩大后再逐步工程化。

三十三、初学者最容易遇到的几个问题

1. 摄像头打不开

检查 cv2.VideoCapture(0),可以尝试 0、1、2,不同电脑摄像头编号可能不同。

2. 找不到 fatiguebest.pt

检查 fatiguebest.pt 是否和 exportonnx.py 处于正确目录。因为代码直接 YOLO("fatiguebest.pt") 读取当前路径下的模型。

3. 找不到 fatiguebest.onnx

说明 exportonnx.py 没有成功导出。先检查 fatiguebest.pt 是否正常,然后重新运行:

复制代码
python exportonnx.py

4. ONNX Runtime 安装问题

可以安装:

复制代码
pip install onnxruntime

然后测试:

复制代码
import onnxruntime

5. CPU 运行比较慢

当前程序明确使用 providers=["CPUExecutionProvider"],因此是 CPU 推理。对于学习和功能验证已经足够。

如果后续追求速度,可以进一步研究 CUDA、TensorRT、OpenVINO、Intel NPU、DirectML,但不建议在第一个项目里一次全部加入。

三十四、从"能运行"到"真正的疲劳驾驶检测"

当前程序是单帧检测,也就是说:当前这一帧有没有疲劳目标。

但真正的疲劳驾驶判断通常不能只看一帧。

例如闭眼 1 帧并不代表驾驶员疲劳,可能只是正常眨眼。

所以更合理的系统应该增加时间维度。例如:

复制代码
连续检测
   ↓
闭眼
   ↓
持续 1 秒
   ↓
持续 2 秒
   ↓
持续 3 秒
   ↓
触发疲劳报警

这样才能降低误报警。

三十五、进一步升级成"疲劳驾驶 AI 系统"

基础 YOLO 项目跑通之后,可以继续升级。

第一阶段:YOLO 目标检测

复制代码
摄像头
  ↓
YOLO
  ↓
疲劳目标

第二阶段:增加连续帧判断

复制代码
YOLO
  ↓
连续N帧
  ↓
状态统计
  ↓
疲劳判断

第三阶段:增加眼睛状态

例如 Open Eye、Closed Eye,通过连续闭眼时间判断疲劳。

第四阶段:增加打哈欠

例如 Normal、Yawn,结合闭眼时间 + 打哈欠次数综合判断。

第五阶段:增加报警

最终:

复制代码
摄像头
  ↓
AI视觉
  ↓
驾驶员状态
  ↓
疲劳评分
  ↓
报警

甚至可以连接蜂鸣器、PLC、MES、车辆控制系统、云平台。

三十六、为什么先使用 ONNX,而不是直接部署复杂 AI 框架?

对于初学者而言,YOLO 训练和模型部署其实是两个不同的问题。

训练阶段:

复制代码
Ultralytics
PyTorch
GPU
数据集

部署阶段:

复制代码
ONNX
ONNX Runtime
OpenCV
CPU/GPU/NPU

将二者分开以后,整个项目结构会非常清晰。也就是:

复制代码
训练框架
  ↓
PT
  ↓
模型转换
  ↓
ONNX
  ↓
部署框架

这也是本文三个 Python 程序的核心设计思路。

三十七、入门学习建议

如果是第一次学习 YOLO,不建议一上来就研究网络结构、Loss、Backbone、Neck、Attention、Transformer、量化、剪枝、TensorRT、NPU。

建议按照下面的路线学习。

第一步:先理解图片 → 标签 → 数据集

第二步:跑通 fatigueTrain.py

第三步:找到 best.pt

第四步:跑通 exportonnx.py

第五步:得到 fatiguebest.onnx

第六步:运行 fatigueRun.py

第七步:理解预处理 → ONNX 推理 → 置信度 → NMS → 坐标转换 → 画框

到这里,就已经完成了一个完整的 AI 视觉项目入门。

三十八、总结

本文通过三个 Python 程序完成了一个完整的疲劳驾驶 AI 视觉检测入门项目。

fatigueTrain.py 负责:

复制代码
数据集
  ↓
YOLOv8训练
  ↓
best.pt

提供的训练程序默认加载 yolov8n.pt,使用 fatigue/fatiguedata.yaml 数据集配置,并设置了 20 个 epoch、640 输入尺寸、batch 16 等参数。

然后 exportonnx.py 完成:

复制代码
best.pt
  ↓
ONNX
  ↓
模型检查

最后 fatigueRun.py 完成:

复制代码
摄像头
  ↓
OpenCV
  ↓
图像预处理
  ↓
ONNX Runtime
  ↓
YOLOv8
  ↓
NMS
  ↓
检测框
  ↓
实时显示

这样就形成了一个非常适合初学者理解的 AI 视觉完整闭环:训练 → 导出 → 推理。

对于第一次学习 YOLO 的开发者来说,与其直接研究复杂的 AI 理论,不如先把这三个程序真正运行起来。

当能够独立完成:

复制代码
训练自己的数据集
  ↓
生成自己的 best.pt
  ↓
导出自己的 ONNX
  ↓
调用摄像头实时检测

基本就已经迈过了 YOLO AI 视觉开发最重要的入门阶段。

后续再逐步加入多帧融合、疲劳时间判断、闭眼检测、打哈欠检测、疲劳评分、声音报警、边缘设备部署、GPU 加速、NPU 加速、工业相机,就可以从一个简单的 YOLO 入门案例,逐渐发展成完整的疲劳驾驶 AI 视觉检测系统。

相关推荐
聪明蛋子哟1 小时前
告别API“翻译”之苦:从OpenAPI到MCP,统一AI与工具集成的桥梁
人工智能
海上小飞龙1 小时前
大模型推理的两阶段:一次 Prefill,加上多次 Decode
人工智能·深度学习·语言模型
hhzz1 小时前
【OpenCV 入门到精通 01】认识 OpenCV 与计算机视觉:从零建立全局认知
人工智能·python·opencv·计算机视觉·开源
m4Rk_1 小时前
【论文阅读】Agent 记忆机制(62):DCM-Agent——用双簇记忆化解优化问题的多范式冲突
论文阅读·人工智能·学习·开源·github
xian_wwq2 小时前
【学习笔记】深度认知系列-第13讲AI Agent时代到来——从“回答问题”到“执行任务”
人工智能·笔记·学习
程序员cxuan2 小时前
GPT - 6 Astra 的使用焚诀
人工智能·后端·程序员
golang学习记2 小时前
Cursor Origin:Cursor要造一个AI时代的Github
人工智能·github·cursor
HugoStudio_SWAN2 小时前
洛谷 P10719 \[GESP202406 五级] 黑白格——暴力美学与图像处理的最小外接矩形
c++·图像处理·人工智能·学习·程序人生·算法·目标跟踪
一马平川的大草原2 小时前
如何实现SVG转Mermaid图
python·svg·格式转换·mermaid