简介：基于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_ext
2.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.cpp，GPU 版本的实现在 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 类损失。摘要里提到的 GIOU（Generalized IoU）和 DIoU（Distance 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 配合余弦退火更稳。具体超参参考这样配：
| 参数 | 推荐值 | 说明 |
|---|---|---|
| optimizer | SGD momentum=0.937 | 比 Adam 稳定，泛化更好 |
| 初始学习率 | 0.01（batch size 对应） | 太大容易 NaN，太小收敛慢 |
| 学习率调度 | 余弦退火 | torch.optim.lr_scheduler.CosineAnnealingLR |
| 输入分辨率 | 1024 或 1280 | 遥感目标小，分辨率低会漏检 |
| batch size | 尽量大，8~16 | 受显存限制时优先降 batch |
| epochs | 150~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 的模板匹配规则更严格；另外项目源码可能按 C++11 标准编写，而 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_normalization=True): """把 (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, dtype=np.float32) def draw_obb(image, obb, color=(0, 255, 0), thickness=2): corners = obb_to_corners(obb) pts = corners.reshape((-1, 1, 2)).astype(np.int32) cv2.polylines(image, [pts], isClosed=True, color=color, thickness=thickness)这里最容易搞错的是dx2, dy2方向：短边方向应垂直于长边，所以用-h/2 * sin_a和h/2 * cos_a。如果展开结果里短边方向旋转了 90 度，说明这里正负号配错了。
我自己最终的排查习惯是：拿一张验证集图像，分别画出标签框和预测框，肉眼检查角度方向是否一致；然后随机抽 20 张密集场景图，统计 NMS 前后的检测框数量，看是否出现“成对删除”的现象。如果大多数情况下旋转框能贴合目标长轴，且没有出现一个目标被两个几乎重叠的框覆盖，就可以认为后处理参数是达标的。
从那以后我每次做旋转目标检测项目，都会强制先单独跑一遍标签转换脚本，生成一张带旋转框的可视化图确认角度约定，再进训练流程。这一步只花五分钟，却避开了后面绝大多数莫名其妙的踩坑。希望帮到你。
本文还有配套的精品资源，点击获取
