简介本资源面向计算机视觉开发者与目标检测学习者提供基于Python与YOLOv5的旋转目标检测完整实现重点解决倾斜、旋转物体难以用传统水平矩形框精确定位的问题适用于遥感影像、工业质检、文字检测等场景。压缩包共150个文件约6.26MB以66个Python脚本和33个YAML配置为主涵盖模型定义、训练与推理流程同时包含CUDA与C扩展源码用于旋转框IoU及NMS的加速计算另有少量Markdown说明、Shell脚本与Dockerfile辅助环境搭建。已有852人学习下载。资源围绕Oriented Bounding Box展开涉及角度回归、旋转数据增强、GIOU/DIoU损失调整及旋转NMS后处理等关键环节读者可据此搭建训练、评估与推理链路理解旋转检测从数据标注到部署的完整思路适合具备一定PyTorch基础、希望深入OBB方向的开发者参考实践。1. 旋转框检测落地从 YOLOv5 到 OBB 的工程化路径做遥感影像、航拍巡检或者工业质检的朋友大概率遇到过这种场景明明目标就在图里水平框也框住了但框里塞了一堆背景后处理一算 IoU 就崩分类头也跟着学歪。问题不在 YOLOv5 本身而在于水平框的表达能力对斜向排列的舰船、飞机、文字行、药板天生不够用。这份基于 Python 的 YOLOv5 旋转目标检测资源核心就是把检测头从(x, y, w, h)扩展成(x, y, w, h, θ)并配套一整套旋转 IoU 计算、旋转 NMS 的 CUDA/CPU 算子。资源里能看到poly_nms.cpp、poly_overlaps.cpp、polyiou.cpp、nms_rotated_cpu.cpp、nms_rotated_ext.cpp、poly_nms_cpu.cpp、poly_overlaps_kernel.cu、poly_nms_kernel.cu、poly_nms_cuda.cu以及setup.cfg说明它不是纯 Python 脚本拼凑而是带 C/CUDA 扩展的完整工程。适合已经跑通过原版 YOLOv5、想把手里的斜框数据集真正训起来的人也适合想搞清楚旋转 NMS 到底怎么算的工程师。2. 旋转检测的数学底座OBB 表示、PolyIoU 与算子编译2.1 为什么水平框不够用OBB 到底多存了什么水平框用四个量描述一个轴对齐矩形前提是目标主轴和图像坐标轴大致平行。一旦目标旋转 45°水平框的面积会比真实目标大出接近一倍框内混入大量背景像素。训练时分类分支会把这些背景当成目标特征学进去回归分支也会因为宽高比失真而震荡。OBB 的做法是额外引入角度 θ常见表示有两种一种是(cx, cy, w, h, θ)θ 定义在[-90°, 0°)或[-45°, 45°)区间另一种是四点多边形(x1,y1,x2,y2,x3,y3,x4,y4)。YOLOv5-OBB 这类实现通常走第一种因为回归更稳定角度只需要一个输出通道。角度定义域是第一个容易翻车的地方。如果 θ 取[-90°, 0°)那么 w 和 h 的角色在边界处会互换导致同一个框有两种等价表示回归目标不唯一loss 会抖。工程上一般把 θ 限制在[-45°, 45°)配合长边定义法让 w 始终对应长边这样表示唯一。资源里的polyiou.cpp就是围绕多边形做 IoU 计算它不依赖角度定义域而是把 OBB 转成四个顶点再算多边形交集这也是后面旋转 NMS 的基础。2.2 从源码文件看算子结构CPU 与 CUDA 双路径把资源里的文件按职责分一下结构其实很清晰文件职责运行位置poly_overlaps.cpp/poly_overlaps_kernel.cu计算两个旋转框的多边形 IoUCPU / CUDApoly_nms.cpp/poly_nms_cpu.cpp/poly_nms_cuda.cu/poly_nms_kernel.cu旋转 NMS 主逻辑与核函数CPU / CUDApolyiou.cpp多边形交并比底层实现CPUnms_rotated_cpu.cpp/nms_rotated_ext.cpp旋转 NMS 的 CPU 实现与扩展入口CPUsetup.cfg扩展编译配置构建期这种 CPU CUDA 双路径的设计是有意为之训练时数据量大、框多走 CUDA 核函数推理部署到没有 GPU 的边缘设备时退回 CPU 版本接口保持一致。nms_rotated_ext.cpp是 Python 调用的入口通过 pybind11 把 C 函数暴露成 Python 可调用的模块。2.3 编译扩展setup.cfg 与构建命令这类带 C/CUDA 扩展的项目最容易卡在编译环节。常见做法是先确认 CUDA 工具链和 PyTorch 版本匹配再执行构建。下面是我一般会走的流程# 确认 CUDA 与 PyTorch 的 CUDA 版本一致 python -c import torch; print(torch.version.cuda) nvcc --version # 进入项目根目录执行扩展编译 python setup.py build_ext --inplace # 验证扩展是否可导入 python -c import poly_nms; print(poly_nms ok)build_ext --inplace会把编译产物直接放到源码目录方便 Python 直接 import。如果报nvcc not found说明 CUDA 没进 PATH如果报架构不匹配需要在编译参数里指定TORCH_CUDA_ARCH_LIST比如export TORCH_CUDA_ARCH_LIST7.5;8.0对应 Turing 和 Ampere。setup.cfg里通常定义了扩展模块名、源文件列表和编译选项改源文件列表时记得同步这里否则新加的.cu不会被编译进去。提示编译前先pip install pybind11很多undefined symbol报错其实是 pybind11 头文件路径没找到。3. 数据与训练DOTA 格式转换、角度回归与超参设置3.1 标注格式转换从 DOTA 八参数到 YOLO-OBB 五参数DOTA 数据集的标注是x1 y1 x2 y2 x3 y3 x4 y4 category difficult而 YOLOv5-OBB 训练需要的是归一化的class cx cy w h θ。转换脚本是绕不开的一步下面这段是我常用的转换逻辑import numpy as np def dota_to_obb(points, img_w, img_h): points: 8 个坐标, 顺序为四个顶点 pts np.array(points, dtypenp.float32).reshape(4, 2) # 计算四条边长度取最长边作为 w edges [np.linalg.norm(pts[i] - pts[(i 1) % 4]) for i in range(4)] w max(edges) h min(edges) # 中心点 cx, cy pts.mean(axis0) # 用最长边方向计算角度限制在 [-45, 45) idx int(np.argmax(edges)) dx pts[(idx 1) % 4][0] - pts[idx][0] dy pts[(idx 1) % 4][1] - pts[idx][1] theta np.degrees(np.arctan2(dy, dx)) while theta 45: theta - 90 w, h h, w while theta -45: theta 90 w, h h, w # 归一化 return cx / img_w, cy / img_h, w / img_w, h / img_h, theta这段代码的关键在角度归一化循环当角度超出[-45, 45)时减 90° 并交换 w、h保证表示唯一。cx / img_w这类归一化是 YOLO 系列的标准做法让回归目标落在 0 到 1 之间训练更稳。转换完还要生成对应的images和labels目录结构以及train.txt、val.txt列表文件。3.2 角度回归的损失设计为什么不能直接回归 θ角度是周期量直接回归 θ 会在-45°和45°边界处产生跳变网络学起来很痛苦。常见做法有两种一种是把角度分类成若干个 bin再回归 bin 内的偏移另一种是预测角度的正弦余弦(sin 2θ, cos 2θ)用这两个连续量间接表示角度。YOLOv5-OBB 的实现里角度分支通常输出一个通道配合专门的角度损失比如在 CIoU 基础上加角度惩罚项。训练配置上hyp.yaml里的box、cls、obj权重需要重新调。旋转检测的框回归比水平框难box权重可以适当调大。学习率方面如果是从水平框预训练权重微调初始 lr 可以设小一点比如0.001配合余弦退火。batch size 受显存限制旋转框的 IoU 计算比水平框慢显存占用也更高8G 显存建议 batch 设 4 到 8。# 启动训练指定 OBB 模型配置和数据配置 python train.py --data data/dota.yaml --cfg models/yolov5s_obb.yaml \ --weights yolov5s.pt --batch-size 8 --epochs 100 --img 1024 \ --hyp data/hyp.obb.yaml--img 1024是因为遥感图像目标小输入分辨率太低会丢细节。--weights加载水平框预训练权重能加速收敛但要注意检测头的形状变化加载时会跳过不匹配的层。3.3 旋转 NMS 的后处理参数推理阶段模型输出的框数量很大必须经过旋转 NMS 去重。资源里的poly_nms就是干这个的。调用时核心参数是conf_thres和iou_thresfrom poly_nms import poly_nms # boxes: N x 6, 格式 [cx, cy, w, h, theta, score] keep poly_nms(boxes, iou_thres0.1)旋转 NMS 的iou_thres通常比水平框设得低因为旋转框之间重叠计算更敏感设 0.1 到 0.3 比较常见。设太高会留下大量重复框设太低会误删相邻目标。这个参数没有万能值得在验证集上试。4. 避坑与排查编译、角度、显存三类高频问题4.1 扩展编译报 undefined symbol现象import poly_nms时报undefined symbol: _ZN...。 原因C 和 CUDA 源文件编译时 ABI 不一致或者 pybind11 版本和 PyTorch 不匹配。 解决统一用同一个编译器版本pip install pybind112.10这类和 PyTorch 匹配的版本重新build_ext --inplace。如果还不行在setup.py里加extra_compile_args[-stdc14]。4.2 角度 loss 不下降mAP 卡在低位现象训练几十轮box loss 正常降但角度相关指标不动mAP 上不去。 原因角度定义域没统一标注转换时 θ 落在[-90, 0)而模型按[-45, 45)学目标冲突。 解决检查转换脚本确保所有 θ 都归一化到[-45, 45)并且 w 对应长边。用几张小图可视化验证框和标注是否对齐。4.3 CUDA out of memory现象训练到一半显存爆掉。 原因旋转 IoU 计算在 CUDA 上会生成中间张量框越多占用越大--img设太大也会放大。 解决降 batch size降--img或者把 NMS 放到 CPU 上跑。训练时可以先关掉旋转 NMS 的 CUDA 路径用 CPU 版本速度慢但省显存。4.4 推理结果框重叠严重现象一张图里同一个目标出好几个框。 原因iou_thres设太高或者conf_thres太低导致低质量框没被滤掉。 解决先把conf_thres提到 0.25 以上再把iou_thres降到 0.1 到 0.2 之间试。旋转 NMS 对阈值比水平框敏感需要单独调。4.5 CPU 版本推理速度极慢现象在没有 GPU 的机器上跑推理一张图要好几秒。 原因polyiou.cpp的多边形交集是逐对计算的CPU 单线程下复杂度高。 解决减少候选框数量先按 score 排序取 top-K 再进 NMS或者用 OpenMP 给 CPU 版本加并行。如果部署环境允许优先走 CUDA 路径。5. 进阶技巧用旋转 IoU 可视化验证你的框到底对不对训练跑起来只是第一步真正省时间的是在训练前就把标注和角度定义验证对。我习惯写一个小脚本把转换后的 OBB 画回原图肉眼确认框和目标是贴合的。下面这段用 OpenCV 画旋转框import cv2 import numpy as np def draw_obb(img, cx, cy, w, h, theta, color(0, 255, 0)): 在图上画一个旋转框输入为归一化前的像素值 rect ((cx, cy), (w, h), theta) box cv2.boxPoints(rect) box np.int0(box) cv2.drawContours(img, [box], 0, color, 2) return imgcv2.boxPoints接收((cx,cy),(w,h),angle)返回四个顶点角度单位是度。把标注和预测都画上去如果预测框和标注框方向一致、贴合紧密说明角度定义和回归都对了。如果预测框总是差 90°那就是角度定义域没对齐回去查转换脚本。另一个验证点是旋转 IoU 本身。可以拿两个已知框手算 IoU再和poly_overlaps的输出对比。比如两个完全重合的框 IoU 应该是 1.0旋转 90° 后如果 w、h 互换IoU 仍应是 1.0。这种边界用例能快速暴露实现问题。从那以后我每次做旋转检测都强制先跑一遍可视化验证再开训练。这个习惯帮我省下了至少两轮白跑的 GPU 时间。希望帮到你。本文还有配套的精品资源点击获取 YOLOv5旋转目标检测OBB工程化：从算子编译到训练避坑
简介本资源面向计算机视觉开发者与目标检测学习者提供基于Python与YOLOv5的旋转目标检测完整实现重点解决倾斜、旋转物体难以用传统水平矩形框精确定位的问题适用于遥感影像、工业质检、文字检测等场景。压缩包共150个文件约6.26MB以66个Python脚本和33个YAML配置为主涵盖模型定义、训练与推理流程同时包含CUDA与C扩展源码用于旋转框IoU及NMS的加速计算另有少量Markdown说明、Shell脚本与Dockerfile辅助环境搭建。已有852人学习下载。资源围绕Oriented Bounding Box展开涉及角度回归、旋转数据增强、GIOU/DIoU损失调整及旋转NMS后处理等关键环节读者可据此搭建训练、评估与推理链路理解旋转检测从数据标注到部署的完整思路适合具备一定PyTorch基础、希望深入OBB方向的开发者参考实践。1. 旋转框检测落地从 YOLOv5 到 OBB 的工程化路径做遥感影像、航拍巡检或者工业质检的朋友大概率遇到过这种场景明明目标就在图里水平框也框住了但框里塞了一堆背景后处理一算 IoU 就崩分类头也跟着学歪。问题不在 YOLOv5 本身而在于水平框的表达能力对斜向排列的舰船、飞机、文字行、药板天生不够用。这份基于 Python 的 YOLOv5 旋转目标检测资源核心就是把检测头从(x, y, w, h)扩展成(x, y, w, h, θ)并配套一整套旋转 IoU 计算、旋转 NMS 的 CUDA/CPU 算子。资源里能看到poly_nms.cpp、poly_overlaps.cpp、polyiou.cpp、nms_rotated_cpu.cpp、nms_rotated_ext.cpp、poly_nms_cpu.cpp、poly_overlaps_kernel.cu、poly_nms_kernel.cu、poly_nms_cuda.cu以及setup.cfg说明它不是纯 Python 脚本拼凑而是带 C/CUDA 扩展的完整工程。适合已经跑通过原版 YOLOv5、想把手里的斜框数据集真正训起来的人也适合想搞清楚旋转 NMS 到底怎么算的工程师。2. 旋转检测的数学底座OBB 表示、PolyIoU 与算子编译2.1 为什么水平框不够用OBB 到底多存了什么水平框用四个量描述一个轴对齐矩形前提是目标主轴和图像坐标轴大致平行。一旦目标旋转 45°水平框的面积会比真实目标大出接近一倍框内混入大量背景像素。训练时分类分支会把这些背景当成目标特征学进去回归分支也会因为宽高比失真而震荡。OBB 的做法是额外引入角度 θ常见表示有两种一种是(cx, cy, w, h, θ)θ 定义在[-90°, 0°)或[-45°, 45°)区间另一种是四点多边形(x1,y1,x2,y2,x3,y3,x4,y4)。YOLOv5-OBB 这类实现通常走第一种因为回归更稳定角度只需要一个输出通道。角度定义域是第一个容易翻车的地方。如果 θ 取[-90°, 0°)那么 w 和 h 的角色在边界处会互换导致同一个框有两种等价表示回归目标不唯一loss 会抖。工程上一般把 θ 限制在[-45°, 45°)配合长边定义法让 w 始终对应长边这样表示唯一。资源里的polyiou.cpp就是围绕多边形做 IoU 计算它不依赖角度定义域而是把 OBB 转成四个顶点再算多边形交集这也是后面旋转 NMS 的基础。2.2 从源码文件看算子结构CPU 与 CUDA 双路径把资源里的文件按职责分一下结构其实很清晰文件职责运行位置poly_overlaps.cpp/poly_overlaps_kernel.cu计算两个旋转框的多边形 IoUCPU / CUDApoly_nms.cpp/poly_nms_cpu.cpp/poly_nms_cuda.cu/poly_nms_kernel.cu旋转 NMS 主逻辑与核函数CPU / CUDApolyiou.cpp多边形交并比底层实现CPUnms_rotated_cpu.cpp/nms_rotated_ext.cpp旋转 NMS 的 CPU 实现与扩展入口CPUsetup.cfg扩展编译配置构建期这种 CPU CUDA 双路径的设计是有意为之训练时数据量大、框多走 CUDA 核函数推理部署到没有 GPU 的边缘设备时退回 CPU 版本接口保持一致。nms_rotated_ext.cpp是 Python 调用的入口通过 pybind11 把 C 函数暴露成 Python 可调用的模块。2.3 编译扩展setup.cfg 与构建命令这类带 C/CUDA 扩展的项目最容易卡在编译环节。常见做法是先确认 CUDA 工具链和 PyTorch 版本匹配再执行构建。下面是我一般会走的流程# 确认 CUDA 与 PyTorch 的 CUDA 版本一致 python -c import torch; print(torch.version.cuda) nvcc --version # 进入项目根目录执行扩展编译 python setup.py build_ext --inplace # 验证扩展是否可导入 python -c import poly_nms; print(poly_nms ok)build_ext --inplace会把编译产物直接放到源码目录方便 Python 直接 import。如果报nvcc not found说明 CUDA 没进 PATH如果报架构不匹配需要在编译参数里指定TORCH_CUDA_ARCH_LIST比如export TORCH_CUDA_ARCH_LIST7.5;8.0对应 Turing 和 Ampere。setup.cfg里通常定义了扩展模块名、源文件列表和编译选项改源文件列表时记得同步这里否则新加的.cu不会被编译进去。提示编译前先pip install pybind11很多undefined symbol报错其实是 pybind11 头文件路径没找到。3. 数据与训练DOTA 格式转换、角度回归与超参设置3.1 标注格式转换从 DOTA 八参数到 YOLO-OBB 五参数DOTA 数据集的标注是x1 y1 x2 y2 x3 y3 x4 y4 category difficult而 YOLOv5-OBB 训练需要的是归一化的class cx cy w h θ。转换脚本是绕不开的一步下面这段是我常用的转换逻辑import numpy as np def dota_to_obb(points, img_w, img_h): points: 8 个坐标, 顺序为四个顶点 pts np.array(points, dtypenp.float32).reshape(4, 2) # 计算四条边长度取最长边作为 w edges [np.linalg.norm(pts[i] - pts[(i 1) % 4]) for i in range(4)] w max(edges) h min(edges) # 中心点 cx, cy pts.mean(axis0) # 用最长边方向计算角度限制在 [-45, 45) idx int(np.argmax(edges)) dx pts[(idx 1) % 4][0] - pts[idx][0] dy pts[(idx 1) % 4][1] - pts[idx][1] theta np.degrees(np.arctan2(dy, dx)) while theta 45: theta - 90 w, h h, w while theta -45: theta 90 w, h h, w # 归一化 return cx / img_w, cy / img_h, w / img_w, h / img_h, theta这段代码的关键在角度归一化循环当角度超出[-45, 45)时减 90° 并交换 w、h保证表示唯一。cx / img_w这类归一化是 YOLO 系列的标准做法让回归目标落在 0 到 1 之间训练更稳。转换完还要生成对应的images和labels目录结构以及train.txt、val.txt列表文件。3.2 角度回归的损失设计为什么不能直接回归 θ角度是周期量直接回归 θ 会在-45°和45°边界处产生跳变网络学起来很痛苦。常见做法有两种一种是把角度分类成若干个 bin再回归 bin 内的偏移另一种是预测角度的正弦余弦(sin 2θ, cos 2θ)用这两个连续量间接表示角度。YOLOv5-OBB 的实现里角度分支通常输出一个通道配合专门的角度损失比如在 CIoU 基础上加角度惩罚项。训练配置上hyp.yaml里的box、cls、obj权重需要重新调。旋转检测的框回归比水平框难box权重可以适当调大。学习率方面如果是从水平框预训练权重微调初始 lr 可以设小一点比如0.001配合余弦退火。batch size 受显存限制旋转框的 IoU 计算比水平框慢显存占用也更高8G 显存建议 batch 设 4 到 8。# 启动训练指定 OBB 模型配置和数据配置 python train.py --data data/dota.yaml --cfg models/yolov5s_obb.yaml \ --weights yolov5s.pt --batch-size 8 --epochs 100 --img 1024 \ --hyp data/hyp.obb.yaml--img 1024是因为遥感图像目标小输入分辨率太低会丢细节。--weights加载水平框预训练权重能加速收敛但要注意检测头的形状变化加载时会跳过不匹配的层。3.3 旋转 NMS 的后处理参数推理阶段模型输出的框数量很大必须经过旋转 NMS 去重。资源里的poly_nms就是干这个的。调用时核心参数是conf_thres和iou_thresfrom poly_nms import poly_nms # boxes: N x 6, 格式 [cx, cy, w, h, theta, score] keep poly_nms(boxes, iou_thres0.1)旋转 NMS 的iou_thres通常比水平框设得低因为旋转框之间重叠计算更敏感设 0.1 到 0.3 比较常见。设太高会留下大量重复框设太低会误删相邻目标。这个参数没有万能值得在验证集上试。4. 避坑与排查编译、角度、显存三类高频问题4.1 扩展编译报 undefined symbol现象import poly_nms时报undefined symbol: _ZN...。 原因C 和 CUDA 源文件编译时 ABI 不一致或者 pybind11 版本和 PyTorch 不匹配。 解决统一用同一个编译器版本pip install pybind112.10这类和 PyTorch 匹配的版本重新build_ext --inplace。如果还不行在setup.py里加extra_compile_args[-stdc14]。4.2 角度 loss 不下降mAP 卡在低位现象训练几十轮box loss 正常降但角度相关指标不动mAP 上不去。 原因角度定义域没统一标注转换时 θ 落在[-90, 0)而模型按[-45, 45)学目标冲突。 解决检查转换脚本确保所有 θ 都归一化到[-45, 45)并且 w 对应长边。用几张小图可视化验证框和标注是否对齐。4.3 CUDA out of memory现象训练到一半显存爆掉。 原因旋转 IoU 计算在 CUDA 上会生成中间张量框越多占用越大--img设太大也会放大。 解决降 batch size降--img或者把 NMS 放到 CPU 上跑。训练时可以先关掉旋转 NMS 的 CUDA 路径用 CPU 版本速度慢但省显存。4.4 推理结果框重叠严重现象一张图里同一个目标出好几个框。 原因iou_thres设太高或者conf_thres太低导致低质量框没被滤掉。 解决先把conf_thres提到 0.25 以上再把iou_thres降到 0.1 到 0.2 之间试。旋转 NMS 对阈值比水平框敏感需要单独调。4.5 CPU 版本推理速度极慢现象在没有 GPU 的机器上跑推理一张图要好几秒。 原因polyiou.cpp的多边形交集是逐对计算的CPU 单线程下复杂度高。 解决减少候选框数量先按 score 排序取 top-K 再进 NMS或者用 OpenMP 给 CPU 版本加并行。如果部署环境允许优先走 CUDA 路径。5. 进阶技巧用旋转 IoU 可视化验证你的框到底对不对训练跑起来只是第一步真正省时间的是在训练前就把标注和角度定义验证对。我习惯写一个小脚本把转换后的 OBB 画回原图肉眼确认框和目标是贴合的。下面这段用 OpenCV 画旋转框import cv2 import numpy as np def draw_obb(img, cx, cy, w, h, theta, color(0, 255, 0)): 在图上画一个旋转框输入为归一化前的像素值 rect ((cx, cy), (w, h), theta) box cv2.boxPoints(rect) box np.int0(box) cv2.drawContours(img, [box], 0, color, 2) return imgcv2.boxPoints接收((cx,cy),(w,h),angle)返回四个顶点角度单位是度。把标注和预测都画上去如果预测框和标注框方向一致、贴合紧密说明角度定义和回归都对了。如果预测框总是差 90°那就是角度定义域没对齐回去查转换脚本。另一个验证点是旋转 IoU 本身。可以拿两个已知框手算 IoU再和poly_overlaps的输出对比。比如两个完全重合的框 IoU 应该是 1.0旋转 90° 后如果 w、h 互换IoU 仍应是 1.0。这种边界用例能快速暴露实现问题。从那以后我每次做旋转检测都强制先跑一遍可视化验证再开训练。这个习惯帮我省下了至少两轮白跑的 GPU 时间。希望帮到你。本文还有配套的精品资源点击获取