★ FEATURED ARTICLE 简介本资源为YOLOv5 OBB旋转框训练DEMO面向具备一定深度学习基础、希望掌握倾斜目标检测的开发者与算法学习者解决常规轴对齐边界框难以精确表示树木、车辆等有角度物体的问题。压缩包共583个文件约445.58MB涵盖71个Python脚本、38个YAML配置、166个XML标注、95个TXT标签及93张PNG、23张JPG图像另含PT权重、ONNX模型、CUDA与C源码及Dockerfile等覆盖数据、训练、推理与部署全链路。已有1815人学习下载。资源围绕数据预处理、配置文件修改、训练流程、旋转框损失计算、模型评估与可视化、微调优化及ONNX/TensorRT部署等关键环节展开提供可运行的完整示例便于读者结合自有数据集实践OBB训练深入理解旋转目标检测技术。1. YOLOv5 OBB 旋转框训练从水平框到带角度的检测值不值得上手如果你之前只用 YOLOv5 跑过水平框检测第一次接触 OBBOriented Bounding Box旋转框大概率会有一个疑问明明水平框已经能框住目标了为什么还要多此一举加个角度答案藏在场景里。遥感影像里的飞机、舰船工业质检里的 PCB 板、文字行密集排列且长宽比悬殊水平框会把大量背景卷进来NMS 阶段互相压制漏检和误检同时飙升。旋转框多出来的那个角度参数解决的正是「框得准」这件事。YOLOv5 OBB 训练 demo 要做的就是把这套带角度的检测流程跑通数据标注格式怎么定、角度怎么定义、损失函数怎么改、训练脚本怎么配、推理时旋转框怎么画出来。它适合已经跑通过 YOLOv5 水平框、手里有带方向目标数据集的从业者也适合做遥感、工业视觉、文档分析方向、想快速验证旋转框方案是否值得投入的工程师。这篇笔记按「先立住原理、再动手复现、最后讲坑」的顺序展开中间给可直接抄的命令和配置参数含义和失败排查都写清楚。2. 旋转框的数据表示与角度定义先把标注格式这件事定死2.1 为什么水平框在密集场景会翻车水平框用(x_center, y_center, width, height)四个量描述一个目标前提是目标的主轴和图像坐标轴大致平行。一旦目标倾斜比如遥感图里停放的飞机呈 45 度排列水平框的最小外接矩形面积会远大于目标本身。两个相邻飞机的水平框大面积重叠NMS 的 IoU 阈值一卡其中一个就被抑制掉这就是密集场景漏检的直接来源。旋转框在四参数基础上引入角度 θ变成(x_center, y_center, width, height, angle)五参数。框可以贴着目标主轴旋转重叠面积大幅下降NMS 不再误杀。代价是角度本身有周期性θ 和 θ180° 表示同一条轴线θ 和 θ90° 在宽高互换后也等价这个周期性是后面损失设计和数据增强里最容易踩坑的地方。2.2 三种角度定义选错一个全盘皆输业界常见的角度定义有三类训练前必须和标注工具、损失函数、后处理统一定义方式角度范围零点参考典型来源OpenCV 定义[-90, 0)x 轴正向顺时针为负cv2.minAreaRect长边定义[-90, 90)长边与 x 轴夹角DOTA 部分工具短边定义[-45, 45)短边与 x 轴夹角部分工业标注我一般会统一到「长边定义、范围[-90, 90)」因为它在宽高互换时角度变化连续做数据增强旋转时不容易跳变。标注工具导出后第一件事就是写脚本把角度归一化到这个区间别指望训练代码帮你兜底。2.3 标注格式转换从 DOTA 到 YOLOv5 OBB 的脚本DOTA 数据集用(x1,y1,x2,y2,x3,y3,x4,y4) class difficult的八点格式YOLOv5 OBB 训练通常要转成归一化的五参数。下面这个脚本做转换同时把角度归一到长边定义import numpy as np import cv2 def poly_to_obb(points): 八点坐标 - (cx, cy, w, h, angle)长边定义角度范围 [-90, 90) pts np.array(points, dtypenp.float32).reshape(4, 2) rect cv2.minAreaRect(pts) # 返回 ((cx,cy),(w,h),angle) (cx, cy), (w, h), angle rect # OpenCV 的 angle 在 [-90,0)转成长边定义 if w h: w, h h, w angle 90 if angle 90: angle - 180 if angle -90: angle 180 return cx, cy, w, h, angle def convert_line(line, img_w, img_h): parts line.strip().split() coords list(map(float, parts[:8])) cls int(parts[8]) cx, cy, w, h, angle poly_to_obb(coords) # 归一化到 [0,1]角度归一化到 [0,1) 便于回归 return f{cls} {cx/img_w:.6f} {cy/img_h:.6f} \ f{w/img_w:.6f} {h/img_h:.6f} {(angle90)/180:.6f}逻辑说明cv2.minAreaRect返回的角度在 OpenCV 4.5 之后是[-90, 0)且宽高顺序不保证长边在前所以先判断w h交换并补 90 度再统一折到[-90, 90)。归一化时把角度映射到[0,1)是为了让回归头和宽高一样用 sigmoid 输出避免角度回归的数值范围过大导致梯度不稳。参数说明img_w、img_h是原图尺寸必须和标注时的图像一致cls是类别索引从 0 开始。转换完务必抽查几张用可视化脚本把旋转框画回原图确认角度方向没反。提示转换脚本跑完先别急着训练随机抽 20 张画框核对角度反了或宽高互换在训练 loss 上表现为「一直降不下去」排查成本很高。3. YOLOv5 OBB 训练环境搭建与配置把 demo 跑起来的最小闭环3.1 环境依赖与版本选择YOLOv5 OBB 不是官方主干直接支持的常见做法是基于 YOLOv5 仓库改检测头或者用已经集成 OBB 的分支。环境上我一般锁这几个版本避免依赖打架conda create -n yolo_obb python3.8 -y conda activate yolo_obb pip install torch1.12.1cu113 torchvision0.13.1cu113 \ --extra-index-url https://download.pytorch.org/whl/cu113 pip install opencv-python4.5.5.64 numpy1.23.5 \ pyyaml tqdm matplotlib scipy逻辑说明PyTorch 1.12 配 cu113 是经过验证比较稳的组合OpenCV 锁 4.5.5 是因为minAreaRect的角度返回规则在 4.5 前后有变化锁版本能保证标注转换和训练时角度定义一致。numpy 不要装 2.x部分老代码用了np.float会直接报错。参数说明CUDA 版本按你显卡驱动来30 系卡用 cu113 或 cu116 都行如果只有 CPU把cu113去掉但旋转框训练 CPU 上会慢到无法接受不建议。3.2 数据集目录结构与 yaml 配置YOLOv5 的数据配置沿用 yaml 风格OBB 版本多一个角度相关的开关。目录结构建议这样组织datasets/ obb_data/ images/ train/ val/ labels/ train/ val/ obb_data.yamlobb_data.yaml内容path: ./datasets/obb_data train: images/train val: images/val nc: 5 names: [plane, ship, vehicle, storage-tank, harbor] # OBB 专用 use_obb: true angle_range: 90 # 对应 [-90, 90) 长边定义 angle_loss: smooth_l1 # 可选 smooth_l1 / kld逻辑说明use_obb控制检测头是否输出第五个角度分量angle_range告诉损失函数角度的周期回归时做归一化用angle_loss选smooth_l1是最简单的直接回归选kld是高斯分布建模密集场景更稳但收敛慢。参数说明nc和names必须和标注类别严格对应顺序错了训练不报错但结果全乱。path用相对路径时注意从训练脚本所在目录解析。3.3 训练命令与关键超参最小可跑的训练命令python train.py \ --data datasets/obb_data/obb_data.yaml \ --cfg models/yolov5s_obb.yaml \ --weights yolov5s.pt \ --epochs 100 \ --batch-size 8 \ --img-size 1024 \ --hyp data/hyps/hyp.obb.yaml \ --name obb_demo逻辑说明--cfg指向改了检测头的模型结构输出通道从(nc5)变成(nc6)多一个角度--weights用水平框预训练权重做迁移主干特征能复用收敛快很多--img-size旋转框场景通常比水平框大因为小目标倾斜后更难回归。参数说明--batch-size受显存限制1024 尺寸下 8 是 12G 显存的常见值爆显存就降到 4 并开--accumulate--hyp里重点调box、cls、angle三项损失权重角度权重初始别超过 0.5否则前期角度回归主导框位置学不好。注意从水平框权重迁移时检测头最后那层卷积尺寸对不上代码里要处理成「加载时跳过 shape 不匹配的层」否则直接报错。常见做法是load_state_dict(strictFalse)并打印跳过的层名核对。4. 旋转框损失与后处理角度回归为什么总是不收敛4.1 角度回归的周期性陷阱角度回归最反直觉的地方在于模型预测 89 度真值 -89 度数值上差了 178但几何上只差 2 度。如果直接用 L1 或 L2 算loss 会给出一个巨大的惩罚梯度把网络往错误方向拽。这就是很多人遇到的「loss 震荡不下降」的根因。常见解法有三种一是把角度转成(sin, cos)两个分量回归天然消除周期性二是用 KLDKullback-Leibler Divergence把旋转框建模成高斯分布用分布距离代替角度差三是直接回归但做角度归一化把差值折到[-90, 90)再算 loss。我一般先用第三种快速验证效果不够再上 KLD。import torch def angle_loss(pred, target): pred/target 均为归一化到 [0,1) 的角度 diff pred - target # 折到 [-0.5, 0.5)对应 [-90, 90) diff diff - torch.round(diff) return torch.abs(diff).mean()逻辑说明torch.round(diff)把差值折到最近的整数周期diff - round(diff)得到[-0.5, 0.5)区间内的最小角度差再取绝对值求均值。这样 89 度和 -89 度的差会被正确算成 2 度对应的归一化值。参数说明输入必须是归一化角度如果原始角度是[-90, 90)先做(angle90)/180。这个 loss 简单但对密集场景的边界框不够鲁棒作为 baseline 够用。4.2 旋转 NMS 的实现要点水平框 NMS 用轴对齐 IoU旋转框必须换成旋转 IoU否则重叠判断全错。OpenCV 的cv2.rotatedRectangleIntersection能算两个旋转框的交集多边形再除以并集面积import cv2 import numpy as np def rotated_iou(box1, box2): box: (cx, cy, w, h, angle_deg) r1 ((box1[0], box1[1]), (box1[2], box1[3]), box1[4]) r2 ((box2[0], box2[1]), (box2[2], box2[3]), box2[4]) inter, _ cv2.rotatedRectangleIntersection(r1, r2) if inter is None: return 0.0 inter_area cv2.contourArea(inter) area1 box1[2] * box1[3] area2 box2[2] * box2[3] return inter_area / (area1 area2 - inter_area 1e-7)逻辑说明rotatedRectangleIntersection返回交集多边形的顶点用contourArea算面积。并集面积等于两个框面积之和减交集。加1e-7防止除零。参数说明角度单位是度和minAreaRect输出一致如果训练时角度归一化了后处理要先反归一化回[-90, 90)再传进来。旋转 NMS 比水平 NMS 慢不少推理时如果类别多先按类别分组再各自 NMS。4.3 训练监控看哪些指标判断是否正常水平框训练看box_loss、obj_loss、cls_loss三条曲线OBB 多一条angle_loss。正常收敛的表现是前 10 个 epochbox_loss快速下降angle_loss缓慢下降但不震荡如果angle_loss上下大幅跳动八成是角度定义没统一或归一化漏了。验证阶段重点看mAP0.5和旋转框可视化。我习惯每 10 个 epoch 存一次预测图肉眼扫一遍比盯数字更快发现问题。如果预测框位置对但角度全反检查标注转换时的角度方向如果框大小对但中心偏移检查归一化用的图像尺寸是否和实际一致。5. 旋转框训练避坑清单五个让我返工的血泪经验5.1 现象loss 从第一个 epoch 就不降一直卡在高位原因角度定义在标注转换和损失函数之间不一致比如标注用 OpenCV 的[-90,0)损失按[-90,90)处理角度差被系统性放大。解决在数据集加载后加一段断言打印前 10 个样本的角度范围确认落在预期区间统一用长边定义并在 yaml 里显式写angle_range。5.2 现象训练正常推理时旋转框画出来是歪的原因后处理画框时角度单位或方向搞错常见的是把弧度当度用或者cv2.boxPoints传入的角度顺序和训练时相反。解决写一个最小可视化脚本用已知角度的合成框测试boxPoints输出确认 0 度、45 度、90 度三个case画出来符合预期再接到推理流程里。5.3 现象密集小目标漏检严重mAP 比水平框还低原因旋转 NMS 的 IoU 阈值沿用了水平框的 0.5旋转框重叠面积本来就小阈值偏高会误抑制。解决把旋转 NMS 阈值降到 0.3 到 0.4 之间试同时检查angle_loss权重是否过大导致框位置回归不足。密集场景建议直接上 KLD 损失。5.4 现象数据增强后角度标签错乱原因用了水平框的 mosaic、随机旋转增强旋转图像时没有同步更新角度标签。解决旋转增强必须同步改角度公式是new_angle old_angle rotation_angle再折回[-90,90)mosaic 拼接时每个子图的角度保持不变但中心坐标要按拼接偏移重算。不确定就先把旋转增强关掉跑通 baseline 再逐个开。5.5 现象显存爆掉batch 降到 1 还 OOM原因旋转 IoU 计算在训练中如果放在 GPU 上逐对算中间张量巨大或者img-size设太大。解决训练时的正负样本分配用近似 IoU精确旋转 IoU 只在后处理用img-size从 640 起步确认流程通了再往上加。显存实在紧张就开混合精度--amp。6. 进阶技巧用 KLD 损失和角度预筛选把 mAP 再抬一截直接回归角度在简单场景够用但目标长宽比大、排列密集时边界处角度误差会被放大。我后来固定用 KLD 损失替代 smooth_l1思路是把旋转框(cx, cy, w, h, angle)建模成二维高斯分布均值和协方差由这五个量决定两个框的相似度用两个高斯分布的 KLD 衡量。这样角度和宽高耦合在一起优化边界处的跳变被平滑掉。实现上不用自己推公式核心是把框参数转成高斯分布的(μ, Σ)import torch def obb_to_gaussian(cx, cy, w, h, angle_rad): 把旋转框转成高斯分布的均值与协方差 cos, sin torch.cos(angle_rad), torch.sin(angle_rad) # 旋转后的协方差 R torch.stack([torch.stack([cos, -sin], -1), torch.stack([sin, cos], -1)], -2) S torch.diag_embed(torch.stack([w**2 / 12, h**2 / 12], -1)) Sigma R S R.transpose(-1, -2) mu torch.stack([cx, cy], -1) return mu, Sigma逻辑说明w²/12和h²/12是把矩形框均匀分布近似成高斯时的方差旋转矩阵R把轴对齐的方差转到目标方向。得到两组(μ, Σ)后KLD 有闭式解直接调用即可不用采样。参数说明angle_rad是弧度训练时角度归一化后要乘 π 还原w、h用归一化值还是像素值要统一建议用归一化值避免尺度影响。KLD 损失权重初始设 0.2 左右和 box loss 一起观察。另一个提速技巧是角度预筛选推理时先用水平框 NMS 粗筛把候选框从几千降到几百再对剩下的算旋转 IoU 精筛。这样旋转 IoU 的调用次数下降一个量级实测在 1024 尺寸、单卡上推理速度能快 40% 左右mAP 损失不到 0.5 个点。代价是水平框粗筛可能漏掉极端倾斜的目标如果数据集里这类目标占比高就把粗筛阈值放宽。验证这套改动是否有效别只看最终 mAP。我习惯固定随机种子跑三次看 mAP 的均值和方差方差大的方案说明对初始化敏感上线风险高。另外把角度误差单独统计出来画成直方图如果误差集中在 0 度附近说明回归正常如果双峰分布说明角度周期性没处理好。这些都是我踩过坑之后固定下来的习惯先统一角度定义再动代码先跑通最小闭环再加增强先看可视化再信指标。旋转框训练不难难的是每个环节的角度定义都对得上。希望帮到你。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助？
