YOLO里的OBB:带旋转角度的目标检测方案

这是忒小众的一个知识点了,小众到只有一个人让我讲一讲!

YOLO大家都很熟悉,以目标检测出名。它可以检测一个画面中,哪个位置有什么类型的物体。

你可能觉得除了精度和性能以外,它没有优化空间了,已经到头了。其实,这类目标检测,很多场景还覆盖不到。就比如更精确地定位物体。

我们想要框出物体,一般都是垂直的矩形,像下面海上的船舶这样。

但是实际上这不准确,我们更想要下面这样。它可以框出物体的实际矩形轮廓,而且紧紧贴合角度的变动。因为这样可以做很多事情,比如它识别船头船尾,判断行进方向,甚至预测他们是否会发生碰撞。

是的,这类包含方向的边界框目标检测,称为"Oriented Bounding Boxes Object Detection",也叫OBB。可能很多人不知道,YOLO中是自带OBB的。以下是YOLOv8的OBB说明文档:

ruby 复制代码
https://github.com/ultralytics/ultralytics/blob/main/docs/en/tasks/obb.md

普通的YOLO目标检测怎么用,OBB也怎么用,只是选择模型文件不一样。

下面是我从谷歌地球上截取了济南遥墙机场的一个角落,我们检测一下飞机的位置:

ini 复制代码
from ultralytics import YOLO
model = YOLO("yolov8n-obb.pt")
results = model("yaoqiang.png",  save=True)

里面的yolov8n-obb.pt怎么来的?那肯定是YOLO官方提供的,你可以直接下载。跟衣服一样,各种尺寸(X、L、M、S)都有,适用于不同设备和场景。

这些预训练权重文件,是训练了什么类型的数据呢?它又能识别哪些物体呢?

YOLO-OBB的训练集是DOTAv1。DOTAv1是DOTA数据集的v1版本,它的主页如下:

bash 复制代码
https://captain-whu.github.io/DOTA/dataset.html

DOTAv1的数据来自航拍卫星图,总共有15种类型:

复制代码
0:飞机
1:轮船
2:储罐
3:棒球场
4:网球场
5:篮球场
6:田径场
7:港口
8:桥梁
9:大型车辆
10:小型车辆
11:直升机
12:环形交叉路口
13:足球场
14:游泳池

要问我是怎么知道的,这些数据哪里下载的?那肯定是官方说明:

ruby 复制代码
https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/DOTAv1.yaml

如果,你的OBB需求包括在YOLO-OBB之内(检测上述15种物体),那么可以直接使用它训练好的模型。如果,你还有其他想法,那么就以它为基础,享受它已有的经验,进行低成本的定制再训练。

想要训练自己的数据,那肯定得先了解它是如何训练DOTAv1的。我们只需要照葫芦画瓢,按照YOLO训练DOTAv1的格式,组织自己的数据即可。

首先,DOTAv1和YOLO没有任何关系,他们都是各自独立的。唯一联系就是YOLO为了检测自己的OBB算法,选取了DOTAv1数据集进行试验。

DOTAv1的组成格式是这样的:

首先由很多训练图片,然后有针对每一张图片的txt标注信息。

标注顶部给出了"采集日期"。"图像源"(GoogleEarth、GF-2 或 JL-1)。以及地面采样距离(GSD)。如果缺少,可以不写。最重要的是后面的标注:

erlang 复制代码
x1, y1, x2, y2, x3, y3, x4, y4, category, difficult
x1, y1, x2, y2, x3, y3, x4, y4, category, difficult
...

x1, y1, x2, y2, x3, y3, x4, y4表示四个顶点的x坐标和y坐标。category表示物体的分类,difficult表示识别难度(1困难,0简单)。以上就是DOTAv1数据集的格式,但并不是YOLO-OBB的训练格式。

YOLO算法有自己的体系,针对OBB训练,它的格式如下:

复制代码
class_index x1 y1 x2 y2 x3 y3 x4 y4

class_index表示分类的索引,其他8项是4个顶点的坐标,不过它的数值采用的是百分比形式(这是YOLO的体系)。最终实际数据像下面这样:

复制代码
0 0.780811 0.743961 0.782371 0.74686 0.777691 0.752174 0.776131 0.749758

对于DOTAv1格式,可以轻松转换为YOLO-OBB格式,官方也有脚本。

java 复制代码
from ultralytics.data.converter import convert_dota_to_yolo_obb
convert_dota_to_yolo_obb("path/to/DOTA")

数据集准备好了之后,就可以进行训练了。训练方式和普通的检测训练一样,也是通用的代码。

python 复制代码
from ultralytics import YOLO
model = YOLO("yolov8n-obb.yaml")
results = model.train(data="DOTAv1.yaml", epochs=100, imgsz=640)

训练完成后,会生成last.pt和best.pt文件,供后续推理使用。

python 复制代码
from ultralytics import YOLO
model = YOLO("best.pt")  
results = model("test.png")

假设我们有自己的数据,如何进行训练呢?

首先,我们需要准备数据,并建立一个yaml文件。我们志向不大,以摩托车和自行车这两类交通工具的OBB检测为例。新建一个dogcart.yaml文件,里面说明训练的关键信息:

yaml 复制代码
path: datasets/dogcart # 数据集的目录
train: images/train # 训练数据
val: images/val # 验证

names:
  0: bicycle # 分类1 自行车
  1: motorcycle # 分类2 摩托车

最终目录如下:

下面就是进行训练,以yolov8n-obb.pt为基础,训练自己的dogcar.yaml:

python 复制代码
from ultralytics import YOLO
model = YOLO("yolov8n-obb.pt")
results = model.train(data="dogcar.yaml", epochs=100, imgsz=640)

训练完成,获取best.pt,试验推理。

这就用YOLO-OBB实现了旋转物体的自训练和使用。

最后,对于OBB训练数据的标记工具,可以采用Labelme,或者给LabelImg添加插件也可以实现。这类工具就太多了,开源不花钱。

相关推荐
Jay Kay5 分钟前
ReLU 新生:从死亡困境到强势回归
人工智能·数据挖掘·回归
Blossom.11815 分钟前
使用Python和Flask构建简单的机器学习API
人工智能·python·深度学习·目标检测·机器学习·数据挖掘·flask
无声旅者1 小时前
AI 模型分类全解:特性与选择指南
人工智能·ai·ai大模型
Grassto1 小时前
Cursor Rules 使用
人工智能
MYH5161 小时前
深度学习在非线性场景中的核心应用领域及向量/张量数据处理案例,结合工业、金融等领域的实际落地场景分析
人工智能·深度学习
Lilith的AI学习日记1 小时前
什么是预训练?深入解读大模型AI的“高考集训”
开发语言·人工智能·深度学习·神经网络·机器学习·ai编程
聚客AI2 小时前
PyTorch玩转CNN:卷积操作可视化+五大经典网络复现+分类项目
人工智能·pytorch·神经网络
程序员岳焱2 小时前
深度剖析:Spring AI 与 LangChain4j,谁才是 Java 程序员的 AI 开发利器?
java·人工智能·后端
柠檬味拥抱2 小时前
AI智能体在金融决策系统中的自主学习与行为建模方法探讨
人工智能
智驱力人工智能2 小时前
智慧零售管理中的客流统计与属性分析
人工智能·算法·边缘计算·零售·智慧零售·聚众识别·人员计数