Hugging Peers

YOLO_APPLE

☆ Save
by rwcjz615 general-recognition Pytorch
1 downloads 0 likes on ModelScope updated

Model card

YOLO_APPLE

面向无人机视角下类球状水果的识别、定位和自动采摘前端感知项目。项目使用 YOLO 进行水果检测，并将检测框中心点转换为相机坐标系下的三维坐标 (X, Y, Z)，便于后续机械臂或采摘执行机构使用。

功能

  • 读取当前仓库中的 YOLO 数据集：shujuji/shujuji/images 和 shujuji/shujuji/labels
  • 训练单类别水果检测模型
  • 对单张无人机图像执行水果识别
  • 输出标注图和三维坐标 CSV
  • 提供 Gradio 可视化演示
  • 提供 ModelScope 打包与上传脚本

环境

建议使用 Python 3.10、3.11 或 3.12。当前工程依赖 PyTorch/Ultralytics，Python 3.14 通常没有稳定的 PyTorch wheel。

python -m venv .venv
.venv\Scripts\activate
pip install -U pip
pip install -r requirements.txt
pip install -e .

数据集

数据集配置位于 configs/dataset.yaml：

path: ../shujuji/shujuji
train: images/train
val: images/val
test: images/test
names:
  0: fruit

当前数据集统计：

  • train: 400 images / 400 labels
  • val: 50 images / 50 labels
  • test: 50 images / 50 labels

训练模型

python scripts/train.py --data configs/dataset.yaml --model yolov8n.pt --epochs 100 --imgsz 640 --batch 16

训练完成后，最佳权重位于：

runs/train/fruit_yolo3d/weights/best.pt

测试评估

python scripts/evaluate.py --weights runs/train/fruit_yolo3d/weights/best.pt --split test

三维坐标推理

python scripts/predict_3d.py ^
  --weights runs/train/fruit_yolo3d/weights/best.pt ^
  --image shujuji/shujuji/images/test/test_0001_vitrine-027-02.jpg ^
  --camera configs/camera.yaml ^
  --output runs/predict_3d

输出：

  • runs/predict_3d/<image>_3d.jpg：带检测框和坐标的标注图
  • runs/predict_3d/<image>_3d.csv：每个水果目标的三维坐标表

三维坐标原理

检测得到水果中心像素 (u, v) 和深度 Z 后，使用针孔相机模型：

X = (u - cx) * Z / fx
Y = (v - cy) * Z / fy
Z = depth

其中 fx, fy, cx, cy 来自相机标定参数，配置在 configs/camera.yaml。

深度来源支持两种方式：

  • diameter：使用类球状水果已知直径和检测框像素直径估计深度
  • constant：使用固定工作距离

单目 RGB 图像不能直接恢复真实绝对深度；如果采摘系统有双目相机、RGB-D 相机、激光测距或无人机高度/姿态信息，建议把真实深度填入该流程以获得更可靠的三维坐标。

Web 演示

python app.py

默认读取 weights/best.pt。如果权重位于训练目录，可在页面中填写实际路径。

ModelScope 打包与上传

先准备上传目录：

python scripts/prepare_modelscope.py --weights runs/train/fruit_yolo3d/weights/best.pt

如果尚未训练，也可以先打包代码：

python scripts/prepare_modelscope.py

上传到目标仓库：

pip install modelscope
set MODELSCOPE_TOKEN=你的ModelScope访问令牌
python scripts/upload_modelscope.py --repo-id rwcjz615/YOLO_APPLE --local-path modelscope_package

也可以直接使用 ModelScope CLI：

modelscope upload rwcjz615/YOLO_APPLE modelscope_package --repo-type model --token 你的ModelScope访问令牌

上传需要 rwcjz615/YOLO_APPLE 仓库的写入权限。