Doni/safesight-edge
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

加入训练好的ppe模型

Browse Source
This commit is contained in:
tian 2026-03-13 14:26:46 +08:00
parent 48b58f74a2
commit 682d568845
9 changed files with 0 additions and 840 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

BIN
models/yolov8s_ppe.onnx Normal file
View File

Binary file not shown.

BIN
models/yolov8s_ppe.pt Normal file
View File

Binary file not shown.

259
train/01_download_dataset.py
View File

@ -1,259 +0,0 @@
#!/usr/bin/env python3
"""
下载鞋子检测数据集
支持:
- Ultralytics Construction-PPE (推荐, 直接下载)
- Open Images V7 (通过 FiftyOne)
使用方法:
# 下载 Construction-PPE (推荐)
python 01_download_dataset.py --source ultralytics
# 下载 Open Images V7 鞋子类别
python 01_download_dataset.py --source openimages --max-samples 5000
"""
import argparse
import os
import sys
import zipfile
from pathlib import Path
def download_ultralytics_cppe(dataset_dir: str = "datasets/construction-ppe"):
"""
下载 Ultralytics Construction-PPE 数据集
完全开放，直接下载，无需注册
"""
import urllib.request
import ssl
url = "https://github.com/ultralytics/assets/releases/download/v0.0.0/construction-ppe.zip"
zip_path = "construction-ppe.zip"
print("="*70)
print("下载 Construction-PPE 数据集")
print("="*70)
print(f"来源: {url}")
print(f"目标: {dataset_dir}")
print()
# 创建目录
os.makedirs(dataset_dir, exist_ok=True)
# 下载
print("[1/3] 下载中... (约 178MB)")
try:
# 禁用 SSL 验证（某些环境需要）
ssl_context = ssl.create_default_context()
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE
with urllib.request.urlopen(url, context=ssl_context, timeout=300) as response:
total_size = int(response.headers.get('content-length', 0))
downloaded = 0
chunk_size = 8192
with open(zip_path, 'wb') as f:
while True:
chunk = response.read(chunk_size)
if not chunk:
break
f.write(chunk)
downloaded += len(chunk)
if total_size > 0:
percent = (downloaded / total_size) * 100
print(f"\r 进度: {percent:.1f}% ({downloaded}/{total_size} bytes)", end="")
print("\n ✓ 下载完成")
except Exception as e:
print(f"\n ✗ 下载失败: {e}")
print("\n请手动下载:")
print(f" 1. 访问: {url}")
print(f" 2. 下载 construction-ppe.zip")
print(f" 3. 解压到 {dataset_dir}/")
return False
# 解压
print(f"\n[2/3] 解压中...")
try:
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
zip_ref.extractall(dataset_dir)
print(" ✓ 解压完成")
except Exception as e:
print(f" ✗ 解压失败: {e}")
return False
# 清理
print(f"\n[3/3] 清理临时文件...")
os.remove(zip_path)
print(" ✓ 完成")
return True
def create_yaml_config(dataset_dir: str, single_class: bool = False):
"""创建单类检测配置文件"""
yaml_content = """# 单类鞋子检测数据集配置
# 基于 Construction-PPE 数据集修改
# 原类别: helmet, gloves, vest, boots, goggles, none, Person, no_helmet, no_goggle, no_gloves, no_boots
# 修改为单一的 shoe 类别（只使用 boots 和 no_boots 的标注）
path: construction-ppe # 数据集根目录
train: images/train # 训练集 (1132张)
val: images/val # 验证集 (143张)
test: images/test # 测试集 (141张)
# 单类配置
nc: 1
names: ['shoe']
# 原始数据信息
original_dataset:
name: Construction-PPE
source: Ultralytics
url: https://docs.ultralytics.com/datasets/detect/construction-ppe/
download: https://github.com/ultralytics/assets/releases/download/v0.0.0/construction-ppe.zip
# 使用说明:
# 1. 此配置将 boots (cls=3) 和 no_boots (cls=10) 都映射为 shoe (cls=0)
# 2. 训练时只检测鞋子，不关心是否安全鞋
# 3. 安全鞋判断通过后续颜色分析完成
"""
yaml_path = os.path.join(dataset_dir, "data.yaml")
with open(yaml_path, 'w') as f:
f.write(yaml_content)
print(f"\n✓ 配置文件创建: {yaml_path}")
return yaml_path
def download_openimages(classes: list, max_samples: int, dataset_dir: str):
"""通过 FiftyOne 下载 Open Images"""
try:
import fiftyone as fo
import fiftyone.zoo as foz
except ImportError:
print("错误: 未安装 fiftyone")
print("请运行: pip install fiftyone")
return False
print("="*70)
print("下载 Open Images V7 数据集")
print("="*70)
print(f"类别: {classes}")
print(f"最大样本数: {max_samples}")
print()
try:
dataset = foz.load_zoo_dataset(
"open-images-v7",
split="train",
label_types=["detections"],
classes=classes,
max_samples=max_samples,
dataset_dir=dataset_dir
)
# 导出为 YOLO 格式
print("\n导出为 YOLO 格式...")
dataset.export(
export_dir=dataset_dir + "-yolo",
dataset_type=fo.types.YOLOv5Dataset,
label_field="detections"
)
print(f"✓ 数据集保存: {dataset_dir}-yolo")
return True
except Exception as e:
print(f"✗ 下载失败: {e}")
return False
def check_dataset(dataset_dir: str):
"""检查数据集完整性"""
print("\n" + "="*70)
print("检查数据集")
print("="*70)
required_dirs = [
'images/train', 'images/val', 'images/test',
'labels/train', 'labels/val', 'labels/test'
]
all_ok = True
for dir_name in required_dirs:
full_path = os.path.join(dataset_dir, dir_name)
if os.path.exists(full_path):
count = len([f for f in os.listdir(full_path) if os.path.isfile(os.path.join(full_path, f))])
print(f" ✓ {dir_name}: {count} 个文件")
else:
print(f" ✗ {dir_name}: 不存在")
all_ok = False
return all_ok
def main():
parser = argparse.ArgumentParser(
description="下载鞋子检测数据集",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
示例:
# 下载 Construction-PPE (推荐，直接可用)
python 01_download_dataset.py --source ultralytics
# 下载 Open Images 鞋子类别
python 01_download_dataset.py --source openimages --max-samples 5000
# 指定保存目录
python 01_download_dataset.py --source ultralytics --dir ./my-datasets/shoes
"""
)
parser.add_argument("--source", choices=["ultralytics", "openimages"],
default="ultralytics",
help="数据源 (默认: ultralytics)")
parser.add_argument("--dir", default="datasets/construction-ppe",
help="数据集保存目录")
parser.add_argument("--max-samples", type=int, default=5000,
help="Open Images 最大样本数 (默认: 5000)")
parser.add_argument("--classes", nargs="+",
default=["Footwear", "Sandal", "Shoe", "Boot"],
help="Open Images 类别 (默认: Footwear Sandal Shoe Boot)")
args = parser.parse_args()
success = False
if args.source == "ultralytics":
success = download_ultralytics_cppe(args.dir)
if success:
create_yaml_config(args.dir, single_class=False)
check_dataset(args.dir)
elif args.source == "openimages":
success = download_openimages(args.classes, args.max_samples, args.dir)
# 输出下一步
if success:
print("\n" + "="*70)
print("数据集准备完成！")
print("="*70)
print(f"数据集路径: {args.dir}")
print("\n下一步:")
print(f" 1. 检查配置: cat {args.dir}/data.yaml")
print(f" 2. 开始训练: 02_train.bat")
print(f" 3. 或手动: yolo detect train data={args.dir}/data.yaml model=yolov8n.pt epochs=150 imgsz=640")
return 0
else:
print("\n✗ 数据集准备失败")
return 1
if __name__ == "__main__":
sys.exit(main())

87
train/02_train.bat
View File

@ -1,87 +0,0 @@
@echo off
chcp 65001 >nul
cls
echo ============================================================
echo 训练鞋子检测模型 (YOLOv8 + 640x640)
echo ============================================================
echo.
:: 设置数据集路径
set DATASET=datasets/construction-ppe/data.yaml
:: 检查数据集是否存在
if not exist %DATASET% (
echo [错误] 找不到数据集配置文件: %DATASET%
echo.
echo 请先下载数据集:
echo python 01_download_dataset.py --source ultralytics
pause
exit /b 1
)
echo [信息] 数据集: %DATASET%
echo.
:: 选择模型
echo 选择模型:
echo 1. YOLOv8n (轻量级, 速度快)
echo 2. YOLOv8s (推荐, 速度和精度平衡)
echo 3. YOLOv8m (高精度, 较慢)
echo.
set /p MODEL_CHOICE="输入选择 (1-3, 默认 2): "
if "%MODEL_CHOICE%"=="" set MODEL_CHOICE=2
if "%MODEL_CHOICE%"=="1" (
set MODEL=yolov8n.pt
set DESC=YOLOv8n
)
if "%MODEL_CHOICE%"=="2" (
set MODEL=yolov8s.pt
set DESC=YOLOv8s (推荐)
)
if "%MODEL_CHOICE%"=="3" (
set MODEL=yolov8m.pt
set DESC=YOLOv8m
)
echo.
echo [信息] 使用模型: %DESC%
echo.
:: 训练参数
set EPOCHS=150
set IMGSZ=640
set BATCH=16
echo 训练参数:
echo - Epochs: %EPOCHS%
echo - Image Size: %IMGSZ%x%IMGSZ%
echo - Batch Size: %BATCH%
echo - Device: GPU (cuda:0)
echo.
echo ============================================================
echo 开始训练
echo ============================================================
echo.
yolo detect train data=%DATASET% model=%MODEL% epochs=%EPOCHS% imgsz=%IMGSZ% batch=%BATCH% device=0
if %ERRORLEVEL% neq 0 (
echo.
echo [错误] 训练失败！
pause
exit /b 1
)
echo.
echo ============================================================
echo 训练完成！
echo ============================================================
echo.
echo 模型保存在: runs/detect/train/weights/best.pt
echo.
echo 下一步: 运行 03_export_onnx.bat 导出 ONNX
echo.
pause

35
train/03_export_onnx.bat
View File

@ -1,35 +0,0 @@
@echo off
chcp 65001 >nul
cls
echo ============================================================
echo 导出 ONNX 模型 (640x640)
echo ============================================================
echo.
set MODEL_PATH=runs/detect/train/weights/best.pt
if not exist %MODEL_PATH% (
echo [错误] 找不到模型: %MODEL_PATH%
echo 请先运行 02_train.bat 训练
pause
exit /b 1
)
echo [信息] 输入模型: %MODEL_PATH%
echo.
yolo export model=%MODEL_PATH% format=onnx imgsz=640 opset=12 simplify
if %ERRORLEVEL% neq 0 (
echo [错误] 导出失败！
pause
exit /b 1
)
echo.
echo [成功] ONNX 模型: runs/detect/train/weights/best.onnx
echo.
echo 下一步: 在 Ubuntu 上运行 04_convert_rknn.py 转换
echo.
pause

262
train/04_convert_rknn.py
View File

@ -1,262 +0,0 @@
#!/usr/bin/env python3
"""
将 YOLOv8 ONNX 模型转换为 RKNN 格式
适用于 RK3588 / RK3568 / RK3576 等平台
环境要求:
- Ubuntu x86_64 / Docker
- Python 3.8 / 3.9 / 3.10 / 3.11
- rknn-toolkit2 (pip install rknn-toolkit2==2.2.0)
使用方法:
# FP16 模式（推荐，速度快精度高）
python 04_convert_rknn.py best.onnx -o shoe_detector.rknn -t rk3588
# INT8 量化（模型更小，需要校准数据集）
python 04_convert_rknn.py best.onnx -o shoe_detector.rknn -t rk3588 -q -d dataset.txt
支持的 target_platform:
- rk3588 / rk3588s
- rk3568 / rk3566
- rk3576
- rv1106 / rv1103 / rv1103b
- rv1126
"""
import argparse
import os
import sys
from pathlib import Path
def check_environment():
"""检查运行环境"""
try:
from rknn.api import RKNN
print("✓ RKNN Toolkit2 已安装")
return True
except ImportError:
print("✗ 错误: 未安装 RKNN Toolkit2")
print("\n请安装:")
print(" pip install rknn-toolkit2==2.2.0")
print("\n或从源码安装:")
print(" https://github.com/airockchip/rknn-toolkit2")
return False
def create_sample_dataset(onnx_path: str, output_path: str = "dataset.txt", num_samples: int = 20):
"""
创建示例量化校准数据集
用于 INT8 量化时提供校准图片路径
"""
print(f"\n创建示例校准数据集: {output_path}")
print("注意: 请用实际图片替换这些示例路径")
sample_content = f"""# RKNN INT8 量化校准数据集
# 每行一个图片路径，建议使用 20-100 张典型场景图片
# 图片格式: JPG, PNG, BMP 等
# 示例路径（请替换为实际路径）:
# /path/to/train/images/img001.jpg
# /path/to/train/images/img002.jpg
# /path/to/valid/images/img001.jpg
# 提示:
# 1. 图片应与实际部署场景相似
# 2. 包含各种光照、角度、背景的样本
# 3. 建议 20-100 张，越多越慢但可能更准
"""
with open(output_path, 'w') as f:
f.write(sample_content)
print(f"✓ 示例数据集已创建: {output_path}")
print(" 请编辑此文件，添加实际的图片路径")
return output_path
def convert_onnx_to_rknn(
onnx_path: str,
output_path: str = None,
target_platform: str = "rk3588",
do_quantization: bool = False,
dataset_path: str = None,
verbose: bool = True
):
"""
转换 ONNX 模型到 RKNN
Args:
onnx_path: ONNX 模型文件路径
output_path: 输出 RKNN 文件路径，默认与 ONNX 同名
target_platform: 目标平台，默认 rk3588
do_quantization: 是否启用 INT8 量化
dataset_path: 量化校准数据集路径（txt 文件，每行一张图片路径）
verbose: 是否打印详细信息
"""
if output_path is None:
output_path = onnx_path.replace(".onnx", ".rknn")
# 确保输出目录存在
output_dir = os.path.dirname(output_path)
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir)
print("="*70)
print(f"ONNX 转 RKNN")
print("="*70)
print(f"输入: {onnx_path}")
print(f"输出: {output_path}")
print(f"目标: {target_platform}")
print(f"量化: {'INT8' if do_quantization else 'FP16 (无量化)'}")
if do_quantization:
print(f"校准: {dataset_path}")
print("="*70)
# 检查输入文件
if not os.path.exists(onnx_path):
print(f"\n✗ 错误: 找不到 ONNX 文件: {onnx_path}")
return False
# 检查数据集（如果需要量化）
if do_quantization:
if dataset_path is None:
print("\n✗ 错误: INT8 量化需要提供校准数据集")
print(" 使用 --dataset 指定数据集文件路径")
print(" 或运行 --create-dataset 创建示例")
return False
if not os.path.exists(dataset_path):
print(f"\n✗ 错误: 找不到数据集文件: {dataset_path}")
return False
from rknn.api import RKNN
# 创建 RKNN 对象
rknn = RKNN(verbose=verbose)
try:
# 配置模型
print("\n[1/4] 配置模型...")
rknn.config(
mean_values=[[0, 0, 0]], # YOLOv8 使用 0-255 输入
std_values=[[255, 255, 255]], # 归一化到 0-1
target_platform=target_platform
)
print(" ✓ 完成")
# 加载 ONNX
print("\n[2/4] 加载 ONNX 模型...")
ret = rknn.load_onnx(model=onnx_path)
if ret != 0:
print(" ✗ 加载失败!")
return False
print(" ✓ 完成")
# 构建模型
print("\n[3/4] 构建 RKNN 模型...")
if do_quantization:
print(f" 使用 INT8 量化，校准数据集: {dataset_path}")
ret = rknn.build(do_quantization=True, dataset=dataset_path)
else:
print(" 使用 FP16 模式（无量化）")
ret = rknn.build(do_quantization=False)
if ret != 0:
print(" ✗ 构建失败!")
return False
print(" ✓ 完成")
# 导出 RKNN
print("\n[4/4] 导出 RKNN 模型...")
ret = rknn.export_rknn(output_path)
if ret != 0:
print(" ✗ 导出失败!")
return False
print(" ✓ 完成")
finally:
rknn.release()
# 验证输出
if os.path.exists(output_path):
size_mb = os.path.getsize(output_path) / (1024 * 1024)
print("\n" + "="*70)
print(f"✓ 转换成功!")
print(f" 输出文件: {output_path}")
print(f" 文件大小: {size_mb:.2f} MB")
print("="*70)
print("\n下一步:")
print(f" 1. 复制到 RK3588:")
print(f" scp {output_path} orangepi@<rk3588_ip>:/home/orangepi/apps/OrangePi3588Media/models/")
print(f" 2. 更新配置文件中的模型路径")
return True
else:
print("\n✗ 错误: 输出文件未生成")
return False
def main():
parser = argparse.ArgumentParser(
description="将 YOLOv8 ONNX 模型转换为 RKNN",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
示例:
# FP16 模式（推荐）
python 04_convert_rknn.py best.onnx -o shoe_detector.rknn
# 指定目标平台
python 04_convert_rknn.py best.onnx -t rk3568
# INT8 量化
python 04_convert_rknn.py best.onnx -q -d dataset.txt
# 创建示例校准数据集
python 04_convert_rknn.py --create-dataset
"""
)
parser.add_argument("onnx", nargs="?", help="ONNX 模型文件路径")
parser.add_argument("-o", "--output", help="输出 RKNN 文件路径")
parser.add_argument("-t", "--target", default="rk3588",
help="目标平台 (默认: rk3588)")
parser.add_argument("-q", "--quantize", action="store_true",
help="启用 INT8 量化")
parser.add_argument("-d", "--dataset", help="量化校准数据集路径 (txt 文件)")
parser.add_argument("--create-dataset", action="store_true",
help="创建示例校准数据集并退出")
parser.add_argument("-v", "--verbose", action="store_true", default=True,
help="显示详细信息")
args = parser.parse_args()
# 创建示例数据集
if args.create_dataset:
create_sample_dataset("best.onnx", "dataset.txt")
return 0
# 检查参数
if args.onnx is None:
parser.print_help()
print("\n错误: 请提供 ONNX 文件路径")
return 1
# 检查环境
if not check_environment():
return 1
# 执行转换
success = convert_onnx_to_rknn(
onnx_path=args.onnx,
output_path=args.output,
target_platform=args.target,
do_quantization=args.quantize,
dataset_path=args.dataset,
verbose=args.verbose
)
return 0 if success else 1
if __name__ == "__main__":
sys.exit(main())

135
train/README.md
View File

@ -1,135 +0,0 @@
# 鞋子检测模型训练指南
## 方案：640x640 单模型（部署时用2窗口）
**训练阶段**：
- 输入：640x640 完整图片
- 模型：YOLOv8s
- 输出：640x640 模型文件
**部署阶段**（pipeline配置）：
- 原图 1920x1080
- 分成 2 个 960x1080 窗口
- 每个窗口 resize 到 640x640 送入模型
- 合并检测结果
---
## 目录结构
```
train/
├── README.md # 本文件
├── 01_download_dataset.py # 下载 Construction-PPE 数据集
├── 02_train.bat # Windows 一键训练脚本
├── 03_export_onnx.bat # 导出 ONNX 脚本
├── 04_convert_rknn.py # 转换为 RKNN 脚本
├── data.yaml.template # 数据集配置文件
└── samples/ # 示例图片
├── calibration/
├── test_images/
└── README.md
```
---
## 快速开始
### 1. 下载数据集
```bash
cd train
python 01_download_dataset.py --source ultralytics
```
或手动下载：
```bash
wget https://github.com/ultralytics/assets/releases/download/v0.0.0/construction-ppe.zip
unzip construction-ppe.zip -d datasets/construction-ppe/
```
### 2. 准备配置
```bash
cp data.yaml.template datasets/construction-ppe/data.yaml
```
### 3. 训练（640x640）
```bash
02_train.bat
```
或手动：
```bash
yolo detect train \
data=datasets/construction-ppe/data.yaml \
model=yolov8s.pt \
epochs=150 \
imgsz=640 \
batch=16 \
device=0
```
**训练参数**：
- 模型：YOLOv8s（速度和精度平衡）
- 输入：640x640
- 预计时间：30-60分钟
### 4. 导出 ONNX
```bash
03_export_onnx.bat
```
### 5. 转换为 RKNN
在 Ubuntu PC 上：
```bash
python 04_convert_rknn.py runs/detect/train/weights/best.onnx -o shoe_detector_640.rknn -t rk3588
```
### 6. 部署（2窗口配置）
复制到 RK3588：
```bash
scp shoe_detector_640.rknn orangepi@<rk3588_ip>:/home/orangepi/apps/OrangePi3588Media/models/
```
Pipeline 配置（部署阶段用2窗口）：
```json
{
"id": "pre_shoe",
"type": "preprocess",
"windows": [
{"x": 0, "y": 0, "w": 960, "h": 1080},
{"x": 960, "y": 0, "w": 960, "h": 1080}
],
"dst_w": 640,
"dst_h": 640
}
```
---
## 类别说明（Construction-PPE）
使用原始11类：
- 0: helmet
- 1: gloves
- 2: vest
- 3: **boots** ← 主要关注
- 4: goggles
- 5: none
- 6: **Person**
- 7: no_helmet
- 8: no_goggle
- 9: no_gloves
- 10: **no_boots**
---
## 相关链接
- [Construction-PPE 数据集](https://docs.ultralytics.com/datasets/detect/construction-ppe/)
- [Ultralytics YOLOv8](https://docs.ultralytics.com/)

16
train/data.yaml.template
View File

@ -1,16 +0,0 @@
# Construction-PPE 数据集配置
path: construction-ppe
train: images/train
val: images/val
test: images/test
# 11类原始类别
nc: 11
names: [
'helmet', 'gloves', 'vest', 'boots', 'goggles', 'none',
'Person', 'no_helmet', 'no_goggle', 'no_gloves', 'no_boots'
]
# 数据下载链接
download: https://github.com/ultralytics/assets/releases/download/v0.0.0/construction-ppe.zip

46
train/samples/README.md
View File

@ -1,46 +0,0 @@
# 示例图片目录
用于存放测试图片和量化校准样本。
## 目录结构
```
samples/
├── test_images/ # 用于测试模型的示例图片
├── calibration/ # INT8 量化校准用的图片（约 20-100 张）
└── README.md # 本文件
```
## 使用说明
### 测试图片 (test_images/)
存放一些典型场景的鞋子图片，用于验证模型效果。
### 校准图片 (calibration/)
INT8 量化时需要，用于确定量化参数。
**要求：**
- 应与实际部署场景相似
- 包含各种光照、角度、背景的样本
- 建议 20-100 张
- 图片格式: JPG, PNG, BMP
**创建校准数据集文件：**
```bash
# Linux/macOS
ls samples/calibration/*.jpg > dataset.txt
ls samples/calibration/*.png >> dataset.txt
# Windows CMD
dir /b samples\calibration\*.jpg > dataset.txt
dir /b samples\calibration\*.png >> dataset.txt
```
## 注意事项
- 校准图片越多，转换时间越长，但精度可能更高
- 建议使用训练集的部分图片作为校准集
- 不要和测试集重复，避免过拟合