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

TF男孩
4,922 阅读5分钟

这是忒小众的一个知识点了，小众到只有一个人让我讲一讲！ 2024-08-08_171548.png

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

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

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

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

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

https://github.com/ultralytics/ultralytics/blob/main/docs/en/tasks/obb.md

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

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

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版本，它的主页如下：

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：游泳池

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

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)。如果缺少，可以不写。最重要的是后面的标注：

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格式，官方也有脚本。

from ultralytics.data.converter import convert_dota_to_yolo_obb
convert_dota_to_yolo_obb("path/to/DOTA")

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

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

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

from ultralytics import YOLO
model = YOLO("best.pt")  
results = model("test.png")

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

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

path: datasets/dogcart # 数据集的目录
train: images/train # 训练数据
val: images/val # 验证

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

最终目录如下：

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

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添加插件也可以实现。这类工具就太多了，开源不花钱。

avatar
文章
阅读
粉丝