Snooker Ball Detection - YOLOv8 on RK3588 NPU

基于 YOLOv8n 的斯诺克目标检测（11 类），训练后导出 ONNX 再量化为 INT8 RKNN，部署到 Rockchip RK3588 NPU。

目录结构

├── direct/              # 全帧直接训练方案
│   ├── direct_train.py
│   ├── direct_export_onnx.py
│   ├── direct_predict_video.py
│   ├── direct_gen_calib.py
│   ├── direct_verify_rknn.py
│   ├── direct_bench_rknn.py
│   └── output/           # 产物目录（.onnx / .rknn / .mp4 / .txt / .jpg）
│
├── pre_handler/         # ROI 预处理方案
│   ├── pre_handler_preprocess.py
│   ├── pre_handler_train.py
│   ├── pre_handler_export_onnx.py
│   ├── pre_handler_predict_video.py
│   ├── pre_handler_gen_calib.py
│   ├── pre_handler_verify_rknn.py
│   ├── pre_handler_bench_rknn.py
│   ├── output_dataset/   # 预处理后的数据集（由 pre_handler_preprocess.py 生成）
│   └── output/           # 产物目录（.onnx / .rknn / .mp4 / .txt / .jpg）
│
├── docker/               # Docker 环境与转换脚本
│   ├── Dockerfile
│   ├── build.sh
│   ├── direct_run.sh     # direct 方案转换 RKNN
│   └── pre_handler_run.sh # pre_handler 方案转换 RKNN
│
├── convert_to_rknn.py    # 核心转换脚本（Docker 容器内调用）
├── input_dataset/        # 原始 Roboflow 数据集（543 张，11 类，CC0）
├── test.mp4              # 测试视频
└── yolov8n.pt            # YOLOv8 nano 预训练权重

两种方案

方案	说明
direct	全帧直接训练，输入 1920x1080 原始图像，不做预处理
pre_handler	先 ROI 裁剪桌面区域（1500x800），过滤桌外标签，再用清洗后的数据训练

direct/ — 全帧直接方案

文件	用途
direct_train.py	训练 YOLOv8n，数据集 input_dataset/，输出到 runs/detect/snooker_score/full_frame_v1/
direct_export_onnx.py	导出 best.pt 为 ONNX，自动移至 output/best.onnx
direct_predict_video.py	PyTorch 模型预测 test.mp4，输出 output/predict_output.mp4
direct_gen_calib.py	从 input_dataset/train/images/ 生成校准集 output/dataset_for_rknn.txt
direct_verify_rknn.py	Orange Pi 单图推理验证，输出 output/rknn_test_output.jpg
direct_bench_rknn.py	全帧推理测速，输出 output/rknn_bench.txt

pre_handler/ — ROI 预处理方案

文件	用途
pre_handler_preprocess.py	从 input_dataset/ 裁剪 ROI + 过滤标签 → output_dataset/
pre_handler_train.py	训练 YOLOv8n，数据集 output_dataset/，输出到 runs/detect/snooker_score/roi_nano_v1/
pre_handler_export_onnx.py	导出 best.pt 为 ONNX，自动移至 output/best.onnx
pre_handler_predict_video.py	ROI 模型预测 test.mp4，输出 output/predict_output.mp4
pre_handler_gen_calib.py	从 output_dataset/train/images/ 生成校准集 output/dataset_for_rknn.txt
pre_handler_verify_rknn.py	Orange Pi 单图推理验证，输出 output/rknn_test_output.jpg
pre_handler_bench_rknn.py	全帧推理测速，输出 output/rknn_bench.txt

direct 方案使用流程

# 1. 训练
cd direct
python direct_train.py

# 2. 导出 ONNX（先修改 direct_export_onnx.py 中的 PT_PATH）
python direct_export_onnx.py

# 3. 构建镜像（仅首次）
bash ../docker/build.sh

# 4. 转换 RKNN
bash ../docker/direct_run.sh

pre_handler 方案使用流程

cd pre_handler

# 0. ROI 数据预处理（自动生成 data.yaml）
python pre_handler_preprocess.py

# 1. 训练
python pre_handler_train.py

# 2. 导出 ONNX（先修改 pre_handler_export_onnx.py 中的 PT_PATH）
python pre_handler_export_onnx.py

# 3. 构建镜像（仅首次）
bash ../docker/build.sh

# 4. 转换 RKNN
bash ../docker/pre_handler_run.sh

类别映射

0: black ball 1: blue ball 2: brown ball 3: green ball 4: pink ball 5: player1 6: player2 7: pocket 8: red ball 9: white ball 10: yellow ball