简介YOLOv5 OBB旋转框训练DEMO是一套面向目标检测开发者和进阶学习者的实操示例针对遥感影像、航拍视频中的倾斜目标如车辆、树木、建筑物检测需求在YOLOv5基础上引入旋转边界框机制弥补常规轴对齐框的不足。资源完整覆盖了OBB训练的数据预处理、配置调整、损失计算、训练验证与推理部署环节适合希望理解旋转框原理并快速跑通流程的读者。压缩包共583个文件大小445.58MB包含166个XML标注文件、95个TXT标签、93张PNG图像、71个Python脚本、38个YAML配置以及若干TensorRT/ONNX相关文件目录结构清晰兼顾数据、代码与训练产物。目前已有1818人学习下载。通过该资源读者可对照实验记录和脚本直观了解OBB数据集的构建方法、旋转框损失函数的实现细节并借助现成的训练、预测、导出脚本迁移到自有数据其中事件日志和推理代码也为排查训练问题、优化模型精度提供了实际参考。1. YOLOv5 OBB旋转框训练Demo先搞懂它解决什么问题水平框检测在遇到细长目标时经常翻车比如遥感影像里的建筑物、工厂传送带上的长条工件、货架上的倾斜标签。你用普通YOLOv5去框一个斜着放的螺丝刀检测框要包住整个目标结果框里一大半是背景两个挨得近的目标还会互相污染特征。这时候就要上旋转框检测OBBOriented Bounding Box让检测框跟着目标角度一起转。这份YOLOv5 OBB旋转框训练Demo就是一个能直接跑起来的旋转框目标检测项目基于YOLOv5的框架扩展了角度回归分支支持带角度标签的训练、验证和推理。适合做遥感目标检测、工业质检、场景文字检测的开发者想在自己数据集上快速落地旋转框方案又不想从头写检测器的人。拿到这份资源你要理解的不只是“跑通训练命令”而是数据集标注格式、模型输出头、损失函数和后处理这几处相对水平框都改在哪。下面按我实际拆解这个项目的顺序一层层讲清楚。2. 资源包结构与环境准备先看清项目骨架再动手2.1 目录结构与关键文件解析解压这个Demo后第一件事不是急着装环境而是先过一遍项目骨架。常见的YOLOv5 OBB分支项目会在原始YOLOv5基础上增加或修改这几个关键路径。models/yolov5s-obb.yaml模型结构配置定义旋转框检测头的输出维度。水平框每个anchor输出5nc个值x,y,w,h,objclasses旋转框版本会在此基础上增加角度分支一般是再加一个或两个输出值。utils/loss.py损失函数。这里会多出角度损失项的计算逻辑常见做法是使用Smooth L1或KLDKullback-Leibler Divergence损失来处理角度回归。utils/metrics.py评估指标。从水平IoU换成旋转IoU计算mAP的统计方式跟着变。utils/datasets.py或utils/augmentations.py数据加载与增强逻辑。Mosaic、随机仿射变换等增强操作会针对旋转框做适配尤其在随机旋转时要同步更新标签的角度值。detect.py / val.py / train.py入口脚本。训练时通过参数指定OBB模式推理输出会多出角度字段画框函数也要改成绘制旋转矩形。我第一次打开这类Demo时犯过一个错直接用pip装好requirements.txt就开跑结果报“ModuleNotFoundError: utils.plots”之类的错。原因就是路径没加对项目里很多模块是通过相对路径互相import的必须从项目根目录执行命令。2.2 Python环境安装与依赖核对环境依赖方面这类OBB Demo通常会要求在PyTorch环境下运行。我一般建议先创建独立的conda环境避免和已有项目打架。conda create -n yolov5-obb python3.8 conda activate yolov5-obb cd /path/to/yolov5-obb pip install -r requirements.txt逻辑说明这条命令创建一个干净的Python 3.8环境然后进入项目目录安装依赖。requirements.txt里一般包含torch、torchvision、opencv-python、numpy、pyyaml、tqdm、matplotlib等基础包。PyTorch版本建议和你的显卡驱动匹配如果是30系及以上显卡装CUDA 11.x配套的torch版本即可。参数说明python3.8不是必须的3.9、3.10也能跑但有些项目依赖的老版本pycocotools在3.10上编译容易出问题用3.8最省事。装完依赖后我先跑一个简单命令验证torch和CUDA是否正常工作python -c import torch; print(torch.__version__, torch.cuda.is_available())逻辑说明如果输出里cuda.is_available()是False说明torch装成了CPU版本或者驱动没对上。遇到这种情况直接去PyTorch官网选对应CUDA版本的安装命令重新装不要在这个问题上耗时间。依赖装好后建议先跑一次官方提供的demo数据训练一遍验证环境和代码流程没问题再上自己的数据集。这个Demo里通常会带一个小的示例数据集下载解压后放在datasets目录下数据结构大概是datasets/ demo_obb/ images/train/xxx.jpg images/val/xxx.jpg labels/train/xxx.txt labels/val/xxx.txt标签文件格式是旋转框的核心下一章详细说。3. 数据集准备从四点标注到可训练的旋转框3.1 旋转框标注格式两种主流表示法旋转框标注最常用的是两种格式一种是OpenCV风格的cx cy w h angle其中angle是角度值单位是度表示框相对水平方向的旋转角度另一种是四角点坐标格式x1 y1 x2 y2 x3 y3 x4 y4直接记录旋转矩形的四个顶点坐标按顺时针或逆时针顺序排列。YOLOv5 OBB分支在数据加载阶段一般会把格式统一成某种内部表示常见的是把四角点转换成中心点加宽高加角度的表示或者反过来。这个Demo里标签文件的组织方式我见过两种变体一种是标签每行格式class_id x1 y1 x2 y2 x3 y3 x4 y4坐标值做了归一化另一种是class_id cx cy w h angle这种格式下角度范围通常是0到180或-90到90取决于作者实现。你在用之前一定要去看dataset.py里怎么解析标签否则标注格式对不上训练出来的模型会“看起来很努力但效果完全不对”。如果你手头标注工具导出的格式不是项目需要的可以写一个转换脚本。下面这个例子演示了如何把四角点格式转换成cx cy w h angle# convert_quad_to_cxywh_angle.py import numpy as np def quad_to_cxywh_angle(points): 将四角点坐标转为 (cx, cy, w, h, angle) points: 8个数值 [x1, y1, x2, y2, x3, y3, x4, y4] 返回: [cx, cy, w, h, angle]角度单位为度 pts np.array(points, dtypenp.float32).reshape(4, 2) # 计算四条边的长度取平均得到宽和高 d1 np.linalg.norm(pts[1] - pts[0]) d2 np.linalg.norm(pts[2] - pts[1]) w (d1 d2) / 2 h (d1 d2) / 2 # 中心点取四个顶点平均值 cx pts[:, 0].mean() cy pts[:, 1].mean() # 角度用第一条边的方向角弧度转角度 angle np.arctan2(pts[1][1] - pts[0][1], pts[1][0] - pts[0][0]) angle_deg angle * 180 / np.pi return [cx, cy, w, h, angle_deg] # 示例一个旋转30度的矩形 quad [100, 100, 186, 150, 166, 236, 80, 186] cx, cy, w, h, angle quad_to_cxywh_angle(quad) print(fcx{cx:.2f}, cy{cy:.2f}, w{w:.2f}, h{h:.2f}, angle{angle:.2f})逻辑说明这个脚本的核心是几何换算。四角点坐标直接来自标注工具通常是有序的顺时针或逆时针四点。宽高通过取两条对边距离的平均来抗噪声角度用第一条边的方向角这样只要顶点顺序稳定角度计算就不会出现跳变。参数说明角度范围很关键。如果后续训练时项目把角度限制在0到180度那这里算出来的负角度就要加180归一化。不同的OBB实现对这个参数的处理逻辑不同有的用180表示半周期有的用90表示四分之一周期你转换脚本里写死之前先确认目标项目的角度编码方式。3.2 数据yaml配置与标签目录组织数据yaml文件在YOLOv5体系中扮演“地图”的角色OBB版本会在此基础上额外增加角度相关的配置项。下面是一个典型的OBB数据配置# data_obb.yaml train: datasets/demo_obb/images/train val: datasets/demo_obb/images/val nc: 1 names: [object] # 角度相关配置 angle: 180逻辑说明train和val指向图片目录labels目录与图片目录同级只是把images换成labels这是YOLOv5默认的标签路径推导规则。nc和names是类别数及名称。angle配置角度范围取180表示模型预测0到180度的旋转角这个值直接决定角度回归分支的输出设计和损失计算方式。参数说明angle: 180适合大部分场景但如果你的数据里目标角度天然有对称性比如所有工件都是水平的或垂直的可以改成angle: 90降低角度回归难度。新手先别动这个参数用默认值跑通再说。目录组织上最容易出错的是图片和标签的命名对应关系。YOLOv5要求图片a.jpg对应的标签是a.txt文件名前缀必须完全一致否则训练时会报“No labels found in ...”。另外标签文件必须是普通文本不能是UTF-8带BOM否则解析第一行标签时会读出一个隐藏字符导致坐标解析错位。我习惯在训练前先写一个检查脚本统计每张图片对应的标签文件是否存在、标签坐标有没有越界# check_labels.py import os from pathlib import Path img_dir Path(datasets/demo_obb/images/train) label_dir Path(datasets/demo_obb/labels/train) missing 0 for img in img_dir.iterdir(): label_path label_dir / (img.stem .txt) if not label_path.exists(): missing 1 print(f缺少标签: {img.name}) else: # 检查标签坐标范围 for line in label_path.read_text().strip().splitlines(): parts line.split() coords list(map(float, parts[1:])) if any(c 0 or c 1 for c in coords): print(f坐标越界: {label_path.name} - {line}) print(f检查完成缺失标签数: {missing})逻辑说明这个小脚本对所有训练图片逐一检查标签文件是否存在并检查归一化坐标是否超出0到1的范围可以有效拦截两种最基础的数据问题。4. 训练与推理核心命令和参数细节4.1 训练命令与OBB关键超参数据准备好、检查脚本跑完没有异常就可以开始训练了。这个Demo的基础训练命令和原版YOLOv5非常接近python train.py \ --data data_obb.yaml \ --weights yolov5s.pt \ --img 640 \ --batch-size 16 \ --epochs 100 \ --device 0 \ --project runs/train_obb逻辑说明--weights可以选择官方预训练权重作为初始权重迁移学习能显著加速收敛。如果不指定模型会从头开始训练收敛慢且效果通常差一些。--device 0指定使用第一张GPU卡。--project指定输出目录训练过程中的权重和日志会保存到这里。参数说明--img 640表示输入图像缩放尺寸如果你的目标比较小可以设成896或1280显存不够就减小batch-size或者用--cache参数把图片加载进内存。这个Demo如果支持多尺度训练还会有一个--multi-scale的选项每10个batch随机改变输入尺寸能提升模型对尺度变化的鲁棒性。角度相关的训练参数一般在hyp.yaml文件里配置OBB版本会多出角度损失权重这一项。常见的角度损失权重设置在0.5到2.0之间权重太大会让模型在早期训练时把大量精力花在角度回归上导致位置和分类学不好权重太小则角度输出不准确推理时旋转框看着歪歪扭扭。# hyp_obb.yaml 关键片段 lr0: 0.01 lrf: 0.2 angle_loss: 0.8逻辑说明lr0是初始学习率lrf是最终学习率系数训练结束时学习率会衰减为lr0*lrf。angle_loss是角度损失相对位置损失的权重0.8表示角度损失在总损失中占八成的位置损失权重。4.2 推理与结果可视化读懂旋转框输出训练完成后推理命令和原版YOLOv5相似python detect.py \ --weights runs/train_obb/exp/weights/best.pt \ --source data/images/ \ --conf-thres 0.25 \ --iou-thres 0.45 \ --device 0逻辑说明--source可以指向图片目录、单张图片或视频文件。推理完的结果保存在runs/detect/exp下旋转框会以带角度的矩形画在图上。这里有一个值得注意的地方水平框模型推理输出的xywh是x_center, y_center, w, h而OBB模型在xywh基础上多出angle字段detect.py在处理预测结果时会多一步角度读取。如果你在自己写的推理脚本里接这个模型要知道输出Tensor的格式通常是这样[ batch_size, num_queries, 6 num_classes ] # 其中6 cx, cy, w, h, angle, obj_conf # 后面是各类别概率从预测结果里取旋转框并画图的逻辑大致是# draw_obb.py import cv2 import numpy as np def draw_obb(image, cx, cy, w, h, angle_deg, color(0, 255, 0)): 根据中心点、宽高和角度画旋转矩形 angle_deg是角度制的旋转角cv2画旋转矩形时用负弧度制 angle_rad -angle_deg * np.pi / 180.0 cos_v np.cos(angle_rad) sin_v np.sin(angle_rad) # 半宽半高 hw w / 2.0 hh h / 2.0 # 旋转矩形的四个顶点相对中心点的偏移 corners np.array([ [-hw, -hh], [hw, -hh], [hw, hh], [-hw, hh] ]) # 旋转矩阵 rotation np.array([ [cos_v, -sin_v], [sin_v, cos_v] ]) rot_corners corners rotation.T np.array([cx, cy]) # 转为整数坐标并按多边形绘制 pts rot_corners.astype(np.int32).reshape(-1, 1, 2) cv2.polylines(image, [pts], isClosedTrue, colorcolor, thickness2) return image逻辑说明旋转矩形本质上就是把四个角点先按宽高摆好然后围绕中心点旋转指定角度。cv2.polylines直接画闭合多边形不需要自己重新计算每条边的方程。这里的角度正负号取决于模型输出的角度定义我拆过的这个Demo里角度是顺时针为正所以画图时要取负号转成OpenCV的逆时针方向。参数说明w和h的定义也很讲究。有的实现里w是长边h是短边角度指向长边方向有的实现则固定用第一个维度作为w。如果推理结果看起来有些框长边和短边反了检查一下模型训练时标签里是长边在前还是短边在前画图时对应调整即可。5. 避坑手册旋转框训练中遇到的五类典型问题5.1 训练loss不下降或震荡角度周期性问题在作祟现象前10个epoch的box loss在下降但总loss波动剧烈训练曲线像锯齿验证集mAP长时间在低值徘徊。原因角度这个量的周期性和坐标不一样。当模型预测的角度值是179度而标签角度是0度时数值上差了179度但实际旋转框几乎重合。多数损失函数不具备角度周期性感知直接用Smooth L1算角度差会让模型在角度边界附近反复横跳。解决优先检查项目是否支持角度范围90度模式把angle从180改成90。这样角度编码区间从0到180压缩到0到90配合数值归一化能有效减少边界问题。如果项目支持使用KLD损失替代Smooth L1做角度回归切换过去通常能缓解震荡。提示KLD损失对角度差做了周期性处理在角度接近边界时梯度更平滑是旋转框检测里常用的角度回归损失方案。5.2 训练时报“Assertiont ≥ 0 t nfailed”坐标越界现象数据集加载到一半直接报错异常发生在读取标签并计算grid索引时个别样本的坐标出现负数或超过1的值。原因标签坐标没有归一化或归一化时用了错误的图片尺寸。常见场景是标注工具导出的是绝对像素坐标你直接放进labels目录没做除以图片宽高的操作。另一种情况是图片存在EXIF旋转信息标注时看的图片和训练时读入的图片方向不一致坐标跟着失效。解决写检查脚本先跑一遍第3章的脚本就能拦截这个问题。如果确认标签本身没问题但报错还在检查图片读取路径下的原图尺寸与标注时是否一致尤其注意标注时是否裁剪过图片。我通常会在转换脚本里把每个标签的坐标极值打印出来一眼定位是不是归一化步骤出了问题。5.3 模型在验证集mAP很高但现场效果差数据增强与角度分布不匹配现象在自己的测试集上mAP有0.85拿到现场拍的照片上检测率直线下降尤其角度接近45度或135度的目标几乎全漏。原因训练数据里的角度分布不均匀大量目标集中在水平或垂直方向45度和135度方向的样本很少。模型对样本稀疏的角度区域天然拟合不足。旋转框检测对角度分布比水平框更敏感因为模型需要学到的不仅是“这里有个东西”还包括“这个东西转了多大幅度”。解决检查训练集的角度分布写个统计脚本画出角度直方图。如果某个区间样本很少可以加大该角度区间的数据增强权重比如在Mosaic之外的随机旋转增强中把旋转角度范围从±15度扩大到±30度。如果项目里支持角度抖动angle jitter增强打开它并设一个合理的值比如20度。5.4 推理时旋转框没有完全贴合目标后处理里缺了一步角度归一化现象检测框大概位置正确但整体旋转角度比实际目标偏移了几度到十几度目标越长偏移越明显。原因模型在训练时角度是在特定区间编码的推理时输出层的激活函数如果没做对应解码输出的角度可能是未经处理的原始值。比如训练时用sigmoid把角度压缩到0到1再乘以180推理时直接拿sigmoid输出乘以180但没减偏移值。解决翻一下models/yolov5s-obb.yaml里角度分支的最终激活函数定义再对照decode逻辑。一般推理结果里角度字段要经过一次缩放和平移公式是angle_pred sigmoid(x) * angle_range - angle_offset具体数值看代码实现。直接在detect.py里打印几组原始输出和解码后的角度值对比标签就能看出解码是否正确。5.5 显存OOM旋转框分支让显存压力变大现象原版YOLOv5用batch-size 16跑得动切到OBB版本同样配置直接报CUDA out of memory。原因OBB模型增加了一个角度输出通道特征图上的输出维度变大同时损失函数里旋转IoU的计算涉及更多矩阵运算显存占用比同体积水平框模型高20%到40%。解决不要和原版参数较劲先降batch-size到8或4跑通。如果必须保持batch-size可以开梯度累积对训练节奏影响较小。另一个实用做法是把--img从640降到544输入尺寸小了显存会空出不少mAP掉得有限。6. 进阶验证技巧合成旋转矩形测试角度预测准确性模型训练完、推理流程跑通以后建议你用合成数据做一个快速验证把角度预测的准确性单独测出来。这个技巧能帮你把“检测框位置准不准”和“角度准不准”两个维度解耦定位问题时特别有用。写一个脚本生成一批已知角度旋转矩形的图片矩形中心位置随机但角度是预设值比如每隔15度一张图# synthetic_test.py import cv2 import numpy as np def create_synthetic_rect(size(640, 640), angle_deg30): 生成一张含单个旋转矩形的合成图 img np.zeros((size[0], size[1], 3), dtypenp.uint8) img[:] (128, 128, 128) cx, cy 320, 320 w, h 200, 60 angle_rad -angle_deg * np.pi / 180.0 corners np.array([ [-w/2, -h/2], [w/2, -h/2], [w/2, h/2], [-w/2, h/2] ]) rot np.array([ [np.cos(angle_rad), -np.sin(angle_rad)], [np.sin(angle_rad), np.cos(angle_rad)] ]) rot_corners corners rot.T np.array([cx, cy]) pts rot_corners.astype(np.int32).reshape(-1, 1, 2) cv2.fillPoly(img, [pts], (0, 255, 0)) return img # 生成一组固定中心点、不同角度的测试图 for angle in range(0, 180, 15): img create_synthetic_rect(angle_degangle) cv2.imwrite(ftest_angle_{angle}.jpg, img)然后用训练好的模型对这批合成图做推理重点看预测的角度值与真实角度的误差。如果误差在5度以内说明角度回归分支工作正常推理时的解码逻辑也对如果误差很大优先检查模型推理输出的角度解码和画图时的角度符号约定。合成数据没有真实场景的纹理和噪声能测出模型的几何感知能力但测不出泛化能力。所以它适合当冒烟测试用在跑新数据集、改角度超参之后花十分钟验证一下角度分支没坏然后再去精调别的参数。我从那以后每次换数据集都强制走一遍“合成矩形验证角度、真实数据验证检测率”的流程能在正式训练前拦截掉大部分标注格式和角度编码的问题。希望帮到你。本文还有配套的精品资源点击获取