# ball_detection

**Repository Path**: oftenlin/ball_detection

## Basic Information

- **Project Name**: ball_detection
- **Description**: ball_detection
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-01-20
- **Last Updated**: 2026-01-20

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 台球检测模型 - YOLOv5

本项目使用YOLOv5实现台球检测功能，并支持模型转换至地平线RDK X5平台。

## 环境准备

确保系统已安装Python环境和Git。

## 1. 获取YOLOv5

```bash
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
git checkout v7.0

# 下载预训练权重
wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt
```

## 2. 数据准备

### 2.1 准备训练数据

```bash
cd data
unzip dataset.zip
python3 ./scripts/convert_annotations.py
```

转换脚本将生成`dataset`目录，其中包含了按YOLOv5格式组织的数据。

```bash
# 将数据集移动到YOLOv5目录
mv dataset yolov5/
```

## 3. 模型训练

```bash
cd yolov5
python train.py --data dataset/billiards.yaml --epochs 200 --weights './yolov5s.pt' --cfg yolov5s.yaml --batch-size 32
```

训练完成后，将最佳权重保存到weights目录：

```bash
mv runs/train/exp/weights/best.pt ./weights/
```

## 4. 模型导出

### 4.1 导出X86平台ONNX模型

```bash
python3 export.py --weights ./weights/best.pt --include onnx --opset 11 --img 640 --device 0
mv ./weights/best.onnx ../weights/billiards_detection_x64.onnx
```

批量预测

```bash
rm -rf ./webapp/data/*
python3 ./scripts/save_predictions_number.py \
  --data ./test/images \
  --weights ./weights/billiards_detection_x64_yolov5m_e500.onnx \
  --output ./webapp/data/predictions_number.csv \
  --device cpu \
  --img-size 640 \
  --conf-thres 0.25 \
  --iou-thres 0.45 \
  --dedup-iou 0.5
```
```bash
python ./scripts/analyze_predictions_number.py \
  --pred-csv ./webapp/data/predictions_number.csv \
  --images-dir ./test/images \
  --labels-dir ./test/labels \
  --output-dir ./webapp/data \
  --iou 0.5
```

```bash
mkdir -p ./webapp/data/images/
cp ./test/train/images/* ./webapp/data/images/
cd ./webapp/
python3 run.py

```

### 4.2 导出RDK X5平台ONNX模型

需要修改YOLOv5的检测层以适配RDK X5：

```bash
# 备份原始模型定义
cp models/yolo.py models/yolo_trainning.py 
cp models/yolo.py models/yolo_quat_rdk.py

# 修改yolo_quat_rdk.py中的Detect forward方法
# 将以下代码添加到Detect类中:
# def forward(self, x):
#     return [self.m[i](x[i]).permute(0,2,3,1).contiguous() for i in range(self.nl)]

# 替换原始文件
rm models/yolo.py
cp models/yolo_quat_rdk.py models/yolo.py 

# 导出RDK X5模型
python3 export.py --weights ./weights/best.pt --include onnx --opset 11 --img 640 --device 0
mv ./weights/best.onnx ../weights/billiards_detection_rdk_x5.onnx
```

## 5. X64平台模型推理测试

```bash
cd ..
rm -rf output/*
python3 scripts/batch_inference.py --weights ./weights/billiards_detection_x64.onnx --output-dir ./output/ --data ./data/test_data/ --output ./output/results.csv
```

## 6. RDK X5平台模型量化

### 6.1 准备量化所需文件

```bash
rm -rf ~/horizon/data
mkdir -p ~/horizon/data

cp -r /test/images ~/horizon/data/billiards_detect_images
cp configs/billiards.yaml ~/horizon/data/
cp weights/billiards_detection_rdk_x5.onnx ~/horizon/data/
```

### 6.2 启动地平线Docker环境

```bash
docker run -it --rm --gpus all --shm-size=15g \
  -v /home/oftenlin/horizon/horizon_x5_open_explorer_v1.2.8:/open_explorer \
  -v /home/oftenlin/horizon/data:/data/horizon_x5/data \
  openexplorer/ai_toolchain_ubuntu_20_x5_gpu:v1.2.8-py310
```

### 6.3 在Docker中进行模型量化

```bash
# 进入数据目录
cd /data/horizon_x5/data/

# 检查模型是否支持量化
hb_mapper checker --model-type onnx --march bayes-e --model billiards_detection_rdk_x5.onnx
```

如果出现error信息，表示模型支持量化。

### 6.4 准备校准数据

```bash
# 找到预处理脚本
find / -name "02_pr*" -print | grep yolo

# 复制测试图片到指定位置
cp -r billiards_detect_images /open_explorer/samples/ai_toolchain/horizon_model_convert_sample/04_detection/08_yolov4/mapper/

# 进入预处理目录
cd /open_explorer/samples/ai_toolchain/horizon_model_convert_sample/04_detection/08_yolov4/mapper/

# 修改预处理脚本参数
# 编辑preprocess.py，将以下参数修改为：
# model_input_height = 640
# model_input_width = 640

# 执行预处理
./02_preprocess.sh

# 移动校准数据
mv calibration_billiards_data_rgb_f32/ /data/horizon_x5/data/
```

### 6.5 执行模型量化

```bash
cd /data/horizon_x5/data/
hb_mapper makertbin --model-type onnx --config billiards.yaml
```

## 项目结构

```
.
├── data/                    # 数据目录
│   ├── dataset.zip          # 原始数据集
│   └── test_data/           # 测试数据
├── scripts/                 # 脚本目录
│   ├── convert_annotations.py  # 数据转换脚本
│   └── batch_inference.py   # 批量推理脚本
├── configs/                 # 配置文件目录
│   └── billiards.yaml       # 地平线量化配置
├── weights/                 # 权重文件目录
│   ├── billiards_detection_x64.onnx    # X86平台模型
│   └── billiards_detection_rdk_x5.onnx # RDK X5平台模型
└── output/                  # 输出目录
    └── results.csv          # 推理结果
```

## 注意事项

1. 训练前确保数据集已按照YOLOv5格式组织
2. 修改RDK X5导出模型时，需正确修改Detect类的forward方法
3. 量化前确保校准数据的预处理参数与模型输入尺寸一致

## 参考

- [YOLOv5 GitHub](https://github.com/ultralytics/yolov5)
- [地平线RDK X5开发文档](https://developer.horizon.ai/) 

python3 ./scripts/batch_inference_eval.py --data ./dataset_numbers/images/val/ --weights ./weights/billiards_number_yolov5s_600epoch_x64.onnx --gt-path ./dataset_numbers/labels/val/ --output-dir output_number --classifier-type number


python ./scripts/batch_inference_eval_modified.py --data ./dataset_numbers/images/val/ --weights ./weights/billiards_number_yolov5s_600epoch.onnx --gt-path ./dataset_numbers/labels/val/ --output-dir output_number_4 --classifier-type number --conf-thres 0.2 --max-labels 3


 python scripts/convert_category_annotations.py --json data_category/category_annotations.json --images data_category/images --output dataset_category --ball-size-ratio 0.02

  python scripts/convert_category_annotations.py --json data_category/category_annotations.json --images data_category/images --output output_sample --ball-size-ratio 0.02


python3 export.py --weights ../weights/billiards_category_yolov5s_200_epoch.pt --include onnx --opset 11 --img 640 --device 0


python3 run_experiments.py  --data dataset_numbers/billiards_numbers.yaml --project runs/experiments_number_2