2026-03-13 11:05:39 +08:00

4.0 KiB

Raw Blame History

鞋子检测模型训练指南

目录结构

train/
├── README.md                 # 本文件
├── 01_download_dataset.py    # 下载数据集脚本
├── 02_train.bat              # Windows 训练脚本
├── 03_export_onnx.bat        # 导出 ONNX 脚本
├── 04_convert_rknn.py        # 转换为 RKNN 脚本
├── data.yaml.template        # 数据集配置文件模板
└── samples/                  # 示例图片（用于测试）
    └── README.md

快速开始

1. 环境准备（Windows + GPU）

# 安装 PyTorch (CUDA 11.8)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 安装 ultralytics
pip install ultralytics

2. 下载数据集

手动下载（推荐）：

访问：https://universe.roboflow.com/nedrick-chandra-gpg1l/safety-shoes-detection-5qgkg
点击 "Download" → 选择 "YOLOv8" 格式
解压到 datasets/safety-shoes-detection/ 目录

或使用脚本（需要 API Key）：

python 01_download_dataset.py --api-key YOUR_API_KEY

3. 准备数据集配置

复制模板并修改路径：

cp data.yaml.template datasets/safety-shoes-detection/data.yaml
# 编辑 data.yaml，确保路径正确

4. 训练模型

一键训练：

02_train.bat

或手动训练：

# YOLOv8n - 轻量级，速度快
yolo detect train data=datasets/safety-shoes-detection/data.yaml model=yolov8n.pt epochs=150 imgsz=640 batch=16 device=0

# YOLOv8s - 精度更高（可选）
# yolo detect train data=datasets/safety-shoes-detection/data.yaml model=yolov8s.pt epochs=150 imgsz=640 batch=16 device=0

训练完成后，模型保存在：runs/detect/train/weights/best.pt

5. 导出 ONNX

03_export_onnx.bat

输出：runs/detect/train/weights/best.onnx

6. 转换为 RKNN

在 Ubuntu PC 上运行：

# 安装 RKNN Toolkit2
pip install rknn-toolkit2==2.2.0

# 转换（FP16 模式 - 推荐）
python 04_convert_rknn.py runs/detect/train/weights/best.onnx -o shoe_detector.rknn -t rk3588

# 或 INT8 量化（需要校准数据集）
# python 04_convert_rknn.py runs/detect/train/weights/best.onnx -o shoe_detector.rknn -t rk3588 -q -d dataset.txt

7. 部署到 RK3588

scp shoe_detector.rknn orangepi@<rk3588_ip>:/home/orangepi/apps/OrangePi3588Media/models/

然后在 configs/full_pipeline_1080p.json 中更新模型路径。

训练参数说明

参数	YOLOv8n	YOLOv8s	说明
模型大小	3.2MB	11MB	文件大小
推理速度	~30-40ms	~50-60ms	RK3588 NPU
mAP	~0.75	~0.82	精度
推荐场景	实时检测	高精度	选择建议

数据集说明

Safety Shoes Detection

来源: Roboflow Universe
类别: safety-shoes / no-safety-shoes
图片数: 约 1000+ 张
场景: 工地安全鞋检测

转换为单类检测

我们将两类合并为单一的 shoe 类别：

检测所有鞋子（安全鞋、运动鞋、布鞋等）
后续通过颜色分析判断是否为劳保鞋

常见问题

Q1: 训练时显存不足？

降低 batch size：

yolo detect train ... batch=8  # 默认 16，改为 8

Q2: 如何提高精度？

增加训练 epoch：epochs=200
使用更大模型：model=yolov8s.pt
增大输入尺寸：imgsz=768
收集更多现场图片 fine-tune

Q3: RKNN 转换失败？

确保使用正确的 opset (12)
使用 simplify=True 导出 ONNX
检查 RKNN Toolkit2 版本与板端驱动匹配

Q4: 检测不到鞋子？

降低置信度阈值：conf=0.15
检查 class_filter 是否正确设置
确认输入图像尺寸与模型匹配

4.0 KiB Raw Blame History Unescape Escape