Python + YOLOv8 疲劳驾驶 AI 视觉检测入门:从模型训练到 ONNX 实时摄像头检测完整实战
零基础入门 Python + YOLOv8 疲劳驾驶 AI 视觉检测:通过 fatigueTrain.py 训练 YOLOv8n 模型、exportonnx.py 导出并校验 ONNX、fatigueRun.py 结合 ONNX Runtime 与 OpenCV 调用电脑摄像头实时检测三个程序,讲透数据集、YOLO 标签、训练参数、ONNX 导出、图像预处理、置信度、NMS、IoU 与坐标还原的完整实战流程。
一、前言
随着人工智能视觉技术的发展,摄像头已经不仅仅用于视频监控,还可以通过深度学习模型自动分析驾驶员状态。
例如:
- 驾驶员闭眼
- 驾驶员打瞌睡
- 驾驶员疲劳状态
- 驾驶员异常状态
- 驾驶员注意力下降
这些任务都可以通过目标检测模型进行识别。
对于刚开始学习人工智能视觉检测的开发者来说,一个完整的 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 模型的输入尺寸。
batch:batch_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.
把像素值从 0255 转换到 01。
然后:
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 文件中,而是:
- fatigueTrain.py:负责训练
- exportonnx.py:负责模型转换
- fatigueRun.py:负责部署推理
这种结构非常容易理解。
以后如果要做工业视觉项目,也可以采用类似结构:
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 视觉检测系统。
