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

相关推荐
再不会python就不礼貌了21 分钟前
Ollama 0.4 发布!支持 Llama 3.2 Vision,实现多模态 RAG
人工智能·学习·机器学习·ai·开源·产品经理·llama
大大大反派24 分钟前
ONLYOFFICE 8.2深度测评:集成PDF编辑、数据可视化与AI功能的强大办公套件
人工智能·信息可视化·pdf
DK2215140 分钟前
机器学习系列-----主成分分析(PCA)
人工智能·算法·机器学习
SmallBambooCode2 小时前
【人工智能】阿里云PAI平台DSW实例一键安装Python脚本
linux·人工智能·python·阿里云·debian·脚本·模型训练
顾京2 小时前
基于扩散模型的表单插补
人工智能·深度学习·算法
NoneCoder2 小时前
AI时代IDE解析
ide·人工智能
狂奔solar2 小时前
yelp数据集上试验SVD,SVDPP,PMF,NMF 推荐算法
人工智能·机器学习·推荐算法
武子康2 小时前
大数据-216 数据挖掘 机器学习理论 - KMeans 基于轮廓系数来选择 n_clusters
大数据·人工智能·机器学习·数据挖掘·回归·scikit-learn·kmeans
liupenglove3 小时前
ElasticSearch向量检索技术方案介绍
大数据·人工智能·深度学习·elasticsearch·搜索引擎·自动驾驶
黄焖鸡能干四碗3 小时前
【系统文档】系统安全保障措施,安全运营保障,系统应急预案,系统验收相关资料(word原件)
大数据·人工智能·需求分析·软件需求·规格说明书