简介基于Python的YOLOv5旋转目标检测实现面向目标检测算法学习者与工业视觉开发者专门解决遥感图像、文档扫描、工业零件等场景中倾斜或旋转物体的精准框定问题。压缩包共150个文件总大小6.26MB主体为Python脚本与YAML配置文件还包含C/CUDA编写的旋转框NMS后处理源码、Dockerfile、模型说明文档等覆盖环境搭建、数据准备、模型训练、评估与推理全流程。项目在标准YOLOv5基础上引入OBB角度回归分支支持旋转框标注、GIOU/DIoU角度损失计算与旋转NMS并附带清晰目录结构和可调超参配置便于读者理解旋转目标检测的数据增强、损失改进和后处理关键环节。目前已有852人学习下载适合需要快速搭建旋转目标检测系统并开展二次开发的算法爱好者与工程人员。1. 用 YOLOv5 做旋转目标检测OBB 到底比普通框难在哪儿转做航拍和遥感图像检测的同行大概率都会撞上同一个痛点用普通的 YOLOv5 检测密集停放的飞机、集装箱或者车辆水平边界框会把几个相邻对象框成一个巨大的重叠矩形后处理一压漏检就出现了。旋转目标检测OBB就是为了解决这个场景提出的改进方向——在 cx、cy、w、h 之外增加一个角度参数让检测框能贴着目标的长轴方向旋转。这个基于 Python 的 YOLOv5 旋转目标检测项目核心改动不在网络结构而在后处理和评估链路它把旋转框的交并比计算和多边形 NMS 做成了 C/CUDA 扩展文件列表里的 poly_overlaps.cpp、nms_rotated_cpu.cpp、poly_nms_cuda.cu 这些就是真正的加速引擎。适合正在做航拍、遥感、工业质检、交通监控里倾斜目标检测的开发者也适合想彻底理解旋转框 IoU 和 NMS 原理的进阶学习者。2. 旋转框的数学基础与扩展实现从 polyiou 到 nms_rotated_ext2.1 五参数表示法cx、cy、w、h、theta 与角度约定要支持旋转目标检测第一步不是改网络而是统一“旋转框怎么表示”。常见方案是五参数表示法中心点 (cx, cy)、宽度 w、高度 h、旋转角 θ。但这个 θ 在不同框架中的定义差别很大这是整个项目里最容易翻车的地方。长边角约定θ 表示长边与图像 x 轴正方向的夹角取值范围通常是 (-π/2, π/2]。这个约定在遥感检测里很常见DOTA 数据集转 OBB 时多数采用它。OpenCV 约定θ 表示矩形绕中心旋转的角度取值范围通常是 [0, π/2)这个定义更贴近 OpenCV 的 minAreaRect 输出。同一个标注数据用不同角度约定解析w 和 h 谁是长边谁是短边也会跟着变。项目里的 polyiou.cpp 和 poly_overlaps.cpp 在计算 IoU 时默认认为输入是“多边形顶点坐标”也就是四个角点而不是五参数。这意味着使用这个扩展之前你必须先把五参数展开成四角点并保证展开方式和你标注时的角度定义一致。我一般会写一个统一转换函数把所有输入先转成“四角点逆时针排列”的内部格式再送入扩展。这个习惯帮我省下了大量排查时间因为后续所有 C/CUDA 代码只认多边形顶点不认角度。2.2 旋转 IoU 计算原理多边形裁剪与鞋带公式旋转 IoU 不能直接套用水平框的公式。水平框 IoU 是矩形面积求交计算量小旋转框则要计算两个任意四边形的相交面积。项目里 poly_overlaps.cpp、polyiou.cpp 做的事正是这个几何计算。具体流程分四步把两个旋转框分别展开成四个顶点得到两个凸四边形。用 Sutherland-Hodgman 多边形裁剪算法求两个四边形的交集多边形。这一步的本质是用其中一个四边形的四条边依次裁剪另一个四边形保留落在内部的顶点并计算新增的交点。对交集多边形用鞋带公式Shoelace Formula求面积。代入公式 IoU 交集面积 / (面积 A 面积 B - 交集面积)。Sutherland-Hodgman 裁剪是旋转框计算的核心。它按顺序处理每条边维护一个“当前多边形”的顶点列表每处理一条边遍历当前多边形的所有边判断顶点是否在裁剪边内侧并在边的交点处插入新顶点。这个逻辑不复杂但顶点顺序错了、浮点精度没控制好结果都会偏离。项目里 poly_overlaps_kernel.cu 做的事完全一样只不过把每个框对的裁剪计算放到了 CUDA 线程里并行执行。批量检测时输入可能是几千个候选框需要计算两两之间的 IoU 矩阵这一步是整个后处理的最大瓶颈。2.3 旋转 NMS 为什么不能用官方 NMS 替代有人会问既然模型已经输出了旋转框我把它转回水平外接矩形再做 NMS 不行吗可以运行但效果退化明显。密集场景下旋转框之间的真实重叠关系用水平外接矩形近似会严重扩大两个并排停放但头尾错开的车辆水平外接矩形之间的 IoU 会虚高NMS 就会误删其中一个框。项目里的 nms_rotated_cpu.cpp 和 nms_rotated_ext.cpp 实现了专门的旋转 NMS。流程是按置信度降序排列所有候选框。取最高置信度框 B。剔除所有与 B 的旋转 IoU 超过阈值的框。在剩余框里重复以上步骤。其中“旋转 IoU 超过阈值”的判断正是调用了上面说的多边形交并比计算。nms_rotated_ext.cpp 是 PyTorch 的 C 扩展入口用torch::Tensor绑定 Python 侧的数据把 GPU 上的检测结果直接传入 CUDA 函数处理省去了数据来回拷贝的开销。对比一下两套方案的计算量官方 NMS 每次 IoU 计算只有几次乘除法旋转 NMS 每次都要做多边形裁剪加面积计算。所以 CPU 单线程跑旋转 NMS 会非常慢项目里才专门写了 poly_nms_kernel.cu 和 poly_nms_cuda.cu 把 NMS 逻辑并行化。2.4 编译扩展setup.cfg 与 nms_rotated_ext 的构建方式这个项目的扩展部分由一组 C 和 CUDA 源文件组成入口为 nms_rotated_ext.cpp。它通过 PyTorch 的torch.utils.cpp_extension编译成 Python 可直接 import 的模块。setup.cfg 负责声明构建后端和包元数据。标准构建方式是执行# 在项目根目录执行构建当前目录下的 C/CUDA 扩展 python setup.py build_ext --inplace如果想把它安装到当前 Python 环境也可以直接用 pip 的 editable 模式pip install -e .构建完成后在训练脚本或推理脚本里导入扩展。按我拆过的类似项目经验导入方式类似这样import torch # 扩展模块编译产物会出现在项目目录下 # 这里以最常见的导入名 nms_rotated_ext 为例具体以仓库实际文件为准 import nms_rotated_extbuild_ext --inplace会在当前目录生成编译好的.so文件Linux/macOS或.pyd文件Windows这样 Python 解释器在 import 时就能直接加载。CPU 版本的实现在 nms_rotated_cpu.cppGPU 版本的实现在 poly_nms_cuda.cu两者共用 nms_rotated_ext.cpp 作为对外接口。编译时如果本机没有 NVIDIA GPU 或没装 CUDA Toolkit项目里的 .cu 文件会编译失败这时可以只保留 CPU 源文件重新配置后续我会在避坑章节详细展开。3. 数据预处理与训练管道OBB 数据集准备和损失函数改造3.1 把 DOTA 四点标注转成 YOLOv5 OBB 标签要用 YOLOv5 训练旋转目标检测标签格式跟水平框完全不同。水平框每个目标存一行“class_id x_center y_center width height”OBB 则需要额外一个角度参数或者在训练脚本里配置成“class_id cx cy w h angle”。DOTA 数据集是遥感检测最常见的来源它给的是四角点坐标 (x1, y1, x2, y2, x3, y3, x4, y4)。转成 YOLOv5 OBB 格式时要按以下方式计算import math def dota_quad_to_obb(points): 将 DOTA 四角点坐标转换为 YOLOv5 OBB 五参数。 points: [x1, y1, x2, y2, x3, y3, x4, y4]四角点顺序为顺时针或逆时针均可 返回: (cx, cy, w, h, angle)angle 为弧度制表示长边与 x 轴夹角 x1, y1, x2, y2, x3, y3, x4, y4 points # 中心点取四个顶点的均值比取对角线交点更稳妥 cx (x1 x2 x3 x4) / 4.0 cy (y1 y2 y3 y4) / 4.0 # 取点1-点2 与 点2-点3 两条边的长度 len_a math.hypot(x2 - x1, y2 - y1) len_b math.hypot(x3 - x2, y3 - y2) # 长边为 w短边为 h if len_a len_b: w, h len_a, len_b angle math.atan2(y2 - y1, x2 - x1) else: w, h len_b, len_a angle math.atan2(y3 - y2, x3 - x2) # 归一化到 [-pi/2, pi/2) if angle math.pi / 2: angle - math.pi elif angle -math.pi / 2: angle math.pi return cx, cy, w, h, angle这段代码的关键在于角度归一化和长边判断。DOTA 四角点如果是从标注工具里导出的点顺序可能顺时针也可能逆时针统一先算两条边的长度再决定 w、h 和角度就不容易出错。角度归一化的目的是让模型回归目标始终落在有限的连续区间内避免 90 度附近的跳变导致训练不收敛。拿到 cx、cy、w、h、angle 后还需要把中心坐标除以图像宽高做归一化和普通 YOLO 标签一样# class_id cx cy w h angle角度归一化后 0 0.5123 0.4821 0.0731 0.0452 -0.3141 0 0.2432 0.7312 0.0523 0.0384 0.8210标签文件和数据文件的关系与标准 YOLOv5 一致每张图像对应一个同名 .txt放在同目录的 labels 文件夹里。3.2 数据增强的旋转一致性不能只转图像不转角旋转目标检测的数据增强有一个隐藏陷阱图像旋转了标注框的角度必须同步更新Mosaic 拼接时各个子图做了不同的缩放和旋转对应的 OBB 也要跟着做相同的仿射变换。训练脚本里做数据增强时我通常这样处理import random import math import cv2 import numpy as np def rotate_image_and_obb(image, obb, angle_deg): image: BGR 图像 obb: [cx, cy, w, h, angle] 归一化坐标 angle_deg: 旋转角度度正数表示逆时针 返回旋转后的图像和新的 obb h_img, w_img image.shape[:2] center (w_img / 2, h_img / 2) # 以图像中心做旋转 matrix cv2.getRotationMatrix2D(center, angle_deg, 1.0) rotated_image cv2.warpAffine(image, matrix, (w_img, h_img)) # 把归一化坐标还原成像素坐标 cx, cy obb[0] * w_img, obb[1] * h_img # 应用旋转矩阵 cos_a math.cos(math.radians(angle_deg)) sin_a math.sin(math.radians(angle_deg)) dx cx - center[0] dy cy - center[1] new_cx center[0] dx * cos_a - dy * sin_a new_cy center[1] dx * sin_a dy * cos_a # 角度直接叠加 new_angle obb[4] math.radians(angle_deg) # 重新归一化 return rotated_image, [new_cx / w_img, new_cy / h_img, obb[2], obb[3], new_angle]这段代码的关键点在于中心点要跟随旋转矩阵做变换角度要线性叠加。如果只旋转图像而不更新角度值模型会在前几个 epoch 不断被不一致的标签干扰损失函数看起来在下降但验证 mAP 上不去。Mosaic 增强的逻辑类似每个子图各自做旋转缩放再拼接到大图上中心点和角度都要重新计算。3.3 损失函数调整角度回归的周期性和 GIOU 的作用旋转目标检测的损失函数比水平框复杂。水平框回归只需对齐位置和尺寸旋转框还要考虑角度。角度是周期性变量直接用 L1 损失会出现一个尴尬的情况预测角度 89.9 度、真实角度 -89.9 度两者实际只差 0.2 度但 L1 损失算出来是 179.8 度梯度方向完全错误。这个项目采用的方案是为旋转框引入 IoU 类损失。摘要里提到的 GIOUGeneralized IoU和 DIoUDistance IoU都是候选。GIOU 不仅考虑重叠区域还照顾到两个框的外包矩形面积当两个框完全不相交时GIOU 仍然有梯度信号这对旋转框训练尤其重要——旋转框角度差得远时交集面积经常为零普通 IoU 损失会梯度消失。一个实际的损失组成大致是loss λ1 * box_loss(旋转 IoU 或 GIOU) λ2 * cls_loss(BCEWithLogits) λ3 * angle_loss其中角度缺失可以通过边界平滑处理比如用两个通道分别预测 sin2θ 和 cos2θ再在解码时还原出真实角度。这种方式能避开角度边界跳变的问题。如果你打开项目里的训练配置看到损失函数里有 sin/cos 相关的计算说明它已经处理了角度周期性。3.4 训练参数与策略SGD、Adam 和余弦退火的取舍训练脚本的参数设置直接影响收敛质量。摘要里提到 SGD 和 Adam 都可以用配学习率调度策略。我拆过几个 OBB 项目一个经验是旋转目标检测的标注噪声通常比水平框大Adam 容易在早期过度拟合噪声SGD Momentum 配合余弦退火更稳。具体超参参考这样配参数推荐值说明optimizerSGD momentum0.937比 Adam 稳定泛化更好初始学习率0.01batch size 对应太大容易 NaN太小收敛慢学习率调度余弦退火torch.optim.lr_scheduler.CosineAnnealingLR输入分辨率1024 或 1280遥感目标小分辨率低会漏检batch size尽量大8~16受显存限制时优先降 batchepochs150~300小数据集跑 200 轮起步角度归一化[-π/2, π/2)与标签预处理保持一致NMS IoU 阈值0.3~0.5密集场景取 0.3稀疏取 0.5训练启动命令格式与普通 YOLOv5 一致python train.py --data data/obb.yaml --weights yolov5m.pt --epochs 200 --batch-size 8 --img 1024.yaml配置文件里需要注意的一点是标准 YOLOv5 的nc类别数直接沿用但anchors可能需要重新聚类。旋转框的长短边比例分布和水平框不同默认 anchor 不匹配会导致收敛慢。项目里如果有自动聚类脚本先用它跑一遍没有的话按实际数据集的 w/h 分布手动改 anchors 配置。训练过程中的 mAP 评估脚本会把 OBB 检测结果转成旋转 IoU 再算这部分用到的还是扩展里的 poly_overlaps。4. 旋转 NMS 与 CUDA 编译避坑五条高频踩坑记录4.1 编译报错MSVC 环境下 C 标准兼容问题现象Windows 上用 MSVC 编译polyiou.cpp或poly_overlaps.cpp报一堆奇怪的模板错误或C2664/C2440类型转换错误。原因有些代码依赖 GCC 特有的编译行为MSVC 的模板匹配规则更严格另外项目源码可能按 C11 标准编写而 MSVC 默认设置或 PyTorch 的扩展构建参数不匹配。解决优先在 Linux 环境编译。我一般用 Docker 镜像pytorch/pytorch:2.0.0-cuda11.8-cudnn8-devel来构建一次成功。如果必须在 Windows 上编译试着先升级 MSVC 到 2022并确保安装了完整的“用于 Windows 的 C CMake 工具”组件然后清理build/目录和*.pyd缓存文件后重新编译。4.2 CUDA 版本和 PyTorch 版本不匹配导致 import 失败现象扩展编译过程中没有报错但 Python 里import nms_rotated_ext时报类似undefined symbol: _ZN2at6detail...的错误或者直接提示找不到.so文件里的某个符号。原因.cu文件是在某个 CUDA 版本下编译的运行时 PyTorch 加载的是另一个 CUDA 版本的 ABI。nms_rotated_ext.cpp里绑定的 torch 类型定义在不同 PyTorch 版本之间有差异。解决严格保持“编译时 PyTorch/CUDA”与“运行时 PyTorch/CUDA”一致。检查方法很简单python -c import torch; print(torch.__version__, torch.version.cuda) nvcc --version两个命令输出的 CUDA 版本必须匹配。不匹配就重建环境不要纠结。4.3 标签转换后 IoU 异常为 0四角点顺序不一致现象训练时 box_loss 一直偏高验证集 mAP 接近 0单独调用扩展算某两个框的交并比得到的 IoU 是 0但可视化看着明明高度重叠。原因旋转框的四角点输入顺序没有统一。多边形裁剪算法要求输入顶点按逆时针或顺时针一致排列混用两种顺序会算错交集多边形。项目里的 poly_overlaps.cpp 可能默认按逆时针处理数据预处理时输出了顺时针裁剪直接得到空集。解决写一个统一的顶点顺序校正函数在送入扩展前强制所有多边形为逆时针import numpy as np def ensure_ccw(poly): # poly: shape (4, 2) 的顶点坐标 # 用鞋带公式判断方向面积为正说明是逆时针 x poly[:, 0] y poly[:, 1] area 0.5 * np.sum(x * np.roll(y, -1) - y * np.roll(x, -1)) if area 0: poly poly[::-1] # 反转顶点顺序 return poly把ensure_ccw挂在所有转换流程的最后一步从此再也没出过这种问题。4.4 推理时大量真实目标被 NMS 误删现象模型输出置信度正常旋转框位置也贴着目标但 NMS 后目标数量明显变少密集区域漏检严重。原因NMS 的 IoU 阈值设置过低。水平框场景下 IoU 阈值 0.5 通常没问题但旋转框在密集场景下相邻目标的真实重叠区域本来就大甚至两个框本身方向一致、位置接近真实 IoU 就是有 0.6。再用 0.5 的阈值就会误删。另外还有一种情况是置信度阈值太低成千上万个低质量框参与 NMS把高分框挤掉了。解决把 NMS IoU 阈值从 0.5 下调到 0.3 或 0.35置信度阈值至少设为 0.25。下调阈值后检测数量会增多再去判断是真实目标还是重复框。用验证集上的一组图片反复试通常 0.3 到 0.4 之间能找到一个平衡点。4.5 角度超出归一化范围导致 loss 为 NaN现象训练跑到第 20 到 50 个 epoch 之间loss 突然变成 NaN之后无法恢复。原因数据增强或标签转换后有代码路径把角度加到了 ±π 范围之外。比如旋转增强时对角度做了累加却忘记重新归一化模型要学的角度分布里出现跳变点回归目标冲突剧烈最终梯度爆炸。解决每次数据增强输出之后强制角度归一化保证在[-π/2, π/2)区间内def normalize_angle(theta): while theta math.pi / 2: theta - math.pi while theta -math.pi / 2: theta math.pi return theta同时把初始学习率从 0.01 降到 0.003 重新跑。先验证一个 batch 的 loss 能下降再放完整数据集。5. 推理验证与部署技巧让旋转检测框落在该落的地方推理链路跟前向传播不一样的地方在于模型输出的是五参数而可视化需要的是四角点。转换逻辑和训练时相同展开时注意角度与长边的对应关系——角度对应的是长边方向展开的第一个顶点要从长边端点开始。import math import cv2 import numpy as np import torch def obb_to_corners(obb, with_normalizationTrue): 把 (cx, cy, w, h, angle) 转成四角点坐标 cx, cy, w, h, angle obb cos_a, sin_a math.cos(angle), math.sin(angle) # 长边方向向量 dx1, dy1 w / 2 * cos_a, w / 2 * sin_a dx2, dy2 -h / 2 * sin_a, h / 2 * cos_a corners [ (cx dx1 dx2, cy dy1 dy2), (cx dx1 - dx2, cy dy1 - dy2), (cx - dx1 - dx2, cy - dy1 - dy2), (cx - dx1 dx2, cy - dy1 dy2), ] return np.array(corners, dtypenp.float32) def draw_obb(image, obb, color(0, 255, 0), thickness2): corners obb_to_corners(obb) pts corners.reshape((-1, 1, 2)).astype(np.int32) cv2.polylines(image, [pts], isClosedTrue, colorcolor, thicknessthickness)这里最容易搞错的是dx2, dy2方向短边方向应垂直于长边所以用-h/2 * sin_a和h/2 * cos_a。如果展开结果里短边方向旋转了 90 度说明这里正负号配错了。我自己最终的排查习惯是拿一张验证集图像分别画出标签框和预测框肉眼检查角度方向是否一致然后随机抽 20 张密集场景图统计 NMS 前后的检测框数量看是否出现“成对删除”的现象。如果大多数情况下旋转框能贴合目标长轴且没有出现一个目标被两个几乎重叠的框覆盖就可以认为后处理参数是达标的。从那以后我每次做旋转目标检测项目都会强制先单独跑一遍标签转换脚本生成一张带旋转框的可视化图确认角度约定再进训练流程。这一步只花五分钟却避开了后面绝大多数莫名其妙的踩坑。希望帮到你。本文还有配套的精品资源点击获取
