本文分享YOLO11的旋转目标检测任务,在原始目标检测中,添加多一个角度预测,实现定向边界框检测。
其中旋转角度算法设计中,通过回归预测实现的。
目录
1、旋转目标检测概述
2、YOLO11中的OBB数据格式
3、分析模型配置参数
4、分析检测头代码
5、编写模型训练
6、开始训练
7、模型推理
1、旋转目标检测概述
旋转目标检测和OBB定向目标检测,其实是表达同一个意思;
旋转目标检测比目标检测更进一步,并引入了一个额外的角度来更准确地定位图像中的对象。
- 它专注于检测图像或视频中具有任意方向的对象。
- 与传统的水平边界框目标检测不同,旋转目标检测能够更精确地定位和描述那些非水平排列的目标,比如倾斜的飞机、船舶、风车等。
YOLO11中的旋转目标检测的特点
- 更精确的定位:通过使用旋转矩形(通常由中心点坐标、宽度、高度和旋转角度定义)来表示目标,可以更准确地界定目标的范围,减少背景干扰。
- 适应性强:对于那些在图像中以各种角度出现的目标,旋转目标检测方法能更好地适应这些变化,提高检测精度。
- 复杂度增加:相较于仅需预测四个坐标的水平边界框,旋转目标需要预测五个参数(x, y, w, h, θ),这增加了模型学习的难度和计算量。
2、YOLO11中的OBB数据格式
YOLO11中的OBB任务格式,通过其四个角点指定边界框,坐标在 0 和 1 之间标准化。
它的数据格式,如下所示:
class_index x1 y1 x2 y2 x3 y3 x4 y4
看一下标签文件示例,如下所示:
0 0.780811 0.743961 0.782371 0.74686 0.777691 0.752174 0.776131 0.749758
但是在YOLO11算法中,用边界框的中心点 (xy)、宽度(w)、高度(h)和旋转角度(r)的格式,处理损失和输出。
其中旋转角度,通过回归预测实现的。
数据参考官网:定向边界框 (OBB) 数据集概述 - Ultralytics YOLO Docs
旋转目标检测任务:OBB - Ultralytics YOLO 文档
3、分析模型配置参数
下面我们,分析YOLO11中旋转目标检测的模型配置参数
文件位置:ultralytics/cfg/models/11/yolo11-obb.yaml
# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLO11 Oriented Bounding Boxes (OBB) 模型,使用 P3-P5 输出。用法示例请参阅 https://docs.ultralytics.com/tasks/obb
# 参数设置
nc: 80 # 类别数量(number of classes)
scales: # 模型复合缩放常量,例如 'model=yolo11n-obb.yaml' 会调用 yolo11-obb.yaml,使用缩放 'n'
# [深度, 宽度, 最大通道数量]
n: [0.50, 0.25, 1024] # 概述:344 层,2695747 参数,2695731 梯度,6.9 GFLOPs
s: [0.50, 0.50, 1024] # 概述:344 层,9744931 参数,9744915 梯度,22.7 GFLOPs
m: [0.50, 1.00, 512] # 概述:434 层,20963523 参数,20963507 梯度,72.2 GFLOPs
l: [1.00, 1.00, 512] # 概述:656 层,26220995 参数,26220979 梯度,91.3 GFLOPs
x: [1.00, 1.50, 512] # 概述:656 层,58875331 参数,58875315 梯度,204.3 GFLOPs
# YOLO11n 主干网络(backbone)
backbone:
# [from, repeats, module, args] # [来自的层,重复次数,模块类型,参数]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 卷积层,输出通道 64,卷积核大小 3x3,步幅 2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 卷积层,输出通道 128,卷积核大小 3x3,步幅 2
- [-1, 2, C3k2, [256, False, 0.25]] # C3 模块,重复 2 次,输出 256 通道
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 卷积层,输出 256 通道,步幅 2
- [-1, 2, C3k2, [512, False, 0.25]] # C3 模块,重复 2 次,输出 512 通道
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 卷积层,输出 512 通道,步幅 2
- [-1, 2, C3k2, [512, True]] # C3 模块,重复 2 次,输出 512 通道,带 shortcut
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 卷积层,输出 1024 通道,步幅 2
- [-1, 2, C3k2, [1024, True]] # C3 模块,重复 2 次,输出 1024 通道,带 shortcut
- [-1, 1, SPPF, [1024, 5]] # SPPF 模块,卷积核大小 5,输出 1024 通道
- [-1, 2, C2PSA, [1024]] # C2PSA 模块,重复 2 次,输出 1024 通道
# YOLO11n 头部(head)
head:
- [-1, 1, nn.Upsample, [None, 2, "nearest"]] # 上采样层,放大 2 倍
- [[-1, 6], 1, Concat, [1]] # 与 backbone P4 拼接
- [-1, 2, C3k2, [512, False]] # C3 模块,重复 2 次,输出 512 通道
- [-1, 1, nn.Upsample, [None, 2, "nearest"]] # 上采样层,放大 2 倍
- [[-1, 4], 1, Concat, [1]] # 与 backbone P3 拼接
- [-1, 2, C3k2, [256, False]] # C3 模块,重复 2 次,输出 256 通道(P3/8-small)
- [-1, 1, Conv, [256, 3, 2]] # 卷积层,步幅 2,输出 256 通道
- [[-1, 13], 1, Concat, [1]] # 与 head P4 拼接
- [-1, 2, C3k2, [512, False]] # C3 模块,重复 2 次,输出 512 通道(P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]] # 卷积层,步幅 2,输出 512 通道
- [[-1, 10], 1, Concat, [1]] # 与 head P5 拼接
- [-1, 2, C3k2, [1024, True]] # C3 模块,重复 2 次,输出 1024 通道(P5/32-large)
- [[16, 19, 22], 1, OBB, [nc, 1]] # OBB 检测层,输入 P3, P4, P5,检测框输出
这里的模型参数,基本和目标检测是一样的,只是检测头有区别:
- 目标检测,使用的检测头:Detect
- 旋转目标检测,使用的检测头:OBB
4、分析检测头代码
从上面的模型参数配置文件,知道检测头是有区别,这里重点分析OBB代码
代码位置:ultralytics/nn/modules/head.py
整体的流程思路:
- 输入特征图:输入Neck特征融合中的多尺度特征图,包含了图像中物体的多尺度信息。
- 计算 OBB 角度:通过
self.cv4
的卷积层,即:三层卷积操作,根据输入特征图计算物体的旋转角度(θ)。 - 该角度值经过 sigmoid 激活函数处理,将范围映射到 [-π/4, 3π/4]。
- 边界框预测:调用目标检测的
Detect
的前向传播方法,预测常规的物体边界框信息(包括位置和类别概率)。 - 训练与推理的分支:训练模式,返回边界框和旋转角度的预测值,供模型优化和损失计算使用。
- 推理模式:将预测的角度与常规的边界框信息结合,得到最终的旋转边界框。
- 解码旋转边界框:在推理过程中,通过
decode_bboxes
函数对边界框进行解码,结合角度信息,输出最终的旋转物体边界框。
下面看看OBB检测头的代码
class OBB(Detect):
"""YOLOv8 OBB 检测头,用于具有旋转模型的检测。"""
def __init__(self, nc=80, ne=1, ch=()):
"""初始化 OBB,包含类别数量 `nc` 和层通道 `ch`。
:param nc: 类别数量(number of classes)
:param ne: 额外参数的数量(number of extra parameters),默认为1,这里特指用于预测旋转角度的参数
:param ch: 层通道(layer channels),输入特征图的通道数列表"""
super().__init__(nc, ch)
self.ne = ne # 额外参数的数量
c4 = max(ch[0] // 4, self.ne)
self.cv4 = nn.ModuleList(nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.ne, 1)) for x in ch)
def forward(self, x):
"""连接并返回预测的边界框和类别概率。
:param x: 输入特征图
:return: 如果是训练模式,返回边界框和角度;如果是推理模式,返回连接了角度的边界框"""
bs = x[0].shape[0] # 获取 batch 大小
angle = torch.cat([self.cv4[i](x[i]).view(bs, self.ne, -1) for i in range(self.nl)], 2) # OBB 角度 logits
angle = (angle.sigmoid() - 0.25) * math.pi # 角度范围 [-pi/4, 3pi/4]
if not self.training:
self.angle = angle # 推理模式下保存角度
x = Detect.forward(self, x) # 调用父类 Detect 的 forward 方法
if self.training:
return x, angle # 训练时返回边界框和角度
return torch.cat([x, angle], 1) if self.export else (torch.cat([x[0], angle], 1), (x[1], angle)) # 推理时返回角度
def decode_bboxes(self, bboxes, anchors):
"""解码旋转的边界框。
:param bboxes: 边界框
:param anchors: 锚框
:return: 解码后的旋转边界框"""
return dist2rbox(bboxes, self.angle, anchors, dim=1) # 使用 `dist2rbox` 函数解码
1)特征图转为角度
首先看输入的Neck特征融合中的多尺度特征图,如何被处理为角度的
self.cv4 = nn.ModuleList(nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.ne, 1)) for x in ch)
这里代码比较难读,看看可读版的:
# 为每个通道数创建一个卷积序列,并将这些序列存入 nn.ModuleList
self.cv4 = nn.ModuleList()
# 遍历每个输入通道 ch,为每个通道创建一个卷积序列
for x in ch:
# 定义一个卷积序列:两个 3x3 卷积和一个 1x1 卷积
conv_sequence = nn.Sequential(
Conv(x, c4, 3), # 第一个 3x3 卷积,输入通道 x,输出通道 c4
Conv(c4, c4, 3), # 第二个 3x3 卷积,输入通道和输出通道均为 c4
nn.Conv2d(c4, self.ne, 1) # 最后一个 1x1 卷积,输出通道为 ne
)
# 将卷积序列添加到 ModuleList
self.cv4.append(conv_sequence)
- 定义了一个卷积层模块列表
self.cv4
,用于处理输入的特征图并计算旋转角度。 - 每个模块由两个 3x3 卷积层和一个 1x1 卷积层组成:
- 第一个
Conv
将输入通道数变为c4
。 - 第二个
Conv
继续保持c4
通道数。 - 最后的
nn.Conv2d
将通道数变为self.ne
,即角度预测的输出维度。
- 第一个
总结一下,通过三层的卷积层,最终输出为角度的。
2)前向传播 forward代码
然后,重点看看前向传播 forward(self, x)
- 目标:结合输入特征图,预测旋转边界框的角度,并返回边界框信息。
- 流程:
- 首先获取输入特征图的 batch size(
bs
)。 - 然后通过卷积层
self.cv4
计算每个特征图的角度值(theta
)。 - 这些角度是经过
sigmoid
函数处理的,并被映射到 [-π/4, 3π/4] 的范围内。 - 判断是否处于训练模式:
- 训练模式:返回边界框预测和角度(
angle
)。 - 推理模式:在预测输出中添加角度信息,返回组合后的边界框和角度。
- 训练模式:返回边界框预测和角度(
- 首先获取输入特征图的 batch size(
重点分析如何处理角度的代码
angle = torch.cat([self.cv4[i](x[i]).view(bs, self.ne, -1) for i in range(self.nl)], 2)
- 使用
self.cv4
中的卷积层对每个输入特征图x[i]
进行处理,计算 OBB 的角度 logits。 - 计算结果的形状调整为 (batch size, ne, 特征图上的位置数) 并沿着第二维度拼接。
angle = (angle.sigmoid() - 0.25) * math.pi
- 对角度 logits 进行
sigmoid
激活,然后减去 0.25 并乘以 π,将角度范围限制在 [-π/4, 3π/4]。
下面代码更有可读性,看看:
# 初始化一个空列表来存储每层的角度 logits
angle_logits_list = []
# 遍历每层的特征图,计算角度 logits,并调整形状
for i in range(self.nl):
feature_map = x[i] # 获取第 i 层的特征图
angle_logits = self.cv4[i](feature_map) # 通过卷积层计算角度 logits
reshaped_angle_logits = angle_logits.view(bs, self.ne, -1) # 调整形状为 (batch_size, ne, num_positions)
angle_logits_list.append(reshaped_angle_logits) # 将处理后的角度 logits 添加到列表中
# 在第二维度上拼接所有层的角度 logits
angle = torch.cat(angle_logits_list, dim=2) # OBB 角度 logits
# 角度范围 [-pi/4, 3pi/4]
angle = (angle.sigmoid() - 0.25) * math.pi
扩展:YOLO11检测头的代码位置,ultralytics/nn/modules/head.py,包含多个任务的检测头
"Detect", "Segment", "Pose", "Classify", "OBB", "RTDETRDecoder", "v10Detect"
5、编写模型训练
新建文件train.py,编写模型训练的代码。
简洁版,参考下面代码:
from ultralytics import YOLO
# 加载预训练权重文件
model = YOLO("weights/yolo11s-obb.pt")
# 使用dota8.yaml数据集进行训练,训练10个epoch,并将图像大小设置为1024像素
results = model.train(data="dota8.yaml", epochs=10, imgsz=1024)
这里图像的分辨率,默认是1024,也可以改为640的
然后开始模型训练啦~
下面是专业版的OBB模型训练,示例代码如下:
from ultralytics import YOLO
# 加载预训练的模型
# model = YOLO("yolo11m-obb.yaml").load("weights/yolo11m-obb.pt")
model = YOLO("weights/yolo11m-obb.pt")
# 定义训练参数,添加默认值、范围和中文注释
train_params = {
'data': "DOTAv1.5.yaml", # 数据集配置文件路径,需要自定义修改 coco128.yaml, coco8.yaml, dota8.yaml
'epochs': 100, # 总训练轮次,默认值 100,范围 >= 1
'imgsz': 1024, # 输入图像大小,默认值 1024,范围 >= 32
'batch': 8, # 批次大小,默认值 16,范围 >= 1
'save': True, # 是否保存训练结果和模型,默认值 True
'save_period': -1, # 模型保存频率,默认值 -1,表示只保存最终结果
'cache': False, # 是否缓存数据集,默认值 False
'device': None, # 训练设备,默认值 None,支持 "cpu", "gpu"(device=0,1), "mps"
'workers': 8, # 数据加载线程数,默认值 8,影响数据预处理速度
'project': None, # 项目名称,保存训练结果的目录,默认值 None
'name': None, # 训练运行的名称,用于创建子目录保存结果,默认值 None
'exist_ok': False, # 是否覆盖已有项目/名称目录,默认值 False
'pretrained': True,
'optimizer': 'auto', # 优化器,默认值 'auto',支持 'SGD', 'Adam', 'AdamW'
'verbose': True, # 是否启用详细日志输出,默认值 False
'seed': 0, # 随机种子,确保结果的可重复性,默认值 0
'deterministic': True, # 是否强制使用确定性算法,默认值 True
'single_cls': False, # 是否将多类别数据集视为单一类别,默认值 False
'rect': False, # 是否启用矩形训练(优化批次图像大小),默认值 False
'cos_lr': False, # 是否使用余弦学习率调度器,默认值 False
'close_mosaic': 10, # 在最后 N 轮次中禁用 Mosaic 数据增强,默认值 10
'resume': False, # 是否从上次保存的检查点继续训练,默认值 False
'amp': True, # 是否启用自动混合精度(AMP)训练,默认值 True
'fraction': 1.0, # 使用数据集的比例,默认值 1.0
'profile': False, # 是否启用 ONNX 或 TensorRT 模型优化分析,默认值 False
'freeze': None, # 冻结模型的前 N 层,默认值 None
'lr0': 0.01, # 初始学习率,默认值 0.01,范围 >= 0
'lrf': 0.01, # 最终学习率与初始学习率的比值,默认值 0.01
'momentum': 0.937, # SGD 或 Adam 的动量因子,默认值 0.937,范围 [0, 1]
'weight_decay': 0.0005, # 权重衰减,防止过拟合,默认值 0.0005
'warmup_epochs': 3.0, # 预热学习率的轮次,默认值 3.0
'warmup_momentum': 0.8, # 预热阶段的初始动量,默认值 0.8
'warmup_bias_lr': 0.1, # 预热阶段的偏置学习率,默认值 0.1
'box': 7.5, # 边框损失的权重,默认值 7.5
'cls': 0.5, # 分类损失的权重,默认值 0.5
'dfl': 1.5, # 分布焦点损失的权重,默认值 1.5
'pose': 12.0, # 姿态损失的权重,默认值 12.0
'kobj': 1.0, # 关键点目标损失的权重,默认值 1.0
'label_smoothing': 0.0, # 标签平滑处理,默认值 0.0
'nbs': 64, # 归一化批次大小,默认值 64
'overlap_mask': True, # 是否在训练期间启用掩码重叠,默认值 True
'mask_ratio': 4, # 掩码下采样比例,默认值 4
'dropout': 0.0, # 随机失活率,用于防止过拟合,默认值 0.0
'val': True, # 是否在训练期间启用验证,默认值 True
'plots': True, # 是否生成训练曲线和验证指标图,默认值 True
# 数据增强相关参数
'hsv_h': 0.2, # 色相变化范围 (0.0 - 1.0),默认值 0.015
'hsv_s': 0.7, # 饱和度变化范围 (0.0 - 1.0),默认值 0.7
'hsv_v': 0.4, # 亮度变化范围 (0.0 - 1.0),默认值 0.4
'degrees': 30.0, # 旋转角度范围 (-180 - 180),默认值 0.0
'translate': 0.1, # 平移范围 (0.0 - 1.0),默认值 0.1
'scale': 0.5, # 缩放比例范围 (>= 0.0),默认值 0.5
'shear': 0.0, # 剪切角度范围 (-180 - 180),默认值 0.0
'perspective': 0.0, # 透视变化范围 (0.0 - 0.001),默认值 0.0
'flipud': 0.0, # 上下翻转概率 (0.0 - 1.0),默认值 0.0
'fliplr': 0.5, # 左右翻转概率 (0.0 - 1.0),默认值 0.5
'bgr': 0.0, # BGR 色彩顺序调整概率 (0.0 - 1.0),默认值 0.0
'mosaic': 0.5, # Mosaic 数据增强 (0.0 - 1.0),默认值 1.0
'mixup': 0.0, # Mixup 数据增强 (0.0 - 1.0),默认值 0.0
'copy_paste': 0.0, # Copy-Paste 数据增强 (0.0 - 1.0),默认值 0.0
'copy_paste_mode': 'flip', # Copy-Paste 增强模式 ('flip' 或 'mixup'),默认值 'flip'
'auto_augment': 'randaugment', # 自动增强策略 ('randaugment', 'autoaugment', 'augmix'),默认值 'randaugment'
'erasing': 0.4, # 随机擦除增强比例 (0.0 - 0.9),默认值 0.4
'crop_fraction': 1.0, # 裁剪比例 (0.1 - 1.0),默认值 1.0
}
# 进行训练
results = model.train(**train_params)
YOLO11模型训练,思路流程:
- 加载模型:使用 YOLO 类指定模型的配置文件,并加载预训练权重
yolo11m-obb.pt
。 - 定义训练参数:通过字典
train_params
定义了一系列训练参数,涵盖了训练过程中可能涉及的配置项,如数据集路径、训练轮数、图像大小、优化器、数据增强等。 - 执行训练:使用
model.train(**train_params)
将定义的训练参数传入模型,开始训练。 - 保存训练结果:训练完成后,结果保存在
results
中,包含损失和精度等信息。
数据集的配置文件路径需要在 train_params
中指定.
例如,代码中的 "DOTAv1.5.yaml"
是一个数据集配置文件,它定义了训练数据的位置、类别等信息。
6、开始训练
直接运行train.py,就开始训练啦
Ultralytics 8.3.7 🚀 Python-3.9.16 torch-1.13.1 CUDA:0 (NVIDIA A30, 24062MiB)
Starting training for 100 epochs...
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
1/100 10.2G 1.24 3.447 1.101 2 640: 100%|██████████| 177/177 [00:31<00:00, 5.63it/s]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 58/58 [00:06<00:00, 8.73it/s]
all 458 69565 0.699 0.268 0.323 0.229
............
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
100/100 9.9G 0.7904 0.4866 1.019 23 640: 100%|██████████| 177/177 [00:21<00:00, 8.12it/s]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 58/58 [00:04<00:00, 13.21it/s]
all 458 69565 0.808 0.449 0.52 0.387
100 epochs completed in 0.786 hours.
Optimizer stripped from runs/obb/train15/weights/last.pt, 42.2MB
Optimizer stripped from runs/obb/train15/weights/best.pt, 42.2MB
训练完成后,可以在runs/obb/train15路径,可以看到保存的权重、训练记录表格和标签信息等
7、模型推理
使用刚才训练好的模型权重,进行模型推理,看看旋转目标检测的效果
示例代码如下所示
from ultralytics import YOLO
# 加载预训练的YOLOv11n模型
model = YOLO(r"runs/obb/train15/weights/best.pt") # weights/yolo11m-obb.pt
# 对指定的图像文件夹进行推理,并设置各种参数
results = model.predict(
source="/mnt/yolo11/ultralytics-main/datasets/DOTAv1.5/images/val/", # 数据来源,可以是文件夹、图片路径、视频、URL,或设备ID(如摄像头)
conf=0.45, # 置信度阈值
iou=0.6, # IoU 阈值
imgsz=1024, # 图像大小
half=False, # 使用半精度推理
device=None, # 使用设备,None 表示自动选择,比如'cpu','0'
max_det=500, # 最大检测数量
vid_stride=1, # 视频帧跳跃设置
stream_buffer=False, # 视频流缓冲
visualize=False, # 可视化模型特征
augment=False, # 启用推理时增强
agnostic_nms=False, # 启用类无关的NMS
classes=None, # 指定要检测的类别
retina_masks=False, # 使用高分辨率分割掩码
embed=None, # 提取特征向量层
show=False, # 是否显示推理图像
save=True, # 保存推理结果
save_frames=False, # 保存视频的帧作为图像
save_txt=True, # 保存检测结果到文本文件
save_conf=False, # 保存置信度到文本文件
save_crop=False, # 保存裁剪的检测对象图像
show_labels=True, # 显示检测的标签
show_conf=True, # 显示检测置信度
show_boxes=True, # 显示检测框
line_width=1 # 设置边界框的线条宽度,比如2,4
)
看一下检测效果:
对于较大的物体,旋转目标检测的效果:
对于细小密集的物体,旋转目标检测的效果:
YOLO11相关文章推荐:
一篇文章快速认识YOLO11 | 关键改进点 | 安装使用 | 模型训练和推理-CSDN博客
一篇文章快速认识 YOLO11 | 实例分割 | 模型训练 | 自定义数据集-CSDN博客
一篇文章快速认识 YOLO11 | 目标检测 | 模型训练 | 自定义数据集-CSDN博客
YOLO11模型推理 | 目标检测与跟踪 | 实例分割 | 关键点估计 | OBB旋转目标检测-CSDN博客
YOLO11模型训练 | 目标检测与跟踪 | 实例分割 | 关键点姿态估计-CSDN博客
YOLO11 实例分割 | 自动标注 | 预标注 | 标签格式转换 | 手动校正标签-CSDN博客
YOLO11 实例分割 | 导出ONNX模型 | ONNX模型推理-CSDN博客
YOLO11 目标检测 | 导出ONNX模型 | ONNX模型推理-CSDN博客
YOLO11 目标检测 | 自动标注 | 预标注 | 标签格式转换 | 手动校正标签_yolo11 标注平台-CSDN博客
分享完成,欢迎大家多多点赞和收藏,谢谢~