简介这份资源面向计算机视觉开发者与目标检测学习者提供基于Python与YOLOv5的旋转目标检测完整实现用于解决倾斜、旋转物体在传统水平边界框下定位不准的问题适用于遥感影像、工业质检、文字检测等场景。压缩包共150个文件约6.26MB以66个Python脚本和33个YAML配置为主涵盖模型训练、推理与参数配置同时包含C、CUDA与Cython源码用于旋转框NMS与多边形IoU的加速计算另有Markdown说明、Shell脚本及Dockerfile辅助环境搭建。已有852人学习下载。资源围绕Oriented Bounding Box展开涉及角度回归、旋转数据增强、GIOU/DIoU损失调整、旋转NMS后处理等关键环节并给出训练、评估与推理的完整流程可帮助读者快速理解并复现旋转目标检测方案适合具备一定深度学习基础、希望深入掌握YOLOv5-OBB的开发者参考。1. 旋转框检测到底难在哪从 yolov5 的轴对齐框说起如果你用 python 和 yolov5 做过常规目标检测大概率遇到过这种场景遥感影像里的飞机、港口里的集装箱、工业质检里的 PCB 板、无人机视角下的锥桶目标本身是斜着的但 yolov5 默认输出的水平框axis-aligned bbox会把大片背景一起框进去。两个相邻的斜向目标水平框重叠率极高NMS 一压就丢一个mAP 直接掉一截。这就是旋转目标检测要解决的问题——让模型直接回归带角度的框oriented bounding boxOBB而不是先检水平框再想办法补救。这个方向适合三类人一是做遥感、航拍、工业视觉的算法工程师目标朝向不固定二是已经跑通过 yolov5 训练自己数据集、想进一步升级检测头的从业者三是做边缘部署比如树莓派、RK3568 这类板子需要轻量旋转检测方案的人。核心思路不复杂保留 yolov5 的 backbone 和 neck把检测头的回归量从 4 个x, y, w, h扩成 5 个x, y, w, h, θ再配套改损失函数、数据标注格式和后处理。但真正落地时角度回归的边界问题、损失函数的周期性、标注工具的配合每一步都有坑。下面按「原理选型 → 数据准备 → 训练调参 → 避坑 → 进阶验证」的顺序拆开讲。2. 旋转检测头的三种改法与 yolov5 的接入点2.1 为什么不能直接在 yolov5 上加一个角度输出最直觉的做法是在 yolov5 检测头的输出通道上加一路让它多回归一个角度 θ。但这样会立刻撞上两个问题。第一角度是有周期性的θ 和 θπ 在几何上是同一个框对于矩形框来说但 L1/L2 损失会把它们当成差异巨大的两个值导致训练震荡。第二yolov5 的 anchor 机制和正负样本分配TaskAlignedAssigner都是基于 IoU 计算的水平框的 IoU 计算方式不能直接套到旋转框上否则正样本匹配会错得离谱。所以正确的接入点不是「加一个通道」这么简单而是要同时改三处检测头的回归维度、正负样本分配时的 IoU 计算、损失函数里的角度项。常见做法是参考 YOLOv5-OBB 或 YOLOv8-OBB 的思路把回归量定义为 (cx, cy, w, h, θ)其中 θ 用弧度表示范围控制在 [-π/2, π/2) 或 [-π/4, 3π/4)具体取哪个区间取决于你的目标长宽比分布。2.2 三种主流角度表示方式的取舍目前工程上常见的角度参数化方式有三种选哪种直接影响训练稳定性和最终精度。表示方式回归量优点缺点适用场景直接回归角度θ弧度实现简单输出直观边界不连续θ 接近 ±π/2 时损失跳变目标朝向集中、长宽比差异小正弦余弦表示(sinθ, cosθ)天然连续无边界问题需要额外归一化推理时要 atan2 还原通用场景推荐首选高斯分布表示(μ, σ)能建模角度不确定性实现复杂调参成本高标注噪声大、角度模糊的目标我一般会优先用正弦余弦表示因为它把角度映射到单位圆上损失函数天然连续不需要额外处理边界。代价是推理时多一步 atan2 计算但这个开销可以忽略。如果你的数据集里目标长宽比接近 1:1比如正方形目标角度本身就没有明确定义这时候用高斯分布或者干脆退回水平框更合理。2.3 在 yolov5 代码里改检测头的具体步骤假设你已经有一份能跑通的 yolov5 代码python 环境配好torch 和 opencv 都装了下面是在检测头里接入角度回归的最小改动。先找到models/yolo.py里的Detect类把输出通道从na * (nc 5)改成na * (nc 6)多出来的那一维就是角度。# models/yolo.py 中 Detect 类的关键改动 class Detect(nn.Module): def __init__(self, nc80, anchors(), ch()): super().__init__() self.nc nc self.no nc 6 # 原来是 nc 5现在多一个角度维度 self.nl len(anchors) self.na len(anchors[0]) // 2 self.grid [torch.zeros(1)] * self.nl self.anchor_grid [torch.zeros(1)] * self.nl self.register_buffer(anchors, torch.tensor(anchors).float().view(self.nl, -1, 2)) self.m nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) def forward(self, x): z [] for i in range(self.nl): x[i] self.m[i](x[i]) bs, _, ny, nx x[i].shape x[i] x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() if not self.training: # 推理时把角度从 sin/cos 还原成弧度 if self.no self.nc 6: sin_val x[i][..., self.nc 4] cos_val x[i][..., self.nc 5] angle torch.atan2(sin_val, cos_val) x[i] torch.cat([x[i][..., :self.nc 4], angle.unsqueeze(-1)], dim-1) z.append(x[i]) return x if self.training else (torch.cat(z, 1),)这段代码的逻辑是训练阶段输出nc 6维类别 cx cy w h sinθ cosθ推理阶段把 sinθ 和 cosθ 用 atan2 还原成角度最终输出nc 5维。参数上要注意self.no必须和损失函数里的维度对齐否则会报 shape mismatch。另外atan2的输出范围是 (-π, π]如果你的标注角度范围是 [-π/2, π/2)需要在后处理里做一次归一化。2.4 正负样本分配要同步改光改检测头还不够yolov5 默认的 TaskAlignedAssigner 用的是水平框 IoU。旋转框的 IoU 计算需要用到多边形相交面积常见做法是用 shapely 库或者自己写一个基于旋转矩形顶点的相交算法。如果这一步不改正样本会匹配到错误的 anchor 上训练 loss 会一直居高不下。import torch import shapely.geometry as sg def rotated_iou(box1, box2): box1, box2: (cx, cy, w, h, angle) 角度为弧度 返回旋转框 IoU def to_polygon(box): cx, cy, w, h, angle box rect sg.box(-w/2, -h/2, w/2, h/2) rect sg.affinity.rotate(rect, angle, origin(0, 0)) rect sg.affinity.translate(rect, cx, cy) return rect poly1 to_polygon(box1) poly2 to_polygon(box2) inter poly1.intersection(poly2).area union poly1.area poly2.area - inter return inter / (union 1e-7)这个函数在样本分配阶段会被频繁调用shapely 的性能在 CPU 上可能成为瓶颈。如果训练速度明显变慢可以考虑用 numba 加速或者换成基于 OpenCV 的cv2.rotatedRectangleIntersection。参数上注意角度单位要统一标注如果是角度制进来之前先转弧度。3. 旋转框数据标注与格式转换的完整链路3.1 标注工具选型与输出格式旋转框标注和水平框标注最大的区别是工具。labelImg 只支持水平框做不了旋转标注。常见做法是用 roLabelImg、X-AnyLabeling 或者 CVAT。roLabelImg 是 labelImg 的旋转框分支输出 XML 格式角度范围是 [-90, 90)适合小规模数据集。X-AnyLabeling 支持更多格式导出包括 DOTA 格式和 COCO 格式的旋转框变体。DOTA 格式是遥感领域最常用的旋转框标注格式每行是x1 y1 x2 y2 x3 y3 x4 y4 category difficult四个点按顺时针排列。COCO 格式的旋转框没有官方标准常见做法是在 annotation 里加一个angle字段或者用segmentation存四个顶点。3.2 从 DOTA 格式转成 yolov5-OBB 训练格式yolov5-OBB 训练时需要的是class cx cy w h angle格式归一化到 [0, 1]。下面这个脚本把 DOTA 的四点格式转成训练格式。import numpy as np import cv2 def dota_to_obb(points, img_w, img_h): points: 四个顶点 [(x1,y1), (x2,y2), (x3,y3), (x4,y4)] 返回归一化的 (cx, cy, w, h, angle) pts np.array(points, dtypenp.float32) rect cv2.minAreaRect(pts) # 返回 ((cx,cy), (w,h), angle) (cx, cy), (w, h), angle rect # OpenCV 的 angle 范围是 (0, 90]转成弧度并归一化到 [-pi/2, pi/2) angle_rad np.deg2rad(angle) if angle_rad np.pi / 2: angle_rad - np.pi if w h: w, h h, w angle_rad np.pi / 2 if angle_rad np.pi / 2: angle_rad - np.pi # 归一化 cx_norm cx / img_w cy_norm cy / img_h w_norm w / img_w h_norm h / img_h return cx_norm, cy_norm, w_norm, h_norm, angle_rad这里有几个关键点。cv2.minAreaRect返回的角度范围是 (0, 90]需要手动映射到 [-π/2, π/2)。另外 w 和 h 的顺序会影响角度定义必须和训练时的约定一致。我一般约定 w 是长边这样角度范围更集中训练更稳。转换完之后建议随机抽几十张可视化检查确认框和角度都对得上。3.3 数据集目录结构与 yaml 配置yolov5 的数据集目录结构不用大改只是 label 文件里每行多一个角度值。yaml 配置文件里需要额外指定obb: True或者类似的标志位具体取决于你用的 OBB 分支版本。# data/rotated_dataset.yaml path: ../datasets/rotated train: images/train val: images/val test: images/test nc: 5 names: [plane, ship, storage-tank, baseball-diamond, tennis-court] # 旋转框标志 obb: True angle_range: [-1.5708, 1.5708] # [-pi/2, pi/2)angle_range这个参数不是 yolov5 原生的需要你在数据加载代码里自己读取并做校验。如果标注角度超出这个范围要么截断要么报错不要让它悄悄传进去否则训练 loss 会出现莫名其妙的尖峰。4. 训练调参与旋转框特有的损失函数设置4.1 损失函数的三个组成部分怎么调旋转框检测的损失函数一般由三部分组成分类损失BCE、框回归损失CIoU 或旋转 IoU、角度损失。角度损失的设计是核心难点。如果用正弦余弦表示角度损失可以写成def angle_loss(pred_sin, pred_cos, target_angle): pred_sin, pred_cos: 模型输出的 sin 和 cos target_angle: 标注角度弧度 target_sin torch.sin(target_angle) target_cos torch.cos(target_angle) # 用 L1 或者 smooth L1 都可以关键是 sin/cos 天然连续 loss F.smooth_l1_loss(pred_sin, target_sin) F.smooth_l1_loss(pred_cos, target_cos) return loss这个损失的好处是不需要处理角度边界sin 和 cos 在单位圆上连续变化。但要注意sin 和 cos 不是独立的预测出来的 (sin, cos) 模长可能不等于 1推理时用 atan2 还原角度不受影响但训练时最好加一个归一化约束或者直接用 atan2 后的角度算损失。我一般会在 loss 里加一项(sin^2 cos^2 - 1)^2的惩罚权重设小一点0.01 左右。4.2 学习率和 anchor 的调整经验旋转框检测的学习率通常比水平框要小。水平框训练时 lr0 设 0.01 很常见但旋转框我一般从 0.001 开始因为角度回归对参数扰动更敏感。如果 loss 在前 10 个 epoch 下降很慢可以适当提高到 0.005但不要超过 0.01。anchor 方面yolov5 默认的 anchor 是基于 COCO 水平框聚类的直接拿来用效果一般。建议用自己的旋转框数据集重新跑一遍 k-means 聚类聚类的距离度量用旋转框的 IoU 而不是宽高比。聚类脚本可以基于 sklearn 的 KMeans 改把距离函数换成1 - rotated_iou。from sklearn.cluster import KMeans import numpy as np def cluster_anchors(boxes, k9): boxes: N x 5 的数组 (cx, cy, w, h, angle) 返回 k 个 anchor 的 (w, h) wh boxes[:, 2:4] # 只用宽高做聚类 kmeans KMeans(n_clustersk, random_state42).fit(wh) anchors kmeans.cluster_centers_ # 按面积排序 anchors anchors[np.argsort(anchors[:, 0] * anchors[:, 1])] return anchors聚类完把 anchor 写进 yaml 配置文件替换掉默认的 COCO anchor。这一步对最终 mAP 的影响通常在 2 到 5 个点值得花时间做。4.3 训练命令与关键超参数假设代码已经改好数据集也准备好了训练命令和标准 yolov5 差不多只是多指定一些旋转框相关的参数。python train.py \ --data data/rotated_dataset.yaml \ --cfg models/yolov5s-obb.yaml \ --weights yolov5s.pt \ --epochs 300 \ --batch-size 16 \ --img-size 1024 \ --lr0 0.001 \ --lrf 0.01 \ --cos-lr \ --obb \ --angle-loss-weight 0.5--obb是自定义参数用来告诉训练脚本启用旋转框分支。--angle-loss-weight控制角度损失在总损失里的权重默认 0.5 是个比较稳的起点。如果发现框的位置回归得不错但角度总是差一点可以把这个权重提到 1.0。--img-size建议设大一点旋转框在小分辨率下角度信息容易丢失1024 是遥感场景的常用值。提示训练前先用少量数据比如 50 张图跑 10 个 epoch确认 loss 能正常下降、可视化结果框的角度方向正确再上全量数据。这一步能省掉很多事后排查的时间。5. 旋转框检测的避坑与排查清单5.1 角度周期性导致的 loss 震荡现象训练 loss 在某个值附近反复跳变降不下去可视化发现部分框的角度预测在 ±90 度附近来回翻转。原因标注角度范围没有统一或者损失函数直接用了角度差而没有处理周期性。比如标注是 [-90, 90)但模型预测出 91 度直接算差值会得到一个巨大的 loss。解决统一用正弦余弦表示角度损失函数在 sin/cos 空间计算。如果必须用角度直接回归确保在 loss 里做角度归一化把差值映射到 [-π/2, π/2) 再算。5.2 旋转 IoU 计算太慢拖垮训练速度现象换成旋转框后每个 epoch 的训练时间从几分钟涨到几十分钟GPU 利用率很低。原因正负样本分配阶段调用了 Python 层面的 shapely 计算旋转 IoU每次前向都要算几千次CPU 成为瓶颈。解决把旋转 IoU 计算移到 GPU 上用张量操作实现或者用cv2.rotatedRectangleIntersection替代 shapely。另一个思路是减少参与 IoU 计算的候选框数量先用水平框 IoU 粗筛一遍再对 top-k 算旋转 IoU。5.3 标注角度方向与代码约定不一致现象训练 loss 正常下降但推理可视化时框的位置对了、角度全反了或者框的长宽对调了。原因标注工具的角度定义比如从 x 轴正方向逆时针和代码里的角度定义比如从 x 轴正方向顺时针不一致或者 w/h 的定义顺序不同。解决在数据加载阶段加一个可视化调试开关随机抽 20 张图把标注框画出来人工确认角度方向。这一步不要省我见过太多因为角度约定不一致白白训了一周的情况。5.4 小目标旋转框的回归精度差现象大目标的旋转框检测效果不错但小目标比如小于 32x32 像素的角度误差很大mAP 比水平框还低。原因小目标本身像素少角度信息在卷积和下采样过程中丢失严重。另外小目标的旋转 IoU 对角度误差更敏感同样的角度偏差在小目标上导致的 IoU 下降更大。解决增大输入分辨率或者用 P2 层更高分辨率的特征图做检测。损失函数里对小目标的角度损失加权让模型更关注小目标的角度回归。如果还是不行考虑对小目标退回水平框检测只对大目标做旋转框。5.5 导出 ONNX 后角度输出不对现象PyTorch 里推理正常导出 ONNX 后用 onnxruntime 推理角度输出全是 0 或者乱码。原因atan2在 ONNX 里的支持取决于 opset 版本低版本 opset 可能不支持。另外如果推理时的后处理逻辑写在 Python 里而没有导出到 ONNX 图中ONNX 模型输出的还是 sin/cos 原始值。解决导出时指定 opset 版本不低于 11确保 atan2 被支持。或者干脆不在 ONNX 里做 atan2让 ONNX 输出 sin/cos在后处理代码里还原角度。导出后用onnxruntime跑一张测试图和 PyTorch 的输出逐元素对比误差应该在 1e-4 以内。6. 旋转框检测的验证方法与一个后处理技巧训练完之后怎么确认模型真的学到了角度信息而不是在「假装」旋转检测我一般会做三个验证。第一把验证集里所有预测框的角度和标注角度画成散点图如果点集中在 yx 对角线附近说明角度回归是有效的如果散成一团说明角度分支没学到东西。第二单独统计不同角度区间的 mAP比如 0-30 度、30-60 度、60-90 度看模型是否在某些角度区间明显偏弱。第三拿几张训练集里没有的角度组合比如 45 度左右的密集目标做可视化人眼确认框的贴合程度。后处理阶段有一个容易被忽略的技巧角度平滑。旋转框的 NMS 和水平框不同两个角度差很小的框可能 IoU 很高但角度差很大的框 IoU 也可能不低因为矩形对称性。我一般会在 NMS 之前先对角度做一次聚类把角度差小于 5 度的框归为一组组内做加权平均然后再跑旋转 NMS。这个操作在密集场景下能减少框的抖动让最终输出更稳定。def angle_smooth(boxes, angle_thresh5.0): boxes: N x 6 (cx, cy, w, h, angle, score) 对角度接近的框做加权平均 boxes boxes[np.argsort(-boxes[:, 5])] # 按 score 降序 keep [] used np.zeros(len(boxes), dtypebool) for i in range(len(boxes)): if used[i]: continue group [i] for j in range(i 1, len(boxes)): if used[j]: continue angle_diff abs(boxes[i, 4] - boxes[j, 4]) angle_diff min(angle_diff, np.pi - angle_diff) # 处理周期性 if angle_diff np.deg2rad(angle_thresh): group.append(j) used[j] True # 加权平均 weights boxes[group, 5] / boxes[group, 5].sum() avg_box np.average(boxes[group, :5], axis0, weightsweights) avg_box np.append(avg_box, boxes[group, 5].max()) keep.append(avg_box) return np.array(keep)这个函数的参数angle_thresh控制角度聚类的粒度默认 5 度。如果目标角度分布很密集可以调到 3 度如果目标角度差异大调到 10 度。注意加权平均时角度不能直接线性平均严格来说应该在 sin/cos 空间做平均再还原但 5 度以内的差异用线性平均误差可以忽略工程上够用。最后说一个我踩过的坑旋转框检测的 mAP 计算和水平框不一样很多现成的评估脚本直接套用水平框的 IoU 计算结果虚高。一定要确认评估代码里用的是旋转 IoU。我一般会自己写一个简单的评估脚本用cv2.rotatedRectangleIntersection算 IoU和训练框架自带的评估结果交叉验证。如果两者差超过 3 个点大概率是其中一方的 IoU 计算有问题。这套方案从改检测头到训练再到后处理完整跑通大概需要两到三周其中数据标注和格式转换占一半时间。如果你的数据集角度分布比较集中一周内能出第一版结果。希望帮到你。本文还有配套的精品资源点击获取