简介YOLOv5 OBB旋转框训练DEMO是一份面向目标检测开发者与深度学习初学者的软件/插件类示例工程重点解决车辆、树木等倾斜或非矩形目标的检测训练问题。它在YOLOv5基础上引入OBB定向边界框可更精确表达物体位置与方向适合需要掌握旋转框检测流程的读者。压缩包共583个文件整体约445.58MB内容以py脚本、yaml/cfg配置、xml/txt标注文件及png/jpg图像为主同时包含pt/onnx模型权重与训练、预测、导出、wts生成等命令脚本能覆盖从数据准备到部署的主要环节。读者可从中了解OBB标注格式转换、模型结构修改、旋转框损失计算、训练验证与推理可视化等关键知识点借助现成的cmd命令和poly_nms.cpp等源码快速跑通demo并基于自有数据集开展迁移学习与精度调优。目前已有1814人学习下载适合中高级深度学习开发者参考实践。1. yolov5 obb旋转框训练demo把水平框检测改成带角度的旋转框值不值得下做工业质检、遥感目标检测、无人机航拍识别时最常见的一个烦恼就是目标本身是斜着躺的用yolov5默认的水平矩形框去框会把大量的背景、相邻目标甚至干扰物一起框进来。比如检测电线塔、集装箱、车辆或者电路板上的元件水平框的mAP看着还行实际落地的定位精度和后续的切割、识别都会吃亏。这个yolov5 obb旋转框训练demo本质就是把yolov5的检测头扩展出一条角度回归分支让每个目标的输出从x、y、w、h变成x、y、w、h、theta从而让框贴合目标轮廓。它能解决的就是目标倾斜导致水平框定位不准的问题适合正在做遥感、工业检测、或研究与毕业设计需要快速上手旋转框检测的人。如果你只是做常规的行人、车辆检测目标基本都是正立的那这个demo未必需要但只要你遇到框子怎么都只能斜着包住目标的场景这份资源能省掉你从零改网络、调loss的大量时间。demo通常包含完整的训练、验证、推理脚本以及一个可以直接跑通的数据集样例。下面我会从旋转框检测的原理讲起拆解demo的文件结构和标签格式然后一步步带你跑通训练再把我实际踩过的坑逐一列出来。最后给出旋转框后处理和部署时最值得记住的几个技巧。2. 旋转框检测的原理与demo结构先弄清楚角度从哪来、往哪去2.1 旋转框和水平框的本质区别yolov5水平框检测的回归目标只有四个值即中心点x、y、宽w、高h。它的先验anchor是二维的每个anchor负责预测一个矩形框但矩形框的边始终和图像坐标轴平行。旋转框检测则不同每个目标需要额外回归一个角度theta。这个角度的定义方式在不同开源项目里并不统一常见的定义有两种一种是OpenCV风格角度范围是[-90,0)表示框的长边和图像x轴正方向的夹角另一种是DOTA数据集的风格角度范围是[0,90)配合检测框的长边和短边来定义。yolov5 obb的demo里一般默认采用第一种或者某种特定编码这一步是后面所有坑的源头。从网络结构上看yolov5 obb分支在原本的检测头输出上增加一个角度分支这个分支和cls、obj分支并行输出维度会多出一个角度通道。角度回归的loss也不同于普通的smooth L1因为在角度这个周期变量上简单的L1 loss会在边界处出现大的突变。常见的做法是用KLDKullback-Leibler Divergence或者GWDGaussian Wasserstein Distance这类基于高斯分布建模的loss把角度和宽高的耦合关系一起考虑进来。demo里一般已经帮你选好了loss不需要你手动改但你要明白这个角度不是直接回归出一个绝对值就完事的它和w、h的组合决定了框的朝向。2.2 demo里的文件结构和标签格式下载这份demo后你会看到类似这样的目录组织结构yolov5_obb/ ├── models/ │ ├── yolo_obb.py │ ├── yolo.py │ └── common_obb.py ├── utils/ │ ├── loss_obb.py │ ├── nms_obb.py │ └── datasets_obb.py ├── train_obb.py ├── val_obb.py ├── detect_obb.py ├── data/ │ └── demo.yaml └── datasets/ └── demo/ ├── images/ │ ├── train/ │ └── val/ └── labels/ ├── train/ └── val/train_obb.py是训练入口val_obb.py是验证入口detect_obb.py是推理入口。models/下的网络定义文件相比yolov5原版增加了角度回归头。utils/loss_obb.py里封装了旋转框专属的lossutils/nms_obb.py里是旋转框的NMS处理逻辑。这些文件之间的调用关系和原版yolov5一致你不需要改动它们的依赖关系。标签文件是txt格式每一行对应图像里的一个目标字段为class_id x_center y_center width height angle其中x_center y_center width height都是相对于图片宽高的归一化坐标范围是0到1angle是角度值具体范围和编码方式要看demo里的datasets_obb.py怎么解析。我先给出一行示例0 0.500000 0.400000 0.300000 0.200000 -30.5上面的含义是类别为0框中心在图像(0.5, 0.4)的位置框的宽占图像宽度的30%高占图像高度的20%角度为-30.5度。注意这里用的是归一化坐标所以和原图尺寸无关。如果你自己标注的标签是像素坐标需要先做归一化再放进来否则训练会直接出问题。提示拿到demo后第一件事不是直接开训而是先检查datasets_obb.py里读取标签的代码确认angle的解析方式和你标注工具导出的角度定义是否一致。这个不一致通常会导致模型能训练但不收敛。3. 环境搭建与数据集准备把demo跑通的三个准备步骤3.1 环境依赖与版本坑我一般会用conda单独建一个环境避免和别的项目互相污染。这个demo基于PyTorch实现依赖项和yolov5原版差不多但额外的旋转框NMS和loss可能需要编译一些自定义算子。先创建一个Python 3.8左右的环境然后安装依赖conda create -n yolo_oob python3.8 conda activate yolo_oob pip install torch1.10.0 torchvision0.11.0 pip install -r requirements.txtrequirements.txt里通常有numpy、opencv-python、pyyaml、tqdm等。如果demo里包含需要编译的算子你会看到setup.py或者build脚本这时需要确保系统装了gcc和cuda toolkit。默认情况下yolov5 obb demo不会强制要求编译任何算子直接用PyTorch算子也可以跑通只是后处理阶段速度稍慢。我建议先看目录下有没有setup.py如果有先执行python setup.py develop把自定义层编译出来。安装完成后用demo自带的数据集做一次val测试确保模型能够正常前向传播。这一步能帮你区分环境问题和代码问题。常见的报错是module torch has no attribute hstack这是因为PyTorch版本太老需要升级到1.9以上还有ImportError: cannot import name intersect_obb from utils.nms_obb说明自定义NMS没编进包需要重新编译或改用纯torch实现的备选版本。3.2 用roLabelImg标注旋转框数据集数据集准备是整个流程中最耗时的部分也是最容易出错的环节。标注旋转框我一般用roLabelImg它是labelImg的一个支持旋转框标注的版本。安装方式是通过pip或从源码编译启动后标注界面上多了一个旋转矩形框按钮。标注时先拖出一个矩形框然后调整角度让框贴合目标。保存的标签是Pascal VOC格式的XML里面每个物体带robndbox字段包含cx cy w h angle四个属性。这个XML需要转换成demo要求的归一化txt格式。我写了个简单的转换脚本可以直接把roLabelImg的XML文件批量转换到yolo obb格式import xml.etree.ElementTree as ET import os import glob def convert_xml_to_txt(xml_path, out_path, class_names): tree ET.parse(xml_path) root tree.getroot() img_w int(root.find(size/width).text) img_h int(root.find(size/height).text) lines [] for obj in root.iter(object): name obj.find(name).text if name not in class_names: continue class_id class_names.index(name) robndbox obj.find(robndbox) cx float(robndbox.find(cx).text) / img_w cy float(robndbox.find(cy).text) / img_h w float(robndbox.find(w).text) / img_w h float(robndbox.find(h).text) / img_h angle float(robndbox.find(angle).text) lines.append(f{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f} {angle:.2f}) with open(out_path, w) as f: f.write(\n.join(lines)) if __name__ __main__: class_names [container, car] # 按你的类别顺序修改 xml_files glob.glob(labels_xml/*.xml) for xml_file in xml_files: out_file os.path.join(labels_txt, os.path.basename(xml_file).replace(.xml, .txt)) convert_xml_to_txt(xml_file, out_file, class_names)上面的脚本逻辑是读取XML里的图像宽高把robndbox的中心点、宽、高都除以图像宽高做归一化angle字段直接保留。注意angle的单位是度和demo里datasets读取时采用的单位必须一致否则需要转换。转换完成后检查几个标签文件用效果代表性的图对比一下原图看框是否贴合目标。3.3 修改数据配置yaml与训练参数demo里自带一个data/demo.yaml你需要参照它创建自己数据集的yaml。典型的配置如下train: datasets/demo/images/train val: datasets/demo/images/val nc: 1 names: [car] # 角度相关配置 angle_deg: True angle_version: le135train和val是图片路径demo代码会根据图片路径自动找同名的txt标签文件所以标签文件必须放在labels目录下且与图片文件名严格同名。nc是类别数names是类别名称列表。angle_deg表示角度单位是度angle_version是角度范围策略常见的有le135左闭右开-90到135和ocOpenCV风格-90到0。这取决于demo代码里对角度回归的处理方式如果你的数据集角度范围是0到90可能需要设置对应的版本并在后面推理阶段做同样的后处理。训练脚本的常用参数如下python train_obb.py \ --data data/demo.yaml \ --cfg models/yolov5s_obb.yaml \ --weights \ --batch-size 8 \ --imgsz 640 \ --epochs 100 \ --device 0--cfg指定网络结构配置文件yolov5s_obb.yaml对应轻量级的配置比yolov5s多了角度输出维度--weights设为空字符串表示从头训练如果你希望加载yolov5s的预训练权重作为初始化可以指定原版yolov5s.pt但这时需确保权重文件里的维度对得上否则会报shape mismatch。--imgsz输入分辨率遥感图像宽度较大的话可以提高到1024但显存占用会明显增加。4. 训练与验证从命令到日志一步一步确认模型真的在学东西4.1 训练命令与超参数选择训练入口是train_obb.py基本的启动方式上面已经给了。根据我的观察第一次跑demo时先用它自带的数据集跑通比直接用自己的数据要节约时间。自带的demo数据集可能只有几百张图跑20个epoch就能看到loss下降趋势。在这个过程中你应该重点关注以下几个指标box_loss框回归的loss包含宽高和中心点。theta_loss角度回归的loss这是obb训练多出来的部分。obj_loss置信度损失衡量有无目标的概率。cls_loss分类损失。如果theta_loss一直不下降大概率是角度编码或GT解析出了问题。我碰到过的情况是angle_version设为oc但标签里的角度是0到90导致网络学到的角度无法表达实际范围loss看起来在收敛但预测结果全部变成水平框。解决方式是回过头检查标签的角度分布把它打印出来确认和配置里的angle_version一致。显存不足是常见的现实问题。你把batch size直接拉到16很可能碰到CUDA out of memory。解决思路不是只降低batch size而是先计算一下需要的显存。以yolov5s_obb为例输入640x640batch8大约需要6~8GB显存。如果不够优先降batch到4同时开启梯度累积python train_obb.py --data data/demo.yaml --cfg models/yolov5s_obb.yaml --batch-size 4 --accumulate 2 --epochs 100--accumulate 2表示每两个step做一次梯度累积等效于batch size8的效果。这样既保住了训练稳定性又不会被显存卡死。4.2 验证预测与可视化训练完成后用验证集跑一遍指标。命令是python val_obb.py --data data/demo.yaml --weights runs/train/exp/weights/best.pt --save-json验证输出会包含mAP0.5、mAP0.5:0.95以及每一类的PR曲线。旋转框的mAP计算方式和水平框不同它用的IoU是旋转IoU你需要确认代码里用的评估程序是否支持旋转框评估。demo里一般会自己实现旋转IoU的计算如果你的类别特别多或目标特别小这个数值会偏低但趋势和水平框检测一致。推理预测用python detect_obb.py --source datasets/demo/images/val/001.jpg --weights runs/train/exp/weights/best.pt --conf-thres 0.3预测结果会保存到runs/detect目录图片上绘制的是旋转矩形框。可视化这一步非常重要务必多去看几张图尤其是目标密集区域和图像边缘的目标。边缘目标的中心点可能超出图像或者角度被预测成相反的朝向这些在可视化里一眼就能看出来。我看到过不少案例训练日志的mAP不算低但可视化结果里框总是歪一个固定角度后来发现是标签的angle单位从度被读成了弧度只差了57倍肉眼上就是旋转从这一点变成了另一侧。提示验证和预测阶段务必确认detect_obb.py里读取角度后是否有后处理转换。训练时角度可能是-90到0推理后要映射回和标注一致的角度范围否则画出来的框会翻转。建议先在单张图上做测试确认无误再批量跑。5. 避坑与常见问题排查五个让训练翻车的细节逐个说透5.1 现象loss乱跳theta_loss居高不下训练一开始theta_loss出现震荡甚至整个box_loss也在抖动模型完全没有稳定的趋势。原因十有八九是学习率设置过大尤其是从一个未训练的随机权重开始。demo默认的初始学习率是0.01但旋转框分支的角度回归对这个数值比水平框更敏感因为角度差一个小数对应的框方向变化就很大。解决方法是把初始学习率调低到0.001同时开启warmup。在train_obb.py里搜索lr0参数或者直接在命令行指定python train_obb.py --lr0 0.001 --warmup-epochs 3 --epochs 200如果你用了我上面所说的半精度训练FP16还可以先关闭AMP再看效果。部分自定义算子在半精度下数值稳定性差角度分支的梯度容易出现NaN。5.2 现象训练能进行但验证时所有预测框都是水平框模型训练完了mAP也不差但可视化结果显示出来的框全部是矩形角度根本没生效。出现这种现象先怀疑推理端的后处理。预测输出的角度是网络直接给出的原始值但画框时需要把它转换到和GT一致的角度范围。如果推理代码里少了一步角度编码的逆转换就会把角度直接当成0处理或者因为角度范围映射错误把-30度映射成了水平框。解决方式是在detect_obb.py里加一行角度转换验证把输出前的原始角度打印出来和标签真实角度对比。常见的转换逻辑是if angle -90: angle 180 elif angle 90: angle - 180具体要看你的角度定义。如果数据集的GT角度是0到90而网络输出是-90到0你需要把负角度加上90得到横平的范围或者加上180让框的朝向翻转。5.3 现象标签没问题但训练集loss比验证集低很多验证mAP极低过拟合是原因之一但往往被你忽视了另一个更隐蔽的问题训练和验证时图像预处理不一致。demo里datasets_obb.py在训练时通常使用Mosaic数据增强Mosaic会把四张图拼在一起。如果验证时不关Mosaic或者关掉了但没调整标注框的坐标映射验证标注和网络看到的图像内容对不上mAP自然崩盘。解决方式确认训练配置里--mosaic默认是打开的而val_obb.py里使用的是不带Mosaic的原始图像。如果你的自定义数据集比较小建议直接关闭Mosaicpython train_obb.py --no-mosaic --augment False关闭后训练每个epoch看到的图像就是原始图像新手在调试阶段往往能更快定位问题等模型稳定后再逐步打开增强。5.4 现象显存够但训练速度极慢训练一个epoch耗时很长专业工具检查GPU利用率只有30%左右。原因可能是数据加载的瓶颈比如标签读取逻辑里有串行文件操作或者图像解码使用了CPU单线程。yolov5自带的数据加载器已经够用但obb demo的datasets_obb.py可能没有做线程优化导致CPU预处理跟不上GPU训练速度。解决方式是用--workers参数增加数据加载线程数比如--workers 8。同时把图像解码换成opencv的imread避免使用PIL的Image.open后者的速度明显更慢。你还可以开启--cache-images把图像和标签直接缓存到内存第一次加载后训练速度会明显提升前提是机器内存够大。5.5 现象模型推理速度比原版yolov5慢了近一倍旋转框NMS需要计算旋转框之间的IoU普通的torchvision.ops.nms不适用因为它只处理水平框。demo自带的旋转NMS如果实现得不好后处理时间会占整个推理时间的大头。比如一张图上目标只有几十个但旋转IoU的计算复杂度高依然会拖慢速度。解决方式是替换成高效的旋转NMS实现。你可以看看demo的utils/nms_obb.py是否支持传入use_cuda参数或者是否实现了一个基于矢量的快速版本。如果只是纯Python实现考虑改用ops库里的旋转NMS或者自己写一个基于凸包的ioU计算。训练阶段用的是另一个loss推理阶段的速度主要卡在NMS这个优化值得做。6. 进阶旋转框NMS替换和模型导出这两个技巧让你的demo真正落地当你的模型训练到稳定阶段接下来要解决的是能不能用的问题。旋转框检测的推理流程中NMS是保证输出质量的关键。普通NMS用水平框IoU在角度偏差较大的情况下会漏掉那些实际重叠的目标或保留大量重复框。我从demo源码里找到它默认用的是基于OpenCV旋转矩形的RotatedRect方法这个方法的精度取决于角度离散化的粗度在多个目标密集排列时容易出错。我一般会替换成skimage库的rotate辅助计算或者直接用mmrotate里提供的旋转NMS实现。这里给出替换nms_obb.py中核心函数的关键代码import numpy as np from shapely.geometry import Polygon def rotated_iou(boxes1, boxes2): def polygon_from_xywhr(box): cx, cy, w, h, angle box cos_a, sin_a np.cos(angle), np.sin(angle) # 四个角点 corners np.array([[0.5*w, -0.5*h], [0.5*w, 0.5*h], [-0.5*w, 0.5*h], [-0.5*w, -0.5*h]]) rot_mat np.array([[cos_a, -sin_a], [sin_a, cos_a]]) return Polygon(corners rot_mat.T np.array([cx, cy])) polys1 [polygon_from_xywhr(b) for b in boxes1] polys2 [polygon_from_xywhr(b) for b in boxes2] ious np.zeros((len(polys1), len(polys2))) for i, p1 in enumerate(polys1): for j, p2 in enumerate(polys2): inter p1.intersection(p2).area union p1.union(p2).area ious[i, j] inter / union if union 0 else 0 return ious这段代码把框拆成四个角点再用shapely计算多边形交并比。虽然shapely不是最快的方式但它的精度高而且容错性强特别适合调试阶段确认旋转IoU的正确性。等确认无误后你可以再切换到编译好的C实现避免Python循环影响速度。替换NMS之前先跑通一个小数据集对比替换前后的mAP和推理耗时不要盲目替换否则效果不升反降。部署是另一个重点。如果要把训练好的模型导出到ONNX供TensorRT或OpenVINO加速需要修改导出脚本让角度分支作为额外的输出头保留。默认的yolov5导出脚本只导出检测坐标和类别角度不会带出来。一种常见做法是python export_obb.py --weights runs/train/exp/weights/best.pt --include onnx在export_obb.py里你需要在输出层增加一维具体位置在模型结构定义中搜索Detect类的self.no改为self.no nc 5 1其中最后的1就是角度维度。导出后用onnxruntime推理时输出张量的shape是[1, num_anchors, 5 nc]这里的5对应cx, cy, w, h, angle。注意onnxruntime推理出的角度单位是度还是弧度最好打印出来和pytorch推理结果对比一下很多人在这一步发现角度被radian转换搞乱了。关于推理置信度阈值我需要多说两句。旋转框检测的置信度分布和水平框不太一样因为角度分支引入了额外的角度误差所以同一个目标在高置信度的同时角度可能偏差5度到10度。在实际使用中我会把conf-thres设到0.25到0.35之间比水平框检测的0.4低一些否则容易漏检。同时利用角度连续性做一次简单的帧间滤波如果视频序列中前后两帧的目标角度突变超过20度多半是预测错误直接用前一帧的角度修正。这个技巧在处理连续帧图像时特别有效。我还养成了一个习惯每次用新数据训练之前先把demo自带的预训练权重在随机一张图上跑一次推理用一张已知角度的目标去验证预处理、标签解析、后处理这整个链路是否一致。如果这一步的输出角度和预期一致再开始正式训练。这个习惯帮我避免了好几次标签坐标归一化出错和角度范围不匹配造成的反复试错。从那以后每次换数据集、换环境、换服务器我都强制走一遍这个验证流程看起来多花了五分钟实际省下的是至少一个下午的排查时间。希望这个习惯对你也有用也希望这份demo在你的项目里顺利跑起来。本文还有配套的精品资源点击获取