简介:这份资源面向具备一定深度学习基础的计算机视觉开发者与算法工程师,提供基于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_classes,COCO 是 85。旋转版要变成6 + num_classes,多出来的一维是角度。但如果你用正弦编码,实际要多两维,变成7 + num_classes,因为要输出 sin 和 cos。
下面是我改 Head 时常用的输出解析片段,基于 YOLOv5 的Detect层思路:
import torch import math class RotateDetectHead(torch.nn.Module): def __init__(self, nc=80, anchors=(), ch=()): super().__init__() self.nc = nc self.no = nc + 5 + 2 # 5 是 xywh+obj,2 是 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 python=3.8 -y conda activate rot_yolo pip install torch==1.10.0 torchvision==0.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的损失权重,以及新加的角度损失权重。
我一般会这样设:
| 参数 | 水平检测默认 | 旋转检测建议 | 原因 |
|---|---|---|---|
| box | 0.05 | 0.05 | 位置回归不变 |
| cls | 0.5 | 0.5 | 分类不变 |
| obj | 1.0 | 1.0 | 目标性不变 |
| angle | 无 | 0.1~0.3 | 太大会压过位置损失,太小角度学不动 |
| lr0 | 0.01 | 0.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_thres=0.5): # boxes: tensor [N, 6] (cx,cy,w,h,angle,conf) order = scores.argsort(descending=True) 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 降了不代表角度对了。有个监控指标,心里有底。
希望帮到你。
本文还有配套的精品资源,点击获取