简介YOLOv5 OBB旋转框训练DEMO是一份面向目标检测学习者的完整示例工程解决倾斜目标如树木、车辆检测中的旋转框训练问题。资源共583个文件约445.58MB包含166个xml标注、95个txt标签、93个png和23个jpg图像素材、71个py脚本、38个yaml配置、6个pt预训练权重以及训练、预测、导出、wts生成等cmd命令覆盖从数据预处理、损失函数配置、模型训练到评估可视化的全流程。工程中内置md说明文档和ipynb交互示例并配套dockerfile与makefile便于环境复现训练环节涉及OBB坐标与角度标注转换、模型结构配置、旋转框损失计算等关键技术点适合已有YOLO基础、希望扩展旋转框检测能力的中高级开发者。已有1812人学习通过该demo可直观理解OBB边界框表示复现训练闭环并借助poly_nms等扩展模块完成后处理、ONNX导出与部署迁移。1. 旋转框检测到底解决了什么问题一个OBB训练demo的价值边界遥感图里停着十几架朝向各异的飞机工业传送带上的零件横七竖八停车场里车辆斜着排成三排——这种场景拿普通YOLOv5的水平框去框结果要么是框得松松垮垮把邻近目标一起包进去要么是两个斜目标因为水平框重叠率过高被NMS误杀一个。旋转框OBBOriented Bounding Box就是用来解决这类“目标本身带角度”的检测问题的。这份YOLOv5 OBB旋转框训练demo是把YOLOv5从水平框检测改造成可输出带角度预测框的完整训练示例包含数据格式转换、网络头改动、训练参数和后处理全链路。适合正在做遥感目标检测、工业质检、无人机视角识别的人也适合想搞懂旋转框标注格式和角度回归原理的开发者。后面所有内容都围绕一份能真正跑通的demo展开讲清楚每一个参数和每一个坑。2. 从DOTA标注到YOLO格式旋转框参数化与转换脚本2.1 为什么不能直接套用水平框的(x, y, w, h)YOLO水平框检测的标注是四个值中心点坐标、宽、高。这个表示在目标不带角度时完全够用因为水平框无论怎么放都是“上平下平”唯一变量只有位置和大小。但旋转目标不一样同样一个长条目标水平放置和斜45度放置如果继续用x、y、w、h表示那么不同角度的同类目标会对应完全不同的宽高比特征模型要额外去学“这个宽高比和那个宽高比其实是同一个目标转了一下”浪费模型容量还学不干净。旋转框的常见参数化方式有三种四点坐标、中心点加宽高加角度、八个参数四条边的直线方程。四点坐标最直观但回归量多且四点之间顺序容易错乱八参数主要用在某些遥感检测框架里数学上更稳但工程实现复杂。YOLOv5这类单阶段检测器改造时几乎都选“中心点、宽、高、角度”的五元组方案因为原有水平框的回归头是现成的只需要在输出维度上增加一个角度通道损失函数多接一个角度分支改动成本最低。要注意角度不是随便定个范围就行的。我见过好几个OBB demo翻车根源都是角度定义没对齐训练代码里角度范围是0到180度但标注转换脚本生成的角度在-90到0度模型输出经过sigmoid后无论如何也不会落到-90度区间loss永远降不下去。所以拿到一个OBB demo第一件事就是去源码里找角度归一化函数看清楚它到底接受什么范围、以哪条边作为参考边。2.2 DOTA转YOLO-OBB按最小外接矩形生成五点标注DOTA数据集是遥感旋转框检测用得最多的标注格式每行是四个顶点的x、y坐标、类别和难易度。但YOLOv5 OBB训练需要的是class_id cx cy w h angle这种紧凑格式。转换思路很简单用OpenCV的cv2.minAreaRect求四个顶点的最小外接矩形拿到中心、宽高和角度。这里面有个容易踩的细节minAreaRect返回的角度是width边与x轴的夹角范围是[-90, 0)而且width不一定是长边。为了统一必须把长边作为w并把角度换算到模型约定的范围。import cv2 import numpy as np def dota_to_yolo_obb(label_path, img_w, img_h, class_map): DOTA标注转YOLO-OBB五点格式 DOTA一行: x1 y1 x2 y2 x3 y3 x4 y4 class_name difficult 输出一行: class_id cx cy w h angle(弧度) results [] with open(label_path, r) as f: for line in f: parts line.strip().split() if len(parts) 9: continue # DOTA四点坐标读进来顺序不影响minAreaRect pts np.array([ [float(parts[0]), float(parts[1])], [float(parts[2]), float(parts[3])], [float(parts[4]), float(parts[5])], [float(parts[6]), float(parts[7])], ], dtypenp.float32) rect cv2.minAreaRect(pts) (cx, cy), (w, h), angle rect # 统一把长边放到w上angle换算到[0, 90) if w h: w, h h, w angle angle 90.0 # OpenCV的angle范围[-90, 0) - [0, 90) if angle 90.0: angle - 90.0 # 归一化到相对坐标 cx_n cx / img_w cy_n cy / img_h w_n w / img_w h_n h / img_h angle_r angle * np.pi / 180.0 cls_id class_map[parts[8]] results.append(f{cls_id} {cx_n:.6f} {cy_n:.6f} {w_n:.6f} {h_n:.6f} {angle_r:.6f}) return results这个脚本的核心逻辑是先由四个顶点求出最小外接矩形再做两步标准化。第一步是交换宽高保证w始终是长边这样模型回归时面对的物理意义是一致的第二步是把OpenCV的负角度加90度转成0到90度配合demo源码里的角度归一化。角度用弧度是因为很多OBB实现的损失函数里直接对弧度值做smooth L1和标注单位保持一致可以省去后处理时反复换算的麻烦。拿到demo代码包后建议先看一下它的label文件样例有的实现是cx cy w h angle class_id有的实现是class_id cx cy w h angle。我这份脚本输出的是后者和YOLOv5原版的class cx cy w h风格一致如果你的demo是前者把输出行的字段顺序调一下就行。2.3 目录结构与数据集划分demo先跑通再谈全量旋转框数据的目录结构不需要特殊设计和原版YOLOv5保持一致即可。常见布局如下路径内容datasets/demo/images/train/训练图片datasets/demo/images/val/验证图片datasets/demo/labels/train/训练集OBB标注txtdatasets/demo/labels/val/验证集OBB标注txtdatasets/demo/train.txt训练图片路径列表每行一个绝对路径datasets/demo/val.txt验证图片路径列表train.txt和val.txt里写绝对路径还是相对路径取决于datasets.py里怎么解析。我一般写绝对路径少一层坑如果你要换机器复现写相对路径更灵活但前提是代码里做了路径拼接。这个细节看起来小实际遇到过因为路径多了一层导致图片全没加载进去、训练了两百轮还在空跑的情况。数据集划分上有一个明确建议第一次跑demo不要拿全量DOTA也别拿所有类别。挑一个类别、几百张图先跑通训练、验证、推理、画图这一整套流程确认角度定义、loss曲线、后处理都没有问题再扩展到全量数据。这样定位问题的时间能缩短一个数量级。另一点是类别序号必须在data.yaml里从0开始连续编号旋转框的类别映射一旦跳过数字会直接索引越界。3. 网络头与训练参数在哪里把角度回归接进去3.1 Detect head新增angle分支后的输出结构原版YOLOv5的Detect head在每一个尺度上输出的通道数是(5 num_classes) * num_anchors5代表边界框的4个回归量加1个objectness。OBB改造后每个anchor需要额外回归一个角度值于是输出通道数变成(6 num_classes) * num_anchors新加的那个维度就是角度。角度回归的处理方式有两种主流方案一种是直接回归角度标量网络输出经过sigmoid后乘以角度范围比如映射到[0, 180)另一种是把角度编码成sin和cos两个值各占一个通道。直接回归标量实现简单、收敛快缺点是有角度周期性跳变的问题——179度和1度在数值上差178但实际只差2度普通L1损失会把这种“其实很近”的样本当成巨大误差。sin/cos编码能消除周期性问题但两个值之间天然存在相关性损失边界上反而可能更抖。demo级别的训练一般用直接回归就够了因为标签角度已经统一到[0, 90)周期跳变的概率被局限在0度附近的小范围。修改完head之后输出的tensor形状会从(batch, anchors, 5 num_classes)变成(batch, anchors, 6 num_classes)。这一步改完后数据加载器、损失函数、后处理三处都要跟着动数据加载器要解析第6个维度损失函数要新增一个angle loss分支后处理要把角度和宽高一起反算成四个顶点。3.2 训练一条命令跑起来参数表与超参注意demo的训练命令和原版YOLOv5几乎完全相同只不过数据集配置换成了OBB版本模型权重也建议改用OBB分支对应的预训练权重或者从原版yolov5s.pt迁移。python train.py \ --img 1024 \ --batch 8 \ --epochs 100 \ --data demo.yaml \ --hyp hyp.scratch-low.yaml \ --weights yolov5s.pt \ --device 0几个常见参数的含义和注意点我按实际经验列一下参数建议初始值说明--img1024遥感目标通常小且密集输入尺寸低于768会丢失大量小目标角度信息--batch尽量大旋转框demo数据量一般不大batch 8到16是常见起步值--epochs100先跑100轮看loss趋势不要一上来就跑300轮--hyp lr00.01如果loss曲线震荡剧烈降到0.001再试--hyp hyp.scratch-low.yaml低增强配置带角度任务不建议直接上高增强hyp后面会说为什么--weights yolov5s.pt迁移权重用COCO预训练权重做迁移绝大多数情况优于随机初始化注意一个细节--hyp指向的配置文件里默认的mosaic和mixup增强都是开的。OBB训练场景下mosaic拼接大图时框的角度标注不会自动跟着图像变换走尤其是水平翻转增强角度标注如果不同步做负号处理等于给模型喂错误标签。所以第一次训练我建议手动把hyp.scratch-low.yaml里的mosaic: 0.0、mixup: 0.0、fliplr: 0.0都关掉等模型稳定收敛之后再逐步打开。3.3 loss权重怎么调angle loss不是越大越好OBB损失函数通常是在原版分类损失和回归损失基础上叠加一个角度回归损失。问题在于角度loss的权重设置很敏感默认值未必适合你的数据集。如果angle loss权重设得太大模型会优先把角度学好但目标分类和定位反而被忽略最终表现为框方向很准、类别置信度一塌糊涂如果设得太小模型会觉得学角度“不划算”最后输出一堆接近水平的长条框。我处理过的一个demo默认angle loss权重是1.0训练后类别mAP掉了8个点。调到0.1之后角度精度基本没变分类和定位都恢复了。这里给一个经验区间不过还是以你的demo源码里默认值作为基准做微调分支权重参考范围观察指标box loss0.05训练集回归误差是否收敛cls loss0.5 ~ 1.0类别mAP是否正常angle loss0.05 ~ 0.2验证集角度偏差和loss曲线判断角度是否学到了不能只看训练loss要在训练中途就用验证脚本画出预测框直观对比预测和真实标注的长边方向。如果角度方向经常和真实框相差90度那多半是长边定义和数据格式没对齐不是权重问题。4. 后处理要不翻车旋转框解码、NMS与评测口径4.1 解码输出boxPoints反算四个顶点模型输出的仍然是(cx, cy, w, h, angle)这种紧凑形式但可视化、计算IoU、提交评测都需要四个顶点坐标。最省事的做法是直接用cv2.boxPoints它接收中心点、宽高和角度返回旋转矩形的四个顶点。import cv2 import numpy as np def decode_obb(pred, angle_in_degreeTrue): pred: (cx, cy, w, h, angle) 返回: 四边形四个顶点坐标是像素值 cx, cy, w, h, angle pred[:5] if not angle_in_degree: angle angle * 180.0 / np.pi rect ((cx, cy), (w, h), angle) box cv2.boxPoints(rect) # 返回4x2 float32 return box.astype(np.int32)这里最需要注意的就是角度单位。训练时如果标签是弧度模型输出的角度分支经过非线性激活后自然也是弧度但cv2.boxPoints只认度数而且旋转方向是顺时针。角度单位搞错的结果就是画出来的框要么全部偏转一个固定角度要么在目标位置附近疯狂打转。我通常会在解码函数里强制加一个单位参数训练和推理都走同一入口避免两个模块各写各的换算逻辑。4.2 OBB-NMS为什么是刚需普通YOLOv5的NMS用水平框IoU作为重叠判断依据。旋转框场景下这个逻辑直接失效两个斜45度放置的长框中心接近、水平范围几乎完全重叠但实际旋转框IoU可能只有0.3。如果按水平IoU做NMS置信度低的那一个会被错误地抑制掉长条目标密集排列时甚至会整排消失。OBB-NMS的核心就是把IoU计算从水平框换成旋转框。实现方式有两种一种是用shapely直接对多边形求交并面积另一种是用旋转框专用IoU算子加速。shapely版本够直观适合验证和调试但速度慢几百个框跑一次NMS要几十毫秒不适合实时推理。from shapely.geometry import Polygon import numpy as np def obb_iou(box_a, box_b): box_a, box_b: 四个顶点坐标, shape (4, 2) poly_a Polygon(box_a) poly_b Polygon(box_b) if not poly_a.is_valid or not poly_b.is_valid: return 0.0 inter poly_a.intersection(poly_b).area union poly_a.area poly_b.area - inter if union 0: return 0.0 return inter / union def obb_nms(dets, iou_threshold0.5): dets: list of dict, 每个包含 vertices(4,2), score, cls_id dets sorted(dets, keylambda x: x[score], reverseTrue) keep [] while dets: best dets.pop(0) keep.append(best) dets [ det for det in dets if not (det[cls_id] best[cls_id] and obb_iou(best[vertices], det[vertices]) iou_threshold) ] return keep这个示例只体现核心逻辑真实工程里会用C或CUDA算子把旋转IoU的计算提速但算法骨架是一样的按置信度排序、逐类计算、IoU超阈值就抑制。如果你的demo代码里实现的还是水平框NMS训练结果再好也白搭务必先确认后处理这一段有没有换成旋转框版本。4.3 评测与可视化本地验证角度偏差评测也不能照搬水平框那套。遥感旋转框的标准评测指标是DOTA mAP判定一个预测框是否算正样本时要求预测框与真实框的旋转IoU超过阈值比如0.5且角度差不能超过设定阈值。后者是关键区别——有的预测框旋转IoU达标但长边方向反了180度按DOTA规则也是错的。本地快速验证建议分两步走。第一步是可视化随机抽20张验证集图片真实框和预测框一起画出目测角度方向是否正确这步能发现角度定义类问题第二步是量化统计预测框与真实框之间的角度绝对差计算偏差小于15度的比例。这个比例如果低于80%说明角度回归还没收敛先去检查增强配置和angle loss权重而不是急着调NMS阈值。5. 避坑实录角度定义、增强与评测的五个坑5.1 角度定义不一致loss降了但输出框旋转方向全乱现象训练loss正常下降验证集mAP看着也不低但可视化时预测框要么整体偏转一个固定角度要么长边方向和真实框完全是垂直的。原因训练代码的角度归一化函数和数据转换脚本的角度定义不一致。举例来说训练代码认为角度0到90度是长边绕x轴顺时针旋转的角度但你的转换脚本生成的是逆时针角度等于标签整体被镜像了一个角度偏移。模型确实学到了一个规律的映射但这个映射和真实标注对不上。解决打开demo源码找到角度归一化函数确认它接受的范围、参考轴方向然后把转换脚本里的角度换算逻辑对齐过去。有一个简单验证方法拿一张只有一个目标的图跑一次前向推理把预测角度和真实角度打印出来对比。如果预测角和真实角呈现一个固定的差值那基本就是这个问题。从那以后我拿到任何一个OBB demo都会先跑这个单目标验证。5.2 翻转和Mosaic增强偷改角度训练集指标与验证集mAP严重分层现象训练集loss非常低准确率接近100%到验证集上暴跌且越训越严重。原因YOLOv5默认开启的水平翻转增强对水平框完全无感水平框左右翻转后宽高不变、类别不变只是中心点的x坐标变了标签不用做任何额外处理。但OBB不一样水平翻转图像后旋转框的顶点顺序会反转角度必须同步取负或者做周期映射。mosaic拼接虽然没有翻转图像但如果拼接过程中有随机旋转或者仿射变换角度标注同样失效。很多demo的datasets.py里没有实现角度同步变换等于一部分训练样本的标注角度是错的模型被迫去拟合脏标签。解决最稳妥的方式是在第一个版本训练时把fliplr、mosaic、mixup全部关掉只保留随机缩放和色彩抖动这类不影响角度的增强。等模型角度回归稳定后再逐项打开并验证。如果非要在增强条件下训练必须确认代码里翻转时对角度做了angle -angle这样的处理。5.3 极端长宽比与极小目标角度回归震荡甚至NaN现象训练过程中angle loss突然跳到极大值或者loss曲线持续震荡然后出现NaN之后整个模型输出全变成垃圾值。原因细长目标的旋转框对角度极其敏感长宽比10比1的目标旋转1度dIoU就能变化几个百分点梯度方向剧烈变化再加上极小目标占的像素少角度信息原本就弱SGD在这种高曲率区域很容易震荡。解决三个手段配合使用。第一把angle loss换成smooth L1降低离群点带来的梯度冲击。第二对极端长宽比样本降权常见做法是按w/h比值设定一个权重系数比值越大权重越小。第三EMA指数移动平均能帮助稳定训练波动如果demo默认没开建议手动开启YOLOv5里对应--ema参数或配置文件里的开关。还有一个保底手段是关闭AMP混合精度FP16下角度梯度的精度损失可能直接导致NaN这个在下一个坑里展开。5.4 角度类别不平衡模型选择“抄近路”把所有框学成水平现象训练结束后预测角度集中在0度附近旋转目标全被框成接近水平的矩形。原因这是我在一个停车场车辆数据集上遇到的典型问题——绝大多数样本是水平停放的车辆斜着停的只有一小部分。模型发现把全部角度预测成0度就能降低大部分loss角度分支反正贡献不了多少收益干脆偷懒。这种情况不是模型坏了是数据分布和loss权重设计共同导致的。解决方向有两个。数据侧对角度做直方图统计如果某个角度区间样本过少对这类样本做在线旋转增强每个训练轮次随机旋转一个角度再喂给模型。loss侧把angle loss权重适当增大让模型意识到不学角度会有明显惩罚。我曾经把angle loss权重从0.1调到0.3斜停车辆的检出率从31%涨到58%水平框方向的指标只掉了1个点。5.5 AMP与显存坑FP16下角度梯度不稳定现象同样的数据和超参开AMP不到几十轮loss变成NaN关掉AMP一路正常。原因原生YOLOv5默认开启AMP混合精度。水平框回归的四维目标量级都在0到1之间FP16能扛住但角度值可能是弧度制范围0到π而部分loss实现里对角度做差后平方数值范围可以放大好几倍半精度浮点数的表示范围很容易溢出。解决训练OBB时优先关掉AMP用--amp false或者代码里配置关闭。显存不够不要靠AMP省应该降batch或者降输入分辨率。如果你必须用AMP那就把angle loss的权重调低并且监控训练日志里angle分支的梯度统计发现异常立刻回退。6. 把demo验证扎实可视化与两阶段训练6.1 先画图后评估GT与预测叠加检查训练到一半我会停下来跑一次推理画一批训练集和验证集的对比图。画图的代码不用复杂import cv2 def draw_obb(img, box_points, color(0, 255, 0), labelNone): cv2.drawContours(img, [box_points.astype(np.int32)], -1, color, 2) if label: x, y box_points[0].astype(int) cv2.putText(img, label, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)把真实框画成绿色预测框画成红色重点检查三件事长边方向是否一致、角度偏差是否稳定、密集场景下NMS有没有误杀。这组图比任何指标都先暴露问题。只要看到预测框整体长边方向和真实框差90度就可以直接判定是角度定义问题不用再去看mAP曲线。6.2 两阶段训练先冻结backbone再解锁精调OBB任务一个不太好直接迁移的点是COCO预训练权重对水平特征很熟练但角度相关的特征完全没有。一种我在实际中反复用且有效的训练节奏是第一阶段冻结backbone只训练head的检测和角度分支第二阶段解锁全部权重用低学习率精调。# 第一阶段冻结前10层只训检测头和角度分支 python train.py \ --data demo.yaml \ --epochs 50 \ --batch 8 \ --img 1024 \ --weights yolov5s.pt \ --freeze 10 # 第二阶段解锁所有层小学习率精调 python train.py \ --data demo.yaml \ --epochs 100 \ --batch 8 \ --img 1024 \ --weights runs/train/exp/weights/best.pt \ --hyp hyp.finetune.yaml第一阶段让角度分支快速找到正确方向第二阶段再让backbone适应旋转特征。这样调出来的模型比一次全量训练更稳尤其是小数据集上能明显减少抖动量级。我最后一次完整跑通一个OBB demo时就是先冻结10层跑了50轮再解锁全量跑了100轮最终验证集角度偏差均值从11.2度降到5.8度密集停放车辆的漏检率也降了一截。从那以后我每次拿到旋转框项目第一轮训练都强制走一遍冻结再解锁的流程先让框学会转再让模型学会看。希望帮到你少走这几个弯路这个demo跑通就只是时间问题了。本文还有配套的精品资源点击获取