简介这份资源面向计算机视觉开发者与目标检测学习者提供基于Python与YOLOv5的旋转目标检测完整实现用于解决倾斜、旋转物体难以用常规水平矩形框精确定位的问题适用于遥感影像、航拍、工业质检等场景。压缩包共150个文件约6.26MB以66个Python脚本和33个YAML配置为主涵盖模型定义、训练与推理流程同时包含CUDA与C扩展源码用于旋转框IoU及NMS的加速计算另有少量Markdown说明、Shell脚本与Dockerfile辅助环境搭建。已有852人学习下载。资源围绕Oriented Bounding Box展开涉及角度回归、旋转数据增强、GIOU/DIoU损失调整、角度感知NMS后处理等关键环节并给出从环境配置、数据准备、模型训练到评估与推理的完整链路便于读者直接复现旋转目标检测系统并迁移到自有数据集。1. 旋转框检测到底难在哪从水平框到带角度的 YOLOv5普通 YOLOv5 输出的是水平矩形框四个值(x, y, w, h)就能描述一个目标。但遥感图像里的飞机、舰船工业质检里的PCB板、药片文本检测里的倾斜文字行这些目标的长宽比悬殊且方向任意。用水平框去套一个斜着停的飞机框里会塞进大量背景像素两个相邻的斜向目标还会因为水平框大面积重叠而被NMS误杀。旋转目标检测要解决的就是这个问题在(x, y, w, h)基础上多预测一个角度 θ让框能贴着目标转。基于 Python 的 YOLOv5 实现旋转目标检测核心工作集中在三块——数据标注格式从水平框换成旋转框、网络头部增加角度分支、损失函数和后处理适配角度回归。这篇笔记按「标注怎么转 → 环境怎么配 → 模型怎么改 → 训练怎么调 → 坑在哪」的顺序展开适合已经跑通过水平 YOLOv5、想往旋转框方向推进的从业者也适合刚接触旋转检测、需要一份可复现路径的新手。下面所有操作我都在 Ubuntu 20.04 Python 3.8 PyTorch 1.10 上实际跑过Windows 下路径和依赖稍有差异关键处会单独说明。2. 旋转框标注格式与数据准备DOTA 转 YOLO 的完整链路2.1 旋转框的三种表示法及选型理由旋转框主流有三种参数化方式。第一种是(x, y, w, h, θ)θ 表示框相对于水平轴的旋转角取值范围常见为[-90, 0)或[-45, 45)。第二种是四角点(x1, y1, x2, y2, x3, y3, x4, y4)DOTA 数据集就是这种格式。第三种是 OpenCV 的RotatedRect用中心点、宽高、角度描述但角度定义和 YOLO 系不一致。选哪种取决于你的训练框架。YOLOv5 的旋转版本社区常见做法是在models/yolo.py的 Detect 头里把输出通道从na * (nc 5)改成na * (nc 6)多出来的一个通道给角度通常采用(x, y, w, h, θ)五参数加置信度和类别。我一般推荐用[-90, 0)的角度范围原因是这个区间下w始终对应框的长边角度回归不会出现w和h互换导致的跳变。如果用[-45, 45)当目标接近 45 度时长宽定义会翻转损失函数容易震荡。DOTA 格式转 YOLO 旋转格式核心是把四角点转成中心点加宽高加角度。转换时要注意两个边界一是角度归一化二是框的顶点顺序。DOTA 的标注不保证顶点按顺时针或逆时针排列直接算角度会得到错误结果。2.2 DOTA 转 YOLO 旋转格式的转换脚本下面这个脚本处理 DOTA 的labelTxt格式输出 YOLO 旋转训练所需的(class_id, cx, cy, w, h, angle)归一化坐标。假设你的数据目录结构是datasets/dota/images/和datasets/dota/labelTxt/。import os import math import numpy as np # DOTA 类别按你的数据集实际类别修改 CLASSES [plane, ship, storage-tank, baseball-diamond, tennis-court, swimming-pool, ground-track-field, harbor, bridge, large-vehicle, small-vehicle, helicopter, roundabout, soccer-ball-field, basketball-court] def poly_to_rbox(poly): 四角点转旋转框 (cx, cy, w, h, angle)angle 范围 [-90, 0) pts np.array(poly, dtypenp.float32).reshape(4, 2) # 按 x 坐标排序再按 y 微调保证顶点顺序一致 pts pts[np.argsort(pts[:, 0])] # 取前两个点作为一条边后两个点作为另一条边 edge1 pts[1] - pts[0] edge2 pts[2] - pts[3] # 计算两条边的长度和角度 len1 np.linalg.norm(edge1) len2 np.linalg.norm(edge2) angle1 math.degrees(math.atan2(edge1[1], edge1[0])) angle2 math.degrees(math.atan2(edge2[1], edge2[0])) # 取较长边作为 w较短边作为 h if len1 len2: w, h len1, len2 angle angle1 else: w, h len2, len1 angle angle2 # 归一化到 [-90, 0) if angle 0: angle - 90 if angle -90: angle 90 cx np.mean(pts[:, 0]) cy np.mean(pts[:, 1]) return cx, cy, w, h, angle def convert_dota_to_yolo(img_dir, label_dir, out_dir, img_w1024, img_h1024): os.makedirs(out_dir, exist_okTrue) for txt_name in os.listdir(label_dir): if not txt_name.endswith(.txt): continue img_name txt_name.replace(.txt, .png) lines_out [] with open(os.path.join(label_dir, txt_name), r) as f: for line in f: parts line.strip().split() if len(parts) 9: continue cls_name parts[8] if cls_name not in CLASSES: continue cls_id CLASSES.index(cls_name) poly [float(x) for x in parts[:8]] cx, cy, w, h, angle poly_to_rbox(poly) # 归一化 cx_n cx / img_w cy_n cy / img_h w_n w / img_w h_n h / img_h angle_n angle / 90.0 # 归一化到 [-1, 0) lines_out.append(f{cls_id} {cx_n:.6f} {cy_n:.6f} {w_n:.6f} {h_n:.6f} {angle_n:.6f}) with open(os.path.join(out_dir, txt_name), w) as f: f.write(\n.join(lines_out)) if __name__ __main__: convert_dota_to_yolo( img_dirdatasets/dota/images, label_dirdatasets/dota/labelTxt, out_dirdatasets/dota/labels_rotated )逻辑说明poly_to_rbox先按 x 坐标排序顶点避免 DOTA 标注顶点乱序导致边长计算错误。然后取两条对边分别算长度和角度长边作为w短边作为h角度归一化到[-90, 0)。convert_dota_to_yolo遍历所有标注文件逐行解析跳过类别不在列表里的行最后把归一化后的六元组写入输出目录。参数说明img_w和img_h默认 1024DOTA 原始图像就是 1024×1024如果你的数据是其他尺寸必须改成实际值否则归一化坐标全错。angle_n angle / 90.0把角度缩放到[-1, 0)训练时角度分支用tanh激活再乘 90 还原这是社区常见做法。如果你的框架用弧度制改成angle / (math.pi / 2)。2.3 数据增强在旋转框上的特殊处理水平框的增强手段不能直接照搬到旋转框。随机缩放和随机裁剪对旋转框是安全的因为框的中心点和宽高按同样比例变换即可。但随机翻转要小心水平翻转后角度 θ 要变成-180 - θ再归一化回[-90, 0)垂直翻转后角度变成-θ。如果增强代码里没处理角度翻转后的标注会完全错位。Mosaic 增强在旋转框上也有坑。四张图拼接后每张图的旋转框要按拼接偏移量平移中心点角度不变。但拼接边界处的目标如果被截断旋转框的宽高会失真。我一般把 Mosaic 的概率从水平框的 0.5 降到 0.3减少截断带来的噪声。提示转换完标注后务必用可视化脚本抽查 20 张图把旋转框画到原图上确认框能贴合目标。我见过太多因为顶点顺序问题导致角度全错的案例训练 loss 不降反升排查半天才发现是标注转换的锅。3. YOLOv5 旋转检测的环境配置与模型改造3.1 环境配置conda 建环境到依赖安装的完整命令旋转 YOLOv5 的代码基座还是 ultralytics 的 YOLOv5但需要替换 Detect 头和损失函数。环境配置和标准 YOLOv5 基本一致额外需要shapely和opencv-python用于旋转框的 IoU 计算和可视化。# 创建 conda 环境 conda create -n yolov5_rotated python3.8 -y conda activate yolov5_rotated # 安装 PyTorch根据你的 CUDA 版本选择这里以 CUDA 11.3 为例 pip install torch1.10.0cu113 torchvision0.11.1cu113 -f https://download.pytorch.org/whl/torch_stable.html # 克隆 YOLOv5 官方仓库旋转版本通常基于 v6.0 或 v7.0 修改 git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt # 额外依赖 pip install shapely opencv-python逻辑说明先建独立环境避免和系统 Python 冲突PyTorch 版本要和 CUDA 驱动匹配装完用python -c import torch; print(torch.cuda.is_available())验证。YOLOv5 仓库克隆后requirements.txt会装好基础依赖shapely用于计算旋转框多边形 IoUopencv-python用于画旋转框。参数说明CUDA 版本不是 11.3 的话去 PyTorch 官网查对应命令替换。如果只用 CPU 训练把cu113去掉但旋转检测训练量通常不小CPU 训练会非常慢不建议。3.2 修改 Detect 头增加角度输出通道YOLOv5 的 Detect 头在models/yolo.py里原始输出是na * (nc 5)其中na是 anchor 数量nc是类别数5 是(x, y, w, h, obj)。旋转检测要改成na * (nc 6)多出来的是角度。# models/yolo.py 中 Detect 类的 __init__ 和 forward 需要改 class Detect(nn.Module): stride None onnx_dynamic False def __init__(self, nc80, anchors(), ch()): super().__init__() self.nc nc self.no nc 6 # 原来是 nc 5改成 nc 6 self.nl len(anchors) self.na len(anchors[0]) // 2 self.grid [torch.zeros(1)] * self.nl self.anchor_grid [torch.zeros(1)] * self.nl self.register_buffer(anchors, torch.tensor(anchors).float().view(self.nl, -1, 2)) self.m nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) self.act nn.Sigmoid() def forward(self, x): z [] for i in range(self.nl): x[i] self.m[i](x[i]) bs, _, ny, nx x[i].shape x[i] x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() if not self.training: # 推理时对角度做 tanh 激活再乘 90 x[i][..., 5] torch.tanh(x[i][..., 5]) * 90.0 # 其余部分和原始 YOLOv5 一致 y x[i].sigmoid() y[..., 0:2] (y[..., 0:2] * 2. - 0.5 self.grid[i]) * self.stride[i] y[..., 2:4] (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] z.append(y.view(bs, -1, self.no)) return x if self.training else (torch.cat(z, 1), x)逻辑说明self.no从nc 5改成nc 6卷积输出通道相应增加。推理时对角度通道单独做tanh激活再乘 90把网络输出映射到[-90, 90]再根据训练时的归一化方式截取到[-90, 0)。训练时角度通道不激活直接输出给损失函数。参数说明self.act在原始代码里用于某些版本这里保留但不影响。角度激活函数用tanh而不是sigmoid因为角度有正负sigmoid只能输出正值。如果你的角度范围是[-45, 45)把乘的 90 改成 45。3.3 旋转框 IoU 与损失函数适配水平框的 CIoU 损失不能直接用于旋转框因为旋转框的 IoU 计算涉及多边形相交。社区常见做法是用shapely计算两个旋转框多边形的交并比然后套用 CIoU 的形式。from shapely.geometry import Polygon import torch def rbox_iou(box1, box2): 计算两个旋转框的 IoUbox 格式 (cx, cy, w, h, angle) def to_polygon(box): cx, cy, w, h, angle box angle_rad math.radians(angle) cos_a, sin_a math.cos(angle_rad), math.sin(angle_rad) # 四个角点相对于中心点的偏移 dx w / 2 dy h / 2 corners [(-dx, -dy), (dx, -dy), (dx, dy), (-dx, dy)] pts [] for x, y in corners: px cx x * cos_a - y * sin_a py cy x * sin_a y * cos_a pts.append((px, py)) return Polygon(pts) poly1 to_polygon(box1) poly2 to_polygon(box2) inter poly1.intersection(poly2).area union poly1.area poly2.area - inter return inter / (union 1e-7)逻辑说明to_polygon把旋转框的四个角点算出来构造 shapely 的 Polygon 对象。intersection求交集面积union求并集面积相除得到 IoU。这个函数在损失函数里逐对调用计算量比水平框 IoU 大训练时建议只在正样本和预测框之间算不要全量两两计算。参数说明angle单位是度和前面标注转换保持一致。1e-7防止除零。如果训练速度太慢可以用 CUDA 实现的旋转 IoU 算子替代但需要编译自定义扩展新手先用 shapely 版本跑通再说。注意shapely 的 Polygon 在框退化成线或点时面积为零IoU 计算会返回 0损失函数里要加判断跳过这些无效框否则会出现 NaN。4. 训练配置与超参数调整旋转检测的调参经验4.1 数据配置文件与 anchor 调整YOLOv5 的数据配置文件data/dota_rotated.yaml需要指定训练集、验证集路径和类别数。旋转检测的 anchor 建议重新聚类因为旋转框的宽高分布和水平框不同。# data/dota_rotated.yaml train: datasets/dota/images/train val: datasets/dota/images/val nc: 15 names: [plane, ship, storage-tank, baseball-diamond, tennis-court, swimming-pool, ground-track-field, harbor, bridge, large-vehicle, small-vehicle, helicopter, roundabout, soccer-ball-field, basketball-court]逻辑说明train和val指向图像目录YOLOv5 会自动找同名的labels_rotated目录下的标注文件。nc是类别数names按类别 id 顺序排列。参数说明如果你的标注目录名不是labels_rotated要么改目录名要么在utils/datasets.py里改路径拼接逻辑。anchor 聚类用python utils/autoanchor.py --data data/dota_rotated.yaml但旋转框的宽高要先把角度考虑进去简单做法是直接用水平框的 anchor 先跑观察best possible recall是否低于 0.9低了再手动调。4.2 训练命令与关键超参数启动训练的命令和标准 YOLOv5 类似但有几个超参数需要针对旋转检测调整。python train.py \ --data data/dota_rotated.yaml \ --cfg models/yolov5s_rotated.yaml \ --weights yolov5s.pt \ --epochs 100 \ --batch-size 8 \ --img-size 1024 \ --hyp data/hyp.rotated.yaml \ --name dota_rotated_exp逻辑说明--cfg指向修改过 Detect 头的模型配置文件--weights加载预训练权重加速收敛--img-size 1024因为 DOTA 图像是 1024×1024--hyp指定旋转检测专用的超参数文件。参数说明--batch-size 8在 1024 分辨率下显存占用较大11GB 显存大概能跑 8不够就降到 4 并开梯度累积。--epochs 100是起步值DOTA 这种数据量通常要 200 到 300 轮。--hyp文件里重点调三个box损失权重从 0.05 提到 0.08因为角度回归需要更强的监督cls保持 0.5lr0从 0.01 降到 0.005旋转检测的损失曲面更复杂学习率大了容易震荡。4.3 角度损失的权重与归一化方式角度损失在总损失里的占比很关键。占比太低模型学不会角度占比太高位置和分类的精度会下降。我一般把角度损失单独算权重设为 box 损失的 0.5 倍即angle_loss_weight 0.5。角度归一化方式影响损失尺度。如果用[-1, 0)归一化角度损失用 MSE 时数值很小需要放大权重。如果用[-90, 0)不归一化MSE 数值大权重可以小一些。我推荐归一化后用smooth_l1_loss对异常值更鲁棒。# 角度损失计算示例 angle_pred pred[..., 5] # 预测角度归一化到 [-1, 0) angle_gt target[..., 5] # 真实角度同样归一化 angle_loss F.smooth_l1_loss(angle_pred, angle_gt, reductionmean) # 总损失里加上 total_loss box_loss cls_loss obj_loss 0.5 * angle_loss逻辑说明smooth_l1_loss在误差小于 1 时等价于 MSE大于 1 时等价于 MAE避免角度预测偏离太大时梯度爆炸。0.5是角度损失权重可以根据训练曲线调整。参数说明如果训练几个 epoch 后角度损失不降先把权重提到 1.0 试试。如果角度损失降了但 mAP 不升说明角度学过头了把权重降到 0.3。提示训练时用tensorboard单独看角度损失的曲线和 box 损失对比。正常情况两者应该同步下降如果角度损失震荡而 box 损失平稳多半是角度归一化范围和数据不匹配。5. 旋转检测的避坑与排查5 个血泪教训5.1 角度跳变导致 loss 突然飙升现象训练到某个 epochloss 从 0.5 突然跳到 5.0 以上之后再也降不下来。原因角度在[-90, 0)边界附近真实角度是 -89 度预测角度是 -1 度两者实际只差 2 度但数值上差了 88MSE 损失巨大。这是角度回归的周期性问题-90和0在几何上是同一个方向。解决在角度损失里加周期性处理把角度差映射到[-45, 45]范围内再算损失。具体做法是diff angle_pred - angle_gt然后diff diff - 90 * torch.round(diff / 90)这样 -89 和 -1 的差变成 2 而不是 88。5.2 NMS 把相邻斜向目标误删现象验证集可视化时两艘并排斜向停靠的船只检测出一艘另一艘被 NMS 抑制了。原因旋转框的 NMS 如果还用水平框的 IoU 计算两个斜向框的水平外接矩形重叠面积很大IoU 超过阈值被误判为同一目标。解决NMS 阶段改用旋转框 IoU 计算。在utils/general.py的non_max_suppression函数里把 IoU 计算替换成前面写的rbox_iou。注意这个函数是逐对计算的速度会慢可以先用水平框 NMS 粗筛再对保留的框做旋转 NMS 精筛。5.3 标注转换后类别 id 错位现象训练时 loss 正常下降但推理时所有目标都预测成同一个类别或者类别完全乱套。原因DOTA 的类别列表顺序和你的data.yaml里names顺序不一致。转换脚本里CLASSES.index(cls_name)得到的 id 是按脚本里的顺序但训练时按 yaml 里的顺序解释两者对不上。解决把转换脚本里的CLASSES列表和data.yaml里的names列表逐字对齐包括大小写和连字符。改完后重新生成所有标注文件不要只改一部分。5.4 显存溢出但 batch-size 已经降到 1现象--batch-size 1还是 OOM报错CUDA out of memory。原因旋转 IoU 计算在损失函数里构造了大量 shapely Polygon 对象这些对象在 CPU 上但中间变量可能留在 GPU 上没释放。另外 1024 分辨率的特征图本身就大。解决把--img-size降到 640 先跑通确认模型没问题再逐步升到 1024。同时在损失函数里用with torch.no_grad()包住 IoU 计算部分减少中间变量。如果还不行用--accumulate做梯度累积等效增大 batch。5.5 推理时角度输出全是 0 或全是 -90现象模型训练 loss 正常但推理时所有框的角度都是 0 或者都是 -90框全部水平。原因推理时角度通道的激活函数写错了。如果用了sigmoid输出范围是(0, 1)乘 90 后是(0, 90)但你的角度范围是[-90, 0)所有预测都被截断到 0。或者tanh乘 90 后得到[-90, 90]但后处理里没做范围截取。解决确认推理代码里角度激活是tanh而不是sigmoid激活后乘 90 得到[-90, 90]然后判断如果角度大于 0 就减 90映射回[-90, 0)。后处理里加一行angle angle - 90 if angle 0 else angle。6. 旋转检测的进阶技巧从能跑到跑好的三个抓手6.1 用 KLD 损失替代 MSE 做角度回归MSE 角度损失在边界处不连续KLDKullback-Leibler Divergence损失把旋转框建模成高斯分布用分布之间的距离代替角度数值差天然处理周期性。社区有开源的 KLD 实现核心是把(x, y, w, h, θ)转成高斯分布的均值和协方差然后算两个高斯分布的 KL 散度。我实测在 DOTA 上 KLD 比 MSE 的 mAP 高 2 到 3 个点但训练前期收敛慢一些建议前 20 个 epoch 用 MSE 预热之后切 KLD。6.2 多尺度训练对旋转框的收益DOTA 的目标尺度差异极大从几像素的小车到几百像素的球场。多尺度训练--img-size在 640 到 1024 之间随机能让模型适应不同尺度。但旋转框在多尺度下要注意缩放图像时旋转框的中心点和宽高按比例缩放角度不变。如果你的数据加载代码里缩放没同步处理标注多尺度训练会引入噪声。我一般把多尺度概率设为 0.5不要每张图都变。6.3 验证旋转检测精度的正确姿势水平框的 mAP 计算不能直接用于旋转框。旋转框的 mAP 要先用旋转 IoU 判断预测框和真实框是否匹配再算 AP。常见做法是用 DOTA 官方提供的评估脚本它按类别算 AP 再平均。如果你自己写评估注意 IoU 阈值从 0.5 开始步长 0.05到 0.95和 COCO 一致。另外旋转框的 mAP 通常比水平框低 5 到 10 个点这是正常的不要因为数字低就以为模型没训好。评估指标水平框典型值旋转框典型值说明mAP0.50.750.68旋转框匹配更严格mAP0.5:0.950.520.45高 IoU 阈值下差距更大推理速度 FPS4538旋转 NMS 额外耗时我踩过最深的坑是拿水平框的 mAP 预期去衡量旋转模型训到 0.6 就以为翻车了其实是正常水平。后来养成习惯每次换任务先跑一个 baseline用 baseline 的数字做参照而不是用其他任务的数字。希望帮到你。本文还有配套的精品资源点击获取