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

相关推荐
松果财经11 分钟前
蓝思科技赋能灵伴科技:AI眼镜产能与供应链双升级
人工智能·科技
青松@FasterAI1 小时前
【NLP算法面经】本科双非,头条+腾讯 NLP 详细面经(★附面题整理★)
人工智能·算法·自然语言处理
萧鼎1 小时前
智能自动化新纪元:AI与UiPath RPA的协同应用场景与技术实践
人工智能·ui·rpa
果冻人工智能1 小时前
去中心化 AI:赋权还是混乱?
人工智能·深度学习·机器学习·架构·去中心化·区块链·ai员工
Landy_Jay3 小时前
深度学习:基于Qwen复现DeepSeek R1的推理能力
人工智能·深度学习
EterNity_TiMe_3 小时前
【人工智能】蓝耘智算平台盛大发布DeepSeek满血版:开创AI推理体验新纪元
人工智能·python·机器学习·deepseek
RFID舜识物联网3 小时前
RFID测温技术:电力设备安全监测的新利器
网络·人工智能·嵌入式硬件·物联网·安全
豪越大豪3 小时前
豪越消防一体化安全管控平台新亮点: AI功能、智能运维以及消防处置知识库
大数据·人工智能·运维开发
9命怪猫3 小时前
AI大模型-提示工程学习笔记13—自动提示工程师 (Automatic Prompt Engineer)
人工智能·ai·大模型·prompt