YOLOv5 OBB旋转框训练实战：从DOTA标注到角度回归与部署-平芜编程栈

简介：面向目标检测开发者与计算机视觉学习者的 YOLOv5 OBB 旋转框训练演示工程，在 YOLOv5 上扩展倾斜框支持，用顶点与旋转角度描述树木、车辆等带方向目标，弥补轴对齐框定位不准的缺陷。压缩包共 583 个文件、约 445.58MB，xml 与 txt 为原始标注及转换后的标签，png、jpg 是样例图与可视化结果，py、yaml、cmd 覆盖预处理、配置、训练与导出脚本，另有 pt 权重、onnx 模型和 cpp、cu 算子供部署参考。已有 1822 人学习下载。工程把标注转换、OBB 配置修改、旋转框损失计算、validate 评估与推理可视化串成可运行链路，并附 dockerfile、makefile 与 TensorBoard 事件记录，便于对照复现训练、观察损失曲线、排查环境与精度问题。

1. 航拍图里斜着停的车，正框永远框不干净

遥感影像里一架斜着停的飞机，轴对齐框要把它和半条跑道一起框进去；产线上歪放着的工件，正框贴不紧边，裁出来的图总带一圈背景；票据版面里略微倾斜的一行字，正框一裁就把上下两行也带进来。这些场景我拿 YOLOv5 都试过，检测本身没毛病，问题出在框的形状上。这份 YOLOv5 OBB 旋转框训练 demo 干的事就一件：把输出从(x, y, w, h)扩成(x, y, w, h, θ)，让模型直接吐带角度的框。它适合已经有 YOLOv5 训练经验、手上又有一批斜目标数据的人，遥感 DOTA、航拍车辆、工业零件都能套。当成一个可直接运行的软件插件改，比自己从零在 anchor 上加角度通道省太多事。

2. 工程骨架与技术栈：从目录结构看一个 batch 怎么流到 loss

拿到压缩包先别急着跑train.py，先花十分钟把目录看一遍，比后面 debug 两小时划算。旋转框版本的 YOLOv5 在骨架上和原版几乎一致，差异集中在三个地方：数据读取时多吐一个角度、检测头多预测一个通道、后处理换了一套 IoU 算法。把这三处找准了，剩下的都是原版逻辑。

2.1 目录骨架与旋转框专属文件

这套 demo 的目录结构大致如下，和官方 YOLOv5 保持兼容，方便直接拿预训练权重做迁移。

目录 / 文件	职责	旋转框改造点
`data/`	数据集 yaml、超参 yaml、样例图	新增`dota.yaml`，路径指向切好的小图
`models/`	网络定义 yaml 与`common.py`、`yolo.py`	`Detect`头换成`DetectOBB`，输出通道 +1
`utils/dataloaders.py`	Dataset、collate、数据增强	标签读取改成五参数，增强要同步转角度
`utils/loss.py`	`ComputeLoss`正负样本匹配与回归	增加角度分支损失
`utils/general.py`	NMS、坐标变换、日志	`non_max_suppression`换成旋转版
`utils/metrics.py`	mAP 计算	`bbox_iou`换成旋转 IoU
`utils/rbox.py`	旋转框工具函数（新增）	`poly2rbox`、`rbox2poly`、`rbox_iou`
`train.py`/`val.py`/`detect.py`/`export.py`	四个入口脚本	基本不动，参数透传

新增的utils/rbox.py是整份代码里被调用最频繁的文件，转换、IoU、画图都从这儿走。我一般先把它的函数签名固定下来，后面所有模块都按这个约定写，避免出现「训练用弧度、推理用角度」这种玄学 bug。

2.2 一个 batch 的完整调用链

从train.py进去，主干链路是这样的：create_dataloader建LoadImagesAndLabels，__getitem__读图、做增强、把标签整理成(cls, cx, cy, w, h, θ)；collate_fn把不同数量的框拼成一个 batch；送进DetectionModel后，三个尺度的特征图分别过DetectOBB头；最后一层输出被ComputeLoss拿去做正负样本匹配和损失累加。

# train.py 主干（节选，保留和旋转框相关的部分） train_loader, dataset = create_dataloader( train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls, hyp=hyp, cache=opt.cache, rect=opt.rect, rank=local_rank, workers=workers, image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '), ) model = DetectionModel(cfg, nc=model.nc, anchors=hyp.get('anchors')).to(device) compute_loss = ComputeLoss(model) # 内部会判断 head 是否带角度通道 for imgs, targets, paths, _ in train_loader: imgs = imgs.to(device).float() / 255.0 pred = model(imgs) # 训练态返回三个尺度的原始输出 loss, loss_items = compute_loss(pred, targets.to(device)) loss.backward()

pred在训练态下不做 decode，直接是(B, na * no, H, W)的张量，no = nc + 6，多出来的那个 1 就是角度。角度放在最后一维，是为了让原版的xywh解码逻辑一行都不用改，索引偏移天然对齐。

提示：第一次跑之前打印一次pred[0].shape，确认最后一维等于nc + 6。如果还是nc + 5，说明models/yolo.py里解析 head 时没走DetectOBB分支，训练会静默丢掉角度，跑完 mAP 低得莫名其妙。

2.3 环境依赖与装法

依赖不复杂，和原版 YOLOv5 基本重合，额外需要的就是 OpenCV 的旋转矩形相交接口，cv2.rotatedRectangleIntersection在opencv-python>=4.5里都有。

# 建议新建虚拟环境，torch 版本跟显卡驱动对齐即可 conda create -n obb python=3.8 -y conda activate obb pip install torch==1.13.1 torchvision==0.14.1 --extra-index-url https://download.pytorch.org/whl/cu117 pip install opencv-python numpy pyyaml tqdm scipy matplotlib seaborn pandas pip install shapely # 可选的旋转框几何工具，画图时方便

这里有个小细节：opencv-python和opencv-python-headless不要同时装，服务器上没有图形界面时装 headless 版本即可，混装会导致cv2.imshow报错但cv2.rotatedRectangleIntersection又正常，排查起来很费时间。

3. 数据准备：DOTA 八点标注怎么变成可训练的五参数标签

旋转框项目里，数据这块的工作量占一半以上。模型的改造其实不难，难的是把八点标注和多边形裁切搞对，角度定义一旦前后不一致，训练曲线看着正常，验证集 mAP 能直接给你掉到 0.01。

3.1 DOTA 原始格式与切图策略

DOTA 的标注文件一行是x1 y1 x2 y2 x3 y3 x4 y4 类别 difficult，坐标是原图上的绝对像素，一张图动辄 8000×8000。直接丢进网络显存扛不住，而且小目标会被缩得只剩几个像素，所以标准做法是先切图再训练。

切图参数一般取size=1024, gap=200。gap是相邻两张子图的重叠区域，作用是防止目标正好被切在边界上。切完之后要按同样的偏移量重新算每个框的坐标，中心点不在子图范围内的直接丢掉。

def slice_image(img, rboxes, size=1024, gap=200): """ img: HxWx3 的原始大图 rboxes: (N, 5) 的 cx cy w h angle(度) 返回子图列表和对应的旋转框 """ H, W = img.shape[:2] step = size - gap xs = list(range(0, max(W - size, 0) + 1, step)) ys = list(range(0, max(H - size, 0) + 1, step)) # 边界补一张，保证右下角也被覆盖 if xs[-1] != W - size: xs.append(max(W - size, 0)) if ys[-1] != H - size: ys.append(max(H - size, 0)) crops = [] for y in ys: for x in xs: sub_img = img[y:y + size, x:x + size] sub_boxes = rboxes.copy() sub_boxes[:, 0] -= x sub_boxes[:, 1] -= y keep = (sub_boxes[:, 0] > 0) & (sub_boxes[:, 0] < size) & \ (sub_boxes[:, 1] > 0) & (sub_boxes[:, 1] < size) crops.append((sub_img, sub_boxes[keep])) return crops

三个参数的含义要记牢：size决定输入分辨率，和目标尺度强相关；gap越大冗余越多，磁盘占用也越大，200 是个经验值；step是实际步长，由前两者算出来，千万别自己另外指定，否则会出现某些区域漏切。

提示：切图脚本一定要做成一次性离线跑完，把子图和标签落盘成images/和labels/两份文件。在线切图会让每个 epoch 的耗时翻好几倍，dataloader 的 worker 也很难做缓存。

3.2 八点转五参数：poly2rbox 的实现

八点到五参数的标准做法是求最小外接旋转矩形，再统一成「长边为 w、角度落在 [-90, 0)」的约定。这个约定必须写死在一个函数里，训练、推理、评估全部调它，不要每个文件各写一份。

import cv2 import numpy as np def poly2rbox_8pts(polys, img_size=None): """ polys: (N, 8)，DOTA 的 x1 y1 x2 y2 x3 y3 x4 y4 img_size: (W, H)，传入时把 cx cy w h 归一化到 [0, 1] 返回: (N, 5)，顺序为 cx cy w h angle(度)，angle ∈ [-90, 0) """ rboxes = np.zeros((len(polys), 5), dtype=np.float32) for i, poly in enumerate(polys): pts = np.asarray(poly, dtype=np.float32).reshape(4, 2) (cx, cy), (w, h), ang = cv2.minAreaRect(pts) # 长边定义法：交换 w/h 时角度补 90 度，保证几何等价 if w < h: w, h = h, w ang += 90.0 # 统一收敛到 [-90, 0) while ang >= 0: ang -= 90.0 while ang < -90: ang += 90.0 rboxes[i] = [cx, cy, w, h, ang] if img_size is not None: W, H = img_size rboxes[:, [0, 2]] /= W # cx、w 按宽归一化 rboxes[:, [1, 3]] /= H # cy、h 按高归一化 return rboxes

cv2.minAreaRect返回的角度本身就在[-90, 0)区间（不同 OpenCV 版本略有差异，4.5 之后统一了），但当你为了「长边为 w」交换w和h时，角度必须同步加 90 度，否则同一个框会表示成两种形式，模型学到的角度分布就是两个峰，回归直接乱掉。

归一化这一步要特别注意：w除以图像宽度、h除以图像高度，而不是都除以 1024。输入如果不是正方形，两者不等价，推理时反归一化就会把框拉变形。

3.3 增强时的角度处理与归一化

水平框的增强（翻转、缩放、裁剪）在旋转框场景下大部分还能用，但每一条都要检查角度是否跟着变。水平翻转时角度取反，垂直翻转时角度同样取反后再偏移，旋转类增强则直接给角度加上旋转量。

def augment_flip_h(rboxes, w_img): """水平翻转：cx 镜像，角度取反后重新收敛到 [-90, 0)""" out = rboxes.copy() out[:, 0] = w_img - out[:, 0] out[:, 4] = -out[:, 4] out[:, 4] = np.where(out[:, 4] >= 0, out[:, 4] - 90, out[:, 4]) return out

颜色抖动、马赛克、mixup 这些不影响几何的增强可以放心用。mosaic 在旋转框上的实现要点是：四张子图拼完之后，每个框的cx cy要按粘贴偏移量重算，角度不变；如果拼接时做了缩放，w h也要同步乘缩放系数。

这里有个容易被忽略的点：角度归一化到网络输出时，一般除以 90 映射到[-1, 0]区间，或者干脆用弧度除以π/2。选哪种不重要，重要的是训练和推理用同一套，并且把常数写成一个模块级变量，别在两处各写一个魔法数字。

4. 核心模块拆解：角度回归头、损失函数与旋转 NMS

骨架和数据都理清之后，剩下三块是真正需要动脑子的：头怎么改、损失怎么算、NMS 怎么做。这三块决定了模型能不能收敛、收敛得好不好、推理结果干不干净。

4.1 回归头改造：多加一个通道的代价

改造的基本思路是最小侵入：保留原版Detect的全部结构，只把输出通道从na * (nc + 5)改成na * (nc + 6)，多出来的那个通道预测角度。

import torch import torch.nn as nn class DetectOBB(nn.Module): def __init__(self, nc=80, anchors=(), ch=()): super().__init__() self.nc = nc self.no = nc + 6 # xywh + obj + cls + theta self.nl = len(anchors) # 检测层数，通常 3 self.na = len(anchors[0]) // 2 self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) self.register_buffer( 'anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) def forward(self, x): out = [self.m[i](x[i]) for i in range(self.nl)] # 三层原始输出 if self.training: return out # 训练态交给 ComputeLoss 解码 return self._decode(out) # 推理态解成 (B, N, 7)

self.no = nc + 6这一行看着简单，却是最容易被改漏的地方。原版代码里no在多处被引用：_decode里的 reshape、ComputeLoss里的切片索引、export.py里的输出维度检查。少改一处，要么 reshape 报错，要么角度通道被当成类别概率，训练 loss 看着在降，实际学出来的是错位的东西。

角度的输出建议先过sigmoid再线性映射回[-90, 0)，这样值域天然有界，不会出现预测出 270 度这种离谱值。代价是梯度在两端会小一些，如果训练前期角度 loss 下降很慢，可以把sigmoid换成tanh试试，值域对到[-1, 1]再乘 90。

4.2 角度损失：连续回归还是圆周分类

角度回归的难点在于周期性：-89°和+89°在几何上只差 2 度，但在数值上差了 178。直接用 L1 或 smooth-L1，模型在接近边界时会来回震荡，loss 曲线像心电图。

常见做法有两种。一种是加一个周期性的惩罚项，把角度差映射到[-45, 45]再算损失；另一种是 CSL（Circular Smooth Label），把角度离散成 180 个 bin，用圆周高斯做软标签，当成分类任务训。

def circular_smooth_label(angle_rad, num_bin=180, sigma=6): """把连续角度编码成圆周高斯软标签，angle_rad 形状 (N,)""" angle_deg = (angle_rad * 180.0 / torch.pi) % 180.0 # 收敛到 [0, 180) centers = torch.arange(num_bin, device=angle_rad.device).float() diff = angle_deg.unsqueeze(1) - centers.unsqueeze(0) # (N, num_bin) # 圆周距离：把差值折回 [-90, 90)，避免 0 和 179 被当成远亲 diff = torch.remainder(diff + num_bin / 2, num_bin) - num_bin / 2 label = torch.exp(-0.5 * (diff / sigma) ** 2) return label / label.sum(dim=1, keepdim=True)

num_bin=180对应 1 度一个 bin，sigma=6控制高斯核宽度，太大就退化成均匀分布学不到东西，太小又和 one-hot 没区别、失去平滑性。我一般从sigma=6起调，如果角度误差卡在 3 度左右下不去，再降到 4 试试。

CSL 的代价是头多出 180 个通道，小模型上显存会紧一些。如果显存吃紧，退回到「回归 + 周期修正」的方案更划算，代码量也少一半：

def angle_loss(pred, target, beta=1 / 9): """pred/target 单位是弧度，先把差值折到 [-pi/2, pi/2] 再算 smooth-l1""" diff = pred - target diff = torch.remainder(diff + torch.pi / 2, torch.pi) - torch.pi / 2 return torch.nn.functional.smooth_l1_loss(diff, torch.zeros_like(diff), beta=beta)

remainder这一行是整个函数的核心，它把任意角度差都折回一个周期内，等价于在圆上量最短弧长。

4.3 旋转框 NMS：把 IoU 换成多边形相交

原版 NMS 用的是轴对齐 IoU，一行(x1, y1, x2, y2)就能算。旋转框得先算两个旋转矩形的交集面积，最稳的做法是调 OpenCV 的多边形相交。

import cv2 import numpy as np def rbox_iou(b1, b2): """b1/b2 格式为 [cx, cy, w, h, angle(度)]""" r1 = ((b1[0], b1[1]), (b1[2], b1[3]), b1[4]) r2 = ((b2[0], b2[1]), (b2[2], b2[3]), b2[4]) ret, pts = cv2.rotatedRectangleIntersection(r1, r2) if ret == 0: # 完全不相交 return 0.0 inter = cv2.contourArea(cv2.convexHull(pts)) union = b1[2] * b1[3] + b2[2] * b2[3] - inter return inter / max(union, 1e-6)

ret有三个取值：0 表示不相交，1 表示部分相交，2 表示一个完全包含另一个。后两种都能正常算面积，所以只判 0 就够了。

工程上不能真的在 NMS 里做 O(N²) 的两两调用，一张图 300 个候选框就是四万多次 OpenCV 调用，推理耗时直接翻倍。我一般先按类别用轴对齐 IoU 粗筛一遍，把 IoU 大于阈值的候选先聚成簇，簇内再做旋转 IoU 精筛。这样调用次数能压到原来的十分之一左右。

提示：验证集上的 mAP 计算同样依赖旋转 IoU，utils/metrics.py里的bbox_iou必须换掉。忘了这一步，训练时 loss 正常降，但 mAP 一直挂在 0 附近，很容易误判成模型没学到东西。

4.4 训练配置与启动：四个参数决定成败

配置文件主要两处：数据 yaml 和模型 yaml。数据 yaml 里nc和names按你的类别填，路径指向切好的小图目录。

# data/dota.yaml train: ../data/dota/images/train val: ../data/dota/images/val nc: 15 names: ['plane', 'ship', 'storage-tank', 'baseball-diamond', 'tennis-court', 'basketball-court', 'ground-track-field', 'harbor', 'bridge', 'large-vehicle', 'small-vehicle', 'helicopter', 'roundabout', 'soccer-ball-field', 'swimming-pool']

启动命令和原版几乎一致，关键是--img和--batch的搭配。1024 输入下，24G 显存跑yolov5s大约能到batch=8，如果显存不够就降 batch、同时按比例调--lrf。

python train.py \ --img 1024 \ --batch 8 \ --epochs 100 \ --data data/dota.yaml \ --cfg models/yolov5s-obb.yaml \ --weights yolov5s.pt \ --hyp data/hyps/hyp.scratch-low.yaml \ --name dota_s_1024 \ --cache ram

--weights yolov5s.pt是拿水平框预训练权重做初始化。加载时会因为最后一层通道数对不上而跳过检测头，这是正常现象，日志里会打印 skipped 的层名，看到model.24之类的跳过就说明流程是对的。

--cache ram在数据集不大（切完几千张 1024 小图）时非常有用，第一次 epoch 之后所有图常驻内存，后续 epoch 的 dataloader 时间能省掉一半以上。

4.5 评估指标与 ONNX 导出

评估直接调val.py，会把预测框和 GT 都转成旋转框算 mAP。看结果时别只盯mAP@0.5，旋转框场景下mAP@0.5:0.95更能反映角度精度，两个数差得多，说明框位置对了但角度偏。

python val.py \ --weights runs/train/dota_s_1024/weights/best.pt \ --data data/dota.yaml \ --img 1024 \ --task val

导出 ONNX 的坑主要在输出维度上。多了一个角度通道，导出后的输出 shape 会变成(B, N, 7)，下游后处理不改就会把角度当成了额外的类别分数。

python export.py \ --weights runs/train/dota_s_1024/weights/best.pt \ --include onnx \ --opset 12 \ --img 1024 \ --simplify

导出后建议用onnxruntime跑一张图，把输出和 PyTorch 的结果对齐，逐元素比一遍。差在 1e-3 以内算正常，差出量级就说明某个算子在图优化时被合错了，通常是角度那个sigmoid被 fusion 掉又重建时没对齐激活函数。

5. 避坑与常见问题排查：五个让我重跑训练的经历

旋转框项目里，报错反而是好事，最怕的是不报错、结果不对。下面这五条都是我实际踩过的，按「现象 → 原因 → 解决」写清楚，遇到了可以直接对号入座。

5.1 数据与标签相关的坑

第一条：训练 loss 一路正常降，验证 mAP 始终在 0 附近徘徊。

现象很迷惑，loss 曲线平滑下降，没有 NaN 也没有震荡，但每隔 5 个 epoch 算一次 mAP，永远是零点几。原因基本锁定在训练和验证用了两套角度定义。训练集标签走的是poly2rbox，角度落在[-90, 0)；验证集的 GT 却是从另一个脚本转出来的，落到了[0, 180)。IoU 计算时两个框看起来重叠，但角度差 90 度，旋转 IoU 直接算成 0。

解决办法是把 GT 生成和预测解码统一收口到一个模块，验证时重新生成一遍 GT 缓存，并用一个断言检查角度值域。这个断言我后来加进了 dataloader 里，任何一批标签越界就立刻抛异常，比事后排查省事得多。

第二条：切完图之后，边角上的目标整批消失。

现象是训练集里某些类别数量对不上，原图上明明有几百个，切完只剩几十个。原因是切图时过滤条件写成了「框的四个角都在子图内」，大目标跨边界时四个角全在外面，被整批过滤掉。正确的判断条件是「框中心点落在子图内」，然后把框坐标平移到子图坐标系，允许框超出边界，让 dataloader 做裁剪填充。

改完之后同类目标数量立刻对齐了。代价是会有部分框只露出一半，这本身也是数据增强的一部分，不用额外处理。

第三条：小目标密集区域，模型几乎全漏检。

遥感数据里小车、小船这类目标，在原图上可能只有二十几个像素，切到 1024 之后也就三四十个像素。原版 YOLOv5 的 anchor 是在 COCO 上聚类的，尺度偏大，小目标落到最小的那个特征层上也匹配不上。

解决办法是拿自己的训练集重新跑一遍 k-means 聚类，把 9 个 anchor 换成贴合数据分布的。utils/autoanchor.py里现成的函数可以直接用，改一下数据路径跑一遍，把输出抄回 yaml 就行。换完之后小目标的召回率能提一截。

5.2 训练过程中的坑

第四条：显存明明够，跑到中途突然 OOM。

现象是前几十个 iteration 都正常，跑到某个 batch 突然报显存不足。原因通常是 mosaic 增强的随机性，四张图拼在一起时，如果某张图的目标特别密集，单张图的框数量能到几百个，拼完就是上千个框，loss 计算时中间张量瞬间膨胀。

解决有两个方向：一是在 dataloader 里给单图框数量设一个上限，超了就重采样；二是把hyp里的mosaic概率从 1.0 降到 0.5 左右，减少密集样本出现的频率。我一般两个都做，前者治本，后者降低触发概率。

提示：训练脚本里加一句torch.cuda.max_memory_allocated()的打印，出问题的时候能直接看到峰值出现在哪个 epoch，比盲猜强。

第五条：学习率看着没问题，但 loss 在 20 个 epoch 后开始发散。

现象是前期收敛正常，中期 loss 突然往上走，然后一直下不来。原因多半是 batch 改小之后没有同步调学习率。原版默认的学习率是按batch=64调的，你降到 8，实际有效学习率就放大了 8 倍。

解决办法是按线性缩放规则把lr0和lrf一起降下来，或者干脆加--warmup-epochs让前期更平缓。经验值是batch每降一半，lr0乘 0.5，lrf保持比例不变。

5.3 推理与导出相关的坑

第六条：ONNX 模型跑出来的框全是轴对齐的。

现象是 PyTorch 推理正常，转成 ONNX 之后所有框的角度都是 0。原因是导出时的后处理被 trace 成了常量。YOLOv5 的export.py默认走Detect.forward的推理分支，那个分支里有大量的 Python 控制流，trace 过程中会被固化成导出那一刻的值。

解决办法是把角度解码从模型里挪出来，放到后处理脚本里做，模型只负责吐原始张量。这样导出的图是纯卷积，没有控制流，跨平台部署时也不会因为算子不支持而失败。

6. 二次开发：加类别、换角度表示、接后端的三条路

这套代码改造完之后，最常见的二次开发需求有三类，我按改动成本从低到高排一下。

第一类是加类别。改data/*.yaml里的nc和names，然后改模型 yaml 里最后一层的nc参数就行。注意models/yolov5s-obb.yaml里的nc有两处，head部分的那一处容易漏，漏了会在DetectionModel初始化时报 shape 不匹配。改完用预训练权重加载，检测头会被跳过，这是预期行为。

第二类是换角度表示。如果你的下游系统用的是四点坐标，可以在utils/rbox.py里加一个rbox2poly的反向函数，推理结束后统一转一次。如果下游用的是[0, 180)的区间，只需要在poly2rbox的最后加一行偏移，其他逻辑不用动。

改造点	需要改的文件	验证方法
新增类别	`data/.yaml`、`models/-obb.yaml`	训练 1 个 epoch，看`nc`打印
换角度区间	`utils/rbox.py`	画 20 张预测图人工看方向
换损失函数	`utils/loss.py`	对比前 10 epoch 的角度 loss 曲线
接 TensorRT	`export.py`+ 后处理	逐元素对齐 PyTorch 输出