safesight-edge/train

鞋子检测模型训练指南

目录结构

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. 下载数据集

手动下载（推荐）：

1. 访问：https://universe.roboflow.com/nedrick-chandra-gpg1l/safety-shoes-detection-5qgkg
2. 点击 "Download" → 选择 "YOLOv8" 格式
3. 解压到 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: 如何提高精度？

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

Q3: RKNN 转换失败？

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

Q4: 检测不到鞋子？

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

---

相关链接

• Ultralytics YOLOv8 文档
• RKNN Toolkit2 文档
• Roboflow Universe - Safety Shoes