简介YOLOv5 OBB旋转框训练DEMO是一份面向目标检测研究者和深度学习开发者的实践资源在YOLOv5框架基础上扩展了旋转边界框Oriented Bounding Box支持可更精确地检测车辆、树木等具有方向性的非矩形目标解决常规轴对齐边界框难以贴合倾斜物体的问题。包体共583个文件大小约445.58MB涵盖Python训练/预测脚本、YAML模型配置、XML标注和TXT标签文件、PyTorch权重pt和ONNX导出模型并包含CUDA/C编写的PolyNMS后处理源码及dockerfile便于从数据处理到推理部署全链路复现。已有1818人学习下载。整套资源覆盖数据格式转换、OBB配置修改、损失函数设计、模型评估与可视化的完整训练闭环附有多组预训练权重和运行脚本适合希望掌握旋转框检测技术、进行模型微调或二次开发的使用者参考与实操。1. 旋转框检测的痛与路为什么 yolov5 obb 值得专门跑一遍如果你做过遥感图像或者工业质检里的目标检测大概率遇到过这个问题飞机、船舶、带角度的小零件在画面里不是横平竖直的用常规的水平框去框一个框里往往塞进好几个目标——你说它检测对了可 IoU 算出来就是难看你说它错了可它明明把目标圈住了只是圈得不够严丝合缝。这就是旋转框检测OBB, Oriented Bounding Box要解决的核心问题。yolov5 obb 旋转框训练 demo 不是一个完整的平台而是一套能让 YOLOv5 输出带角度预测框的训练链路适合做遥感、航拍、工业视觉的人拿来当起点。这个 demo 的价值在于省事它把数据格式转换、角度编码方式、损失函数适配、检测头修改这几件最容易卡壳的事都串起来了你不用自己从零去翻论文复现。你只需要准备好标注好的旋转框数据按它的目录结构放进去改几个配置文件就能开始训练。我的建议是不要拿它直接上生产先跑通一个小数据集把整个流程看明白了再考虑怎么往自己的业务上引。2. 技术栈与整体结构先把 OBB 这层窗户纸捅破2.1 从水平框到旋转框YOLOv5 到底改了什么常规 YOLOv5 的检测头输出是x, y, w, h, conf, cls其中w和h是轴向对齐的宽高。OBB 版本在这里加了一个参数常见做法是在每个 anchor 的回归分支里追加angle角度形成x, y, w, h, angle, conf, cls。这个角度并不是随便加的它的编码方式直接决定了网络好不好收敛。demo 里最常见的角度定义是 OpenCV 风格角度范围从-90°到0°其中w始终是长边旋转角绝对值较大的那条边对应的方向。这个定义有一个好处它和大多数标注工具导出的格式一致。比如说你在标注软件里旋转一个矩形框导出的 JSON 里带一个theta字段通常就是这个范围。如果你用的是0°到90°或者-180°到180°训练前必须统一。我拆过这个 demo 的loss.py发现它处理角度回归的方式比较直接角度作为普通回归目标参与CIoU或SmoothL1损失计算并没有做圆形平滑标签CSL那种复杂的编码。这意味着你要自己注意角度周期性的问题——89°和-89°只差2°但如果网络直接回归数值Loss 会认为它差了178°。解决这种问题更专业的做法是引入 CSL 或环形平滑demo 里面用了一个取巧的方式在 loss 计算的时候把角度差映射回[-90, 90]区间再做损失计算。2.2 数据流向标注文件怎么一步步变成训练样本整个工程的数据流可以概括成标注文件 - 数据集清单 - 训练样本 - 模型输出 - 后处理五步。你拿到手的.txt标注文件格式大概是这样的class_id x1 y1 x2 y2 x3 y3 x4 y4这是四个角点的绝对坐标或归一化坐标具体看 demo 的datasets.py怎么读。常见做法是存归一化坐标也就是相对于图片宽高的比例值范围在0到1之间。demo 里通常还会有一行difficult标志位用于忽略某些难例。如果你的数据是从 DOTA 数据集转过来的注意它原始格式是x1 y1 x2 y2 ... x4 y4的多边形点坐标demo 里会有一段预处理代码去解析并转成 OBB 表示。接下来datasets.py里做的事情是读取图片路径和对应标注文件通过索引把它们配对然后在__getitem__里完成图像增强和标注同步变换。旋转框增强是一个很容易写错的地方你对图片做了随机旋转那么标注框的四个角点坐标也必须跟着做同样的旋转矩阵变换同时图片旋转后的黑边填充区域要扣除。2.3 完整目录结构拆解每个文件负责什么demo 工程目录大体长这样. ├── data/ # 数据集配置 │ └── obb.yaml # 数据集路径、类别数、类别名 ├── datasets/ # 数据集存放目录 │ ├── images/ # 训练/验证图片 │ └── labels/ # 旋转框标注 txt ├── models/ │ ├── yolo5s-obb.yaml # 网络结构定义 │ └── common.py # OBB 相关模块 ├── utils/ │ ├── datasets.py # 数据加载与增强 │ ├── loss.py # 损失函数含角度回归 │ └── plots.py # 可视化推理结果 ├── train.py # 训练入口 ├── detect.py # 推理脚本 └── export.py # 导出 ONNX / TensorRT重点看三个文件models/yolo5s-obb.yaml定义了检测头的输出维度最关键的是nc类别数和no每个 anchor 的输出维度。如果类别数是4那么no 5 1 4 104 个坐标 1 个角度 1 个置信度 4 个类别概率。想不起来这个公式后面改模型一定会翻车。同样地utils/loss.py的build_targets函数里要处理角度匹配逻辑这里不能照搬原版 YOLOv5必须把角度信息一路传到损失计算的阶段。3. 从零跑通训练环境搭建、数据准备与命令实战3.1 环境安装避开版本地狱的稳妥选择这个 demo 对环境的要求不算苛刻Python 版本建议装3.8到3.10。我建议你新建一个独立的 conda 环境避免把系统 Python 环境搞乱。核心依赖就这几个conda create -n obb python3.9 conda activate obb pip install torch1.13.1 torchvision0.14.1 --index-url https://download.pytorch.org/whl/cu117 pip install opencv-python pyyaml tqdm matplotlib seaborn这里有个血泪经验torch的版本继承和你本地 CUDA 驱动有非常强的绑定关系。如果你显卡驱动是新的装老版本 CUDA 工具包反而会出问题如果驱动太老装新版本 PyTorch 会直接报CUDA driver version is insufficient。你可以先用nvidia-smi查看驱动支持的 CUDA 版本再决定装哪个 cu 后缀的 torch 包。装完以后验证 torch 能用 GPUimport torch print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果输出是True且能看到显卡型号环境就过了。不需要安装额外的 DOTA 解析库demo 自带的代码已经处理了格式转换。OpenCV 要确认是opencv-python不是opencv-contrib-python这俩在某些版本下会打架装一个就行。3.2 准备数据集从你的标注文件开始改造把你的数据整理成下面这个结构。图片和标注分开放在images和labels目录下文件名一一对应不要求后缀一致但主名必须相同datasets/ ├── images/ │ ├── train/ │ │ ├── 0001.jpg │ │ └── 0002.jpg │ └── val/ │ ├── 0003.jpg │ └── 0004.jpg └── labels/ ├── train/ │ ├── 0001.txt │ └── 0002.txt └── val/ ├── 0003.txt └── 0004.txt标注内容示例0 0.234 0.512 0.301 0.567 0.287 0.618 0.220 0.563每个数字代表的含义是第一个数是类别 id从0开始后面八个数按顺序是四个角点的(x, y)坐标对这里存的是归一化比例取值范围是0到1。注意四个角点的顺序必须是顺时针或逆时针连续排列也就是要么A-B-C-D要么D-C-B-A不能跳着来否则计算外接旋转矩形时会出现角度偏差。之后修改data/obb.yamltrain: /absolute/path/to/datasets/images/train val: /absolute/path/to/datasets/images/val nc: 1 names: [aircraft]这里train和val指向的是图片目录不是数据集的根目录。labels 目录会自动通过替换images为labels来推导前提是两者在同级路径下。如果你把 labels 放在别的地方就得在datasets.py里手动改label_dir的逻辑。想快速验证标注格式是否正确写一个小脚本把 label 里的框画出来看看import cv2 import numpy as np img cv2.imread(datasets/images/train/0001.jpg) with open(datasets/labels/train/0001.txt) as f: for line in f: vals list(map(float, line.strip().split())) pts np.array(vals[1:9], dtypenp.float32).reshape(4, 2) pts[:, 0] * img.shape[1] pts[:, 1] * img.shape[0] pts pts.astype(np.int32) cv2.polylines(img, [pts], True, (0, 255, 0), 2) cv2.imwrite(check.jpg, img)跑完看看check.jpg如果框和精确目标贴合得很好说明数据格式没问题。这里最常见的翻车情况是角点顺序乱了画出来的矩形变成了交叉线或者蝴蝶形状那就要回到标注工具里检查导出设置。3.3 训练启动参数怎么传、模型怎么选一切都准备好以后训练命令大概是这个形式python train.py --data data/obb.yaml --cfg models/yolo5s-obb.yaml \ --weights yolov5s.pt --epochs 100 --batch-size 16 \ --img 640 --device 0这里各参数的含义是--data指定数据集配置文件--cfg指定网络结构--weights是预训练权重建议用官方yolov5s.pt做初始化不用旋转框相关的预训练也能跑起来因为前面几层提取的特征是通用的--img是训练时的输入尺寸如果图片是1024x1024的遥感图实验建议先用800或640显存不够再加--batch-size减半。训练过程中重点看两个指标Box损失和Angle损失的下降曲线。如果Angle损失一直不降大概率是角度编码有问题多数情况下问题出在你标注的旋转框长边方向不统一。举个例子同一个目标一张图里标注时w对应x方向另一张图标注时w对应y方向网络就糊涂了——它学到了同一种目标有两种角度分布自然学不下去。3.4 推理验证用文本测试一个修改后的推理脚本训练完后用detect.py验证单张图片python detect.py --weights runs/train/exp/weights/best.pt \ --source datasets/images/val/0003.jpg --img 640demo 的detect.py会在图上画出旋转框把角度信息显示在标签后面。这里要留意一个细节输出图片上框的角点顺序可能看起来是乱的因为后处理还原的时候是按左上角、右上角、右下角、左下角的固定顺序排的不一定和你标注时的顺序一致但几何上仍是同一个四边形。这不影响评估。4. 核心代码拆解数据加载、Loss 角度回归与后处理解码4.1 数据加载层的旋转框标注解析数据读取是 OBB 和普通检测最不一样的地方。普通检测只需要读四个数(x, y, w, h)这里需要解析八个坐标并转换成内部表示。demo 的datasets.py里有一个函数专门做这个转换def load_obb_label(path, img_shape): 加载 OBB 标注并转换为 (cx, cy, w, h, angle) 表示 angle 范围 [-90, 0)w 始终是长边 labels [] with open(path, r) as f: for line in f: parts line.strip().split() cls int(parts[0]) pts np.array(parts[1:9], dtypenp.float64).reshape(4, 2) # 归一化坐标乘以图像宽高得到像素坐标 pts[:, 0] * img_shape[1] pts[:, 1] * img_shape[0] # 用 OpenCV 求最小外接旋转矩形 rect cv2.minAreaRect(pts.astype(np.float32)) (cx, cy), (w, h), angle rect # 统一长边为 wangle 范围对齐 if w h: w, h h, w angle 90 labels.append([cls, cx / img_shape[1], cy / img_shape[0], w / img_shape[1], h / img_shape[0], angle]) return labels这段代码的关键在最后三步cv2.minAreaRect拿到的是浮点精度的最小外接矩形角度范围是[-90, 0)对应 OpenCV 的惯例。当w h时把长边统一到w角度相应地加90°转换到同一量纲。这不是表面功夫它是为了后面计算损失时角度语义一致。如果不做这段统一网络可能在训练前期反复横跳。另外一个容易忽略的点这里的cx,cy再归一化一次返回的相对坐标参与网络训练。cv2.minAreaRect本身对浮点输入是稳定的但如果你从txt读出来的坐标精度不够比如只保留两位小数最小外接矩形可能出现微小旋转在训练时表现为角度抖动。建议你的标注文件保留至少六位小数。4.2 Loss 函数角度回归的核心写法loss.py里角度损失的实现是理解整个 OBB 训练的关键。普通的 YOLOv5 只有三个损失项box_loss坐标、obj_loss置信度、cls_loss分类。OBB 版本多了一个angle_loss它参与box_loss的合并或单独计算取决于具体实现。我拆过的这个版本是单独计算的def angle_loss(pred_angle, target_angle): 计算角度回归损失 pred_angle: 网络预测的角度, shape (N,) target_angle: 目标角度 # 将角度差映射回 [-90, 90]避免周期性跳变 diff pred_angle - target_angle while diff 90: diff - 180 while diff -90: diff 180 loss torch.abs(diff) * 0.01 # 缩放系数防止角度 Loss 数值过大 return loss.mean()这里的核心是周期映射。直接让网络去回归-90到0的数值在边界处损失会出现巨大的跳变——比如目标角度是-89.5°预测是-0.5°数值上相差89°但实际几何意义只差0.5°因为-0.5°和-90.5°等价。映射操作通过把差值折叠到[-90, 90]区间内解决了这个不连续问题。但直接算绝对值还有一个问题梯度的幅度恒定在接近 0 时不会变小。更好的做法是换用 SmoothL1 或者把角度编码成两个正弦余弦值再计算损失。demo 用了简化方案训练效果在小数据集上够用。如果你在训练中遇到角度 Loss 后期降不下去可以考虑改成smooth_l1_lossloss F.smooth_l1_loss(pred_angle, target_angle, beta1.0)4.3 后处理从回归值到旋转矩形推理阶段网络输出的(x, y, w, h, angle)需要解码回四角点坐标渲染到图上。demo 的plots.py里做了解码操作def decode_obb(cx, cy, w, h, angle): 将 OBB 表示解码为四角点 angle 单位: 度 theta np.deg2rad(angle) cos_t, sin_t 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]) # 旋转矩阵作用在每个角点上 pts_x cx dx * cos_t - dy * sin_t pts_y cy dx * sin_t dy * cos_t return np.stack([pts_x, pts_y], axis1)这个解码公式是标准旋转矩阵的应用先构造未旋转的矩形以中心点对称的四角点再用angle做旋转。angle正负决定旋转方向不同标注工具可能方向相反。如果你跑完推理后发现框顺时针偏了而不是逆时针偏了把这里sin_t的符号取反就能修正这是一个值得记住的调试点。4.4 Anchor 匹配和原版 YOLOv5 的区别OBB 版的 anchor 匹配策略基本沿用 YOLOv5 的思路对每个 prediction 层的 anchor计算和 GT 的 IoU高于阈值就匹配为正样本。不同点在于 IoU 计算方式——水平框直接用w, h相交区域OBB 需要先还原出旋转矩形的四角点再计算任意四边形之间的 IoU。这一步计算量大demo 里用了简化近似把角度带入后先计算外接水平矩形的 IoU再乘以角度差接近度的惩罚系数。这算是一个工程取巧如果你追求更准的效果可以用gIoU或skew_iou代替。但要注意计算开销skew_iou需要逐对计算多边形相交面积在 batch size 比较大时会成为训练瓶颈。我的建议是前期用简化方案跑通流程后期再上精确版本。5. 避坑指南数据标注、角度编码与训练收敛的常见问题5.1 标注角点顺序不一致画出来的框交叉现象训练前可视化标注框时部分框显示成蝴蝶交叉形或者角点连线穿过了矩形内部。原因标注工具导出的坐标不是严格按顺时针或逆时针排的。有的工具按左上、右上、右下、左下导出有的按左上、左下、右下、右上导出。你的代码如果默认按一种顺序读取对另一种顺序的数据就会错乱。解决加载标签时统一做一次多边形重排序。常见做法是取几何中心再按角点相对中心的角度排序保证环绕方向一致def order_points(pts): center pts.mean(axis0) angles np.arctan2(pts[:, 1] - center[1], pts[:, 0] - center[0]) return pts[np.argsort(angles)]经过这个排序后四角点一定是顺时针或逆时针连续排列。5.2 训练时 Angle Loss 不下降或剧烈震荡现象训练了十几个 epochBox Loss 在降但 Angle Loss 几乎不降甚至抖动越来越大。原因最常见的是长边方向不一致。同一个类别有的标注w是水平方向有的标注w是垂直方向因为标注工具在自动生成旋转框时没有统一规则。网络同时学习两种角度分布结果只能取平均。解决加载标签时统一规则核心是if w h: swap(w, h)那段逻辑。另外检查你的数据集里是否存在大量退化框几乎变成正方形的目标这种目标的角度本来就无意义可以设置一个长宽比阈值小于该阈值的目标忽略角度损失。5.3 图片增强导致标注错位现象训练开启 Mosaic 或随机旋转增强后Loss 开始正常下降但中途突然变大或者验证集精度时好时坏。原因Mosaic 拼接时四张图的标注框坐标合并在同一张图上但旋转框的角点变换没有正确应用。比如你对整个 Mosaic 图做了水平翻转但没有把角点坐标做相应变换。解决检查数据增强函数是否有独立的 OBB 分支。最基本的翻转时要同时更新四个角点坐标和目标角度。水平翻转角度变为-angle垂直翻转角度不变但所有x坐标要取反。如果不想多写代码可以先关闭增强训练几轮确认 Loss 正常再逐步打开。5.4 推理结果整体旋转了 90 度现象预测框的位置是对的但方向整体偏转 90 度长的变成短的短的变成长的。原因推理解码时角度符号的方向与训练时的标注约定不一致。openCV的minAreaRect角度是按 x 轴逆时针方向定义的但网络回归的角度经过解码可能变成了顺时针。解决验证decode_obb函数时需要先跑一个单样本测试对比一张图上 GT 和推理输出。如果差 90 度翻解码函数里的旋转矩阵符号即可解决。5.5 训练时报形状不匹配的错误现象训练刚开始loss.py里出现RuntimeError: The size of tensor a (10) must match the size of tensor b (12)。原因models/yolo5s-obb.yaml里的输出维度没有改。比如你的类别数是5但定义的no还是默认的11导致标签数量与预测数量对不上。解决记住公式no 5 nc其中5是(x, y, w, h, angle)五个回归参数加一个置信度其实是6 nc。准确说 YOLOv5 的no应该是5 1 ncOBB 版是6 nc。类别改了就必须同步改yaml里nc字段有时候nc还出现在anchors配置里都要一一对应。6. 二次开发把 demo 改造成你的专用检测器6.1 切换到自己的数据集三个必改点第一步是换数据集。除了改data/obb.yaml里的路径和类别名之外还有三个容易漏掉的地方utils/datasets.py里的class_names要改否则可视化图上标签名和实际类别对不上models/yolo5s-obb.yaml里的nc要改成你的实际类别数这个前面踩过坑还有train.py里的--classes参数如果传了固定值也要同步调整。这些改完以后建议用小数据集先跑通一个 epoch确认没有报错再放大规模。我一般会用每类 20 张左右的图片做一个冒烟测试目的是验证数据链路而不是追求精度。冒烟测试通过后再把完整数据挂上去跑。6.2 嵌入自己的预处理图像降采样与归一化遥感图往往是大尺寸原图直接训练会 OOM。demo 默认用--img 640做训练输入这意味着大图被直接缩放。如果目标特别小比如几十个像素的汽车直接缩放会抹掉细节。常见的做法是在datasets.py的__getitem__里先对原图做滑窗裁剪把每个裁剪块作为独立样本输入训练。这个改造不复杂但是要小心标注坐标跟随裁剪偏移def crop_obb(img, labels, crop_size, stride): h, w img.shape[:2] for y in range(0, h - crop_size 1, stride): for x in range(0, w - crop_size 1, stride): crop img[y:ycrop_size, x:xcrop_size] new_labels [] for label in labels: # 检查标注中心点是否落在裁剪块内 cx, cy label[1] * w, label[2] * h if x cx x crop_size and y cy y crop_size: # 坐标平移并重新归一化到 crop_size new_label label.copy() new_label[1] (cx - x) / crop_size new_label[2] (cy - y) / crop_size new_labels.append(new_label) yield crop, new_labels这里要注意裁剪块大小最好和--img保持一致否则又要做一次缩放变换。裁剪后的标注坐标要重新归一化同时裁到边缘的目标如果中心点不在块内就丢掉。这样处理以后模型的召回率会明显提升因为小目标有了更大的有效分辨率。6.3 角度精度的进一步提升替换损失函数demo 的角度损失是简单的绝对值映射精度有限。如果目标是遥感里的飞机船舶需要高精度角度才能满足业务要求可以考虑把角度编码改成双正弦形式。做法是把angle拆成(sin(2*angle), cos(2*angle))两个值参与回归解码时用atan2还原。这样网络学的不再是单点值而是一个周期函数的拟合角度误差会大幅减小。不过改动面比较大涉及数据加载的输出维度、损失函数、解码函数三处需要统一修改。6.4 部署时的加速导出 ONNX 的坑demo 里提供的export.py可以导出 ONNX 模型做 TensorRT 加速。实际部署时有一个需要注意的地方ONNX 导出后的输出格式是(batch, anchors, 6 nc)的原始张量需要你在部署代码里写一遍 NMS 和后处理。demo 自带的detect.py里的后处理逻辑依赖 PyTorch 的算子ONNX 导出时有些算子不支持。我建议二次开发时把后处理函数改写成纯 NumPy 或纯 C 可实现的算子确保导出时不被卡住。我自己的习惯是每次改完数据增强或者解码逻辑都先跑一遍可视化比对脚本把 GT 框和推理框画在同一张图上对比。这个脚本是我从第一次做 OBB 项目时就保留至今的土办法它帮我发现了至少三次角点顺序、角度符号、尺寸归一化的问题。每次改完必跑希望这个习惯也能帮到你。本文还有配套的精品资源点击获取