简介这份资源面向具备一定深度学习基础的计算机视觉开发者与算法工程师提供基于Python与YOLOv5的旋转目标检测完整实现用于解决传统水平边界框难以精确框定倾斜、旋转物体的痛点可应用于遥感影像、航拍、工业质检等场景。压缩包共150个文件约6.26MB以66个Python脚本和33个YAML配置为主涵盖模型定义、训练与推理流程同时包含CUDA与C源码用于旋转框NMS及多边形IoU的加速计算另有少量文档、Shell脚本与权重说明文件结构清晰便于二次开发。资源围绕Oriented Bounding Box展开涉及角度回归、旋转数据增强、GIOU/DIoU损失调整及角度感知NMS等关键环节读者可据此搭建训练、评估到推理的完整链路。目前已有852人学习下载适合希望将YOLOv5迁移到旋转检测任务的中高级开发者参考。1. 旋转框检测当 YOLOv5 遇上带角度的目标做遥感影像或者工业质检的同行大概率遇到过这个场景明明标注框把舰船、飞机、文字条包得严严实实模型训练 loss 也降下去了可一到推理密集排列的细长目标就开始互相压框NMS 一过要么漏检要么框歪。水平框的先天缺陷就在这里——它假设目标轴向对齐而现实里大量目标是有朝向的。基于 python 的 yolov5 实现的旋转目标检测本质就是给 YOLOv5 的检测头加一个角度回归分支让输出从 (x, y, w, h) 变成 (x, y, w, h, θ)再用旋转 NMS 替代普通 NMS。这套方案适合已经跑通过 YOLOv5 水平检测、手里有 DOTA 或自建旋转标注数据、想用最小改动拿到可用旋转框的工程师。它不推翻 YOLOv5 的主干和 Neck改动集中在 Head、损失和标签分配三处这也是它比换一套全新框架更值得投入的原因。2. 角度回归到底加在哪从水平头到旋转头的改造逻辑2.1 为什么不是简单多输出一个数很多人第一反应是在 YOLOv5 的 85 维输出后面再拼一维角度不就完了。真这么做训练初期就会翻车。原因在于角度的周期性——179° 和 -179° 在几何上几乎一样但在 L1 或 MSE 损失里差了 358梯度会把回归头带偏。所以旋转检测的角度处理必须解决两个问题周期性和边界突变。常见做法有三类。第一类是把角度当分类问题切成 180 个 bin 做交叉熵缺点是精度受 bin 宽度限制。第二类是用正弦编码把 θ 映射成 (sinθ, cosθ) 两个连续值回归天然消除周期性这也是我一般会推荐给新手的方案改动小、稳定。第三类是基于高斯分布的表示把旋转框建模成二维高斯用 KL 散度做损失精度高但实现复杂适合有经验后再上。选型上如果你的数据里目标长宽比普遍大于 3、角度分布连续正弦编码足够用如果目标接近正方形角度本身定义就模糊这时候硬回归角度反而有害得考虑换成基于点的表示。2.2 检测头输出的维度变化以 YOLOv5s 为例原版每个 anchor 输出5 num_classesCOCO 是 85。旋转版要变成6 num_classes多出来的一维是角度。但如果你用正弦编码实际要多两维变成7 num_classes因为要输出 sin 和 cos。下面是我改 Head 时常用的输出解析片段基于 YOLOv5 的Detect层思路import torch import math class RotateDetectHead(torch.nn.Module): def __init__(self, nc80, anchors(), ch()): super().__init__() self.nc nc self.no nc 5 2 # 5 是 xywhobj2 是 sin/cos self.nl len(anchors) self.na len(anchors[0]) // 2 self.m torch.nn.ModuleList( torch.nn.Conv2d(x, self.no * self.na, 1) for x in ch ) def forward(self, x): for i in range(self.nl): x[i] self.m[i](x[i]) bs, _, ny, nx x[i].shape # 重排成 [bs, na, ny, nx, no] x[i] x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2) return x逻辑说明self.no从 85 变成 87多出的两维就是 sin/cos。重排是为了后续 decode 时按 anchor 维度取用。参数上anchors沿用 YOLOv5 默认的 9 个聚类框即可旋转检测对 anchor 形状不敏感因为角度分支会补偿朝向但如果你数据里目标尺度跨度极大建议重新用 kmeans 聚一次把长宽比纳入距离度量。2.3 标签分配要跟着改YOLOv5 用的是跨网格的标签分配正样本判定基于宽高比和中心距离。旋转框下中心点没变但 IoU 计算要换成旋转 IoU。如果你还用水平 IoU 做分配密集场景里两个朝向不同的框会被判成高重叠正样本互相污染。我一般会保留 YOLOv5 的分配框架只把 IoU 计算替换成旋转版本。旋转 IoU 没有闭式解工程上用近似把两个旋转框各采样成多边形调cv2.rotatedRectangleIntersection求交再算面积比。这个函数在 OpenCV 里现成速度够用训练时每个 batch 多花的时间在可接受范围。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)参数说明角度单位统一用度OpenCV 的rotatedRectangleIntersection要求角度制。inter返回的是交集多边形顶点用contourArea算面积。注意这个函数在框完全包含或完全分离时有边界行为加个1e-7防止除零。如果你的数据量很大这个逐对计算会成为瓶颈可以先用水平 IoU 粗筛只对水平 IoU 大于 0.1 的框对算旋转 IoU。3. 用 python 把旋转 YOLOv5 在本地跑通的最小路径3.1 环境与依赖的确定热词里 python 安装、python 环境变量配置、vscode python 环境配置出现频率很高说明不少人是卡在环境上的。旋转检测比普通检测多一个依赖OpenCV 的 contrib 版本因为rotatedRectangleIntersection在标准opencv-python里也有但如果你要用cv2.boxPoints做可视化contrib 更稳。我一般用 conda 建环境避免和系统 python 打架conda create -n rot_yolo python3.8 -y conda activate rot_yolo pip install torch1.10.0 torchvision0.11.0 -f https://download.pytorch.org/whl/torch_stable.html pip install opencv-python numpy tqdm pyyaml matplotlib逻辑说明python 3.8 是 YOLOv5 官方长期验证的版本torch 1.10 对应 CUDA 11.3兼容性最好。-f指定官方 wheel 源避免 pip 默认源拉到 CPU 版。装完后用python -c import torch; print(torch.cuda.is_available())验证返回 True 才算环境通了。这一步不过后面训练全是 CPU 在跑一个 epoch 能等到天亮。3.2 数据标注格式的转换旋转检测的标注格式没有统一标准DOTA 用(x1,y1,x2,y2,x3,y3,x4,y4) class有些工具导出(cx,cy,w,h,angle) class。YOLOv5 训练需要归一化的(cls, cx, cy, w, h, angle)每行一个目标角度用弧度。下面是我常用的 DOTA 转 YOLO 旋转格式脚本import os import math import numpy as np def dota_to_rotated_yolo(txt_path, img_w, img_h, out_path): with open(txt_path, r) as f: lines f.readlines() out_lines [] for line in lines: parts line.strip().split() if len(parts) 9: continue coords list(map(float, parts[:8])) cls parts[8] pts np.array(coords).reshape(4, 2) # 用最小外接矩形求 cx,cy,w,h,angle rect cv2.minAreaRect(pts.astype(np.float32)) (cx, cy), (w, h), angle rect # 归一化 cx / img_w cy / img_h w / img_w h / img_h # 角度转弧度并归一到 [-pi/2, pi/2) angle math.radians(angle) if angle math.pi / 2: angle - math.pi out_lines.append(f{cls} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f} {angle:.6f}) with open(out_path, w) as f: f.write(\n.join(out_lines))逻辑说明cv2.minAreaRect返回的角度范围是 [-90, 0)不同 OpenCV 版本行为略有差异所以后面做了归一化到 [-π/2, π/2)。参数上img_w和img_h必须和实际送入网络的尺寸一致如果你训练时用 letterbox 缩放这里要用缩放后的尺寸否则归一化坐标会偏。常见坑是 DOTA 原图很大直接归一化后 w、h 很小回归时数值不稳定建议先把大图裁成 1024×1024 的 patch 再转。3.3 训练配置的关键参数YOLOv5 的hyp.scratch.yaml里有一堆超参旋转检测要动的主要是三个box、cls、obj的损失权重以及新加的角度损失权重。我一般会这样设参数水平检测默认旋转检测建议原因box0.050.05位置回归不变cls0.50.5分类不变obj1.01.0目标性不变angle无0.1~0.3太大会压过位置损失太小角度学不动lr00.010.005多一个回归分支初始学习率降一半更稳角度损失用正弦编码时损失函数是MSE(sin_pred, sin_gt) MSE(cos_pred, cos_gt)不是直接对角度做回归。这样写的好处是梯度有界不会因为角度跳变产生大梯度。训练命令和原版 YOLOv5 一致python train.py --data data/rotated.yaml --cfg models/yolov5s-rotated.yaml --weights yolov5s.pt --batch-size 8 --epochs 100 --hyp data/hyp.rotated.yaml参数说明--weights加载预训练权重能加速收敛但注意预训练权重的 Head 是 85 维你的新 Head 是 87 维加载时会跳过不匹配的层这是正常的。--batch-size根据显存调旋转 IoU 计算吃内存8G 显存建议不超过 8。--epochs100 起步旋转检测收敛比水平检测慢因为角度分支要从头学。4. 旋转 NMS 与后处理推理阶段最容易翻车的地方4.1 普通 NMS 为什么在旋转框上失效水平 NMS 用 IoU 阈值抑制重叠框但旋转框的 IoU 计算本身就要用旋转版本。如果你推理时忘了换用水平 IoU 去抑制密集排列的舰船会被误删——因为两个朝向不同的船水平外接框重叠度很高但实际不重叠。这就是热词里 yolov5 后处理被频繁搜索的原因很多人训练 loss 正常推理结果一塌糊涂问题就出在这。旋转 NMS 的实现思路先按置信度排序取最高分框计算它和剩余框的旋转 IoU超过阈值的抑制掉循环直到处理完。def rotated_nms(boxes, scores, iou_thres0.5): # boxes: tensor [N, 6] (cx,cy,w,h,angle,conf) order scores.argsort(descendingTrue) keep [] while order.numel() 0: i order[0].item() keep.append(i) if order.numel() 1: break ious torch.tensor([ rotated_iou(boxes[i].tolist(), boxes[j].tolist()) for j in order[1:] ]) mask ious iou_thres order order[1:][mask] return keep逻辑说明rotated_iou复用前面的多边形求交。参数上iou_thres旋转检测一般设 0.3~0.5比水平检测低因为旋转框本身贴合更紧阈值太高会漏抑制。这个实现是逐框循环速度慢生产环境建议用 CUDA 加速版本或者把框转成多边形后用torchvision.ops.nms的变体。4.2 角度解码的边界处理推理时网络输出的是 sin 和 cos要还原成角度用atan2。但atan2返回 [-π, π]而训练时角度归一化到了 [-π/2, π/2)这里有个不一致。我一般会在解码后统一把角度映射回 [-π/2, π/2)因为旋转框的 w 和 h 可以互换角度加 π/2 等价于交换宽高。def decode_angle(sin_val, cos_val): angle math.atan2(sin_val, cos_val) if angle math.pi / 2: angle - math.pi elif angle -math.pi / 2: angle math.pi return angle参数说明sin_val和cos_val是网络直接输出不需要额外归一化因为正弦编码本身就在 [-1,1]。这个解码逻辑要和训练时的角度归一化严格对应否则会出现训练时角度对、推理时角度差 90° 的玄学问题。5. 避坑与排查旋转 YOLOv5 训练中真实踩过的坑5.1 现象loss 震荡不收敛角度分支梯度爆炸原因角度损失权重设太大或者用了直接角度回归而不是正弦编码。直接回归时179° 和 -179° 的 loss 是 358梯度瞬间把回归头带飞。解决换成正弦编码角度损失权重从 0.1 开始试观察 loss 曲线如果角度 loss 占比超过总 loss 的 30%就往下调。5.2 现象推理结果框位置对但角度全错原因训练时角度归一化范围和推理解码范围不一致。比如训练用了 [0, π)推理按 [-π/2, π/2) 解码差了一个象限。解决把训练和推理的角度处理写成一个函数两边共用别手写两套。我一般会在utils/angle_utils.py里放normalize_angle和decode_angle训练脚本和推理脚本都 import 它。5.3 现象密集小目标漏检严重原因旋转 IoU 计算在目标很小时数值不稳定contourArea返回的面积接近 0正样本分配时被过滤掉。解决对小于 8×8 像素的目标先用水平 IoU 做分配只对中等以上目标用旋转 IoU。或者在标签分配阶段放宽正样本阈值让更多小目标进入正样本。5.4 现象训练速度比水平检测慢一倍以上原因旋转 IoU 逐对计算Python 循环开销大。解决把旋转 IoU 计算放到 GPU 上用多边形面积公式向量化实现或者减少参与旋转 IoU 计算的框对数量先用中心距离粗筛只对距离小于两个框对角线之和的框对算旋转 IoU。5.5 现象模型在验证集上 mAP 正常但可视化时框画歪了原因可视化代码用的cv2.boxPoints角度单位是度而网络输出是弧度没转换。解决可视化前统一转成度并且注意cv2.boxPoints返回的四个点顺序画线时要按顺序连否则框会交叉。这个坑很隐蔽因为 mAP 计算用的是坐标数值不受可视化影响所以指标正常但图不对。6. 把旋转检测推到可用验证方法与一个提点技巧训练完一个旋转 YOLOv5别只看 mAP。旋转检测的 mAP 计算本身就有坑很多开源实现用的是水平 IoU 算 AP指标虚高。验证时我一般做三件事。第一用旋转 IoU 重新算一遍 AP对比水平 IoU 的 AP如果差距超过 10 个点说明你的模型在角度上还没学好水平框碰巧重叠而已。第二挑密集场景的可视化图人眼过一遍。旋转检测的最终裁判是下游任务比如 OCR 里的文字条检测框歪 5° 可能不影响识别歪 15° 就切不干净了。所以我会把可视化图按目标密度排序先看最密的那几张。第三测推理速度。旋转 NMS 是后处理瓶颈用torch.cuda.Event计时把前向和 NMS 分开测。如果 NMS 占了总时间的一半以上就得换向量化实现。提点技巧上我踩过最值的坑是在角度分支上加一个辅助损失让网络同时预测角度的 sin 和 cos 的符号但不参与反向传播只用来监控。这样训练时能提前发现角度学反了的情况不用等 100 个 epoch 跑完才看出来。具体做法是在compute_loss里加一个angle_metric只记录不反传。# 只监控不反传的角度符号准确率 with torch.no_grad(): pred_sign torch.sign(pred_angle) gt_sign torch.sign(gt_angle) angle_acc (pred_sign gt_sign).float().mean() # 写入日志 mloss.append(angle_acc)这个习惯帮我省了很多后悔药。旋转检测的调试周期比水平检测长因为角度是隐变量loss 降了不代表角度对了。有个监控指标心里有底。希望帮到你。本文还有配套的精品资源点击获取
