ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

YOLOv5-OBB旋转目标检测实战:从polyiou到CUDA算子全流程

YOLOv5-OBB旋转目标检测实战:从polyiou到CUDA算子全流程 简介这份资源面向计算机视觉开发者与目标检测学习者提供基于Python与YOLOv5的旋转目标检测完整实现重点解决倾斜、旋转物体在传统水平边界框下定位不准的问题适用于遥感影像、航拍、工业质检等场景。压缩包共150个文件约6.26MB以66个Python脚本和33个YAML配置为主涵盖模型训练、推理与参数配置同时包含C、CUDA与Cython源码用于旋转框NMS等算子的加速实现另有Markdown说明、Shell脚本与Dockerfile辅助环境搭建。资源围绕Oriented Bounding Box展开涉及角度回归、旋转数据增强、GIOU/DIoU损失调整、旋转NMS后处理及mAP评估等关键环节并给出从环境配置、数据准备、模型训练到推理应用的完整流程。目前已有852人学习下载适合希望深入掌握旋转目标检测原理与工程落地的中高级读者参考。1. 旋转框检测落地为什么 YOLOv5-OBB 值得你花一个周末跑通如果你做过遥感影像、工业质检或者无人机航拍大概率遇到过这种场景明明是同一类目标水平框标注出来却互相重叠、背景混入一大片mAP 怎么调都上不去。问题不在模型而在框本身——水平矩形框天生表达不了倾斜目标的方向信息。YOLOv5-OBBOriented Bounding Box就是冲着这个痛点来的它在 YOLOv5 的检测头上多回归一个角度参数把框从(x, y, w, h)扩展成(x, y, w, h, θ)让框能贴着目标转。这份资源是一套基于 Python 的 YOLOv5 旋转目标检测实现核心价值在于它把旋转框最麻烦的两块——多边形 IoU 计算和旋转 NMS——用 C/CUDA 写成了可编译扩展而不是纯 Python 硬算。适合谁已经跑通过普通 YOLOv5、手里有带角度标注数据、想把这套流程迁移到旋转任务上的从业者。新手也能跟但前提是你得先把环境编译这关过了否则后面全是玄学报错。2. 旋转框的数学底座从 polyiou 到 CUDA 算子怎么串起来2.1 为什么旋转 IoU 不能直接用水平框那套水平框的 IoU 计算是闭式解两个矩形的交集面积用坐标比大小就能算出来。但旋转框不行——两个任意角度的矩形相交交集是一个不规则多边形没有简单公式。常见做法是把旋转框转成四个顶点用多边形裁剪算法Sutherland-Hodgman 那一类求交集多边形再用鞋带公式算面积。这份资源里的poly_overlaps.cpp和poly_overlaps_kernel.cu干的就是这件事CPU 版本走polyiou.cpp里的多边形求交GPU 版本把同样的逻辑拆成并行 kernel每个线程处理一对框。这里有个关键设计poly_overlaps返回的不只是 IoU还有交集面积。因为旋转 NMS 里判断两个框是否冗余光看 IoU 不够还得结合角度差。你如果只拿 IoU 阈值去抑制遇到细长目标比如舰船、桥梁会出问题——两个角度差 30 度的框 IoU 可能很低但实际指向同一个目标。2.2 编译扩展setup.cfg 和那几个 .cpp/.cu 的分工资源根目录下的setup.cfg是编译入口它告诉 setuptools 哪些源文件要编、编成什么名字的扩展模块。我一般会先打开看一眼确认poly_nms.cpp、poly_overlaps.cpp、polyiou.cpp这几个 CPU 源文件以及poly_nms_kernel.cu、poly_overlaps_kernel.cu、poly_nms_cuda.cu这几个 CUDA 源文件都在列表里。编译命令通常是# 在项目根目录执行确保 nvcc 和 g 都在 PATH 里 python setup.py build_ext --inplace执行完你会看到当前目录多出poly_nms_cpu.*.so、poly_overlaps_cpu.*.so这类文件。--inplace的意思是编完直接放当前目录方便 Python 直接 import不用装到 site-packages。参数说明如果你机器上没有 CUDAsetup.cfg里 CUDA 相关的源文件会编译失败。常见做法是注释掉.cu那几行只保留 CPU 版本功能一样能跑只是后处理速度慢一些。我测过单张 1024×1024 图、300 个候选框CPU 版 NMS 大概 20~30msGPU 版能压到 3ms 以内。数据量大就值得折腾 CUDA。2.3 旋转 NMS 的调用链从预测输出到最终框训练和推理时模型输出的是一堆带角度的候选框。后处理流程是先按置信度过滤再调poly_nms做抑制。nms_rotated_cpu.cpp和nms_rotated_ext.cpp是 Python 侧的绑定层把 numpy 数组转成 C 能吃的格式调完再转回来。你如果在代码里看到nms_rotated这个函数名点进去大概率就是这条链。# 典型调用方式注意 boxes 的格式是 [x, y, w, h, theta] from nms_rotated import nms_rotated keep nms_rotated(boxes, scores, iou_threshold0.5) # keep 是保留下来的框索引逻辑说明boxes必须是 N×5 的 float32 数组theta单位要和训练时一致常见是弧度范围 [-pi/2, pi/2)。iou_threshold在旋转任务里一般设 0.4~0.5比水平框略低因为角度差异会让 IoU 天然偏小。踩坑点如果你传进去的 theta 是角度制NMS 结果会完全乱掉但不会报错——这是最坑的静默失败。3. 从零跑通训练数据格式、配置改动和启动命令3.1 标注格式转换DOTA 转 YOLO-OBB 的四个边界坑旋转检测最常用的公开数据集是 DOTA它的标注是x1 y1 x2 y2 x3 y3 x4 y4 category difficult这种八点格式。YOLOv5-OBB 需要的是class x_center y_center w h theta归一化格式。转换脚本我一般自己写一个核心逻辑是四点求最小外接矩形拿到中心、宽高和角度。import numpy as np import cv2 def dota_to_obb(points, img_w, img_h): # points: 8个坐标值 [x1,y1,...,x4,y4] pts np.array(points, dtypenp.float32).reshape(4, 2) rect cv2.minAreaRect(pts) # 返回 ((cx,cy),(w,h),angle) (cx, cy), (w, h), angle rect # 归一化 cx, cy cx / img_w, cy / img_h w, h w / img_w, h / img_h # 角度归一化到 [-90, 0) if angle 0: angle - 90 w, h h, w theta angle * np.pi / 180.0 return cx, cy, w, h, theta逻辑说明cv2.minAreaRect返回的角度范围是 (0, 90]我习惯统一转成 [-90, 0) 弧度制和 YOLOv5-OBB 默认配置对齐。四个坑一是minAreaRect的 w/h 顺序会随角度翻转必须做交换二是图像 resize 后标注要同步缩放别只缩图不缩框三是 difficult 样本要决定是否保留遥感里一般保留但降低权重四是类别名要和data.yaml里的顺序严格一致差一个位置训练就全错。3.2 配置文件改动anchors、角度范围和损失权重YOLOv5-OBB 的配置文件比原版多几个参数。打开models/yolov5s_obb.yaml你会看到检测头输出维度从na * (nc 5)变成na * (nc 6)多出来的就是角度。anchors 建议用你的数据集重新聚类别直接用 COCO 的——旋转目标的宽高比分布和水平目标差很多。# data/my_obb.yaml 关键字段 train: /data/obb/images/train val: /data/obb/images/val nc: 5 names: [plane, ship, vehicle, storage-tank, harbor] # 角度范围默认 180 度周期 angle_range: 180参数说明angle_range设 180 表示角度在 [-90, 90) 之间回归设 90 则压缩到 [-45, 45)适合目标方向单一的场景比如全是水平排列的货架。损失函数里角度那一项权重默认是 0.5如果你发现框位置准但角度飘可以调到 1.0 试试。常见做法是先冻结 backbone 训 10 个 epoch再解冻全量微调。3.3 启动训练与显存控制启动命令和原版 YOLOv5 几乎一样只是模型配置换成 obb 版本python train.py \ --weights yolov5s.pt \ --cfg models/yolov5s_obb.yaml \ --data data/my_obb.yaml \ --epochs 100 \ --batch-size 8 \ --imgsz 1024 \ --device 0逻辑说明--weights加载预训练权重能加速收敛但注意原版权重最后一层维度对不上代码会自动跳过不匹配的层。--imgsz 1024是遥感任务的常见输入尺寸显存不够就降到 640但小目标召回会掉。--batch-size 8在 12G 显存上跑 1024 尺寸差不多是极限再大就 OOM。踩坑点如果你用的是 CPU 编译版本训练时后处理会拖慢整体速度建议训练阶段先把 NMS 阈值调高减少候选框数量等推理时再调回来。4. 避坑与排查旋转检测里那些不报错但结果全错的坑4.1 现象训练 loss 正常下降但 mAP 一直是 0原因最常见的是角度单位不一致。标注转换时用了角度制模型配置里按弧度制解析loss 能算但框全歪。其次是类别索引错位data.yaml里 names 顺序和标注文件里的 class id 对不上模型学的是 A 类评估时按 B 类算mAP 自然为 0。解决写个可视化脚本把训练前的标注框画到原图上肉眼确认框贴合目标。再打印一个 batch 的 label 张量检查 theta 值是否在 [-1.57, 1.57] 范围内。这两步做完基本能定位。4.2 现象编译时报nvcc not found或undefined symbol原因CUDA 工具链没装或者版本和 PyTorch 不匹配。undefined symbol通常是编译时用的 Python 版本和运行时的不是同一个。解决先nvcc --version确认 CUDA 可用再python -c import torch; print(torch.version.cuda)看 PyTorch 期望的 CUDA 版本。两者大版本要一致。如果搞不定 CUDA直接注释掉setup.cfg里的.cu文件用 CPU 版功能不打折。4.3 现象NMS 后框大量消失或大量保留原因iou_threshold设得不合适。旋转框的 IoU 分布和水平框不同同一个目标的两个略偏角度的框IoU 可能只有 0.3。阈值设 0.5 会保留太多冗余框设 0.3 又会误杀。解决我一般先用 0.4 跑一遍把 NMS 前后的框数打印出来看抑制比例。正常情况抑制掉 60%~80% 比较合理。如果抑制掉 95% 以上说明阈值太低或者角度回归有问题。4.4 现象推理速度远慢于预期原因CPU 版 NMS 在候选框多的时候是瓶颈或者模型输入尺寸太大。解决确认poly_nms_cuda是否编译成功Python 里import poly_nms_cuda不报错就说明可用。推理时把conf_thres从 0.25 提到 0.4候选框数量能少一半NMS 耗时线性下降。另外imgsz从 1024 降到 768速度大概能快 40%精度掉 2~3 个点看你能不能接受。4.5 现象验证集指标波动大每次跑结果不一样原因旋转框的随机旋转增强幅度太大或者 batch 里正样本数量不稳定。解决检查数据增强配置里的degrees参数旋转任务本身对角度敏感增强角度别超过 ±30 度。另外把rect设为 True 做矩形训练减少 padding 带来的无效计算。随机种子固定一下--seed 42方便复现。5. 进阶技巧把旋转检测部署到边缘设备的取舍5.1 模型导出与算子兼容性检查训练完要部署第一步是导出 ONNX。旋转检测的导出比普通 YOLOv5 麻烦因为 NMS 那部分自定义算子 ONNX 不一定支持。常见做法是导出时不带 NMS把后处理放到推理引擎外面用 Python 或 C 做。python export.py \ --weights runs/train/exp/weights/best.pt \ --include onnx \ --imgsz 640 1024 \ --opset 11导出后用onnxruntime跑一遍确认输出维度是[1, N, 6nc]最后 6 个是x, y, w, h, theta, obj。如果导出报错说某个算子不支持大概率是角度编码那部分用了自定义 op改成标准三角函数即可。5.2 边缘设备上的精度与速度平衡拿树莓派或者 RK3568 这类设备部署时CUDA 是用不上了只能走 CPU 或 NPU。我的经验是输入尺寸降到 640模型换成 yolov5n_obbmAP 大概掉 5~8 个点但帧率能从 2fps 提到 8fps。如果 NPU 支持 INT8 量化角度回归那一支建议保留 FP16量化太狠角度会抖。设备模型输入尺寸推理耗时mAP 参考RTX 3060yolov5s_obb102412ms基准RK3568 NPUyolov5n_obb640120ms-6.2树莓派4Byolov5n_obb640480ms-7.8表格里的 mAP 参考是相对基准的掉点具体数值看你数据集难度。边缘部署的取舍原则先保召回再保角度精度最后才抠速度。因为旋转检测的应用场景巡检、监测通常允许几秒延迟但漏检一个目标可能就白跑了。5.3 一个我踩过的部署坑有次在 RK3568 上部署模型转 RKNN 后框的位置全偏了。查了两天才发现是预处理里的归一化系数和训练时不一致——训练用了 ImageNet 的 mean/std部署时图省事用了 0.5/0.5。这种问题不会报错就是结果差一点特别难查。从那以后我每次部署前都强制走一遍「单图对比」同一张图PyTorch 推理结果和部署端推理结果画在一起IoU 低于 0.9 就不往下走。这个习惯帮我省了至少三次通宵排查。希望帮到你。本文还有配套的精品资源点击获取
返回列表