ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

YOLOv5多任务改造:旋转框检测与语义分割的工业实战

YOLOv5多任务改造:旋转框检测与语义分割的工业实战 简介本资源是一套基于YOLOv5框架、面向WoodScape鱼眼车载数据集的旋转框目标检测与语义分割联合任务完整实现方案适用于计算机视觉方向的本科生、研究生及初入AI工程领域的开发者尤其适合作为课程设计、毕业设计或科研原型快速验证项目。压缩包共76个文件6.14MB涵盖46个Python核心脚本含训练/推理/数据预处理/可视化模块、8个YOLO配置yaml文件支持旋转框检测与分割双分支结构、5个JSON标注转换工具及图像/文档类辅助文件目录组织清晰模块职责明确如segbranch.jpg、lossweight.jpg等图示直观呈现分割分支设计与损失权重策略。已有126人下载学习配套README.md详述环境配置、数据准备与运行流程并提供实测可用的训练日志、可视化效果图及关键代码注释支持在此基础上拓展多任务学习或部署优化。1. 项目缘起当旋转框遇上语义分割一个被低估的实战场景最近在整理一些计算机视觉的实战项目时我发现一个很有意思的现象很多朋友在接触目标检测和语义分割时往往是分开学习的。要么是拿YOLOv5跑个COCO数据集检测一下行人车辆要么是用U-Net、DeepLabV3做个简单的图像分割。但实际工业场景中尤其是像自动驾驶、遥感影像分析、工业质检这些领域问题往往是复合型的——你不仅需要知道目标在哪里检测还需要精确地知道它的轮廓和类别分割甚至目标本身可能还是倾斜的旋转框。这就引出了我们今天要拆解的这个实战项目基于YOLOv5和WoodScape数据集的旋转框目标检测与语义分割。乍一看标题有点长技术栈也似乎混搭但它恰恰戳中了一个非常实际的痛点如何在一个统一的框架下同时处理倾斜目标的精确定位和像素级分类WoodScape这个数据集可能很多人不熟悉它是一个专注于自动驾驶周边感知的多任务数据集包含了环视鱼眼图像里面的目标比如车辆、行人由于视角原因常常是倾斜的用水平框axis-aligned bounding box会引入大量背景噪声效果大打折扣。所以这个项目的核心价值不在于用了多新的模型而在于它提供了一套务实的技术缝合方案。它没有去追逐最新的YOLOv8或YOLOv9而是基于成熟、社区资源极其丰富的YOLOv5进行改造增加了旋转框检测头并并行接入了语义分割分支。对于想深入理解多任务学习、模型改造以及处理复杂视觉任务的朋友来说这是一个绝佳的练手项目。它回答了这样一个问题当你的数据既有倾斜目标需要精确定位又需要对场景进行像素级理解时该怎么搭建你的训练Pipeline2. WoodScape数据集深度解析为什么它适合做多任务学习在动手写代码之前我们得先吃透数据。WoodScape数据集是这个项目的基石它的特性直接决定了我们模型架构的设计方向。WoodScape是一个为自动驾驶车辆环视系统Surround View设计的大规模数据集。它最大的特点在于其多任务标注和鱼眼图像。与KITTI、Cityscapes等常见数据集不同WoodScape的图像来自车辆的四个鱼眼摄像头形成了360度的环视视角。这种图像会产生严重的径向畸变目标尤其是靠近图像边缘的车辆、行人会呈现明显的倾斜和变形。2.1 数据标注格式与挑战WoodScape提供了四大类任务的标注语义分割像素级的类别标签包括道路、人行道、车辆、行人、天空等数十个类别。实例分割为每个独立的物体实例提供像素级掩码。2D目标检测传统的水平边界框。旋转目标检测这正是本项目关注的重点。标注以旋转矩形框Rotated Bounding Box的形式提供通常用(cx, cy, w, h, angle)或(x1, y1, x2, y2, x3, y3, x4, y4)来表示一个倾斜的矩形。对于旋转框(cx, cy, w, h, angle)是更紧凑的表示法其中angle的定义是关键。不同库如OpenCV, DOTA_devkit对旋转角度的定义可能不同例如是相对于x轴的角度还是矩形第一条边的角度角度范围是[0, 90°]还是[0, 180°]。WoodScape数据集通常采用与OpenCV一致的表示法angle表示矩形第一条边宽度w对应的边与x轴正方向的夹角范围在(-90, 0]度之间有时是[0, 90)。在数据预处理时必须首先确认并统一角度定义否则后续的损失计算会完全错误。2.2 鱼眼图像的处理策略鱼眼图像给模型带来了额外的挑战畸变传统的CNN是在透视图像上训练出来的直接输入鱼眼图模型需要额外学习畸变不变性这增加了学习难度。目标形态边缘的目标被严重拉伸和弯曲水平框会包含大量无关区域。因此常见的预处理策略有两种去畸变利用相机标定参数将鱼眼图像校正为透视图像。这样做的好处是可以让模型专注于学习通用的特征但会损失部分图像信息特别是边缘部分会被拉伸或产生黑边。直接使用不进行去畸变让模型和数据增强来学习这种畸变模式。这要求数据增强必须足够强大并且模型容量要足够。对于WoodScape由于其环视特性保留畸变信息有时对理解场景布局更有帮助。在本项目中为了简化流程并专注于模型改造我建议采用去畸变后的图像作为输入。这样我们可以直接利用在ImageNet等大数据集上预训练的主干网络Backbone权重加速收敛。去畸变的代码需要用到数据集中提供的相机内参intrinsics和畸变系数distortion coefficients通常使用OpenCV的cv2.fisheye.undistort函数即可完成。3. YOLOv5架构改造嵌入旋转框检测与语义分割分支原版YOLOv5是一个高效的单阶段目标检测器输出的是水平框(x_center, y_center, width, height, obj_conf, class_conf...)。我们的目标是让它同时输出旋转框和分割掩码。3.1 旋转框检测头的设计这是改造的核心。我们不能简单地在回归层增加一个角度参数因为角度的周期性0度等于180度会导致回归损失函数的不连续例如预测179度与真实值-179度实际上只差2度但L1/L2损失会计算为很大的误差。主流解决方案是使用“角度分类”或“高斯Wasserstein距离GWD”损失。方案一角度分类本项目推荐将连续的角度离散化。例如将[-90° 90°)的范围均匀划分为180个bins每度一个bin。这样角度预测就变成了一个分类问题。模型需要输出t_x, t_y, t_w, t_h中心点和宽高的偏移量与YOLO原版一致。t_theta一个180维的向量表示角度属于每个bin的概率。obj_conf目标置信度。class_conf类别置信度。在推理时通过softmax得到角度分布取argmax得到离散角度或者通过期望计算得到更精细的角度值。需要对YOLOv5的Head部分进行如下修改在models/yolo.py中Detect类对应的输出卷积层增加角度的输出通道。例如原来每个anchor的输出通道是(5 num_classes)现在需要改为(5 angle_bins num_classes)。在loss.py中需要修改损失函数compute_loss。框回归损失GIoU Loss需要替换为适用于旋转框的损失如旋转IoURotated IoU损失或SkewIoU损失。同时增加一个角度分类的交叉熵损失。旋转IoU的计算比水平IoU复杂得多推荐使用成熟的库如rbox_iou_opsCUDA实现或shapelyPython几何计算库较慢但易于调试。在训练初期可以先用水平框损失预热再引入旋转框损失。方案二高斯Wasserstein距离损失将旋转框建模为一个二维高斯分布用两个分布的Wasserstein距离作为损失。这种方法能天然地处理旋转对称性问题如正方形旋转90度后不变理论更优美但实现和理解起来更复杂。对于初次改造强烈建议采用方案一角度分类因为它更直观易于调试且与YOLOv5原有的分类分支结构相似。3.2 语义分割分支的添加我们需要在YOLOv5的架构上添加一个并行的分割解码器Decoder。一个经典且高效的做法是借鉴U-Net或FPN的思想利用主干网络的不同阶段特征图进行融合。具体改造步骤特征提取YOLOv5的BackboneCSPDarknet会产生三个不同尺度的特征图P3, P4, P5分别对应浅层细节和深层语义信息。分割头设计新增一个分割分支。通常我们将深层语义信息丰富的P5特征图上采样然后与空间细节丰富的P4、P3特征图进行逐元素相加Add或通道拼接Concat。特征融合可以设计一个轻量级的FPN结构。例如对P5进行2倍上采样与P4融合。对融合后的特征再进行2倍上采样与P3融合。最终对融合了多尺度信息的特征图进行一个简单的卷积层将通道数映射到分割类别数num_seg_classes。输出分割头的输出是一个[B, num_seg_classes, H, W]的张量其中H, W是输入图像的下采样尺寸如原图640x640下采样4倍或8倍。训练时使用交叉熵损失或Dice损失对类别不平衡问题更鲁棒。代码修改点在models/common.py中定义新的分割模块如SegmentationHead。在models/yolo.py的Model类中初始化这个分割头并在前向传播forward方法中返回检测结果和分割结果。在loss.py中增加分割损失的计算。3.3 多任务损失的平衡现在模型有三个主要损失检测损失框回归角度分类目标置信度类别分类、分割损失。直接简单相加 (L_total L_det L_seg) 可能会导致一个任务主导训练。需要引入动态权重平衡常见方法有不确定性加权为每个任务损失学习一个可训练的参数log方差自动调整权重。L_total sum(exp(-s_i) * L_i s_i)其中s_i是任务i的待学习参数。GradNorm动态调整权重使得不同任务的梯度幅度相近。手动调参初期实用先固定一个任务如检测的权重为1.0然后以0.1为步长调整另一个任务分割的权重观察验证集上两个任务指标的变化找到一个平衡点。在项目初期建议从手动调参开始快速验证架构的可行性。4. 项目实战从环境搭建到训练调优全流程假设我们已经拿到了WoodScape数据集可能需要申请并完成了去畸变预处理将其整理成了YOLO格式的旋转框标注和分割掩码图。4.1 环境配置与代码结构# 基础环境建议使用Python 3.8 PyTorch 1.7 git clone 本项目源码仓库 cd rotated_yolov5_seg pip install -r requirements.txt # 需要额外安装旋转框IoU计算库如 rotated_iou 或 box_iou_rotated # 项目目录结构建议 rotated_yolov5_seg/ ├── data/ │ ├── woodscape.yaml # 数据集配置文件 │ ├── images/ # 训练/验证图像 │ ├── labels_rotated/ # 旋转框标注 (YOLO格式: cls cx cy w h angle) │ └── masks/ # 语义分割标注图 (单通道PNG像素值为类别ID) ├── models/ │ ├── common.py # 修改添加分割头等模块 │ ├── yolo.py # 修改整合检测与分割头 │ └── rotated_yolov5s.yaml # 新的模型配置文件 ├── utils/ │ ├── datasets.py # 修改支持同时加载图像、旋转框标签、分割掩码 │ ├── loss.py # 重写包含旋转框损失和分割损失 │ └── metrics.py # 扩展评估旋转框mAP和分割mIoU ├── train_multi_task.py # 新的多任务训练脚本 └── detect_seg.py # 新的推理脚本可同时输出检测框和分割图关键修改点详解data/woodscape.yaml:path: ../datasets/woodscape train: images/train val: images/val # 类别数 nc: 10 # 检测类别数如 car, pedestrian, cyclist... # 分割类别数 (通常比检测类别多包含背景、道路等) seg_nc: 19 # 类别名称 names: [car, pedestrian, ...] seg_names: [road, sidewalk, car, ...] # 旋转框角度bins angle_bins: 180utils/datasets.py在LoadImagesAndLabels类的__getitem__方法中需要同时加载图像、旋转框标签文件和对应的分割掩码文件。返回的字典应包含img,rotated_labels,masks。train_multi_task.py主训练循环需要处理两个损失。在每个batch中# 前向传播 pred, seg_out model(imgs) # 计算损失 loss, loss_items compute_loss((pred, seg_out), (rotated_targets, masks)) # loss_items 应包含 [box_loss, angle_loss, obj_loss, cls_loss, seg_loss]优化器可以对所有参数一起优化也可以为不同部分设置不同的学习率如分割头可以设置更大的初始lr。4.2 训练技巧与参数调优预热训练不要一开始就上多任务。可以先冻结分割头只训练旋转框检测任务几轮让检测部分先收敛到一个较好的状态。然后再解冻分割头用较小的学习率开始多任务训练。数据增强针对旋转框旋转增强需要特别小心。对图像和分割掩码进行旋转时旋转框的坐标和角度必须同步、正确地变换。Mosaic和MixUp增强也需要适配旋转框。学习率策略使用Cosine退火或带热重启的Cosine退火。多任务训练时损失曲面可能更复杂适当降低初始学习率如lr01e-3降为3e-4。超参数hyp.scratch.yaml中的超参数需要调整。特别是分类和角度分类损失的权重cls_pw,angle_pw以及分割损失的权重seg_pw。建议从seg_pw0.5开始尝试。验证指标需要同时监控检测指标mAP0.5mAP0.5:0.95需使用旋转框mAP计算脚本。分割指标mIoU交并比Pixel Accuracy。训练损失观察各个损失项是否平稳下降有无剧烈震荡。4.3 推理与可视化推理脚本需要同时运行两个头# 模型加载 model attempt_load(best.pt, map_locationdevice) model.eval() # 推理 with torch.no_grad(): det_out, seg_out model(img) # det_out: [1, num_boxes, (5 angle_bins num_det_classes)] # seg_out: [1, num_seg_classes, H, W] # 后处理 # 1. 处理检测输出NMS需要支持旋转框NMS rotated_boxes, angles, confs, cls_ids non_max_suppression_rotated(det_out, ...) # 2. 处理分割输出取argmax得到每个像素的类别ID seg_mask torch.argmax(seg_out[0], dim0).cpu().numpy() # [H, W] # 可视化 plot_rotated_boxes_on_image(original_img, rotated_boxes, angles) overlay_seg_mask(original_img, seg_mask, alpha0.5)可视化时可以将旋转框和分割结果叠加在同一张图上直观评估模型性能。5. 避坑指南与常见问题排查在实际操作中你几乎一定会遇到下面这些问题。5.1 旋转框损失震荡或为NaN原因1角度定义不一致。数据标注、数据增强、损失计算、NMS等环节对旋转角度的定义范围、起始边必须完全一致。解决方案在项目开始时就定义一个全局的angle_utils.py所有角度操作都通过这个工具函数进行并编写单元测试验证。原因2旋转IoU计算不稳定。当两个框几乎平行或重合时几何计算可能出现奇点。解决方案使用成熟的、经过数值稳定性处理的库如mmrotate中的rbox_iou。在损失函数中加入极小值eps防止除零。原因3梯度爆炸。多任务模型可能更不稳定。解决方案使用梯度裁剪torch.nn.utils.clip_grad_norm_并降低初始学习率。5.2 分割结果全为某一类如背景原因严重的类别不平衡。WoodScape中“道路”和“天空”等类别像素占比极大。解决方案使用Dice Loss或Focal Loss替代标准的交叉熵损失它们能更好地处理不平衡问题。在损失函数中为每个类别设置不同的权重权重与类别像素频率的倒数成正比。在数据增强中有针对性地对少数类别区域进行过采样或复制-粘贴。5.3 检测与分割任务互相拖累现象同时训练时一个任务指标上升另一个任务指标下降。解决方案检查共享主干两个任务共享Backbone可能特征存在冲突。可以尝试在Backbone的较高层如C4或C5阶段就将特征图分叉让两个任务拥有相对独立的低层特征提取路径。调整损失权重使用前面提到的“不确定性加权”方法进行自动化调优。课程学习先训练一个任务较多轮次再引入第二个任务进行微调。5.4 模型速度严重下降原因分割头增加了大量计算特别是上采样和特征融合操作。优化方案轻量化分割头减少分割头中的通道数。例如将FPN融合层的通道数从256减至128或64。降低分割输出分辨率模型内部计算可以用1/8或1/16的下采样率只在最终输出前上采样到1/4。推理时分割图可以通过插值快速放大到原图尺寸。知识蒸馏训练一个大型的教师网络然后用它来蒸馏一个小型的学生网络在精度和速度间取得平衡。这个项目就像搭积木把YOLOv5、旋转框、语义分割这几个成熟的“积木块”以正确的方式组合起来。最大的收获不是最终模型的效果而是在这个缝合过程中你对数据流、损失函数、多任务平衡、调试技巧的深刻理解。我自己的经验是先从最简单的部分跑通比如只改旋转框再加分割每一步都做好验证和可视化这样出了问题也能快速定位。WoodScape数据集的复杂性迫使你去思考更底层的问题这比在标准数据集上刷分数要有价值得多。本文还有配套的精品资源点击获取
返回列表