yolov5_obb: 工业级旋转目标检测从数据制作到模型部署全流程

导读

yolov5_obb 是基于 yolov5 目标检测框架开发的一个项目，支持旋转目标检测任务（Oriented Bounding Boxes，OBB），旨在生成更好拟合具有角度位置的物体预测结果，广泛适用于机器人感知，遥感等领域。此前经过调研发现目前全网上关于此方面的资料相对较少，鱼龙混杂，不是比较老旧、乱七八糟，就是一言不合就付费查看，看个寂寞，实在是不忍直视。因此，本文旨在提供一个从数据集制作、划分、安装、训练、验证、部署保姆级教程，帮助大家从0到1快速完成项目的上手和开发，满足日常工作和学习的需求。

数据集制作

工具介绍

X-AnyLabeling[1]不仅是一个标注工具，它还是自动数据标注未来的一大飞跃。它的设计不仅简化了标注过程，还集成了尖端的AI模型，融入了当前最先进的视觉大模型和多模态技术，以获得更出色的体验。X-AnyLabeling 专注于实际应用，力求提供一个工业级的、功能丰富的工具，可帮助开发人员自动标注和处理各种复杂任务的数据。

在最新的 v1.0.0 版本中，X-AnyLabeling 已经完美支持了旋转框的标注，同时提供全方位的角度显示，实操体验方面也与经典的 roLabelImg 基本保持一致，极大降低用户的学习成本，且支持直接导出 DOTA 格式的标签文件，无需转换一键训练。最后，工具中还提供了 yolov5_obb 预训练，并支持加载自定义模型，可极大提升数据标注效率，形成快速闭环。

工具安装

git clone https://github.com/CVHub520/X-AnyLabeling  
cd X-AnyLabeling  
pip install -r requirements.txt  
# pip install -r requirements-gpu.txt  
python anylabeling/app.py  

使用教程

• 在标注之前，准备一个预定义的类别标签文件(*.txt)，格式如下：

vehicle  
boat  

每一行代表一个类别，根据自己的具体任务填写。

• 在顶部菜单栏中点击Format选项，选择DOTA，并导入上一步准备好的标签文件。

基本用法

• 按下快捷键 "O" 来创建一个旋转形状。

• 打开编辑模式（快捷键："Ctrl+J"）并单击选择旋转框。

• 通过快捷键 "zxcv" 旋转所选框，其中：

• z：大角度逆时针旋转
• x：小角度逆时针旋转
• c：小角度顺时针旋转
• v：大角度顺时针旋转

高级用法

此外，您可以使用训练好的模型批量预标记当前数据集。

• 按下快捷键 "Ctrl+A" 打开自动标注模式；
• 选择一个适当的模型或加载本地训练完毕的自定义模型。
• 按下快捷键 "Ctrl+M" 一次运行所有图像。

如果你手头上刚好有一批标注好的 DOTA 格式数据集，同样也可以使用标签转换脚本一键转换成 X-AnyLabeling 的自定义格式导入到工具中进一步查看和修改：

python tools/label_converter.py --task rotation --src_path dota_label_folder --img_path dota_image_folder --mode dota2custom  

详细信息，请参考[2]。

入门指南

准备工作

• 环境要求

• Python 3.7+
• PyTorch ≥ 1.7
• CUDA 9.0或更高版本
• Ubuntu 16.04/18.04

• 安装开始

a. 创建一个conda虚拟环境并激活它：

conda create -n yolov5_obb python=3.8 -y   
source activate yolov5_obb  

b. 确保您的CUDA运行时API版本≤CUDA驱动程序版本。（例如11.3 ≤ 11.4）

nvcc -V  
nvidia-smi  

c. 根据您的机器环境，根据官方说明安装PyTorch和torchvision，并确保cudatoolkit版本与CUDA运行时API版本相同，例如：

pip install torch==1.12.0+cu116 torchvision==0.13.0+cu116 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu116  
nvcc -V  
python  
>>> import torch  
>>> torch.version.cuda  
>>> exit()  

d. 克隆最新版本的YOLOv5_OBB存储库。

git clone https://github.com/CVHub520/yolov5_obb.git  
cd yolov5_obb  

e. 安装yolov5-obb。

pip install -r requirements.txt  
cd utils/nms_rotated  
python setup.py develop  # 或者 "pip install -v -e ."  

注意：

1. 对于Windows用户，请参考[3]，如果在你在生成 utils/nms_rotated_ext.cpython-XX-XX-XX-XX.so方面遇到困难。
2. 需要注意的是，笔者这里对 `poly_nms_cuda` CUDA实现重构了一遍，如果你使用的是 hukaixuan19970627 实现的版本，大概率会因版本问题编译不通过，可参考着修改下，但建议直接使用此修改后的版本，避免冲突。

• DOTA_devkit [可选]

如果您需要切分高分辨率图像并进行评估，建议使用以下工具：

cd yolov5_obb/DOTA_devkit  
sudo apt-get install swig  
swig -c++ -python polyiou.i  
python setup.py build_ext --inplace  

数据集划分

准备自定义数据集文件

注意：确保标签格式为[polygon classname difficulty]，例如，您可以默认将所有difficulty参数重置0，除非另有用途。

  x1      y1       x2        y2       x3       y3       x4       y4       classname     diffcult  
  
1686.0   1517.0   1695.0   1511.0   1711.0   1535.0   1700.0   1541.0   large-vehicle      0  

然后，如果不需要对高分辨率图像进行切分处理的话，可直接修改路径参数并运行当前工程目录下的对应[脚本] (./divide.py)。注意修改下文件里面对应的路径参数等。

否则，您可以按以下步骤操作。

cd yolov5_obb  
python DOTA_devkit/ImgSplit_multi_process.py  

确保您的数据集组织在如下所示的目录结构中：

.  
└── dataset_demo  
    ├── images  
    │   └── P0032.png  
    └── labelTxt  
         
  
 └── P0032.txt  

最后，您可以创建一个自定义的数据yaml文件，参照[yolov5obb_demo.yaml] (./data/yolov5obb_demo.yaml)，同yolov5类似。

注意：

• 对于像DOTA这种超大分辨率的图像数据集，建议在训练/测试之前进行切分处理，以获得更好的性能。
• 对于单类问题，建议添加一个"None"类，实际上将其变为一个2类任务，例如DroneVehicle_poly.yaml配置文件中所示，有利于训练更加稳定。

训练/验证/检测

在正式开始训练任务之前，请遵循以下建议：

1. 确保将输入分辨率设置为32的倍数。
2. 默认情况下，将批处理大小设置为8。如果将其增加到16或更大，请调整框丢失的缩放因子，以帮助收敛theta。

• 要在多个GPU上使用分布式数据并行（DDP）模式进行训练，请参考此shell[脚本] (./sh/ddp_train.sh)。
• 如果要训练原始数据集而不需要切分数据集，可请参考以下命令快速开始🔜：

python train.py \  
  --weights weights/yolov5n.pt \  
  --data data/task.yaml \  
  --hyp data/hyps/obb/hyp.finetune_dota.yaml \  
  --epochs 300 \  
  --batch-size 1 \  
  --img 1024 \  
  --device 0 \  
  --name /path/to/save_dir  

• 要检测图像文件/文件夹/视频，请参考以下命令：

python detect.py \  
    --weights /path/to/*.pt \  
    --source /path/to/image \  
    --img 1024 \  
    --device 0 \  
    --conf-thres 0.25 \  
    --iou-thres 0.2 \  
    --name /path/to/save_dir  

注意：有关更多详细信息，请参考[此文档] (./docs/GetStart.md)。

部署

• 导出*.onnx文件：

python export.py \  
    --weights runs/train/task/weights/best.pt \  
    --data data/task.yaml \  
    --imgsz 1024 \  
    --simplify \  
    --opset 12 \  
    --include onnx  

Python

• 基于onnxruntime后端进行推理：

python deploy/onnxruntime/python/main.py \  
    --model /path/to/*.onnx \  
    --image /path/to/image  

C++

• 进入目录：

cd opencv/cpp  
  
.  
└── cpp  
    ├── CMakeLists.txt  
    ├── build  
    ├── image  
    │   ├── demo.jpg  
    ├── main.cpp  
    ├── model  
    │   └── yolov5m_obb_csl_dotav15.onnx  
    └── obb  
        ├── include  
        └── src  

• 安装OpenCV和Eigen3库。

注意，建议使用OpenCV版本4.6.0或更新，v4.7.0已成功测试。
Eigen是一个高性能的C++模板库，专用于线性代数运算，支持矩阵和向量操作，具有自然的C++语法、无依赖性和跨平台特性，广泛应用于科学计算、图形学和机器学习等领域。

• 将图像和模型文件放在指定目录中。
• 根据您的特定要求和用例修改CMakeLists.txt、main.cpp和yolo_obb.h文件的内容。
• 基于opencv-dnn后端进行推理：

mkdir build && cd build  
cmake ..  
make  

对于 OpenVINO 或者 TensorRT 等其它后端的部署推理，原理是一样的，只需要改下 engine 那块即可。

模型库

此外，本项目还同时提供了 *.pt 和 *.onnx 文件，方便大家下载使用，快速复现。

总结

今天为大家详细介绍了构建旋转目标检测任务的整个pipeline，如果你在做项目的过程中碰到任何问题，可直接在 github issue 反馈，如若需要进一步交流学习或者对X-AnyLabeling工程感兴趣的欢迎添加微信：ww10874，备注yolov5_obb，非诚勿扰。

References

1.X-AnyLabeling: https://github.com/CVHub520/X-AnyLabeling

2.X-AnyLabeling文档: https://medium.com/@CVHub520/x-anylabeling-pioneering-the-annotation-revolution-eed0ae788f7d

3.https://github.com/hukaixuan19970627/yolov5_obb/issues/224

公众号后台回复“数据集”获取100+深度学习各方向资源整理