简介面向计算机视觉开发者与目标检测研究者的YOLOv5旋转目标检测实现资源基于Python和PyTorch搭建聚焦OBB倾斜边界框与角度回归原理涵盖数据预处理、数据增强、损失函数调整、NMS后处理等关键改进适用于遥感图像、航拍影像、工业质检等倾斜物体识别场景可帮助学习者快速搭建、训练并部署旋转目标检测模型。压缩包共150个文件以66个Python脚本为核心辅以33个YAML配置、8个Markdown说明、6个C源码、4个CUDA核函数另含Dockerfile、makefile、ipynb示例等整包仅6.26MB便于快速下载。已有852人学习/下载。资源提供完整的工程目录与可运行代码内置旋转框标注、训练评估、推理模块以及配置文件支持GIOU/DIoU等损失函数扩展适合具备一定深度学习基础、希望深入YOLOv5旋转检测实现细节的开发者参考使用便于二次开发与算法对比。1. 旋转目标检测为什么 yolov5 的常规输出头在航拍图上一碰就碎做遥感目标检测的人基本都遇到过同一个问题拿 yolov5 默认权重跑一张无人机俯拍的停车场一辆斜着停放的车水平框画出来不仅把旁边车道的车也框进去一半两个邻近的检测框还因为重叠太大被 NMS 直接干掉。场景从水平拍摄换成俯视拍摄之后目标的朝向变得任意而 yolov5 的预测框只能输出 x、y、w、h 四个量它天生没有表达角度的能力这就是旋转目标检测要解决的缺口。本文要拆的就是在 python 的 yolov5 框架上加一路角度回归分支把普通水平框检测改成带角度的旋转框检测。适合手里已经跑通过 yolov5 目标检测、现在被航拍图或者工厂质检图里密集斜向目标折磨的从业者读完你至少知道核心代码改哪里、数据集怎么准备、训练翻车了先查哪几个地方。2. 旋转框的表示与角度约定先看懂数据集标注再谈改网络2.1 两种主流角度定义OpenCV 定义与长边定义旋转目标检测在网络层面其实只比水平检测多了一个自由度角度。但这个角度怎么定义直接决定后面损失函数、anchor 匹配和后处理要不要跟着改。目前从业者用的主要是两套约定第一套叫 OpenCV 定义角度范围是 [-90°, 0°)框由旋转矩形的中心点 (cx, cy)、宽 w、高 h 和角度 theta 组成其中 w 永远是与 x 轴夹角较小的那条边h 是另一条边。第二套叫长边定义角度范围是 (-90°, 90°]框的 w 取长边、h 取短边theta 表示长边与 x 轴正方向的夹角。这两套定义本身没有对错之分但如果你混着用代码一跑就会出现一个非常诡异的画面训练时 loss 降得很漂亮验证时画出来的框却整体歪掉。原因在于同一组多边形坐标用两种方式反推 w、h 和 theta结果可能差 90°。我一般建议新手直接跟长边定义走因为它的几何含义更直观后面接旋转 NMS 和可视化时不容易把自己绕晕。要注意的是很多公开遥感数据集 DOTA 的标注格式是四边形的八个坐标点 (x1, y1, x2, y2, x3, y3, x4, y4)它不直接给你角度你需要先从点坐标反推出旋转矩形。2.2 从多边形 8 个点转成 (cx, cy, w, h, theta)坐标变换细节从 DOTA 的四个点还原旋转矩形核心思路是先按点的顺序重建一个多边形然后求它的最小外接旋转矩形。常见做法是遍历四条边分别求这条边作为“底边”时整个多边形在底边方向上的投影长度作为宽、在法线方向上的投影长度作为高取面积最小的那个作为结果。更省事的办法是直接用 cv2.minAreaRect它输入四个点就能算出 (cx, cy), (w, h), angle但这个 angle 是 cv2 自己的约定范围是 [-90°, 0)如果你要转成长边定义得再换算一下。import cv2 import numpy as np def polygon_to_cxcywha(points, angle_formatlongside): 将四个顶点坐标转成 (cx, cy, w, h, theta) points: 长度 8 的数组按 (x1,y1,x2,y2,x3,y3,x4,y4) 顺序传入 angle_format: longside 表示长边定义, opencv 表示 OpenCV 定义 pts np.array(points, dtypenp.float32).reshape(4, 2) rect cv2.minAreaRect(pts) # ((cx, cy), (w, h), angle) (cx, cy), (w, h), angle rect # cv2.minAreaRect 返回的角度范围是 [-90, 0)w 是与 x 轴夹角较小的边 if angle_format longside: theta angle / 180.0 * np.pi # 转弧度 # 长边定义要求 w h如果 OpenCV 给出的是短边就交换 if h w: w, h h, w theta np.pi / 2 # 把角度归一化到 (-pi/2, pi/2] if theta np.pi / 2: theta - np.pi if theta -np.pi / 2: theta np.pi else: theta angle / 180.0 * np.pi return cx, cy, w, h, theta这里有个容易踩的细节很多人直接用 cv2.minAreaRect 返回的 angle 当长边定义用发现训练前期 loss 反复震荡。原因是 OpenCV 的 angle 是相对于 x 轴正方向的夹角范围只有 90°它返回的 w 和 h 的顺序会随目标朝向变化而互换一旦 w、h 互换宽高比这个特征就变了网络等于同时学两个不一致的目标。上面代码里做了一个处理长边定义下如果 h 大于 w 就交换两者并给角度加 90°这样 w 永远是长边目标的语义才稳定。2.3 角度周期性与 IoU为什么 89° 和 -91° 是同一个框角度回归和坐标回归有一个本质区别坐标的数值范围是开的角度的数值范围是闭合且周期性的。长边定义下 89° 和 -91° 其实指向同一个朝向但如果你直接算数值差损失是 180°同理89.5° 和 -90.5° 只差 1°数值上却差了 180°。这就是旋转检测训练中非常经典的“角度周期性”问题也是很多新人第一次跑旋转检测时 loss 下降一半就不再动的元凶之一。处理周期性常见有三种手段。第一种是直接回归连续角度并加一个周期罚项实现简单但容易收敛到局部最优。第二种是把角度按区间切分成分类任务比如每 5° 一个 bin一共 36 类用分类损失去优化规避了连续回归的周期问题代价是角度精度受限。第三种是我个人比较推荐的做法回归角度之余额外计算一个“翻转折损”——把预测角度加减 90° 之后重新算一遍损失取最小值作为最终损失这等于让网络在两条可能的路径里自己选更快的那条。至于旋转框的 IoU没法用水平框那个简单的公式算通常要先把两个旋转框转成多边形再用 Shoelace 公式求交集面积。纯 Python 实现这个计算很慢所以一般会引入 shapely 的几何库做多边形裁剪或者写一个简化版训练时用近似 IoU只在验证时用精确 IoU。3. 在 yolov5 上改出旋转检测头核心代码改动与网络结构3.1 输出头改造从 5 个量变成 6 个量角度怎么回归yolov5 原本的输出头对每个 anchor 预测 5 个量tx、ty、tw、th 和 objectness加上类别数就是 5 num_classes。旋转检测的做法是再加一个通道用来回归角度常见的是在 Detect 层的输出维度里把 5 改成 6。具体的改法每份开源实现略有差异但核心动作一致从 self.no nc 5 改成 self.no nc 6然后修改 forward 里对输出张量的拆分逻辑。# yolov5/models/yolo.py 中 Detect 类的修改示意 class Detect(nn.Module): def __init__(self, nc80, anchors(), ch(), inplaceTrue): super().__init__() self.nc nc # 原来这里是 self.no nc 5 # number of outputs per anchor self.no nc 6 # 5 个基础量 1 个角度 self.nl len(anchors) # 检测层数量通常是 3 self.na len(anchors[0]) // 2 # 每层 anchor 数量 def forward(self, x): # x 是三个不同尺度特征图的列表 for i in range(self.nl): x[i] self.m[i](x[i]) # 卷积输出维度变为 bs, no, h, w # 后续 reshape 时最后一维 no 自动变成 6 nc return x角度分支的具体回归目标是什么值得单独说一下。我们不是让网络直接输出绝对角度而是输出一个相对与 anchor 的偏移量。每个 anchor 本身可以带一个初始角度常见做法是把 anchor 的角度初始化为 0°然后网络预测一个角度残差解码时用 anchor 角度加残差得到最终角度。这样设计的好处是网络的预测范围不需要覆盖全角度只需要预测一个相对较小的偏移学起来更稳定。还有个做法是给每个位置预设多个不同角度的 anchor等于把角度也做成 anchor 匹配的一部分但训练成本会成倍上升小数据集上收益一般技术选型时不建议一开始就上。3.2 损失函数的适配Smooth L1 还是角度分类输出头加了角度通道之后损失函数也必须跟着改。yolov5 原本的回归损失是 CIoU它要求两个框代表同样的几何体而旋转框的 IoU 计算起来太耗时训练时基本不会用旋转 IoU 作为回归损失。比较常见的替代方案是坐标和宽高继续用 CIoU角度单独用 Smooth L1 损失两个损失按权重叠加。Smooth L1 对异常值比 L2 更稳这在角度回归的场景里特别重要因为角度周期性导致的偶发大误差不应该一次把整个梯度带偏。# 角度损失计算示意结合周期性的 Smooth L1 import math import torch def angle_loss(pred_angle, target_angle, reductionmean): pred_angle, target_angle: 弧度制形状一致 考虑了 90° 周期性真实值与预测值差超过 90° 时翻转后再算损失 # 原始差值 diff pred_angle - target_angle # 周期性适配对差值做等价变换把它映射到 [-pi/2, pi/2] 之间 diff torch.remainder(diff, math.pi) diff torch.where(diff math.pi / 2, diff - math.pi, diff) diff torch.where(diff -math.pi / 2, diff math.pi, diff) # Smooth L1 (Huber) 损失delta1.0 abs_diff torch.abs(diff) quadratic torch.clamp(abs_diff, max1.0) linear abs_diff - quadratic loss 0.5 * quadratic**2 linear if reduction mean: return loss.mean() return loss.sum()这个损失函数值得说明一下torch.remainder(diff, math.pi) 先把差值折到 [-π, π]再用两个 where 把它进一步折到 [-π/2, π/2]也就是“翻转等价”的数学表达——预测角度 89° 和 -91°与真实角度的差值会被映射到同一个位置。这样损失函数不会在一个 epoch 里忽高忽低。实际训练时角度损失的权重我一般设在 0.5 到 1.0 之间和 CIoU 损失的权重调成同一个数量级过小会导致角度学不到位过大会让网络优先学角度而忽略框的定位。3.3 旋转 NMS 的替代方案没有 cuda 加速也能先跑起来改了输出头和损失之后推理阶段还有一个绕不开的坎后处理 NMS。水平框的 NMS 只需要算 IoU四个顶点一减一除就出来了旋转框的 IoU 需要求多边形交集纯 CPU 算非常慢。社区里常见的做法是对预测框做“松 NMS”先把旋转框外接成水平框用水平 IoU 把大部分冗余框过滤掉只对剩余少量候选做精确的旋转 IoU。这样既保证了精度不损失太多又能在没有旋转 NMS cuda 扩展的环境里先让推理流程跑通。from shapely.geometry import Polygon def rbox_iou(box1, box2): 计算两个旋转框的 IoU box: (cx, cy, w, h, theta)theta 为弧度长边定义 返回 0~1 的 IoU 值 def rect_to_polygon(box): cx, cy, w, h, theta box cos_v, sin_v math.cos(theta), math.sin(theta) # 四个角相对中心的偏移按逆时针顺序排列 dx [w / 2, -w / 2, -w / 2, w / 2] dy [h / 2, h / 2, -h / 2, -h / 2] points [] for i in range(4): x cx dx[i] * cos_v - dy[i] * sin_v y cy dx[i] * sin_v dy[i] * cos_v points.append((x, y)) return Polygon(points) poly1 rect_to_polygon(box1) poly2 rect_to_polygon(box2) if not poly1.is_valid or not poly2.is_valid: return 0.0 inter_area poly1.intersection(poly2).area union_area poly1.area poly2.area - inter_area return inter_area / union_area if union_area 0 else 0.0这段代码里的 rect_to_polygon 是把 (cx, cy, w, h, theta) 转成 shapely 的多边形对象注意四个角点的顺序必须保证是逆时针或者顺时针一致否则 shapely 建出来的多边形可能是自相交的计算面积会出错。实际做批量推理时建议先用外接水平框做一次初筛把候选框数量压到几百个以内再对这个量级的候选做精确旋转 IoU速度能达到实时性的边缘。我自己在测试机上单张 640×640 的图片候选框约 500 个时旋转 NMS 的耗时大概在 20 毫秒左右可以接受。4. 用旋转数据集训练自己的模型标注、超参数与训练流程4.1 标注用什么工具标注旋转框输出什么格式旋转目标检测的数据集标注比水平框麻烦一些因为你要标的是带角度的斜框。最常见的免费标注工具是 roLabelImg它是 LabelImg 的旋转框分支版本快捷键基本沿用 LabelImg用起来几乎没有学习成本输出的是 XML 文件每个目标除了 xmin、ymin、xmax、ymax 之外还带一个 robndbox 节点里面存的是 cx、cy、w、h、thetatheta 单位是角度。如果项目团队习惯用开源社区里更活跃的工具X-AnyLabeling 也支持旋转框标注而且可以输出 COCO 格式的 segmentation 多边形。注意一点不同工具输出的旋转框格式差别很大有的存中心点加宽高加角度有的存四个顶点坐标选定工具之后中间转换脚本基本是少不了的。标注旋转框的时候有个操作习惯建议尽早养成先按目标的“主轴”方向拉框再调整角度贴齐。如果是舰船、飞机这类有明确首尾方向的目标尽量统一用船头/机头方向作为角度参考手工标注偏差控制在 3° 以内对最终精度影响不大但超过 10° 的话损失函数学起来会非常痛苦同一个类别里角度分布被标得七零八散网络很难收敛。4.2 数据集格式适配labelme / DOTA 转成 yolov5_obb 的格式yolov5 原版训练时读的是 txt 文件每行一个目标class_id, x_center, y_center, width, height坐标是归一化到 0~1 的。旋转检测的 txt 格式和它类似只是宽高后面多了个角度class_id, xc, yc, w, h, theta。如果标注工具输出的是四个顶点的坐标比如 labelme 的 segmentation 字段就需要把多边形转成中心点宽高加角度。第 2 章里写的 polygon_to_cxcywha 函数在这里直接用上。# labelme 标注 json 转 yolov5_obb 训练 txt 的脚本核心片段 import json import os import numpy as np def labelme_json_to_obb_txt(json_path, out_txt_path, class_names): labelme 导出的 json 里shape 的 points 存的是四个顶点坐标 转成 yolov5_obb 的 txt 格式cls_id, xc, yc, w, h, theta 坐标需要除以图片宽高做归一化 with open(json_path, encodingutf-8) as f: data json.load(f) img_w, img_h data[imageWidth], data[imageHeight] lines [] for shape in data[shapes]: label shape[label] if label not in class_names: continue cls_id class_names.index(label) points np.array(shape[points], dtypenp.float32).flatten() # 利用第 2 章的函数多边形转旋转矩形 cx, cy, w, h, theta polygon_to_cxcywha(points, angle_formatlongside) # 归一化到 0-1 cx / img_w cy / img_h w / img_w h / img_h lines.append(f{cls_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f} {theta:.6f}) with open(out_txt_path, w, encodingutf-8) as f: f.write(\n.join(lines))这段转换脚本里最容易出错的地方是归一化。宽和高除以图片宽高没有问题但角度不需要归一化直接保留弧度值就行。有些人在转换时顺手把角度也做了归一化训练时解码忘记乘回来画框就全歪了。还有一个隐蔽的坑labelme 标注的“旋转框”实际上是用多边形工具画出来的四个顶点顶点顺序可能是顺时针也可能是逆时针polygon_to_cxcywha 用了 cv2.minAreaRect 去拟合它不依赖顶点顺序所以这个环节反而安全。4.3 训练超参数学习率、batch、anchors 要不要重算旋转检测训练的超参数整体和水平检测类似但有几个地方需要单独调。首先是 anchoryolov5 默认会在训练前用 k-means 重算 anchor这个对旋转框同样适用但注意旋转框的 anchor 匹配计算量更大如果你的数据集目标尺寸很均匀比如全是长宽比接近的舰船可以直接沿用默认 anchor 跳过重算训练反而更稳定。其次是角度损失的权重我前面建议 0.5~1.0如果发现验证集上框的位置已经准了但角度偏了就上调权重如果框的大小都对但位置飘飘忽忽就是角度权重太高压过了 CIoU。学习率和 batch size 的建议单卡 3080 级别batch size 8初始学习率 0.01 是大多数开源仓库的默认配置不用刻意改。但 warmup 的迭代数建议从默认的 3 个 epoch 加到 5 个因为旋转检测前几个 epoch 角度分支的学习信号比较混乱warmup 太长会拖慢收敛太短容易在前期把角度分支的权重带飞。另外训练轮数上公开 DOTA 数据集常见的是 150~200 epoch自建小数据集 100 epoch 左右就能看到明显的收敛拐点后续加了旋转增强之后再继续训练。4.4 配一个能跑起来的 python 环境torch、cuda 和 opencv 的版本关系环境配置这块看着琐碎实际上很多人卡在这一步。旋转检测在 yolov5 的基础上加了新的依赖比如 shapely 用于几何计算opencv-python 用于 minAreaRect 和图像读取。Python 版本建议 3.8 到 3.10torch 用 1.13 或 2.x 都可以但要注意 torchvision 版本必须和 torch 匹配装错了 import 直接报错。在 vscode 里配 python 环境时我一般习惯用 conda 创建独立环境再装依赖避免和系统 python 打架。# conda 创建环境并安装依赖建议按顺序执行 conda create -n yolo_obb python3.9 -y conda activate yolo_obb # 先装 torch 再装其他依赖 pip install torch2.1.0 torchvision0.16.0 --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python shapely pyyaml tqdm tensorboard这里的核心逻辑是先确认自己的显卡驱动支持的 CUDA 版本再选对应版本的 torch 安装。比如 cuda 11.8 的驱动就用 cu118 后缀的轮子cuda 12.x 的驱动装 cu121。如果装错了 torch 版本最常见的报错是调用 cuda 时报 no kernel image available或者训练时提示 torch.cuda.is_available() 返回 False。装好后可以用 python -c import torch, torchvision; print(torch.version, torchvision.version) 快速验证一下跑通之后再进来改代码。5. 旋转目标检测的常见问题排查角度炸了、mAP 不涨、NMS 框错乱5.1 角度周期性导致损失震荡不收敛现象训练 loss 前期下降正常到某个阶段后开始有规律地上下震荡验证集 mAP 卡在某个值反复横跳。原因是角度损失没有处理 90° 翻转等价性模型预测 89° 而标签是 -91° 时数值上差了 180°梯度把参数往错误方向推。解决方法是换成第 3.2 节里带周期性映射的角度损失函数——先对差值取余数再折到 [-90°, 90°] 区间内确保物理上同一个朝向不会被网络当作两个完全不同的目标。5.2 训练集标注混用了两套角度定义现象训练后验证集画出的框一部分是正的一部分是歪了 90° 的可视化结果非常诡异。原因是你的数据集里一部分标注是 OpenCV 定义的短边角另一部分是长边角网络学到的角度分布是双峰的。排查时要先统计训练集里所有标注的角度直方图如果在 0° 和 ±90° 附近都出现明显峰值基本可以断定定义混杂。解决方法是写一个脚本统一检查所有 txt 标注手动确认每张图的目标方向与角度值的对应关系然后统一转换到同一种角度定义。5.3 Mosaic 增强把旋转框坐标整错了现象开启 Mosaic 增强后验证集 mAP 比不开还低可视化时发现部分检测框的位置是对的但角度偏了大约 45°。原因是 Mosaic 拼接图片时做了旋转和缩放变换但数据加载器里变换的是图片像素标注的五个量 (cx, cy, w, h, theta) 没有同步做旋转矩阵变换。解决方法是如果用的是现成旋转检测仓库检查增强代码里是否对 polygon 或 rbox 做了同步变换如果是自研的增强管线优先关掉 Mosaic换成普通随机缩放、平移和水平翻转——水平翻转时要记得把角度取反。5.4 旋转 NMS 阈值设得太低把重叠目标干掉现象推理时密集排列的车辆或飞机一排几个目标被合并成了一个框。原因是旋转 NMS 的阈值沿用水平检测的 0.45 不合理遥感俯视图里相邻的旋转目标如果朝向一致旋转框的交并比天然比水平框高。解决方法是把 NMS 阈值从 0.45 微调到 0.3 左右同时配合置信度阈值过滤把低分框提前滤掉这样密集场景的漏检会明显改善。5.5 小目标角度预测大偏差但坐标框位置正常现象验证集上小目标的 mAP 偏低可视化发现小目标的框位置贴得挺准就是方向转得不对。原因是小目标本身像素少卷积特征图里角度信息本来就弱再加上角度损失在整体损失中占比不够模型选择了把容量优先分配给收益更明显的坐标回归。解决方法是给角度损失加一个和物体尺寸相关的缩放系数小目标的损失权重适当放大 1.2~1.5 倍或者把特征图上的小目标分支在损失计算时单独加权。还有一个工程技巧推理时对小目标使用更大的测试时增强多尺度水平翻转角度预测的稳定性会有肉眼可见的提升。6. 结果评估与落地部署画框可视化、精度指标和导出技巧旋转检测训练完第一件事不是看 mAP而是把验证集的预测结果可视化出来逐张检查角度是不是真的贴合目标主轴。画旋转框不能沿用 yolov5 那个 matplotlib Rectangle需要自己按角度计算四个顶点坐标再画。下面这段代码是推理可视化最常用的函数直接把它放在检测脚本里调用即可。import cv2 import numpy as np def draw_rbox(img, box, color(0, 255, 0), thickness2): 在图上画旋转框 box: (cx, cy, w, h, theta)theta 为弧度 img 会被原地修改 cx, cy, w, h, theta box cos_v, sin_v np.cos(theta), np.sin(theta) dx np.array([w / 2, -w / 2, -w / 2, w / 2]) dy np.array([h / 2, h / 2, -h / 2, -h / 2]) points np.zeros((4, 2), dtypenp.int32) for i in range(4): x cx dx[i] * cos_v - dy[i] * sin_v y cy dx[i] * sin_v dy[i] * cos_v points[i] [int(x), int(y)] cv2.polylines(img, [points], isClosedTrue, colorcolor, thicknessthickness) return img可视化通过之后再进入评估环节。旋转检测的 mAP 计算不能用 yolov5 里现成的 evaluate 脚本因为标准的 mAP 计算依赖水平 IoU需要替换成旋转 IoU 版本。具体做法是在评估脚本里把人脸框的 IoU 计算函数换成第 3.3 节的 rbox_iou并且把 IoU 阈值的判断逻辑保持不变。注意 rbox_iou 是逐个框计算的如果验证集图片多评估速度会很慢建议先把每张图片的预测结果和标注都缓存下来再一次性计算所有 IoU 矩阵做排序。部署方面如果是把模型导成 TensorRT 或者 ONNX 跑边缘设备yolov5 的导出命令基本可以直接复用但要注意角度分支的输入输出形状。导出 ONNX 时需要把后处理里的旋转 NMS 剥离掉只保留网络本身的输出在推理代码里自己做旋转 NMS。这个和水平检测一样是常规操作——把后处理留到部署框架外部处理灵活性最好。如果目标是树莓派 5 之类的低算力设备建议把模型输入尺寸从 640 降到 512并且把模型宽度缩放因子从 0.33 调低到 0.25角度预测的精度损失相对可控但推理速度能提升将近一倍。回想我自己第一次把旋转检测跑通到真实验收场景时最深的一个教训是不要一上来就追求复杂的角度分类或者旋转 anchor 方案先把连续回归加周期损失这条路走通把数据结构、标注转换和可视化链路理顺再考虑换更精细的角度预测方式。很多模型精度不够不是网络不行而是标签格式、角度定义和 NMS 阈值这些底层环节有偏差。这套方案踩稳之后再迁移到任意旋转检测场景都只是换数据集和调参的问题。希望这篇笔记能帮你少绕几个弯。本文还有配套的精品资源点击获取
