ZeroOne AI
← 返回文章列表

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

👁 1
分类:AI视觉

零基础入门 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

其中:

五、项目目录

建议建立下面的目录:

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

其中:

六、准备疲劳驾驶数据集

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

其中:

二十、打开电脑摄像头

程序:

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 文件中,而是:

这种结构非常容易理解。

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

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 视觉检测系统。

评论(0