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

相关推荐
深圳南柯电子7 分钟前
深圳南柯电子|电子设备EMC测试整改:常见问题与解决方案
人工智能
Kai HVZ7 分钟前
《OpenCV计算机视觉》--介绍及基础操作
人工智能·opencv·计算机视觉
biter008812 分钟前
opencv(15) OpenCV背景减除器(Background Subtractors)学习
人工智能·opencv·学习
吃个糖糖18 分钟前
35 Opencv 亚像素角点检测
人工智能·opencv·计算机视觉
IT古董1 小时前
【漫话机器学习系列】017.大O算法(Big-O Notation)
人工智能·机器学习
凯哥是个大帅比1 小时前
人工智能ACA(五)--深度学习基础
人工智能·深度学习
m0_748232921 小时前
DALL-M:基于大语言模型的上下文感知临床数据增强方法 ,补充
人工智能·语言模型·自然语言处理
szxinmai主板定制专家2 小时前
【国产NI替代】基于FPGA的32通道(24bits)高精度终端采集核心板卡
大数据·人工智能·fpga开发
海棠AI实验室2 小时前
AI的进阶之路:从机器学习到深度学习的演变(三)
人工智能·深度学习·机器学习
机器懒得学习2 小时前
基于YOLOv5的智能水域监测系统:从目标检测到自动报告生成
人工智能·yolo·目标检测