工业级旋转目标检测从数据制作到模型部署 ,挖来的大佬的哈~~ 感谢
yolov5_obb
是基于 yolov5
目标检测框架开发的一个项目,支持旋转目标检测任务(Oriented Bounding Boxes,OBB),旨在生成更好拟合具有角度位置的物体预测结果,广泛适用于机器人感知,遥感等领域。此前经过调研发现目前全网上关于此方面的资料相对较少,鱼龙混杂,不是比较老旧、乱七八糟,就是一言不合就付费查看,看个寂寞,实在是不忍直视。因此,本文旨在提供一个从数据集制作、划分、安装、训练、验证、部署保姆级教程,帮助大家从0到1快速完成项目的上手和开发,满足日常工作和学习的需求。
#数据集制作
工具介绍
X-AnyLabeling[1]不仅是一个标注工具,它还是自动数据标注未来的一大飞跃。它的设计不仅简化了标注过程,还集成了尖端的AI模型,融入了当前最先进的视觉大模型和多模态技术,以获得更出色的体验。X-AnyLabeling 专注于实际应用,力求提供一个工业级的、功能丰富的工具,可帮助开发人员自动标注和处理各种复杂任务的数据。
在最新的 v1.0.0 版本中,X-AnyLabeling
已经完美支持了旋转框的标注,同时提供全方位的角度显示,实操体验方面也与经典的 roLabelImg
基本保持一致,极大降低用户的学习成本,且支持直接导出 DOTA 格式的标签文件,无需转换一键训练。最后,工具中还提供了 yolov5_obb
预训练,并支持加载自定义模型,可极大提升数据标注效率,形成快速闭环。
喜欢的小伙伴欢迎点个
star
,持续关注后续更丰富的功能更新。
工具安装
使用教程
- 在标注之前,准备一个预定义的类别标签文件(*.txt),格式如下:
每一行代表一个类别,根据自己的具体任务填写。
- 在顶部菜单栏中点击
Format
选项,选择DOTA
,并导入上一步准备好的标签文件。
基本用法
- 按下快捷键 "O" 来创建一个旋转形状。
- 打开编辑模式(快捷键:"Ctrl+J")并单击选择旋转框。
- 通过快捷键 "zxcv" 旋转所选框,其中:
- z:大角度逆时针旋转
- x:小角度逆时针旋转
- c:小角度顺时针旋转
- v:大角度顺时针旋转
高级用法
此外,您可以使用训练好的模型批量预标记当前数据集。
- 按下快捷键 "Ctrl+A" 打开自动标注模式;
- 选择一个适当的模型或加载本地训练完毕的自定义模型。
- 按下快捷键 "Ctrl+M" 一次运行所有图像。
如果你手头上刚好有一批标注好的 DOTA 格式数据集,同样也可以使用标签转换脚本一键转换成 X-AnyLabeling
的自定义格式导入到工具中进一步查看和修改:
详细信息,请参考此文档[2]。
入门指南
准备工作
- 环境要求
- Python 3.7+
- PyTorch ≥ 1.7
- CUDA 9.0或更高版本
- Ubuntu 16.04/18.04
- 安装开始
a. 创建一个conda虚拟环境并激活它:
b. 确保您的CUDA运行时API版本≤CUDA驱动程序版本。(例如11.3 ≤ 11.4)
c. 根据您的机器环境,根据官方说明安装PyTorch和torchvision,并确保cudatoolkit版本与CUDA运行时API版本相同,例如:
d. 克隆最新版本的YOLOv5_OBB存储库。
e. 安装yolov5-obb。
注意:
- 对于Windows用户,请参考此问题[3],如果在你在生成 utils/nms_rotated_ext.cpython-XX-XX-XX-XX.so方面遇到困难。
- 需要注意的是,笔者这里对 `poly_nms_cuda` CUDA实现重构了一遍,如果你使用的是 hukaixuan19970627 实现的版本,大概率会因版本问题编译不通过,可参考着修改下,但建议直接使用此修改后的版本,避免冲突。
- DOTA_devkit [可选]
如果您需要切分高分辨率图像并进行评估,建议使用以下工具:
数据集划分
准备自定义数据集文件
注意:确保标签格式为[polygon classname difficulty],例如,您可以默认将所有difficulty参数重置0,除非另有用途。
然后,如果不需要对高分辨率图像进行切分处理的话,可直接修改路径参数并运行当前工程目录下的对应[脚本] (./divide.py)。注意修改下文件里面对应的路径参数等。
否则,您可以按以下步骤操作。
确保您的数据集组织在如下所示的目录结构中:
最后,您可以创建一个自定义的数据yaml文件,参照[yolov5obb_demo.yaml] (./data/yolov5obb_demo.yaml),同yolov5类似。
注意:
- 对于像DOTA这种超大分辨率的图像数据集,建议在训练/测试之前进行切分处理,以获得更好的性能。
- 对于单类问题,建议添加一个"None"类,实际上将其变为一个2类任务,例如DroneVehicle_poly.yaml配置文件中所示,有利于训练更加稳定。
训练/验证/检测
在正式开始训练任务之前,请遵循以下建议:
- 确保将输入分辨率设置为32的倍数。
- 默认情况下,将批处理大小设置为8。如果将其增加到16或更大,请调整框丢失的缩放因子,以帮助收敛
theta
。
- 要在多个GPU上使用分布式数据并行(DDP)模式进行训练,请参考此shell[脚本] (./sh/ddp_train.sh)。
- 如果要训练原始数据集而不需要切分数据集,可请参考以下命令快速开始🔜:
- 要检测图像文件/文件夹/视频,请参考以下命令:
注意:有关更多详细信息,请参考[此文档] (./docs/GetStart.md)。
#部署
- 导出*.onnx文件:
Python
- 基于onnxruntime后端进行推理:
C++
- 进入目录:
- 安装OpenCV和Eigen3库。
注意,建议使用OpenCV版本4.6.0或更新,v4.7.0已成功测试。
Eigen是一个高性能的C++模板库,专用于线性代数运算,支持矩阵和向量操作,具有自然的C++语法、无依赖性和跨平台特性,广泛应用于科学计算、图形学和机器学习等领域。
- 将图像和模型文件放在指定目录中。
- 根据您的特定要求和用例修改CMakeLists.txt、main.cpp和yolo_obb.h文件的内容。
- 基于opencv-dnn后端进行推理:
对于 OpenVINO
或者 TensorRT
等其它后端的部署推理,原理是一样的,只需要改下 engine 那块即可。
#模型库
此外,本项目还同时提供了 *.pt 和 *.onnx 文件,方便大家下载使用,快速复现。
#总结
这里详细介绍了构建旋转目标检测任务的整个pipeline
,如果你在做项目的过程中碰到任何问题,可直接在 github issue 反馈