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进行改造，增加了旋转框检测头，并并行接入了语义分割分支。对于想深入理解多任务学习、模型改造以及处理复杂视觉任务的朋友来说，这是一个绝佳的练手项目。它回答了这样一个问题：当你的数据既有倾斜目标需要精确定位，又需要对场景进行像素级理解时，该怎么搭建你的训练Pipeline？

2. WoodScape数据集深度解析：为什么它适合做多任务学习

在动手写代码之前，我们得先吃透数据。WoodScape数据集是这个项目的基石，它的特性直接决定了我们模型架构的设计方向。

WoodScape是一个为自动驾驶车辆环视系统（Surround View）设计的大规模数据集。它最大的特点在于其多任务标注和鱼眼图像。与KITTI、Cityscapes等常见数据集不同，WoodScape的图像来自车辆的四个鱼眼摄像头，形成了360度的环视视角。这种图像会产生严重的径向畸变，目标（尤其是靠近图像边缘的车辆、行人）会呈现明显的倾斜和变形。

2.1 数据标注格式与挑战

WoodScape提供了四大类任务的标注：

1. 语义分割：像素级的类别标签，包括道路、人行道、车辆、行人、天空等数十个类别。
2. 实例分割：为每个独立的物体实例提供像素级掩码。
3. 2D目标检测：传统的水平边界框。
4. 旋转目标检测：这正是本项目关注的重点。标注以旋转矩形框（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部分进行如下修改：

1. 在models/yolo.py中Detect类对应的输出卷积层，增加角度的输出通道。例如，原来每个anchor的输出通道是(5 + num_classes)，现在需要改为(5 + angle_bins + num_classes)。
2. 在loss.py中，需要修改损失函数compute_loss。框回归损失（GIoU Loss）需要替换为适用于旋转框的损失，如旋转IoU（Rotated IoU）损失或SkewIoU损失。同时，增加一个角度分类的交叉熵损失。
3. 旋转IoU的计算比水平IoU复杂得多，推荐使用成熟的库，如rbox_iou_ops（CUDA实现）或shapely（Python几何计算库，较慢但易于调试）。在训练初期，可以先用水平框损失预热，再引入旋转框损失。

方案二：高斯Wasserstein距离损失将旋转框建模为一个二维高斯分布，用两个分布的Wasserstein距离作为损失。这种方法能天然地处理旋转对称性问题（如正方形旋转90度后不变），理论更优美，但实现和理解起来更复杂。

对于初次改造，强烈建议采用方案一（角度分类），因为它更直观，易于调试，且与YOLOv5原有的分类分支结构相似。

3.2 语义分割分支的添加

我们需要在YOLOv5的架构上添加一个并行的分割解码器（Decoder）。一个经典且高效的做法是借鉴U-Net或FPN的思想，利用主干网络的不同阶段特征图进行融合。

具体改造步骤：

1. 特征提取：YOLOv5的Backbone（CSPDarknet）会产生三个不同尺度的特征图（P3, P4, P5），分别对应浅层细节和深层语义信息。
2. 分割头设计：新增一个分割分支。通常，我们将深层语义信息丰富的P5特征图上采样，然后与空间细节丰富的P4、P3特征图进行逐元素相加（Add）或通道拼接（Concat）。
3. 特征融合：可以设计一个轻量级的FPN结构。例如：
   • 对P5进行2倍上采样，与P4融合。
   • 对融合后的特征再进行2倍上采样，与P3融合。
   • 最终，对融合了多尺度信息的特征图进行一个简单的卷积层，将通道数映射到分割类别数num_seg_classes。
4. 输出：分割头的输出是一个[B, num_seg_classes, H, W]的张量，其中H, W是输入图像的下采样尺寸（如原图640x640下采样4倍或8倍）。训练时使用交叉熵损失或Dice损失（对类别不平衡问题更鲁棒）。
5. 代码修改点：
   • 在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 # 新的推理脚本，可同时输出检测框和分割图

关键修改点详解：

1. 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: 180

2. utils/datasets.py：在LoadImagesAndLabels类的__getitem__方法中，需要同时加载图像、旋转框标签文件和对应的分割掩码文件。返回的字典应包含img,rotated_labels,masks。

3. 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退火。多任务训练时，损失曲面可能更复杂，适当降低初始学习率（如lr0=1e-3降为3e-4）。
• 超参数：hyp.scratch.yaml中的超参数需要调整。特别是分类和角度分类损失的权重cls_pw,angle_pw，以及分割损失的权重seg_pw。建议从seg_pw=0.5开始尝试。
• 验证指标：需要同时监控：
  • 检测指标：mAP@0.5，mAP@0.5:0.95（需使用旋转框mAP计算脚本）。
  • 分割指标：mIoU（交并比），Pixel Accuracy。
  • 训练损失：观察各个损失项是否平稳下降，有无剧烈震荡。

4.3 推理与可视化

推理脚本需要同时运行两个头：

# 模型加载 model = attempt_load('best.pt', map_location=device) 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], dim=0).cpu().numpy() # [H, W] # 可视化 plot_rotated_boxes_on_image(original_img, rotated_boxes, angles) overlay_seg_mask(original_img, seg_mask, alpha=0.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中，“道路”和“天空”等类别像素占比极大。解决方案：
  1. 使用Dice Loss或Focal Loss替代标准的交叉熵损失，它们能更好地处理不平衡问题。
  2. 在损失函数中为每个类别设置不同的权重，权重与类别像素频率的倒数成正比。
  3. 在数据增强中，有针对性地对少数类别区域进行过采样或复制-粘贴。

5.3 检测与分割任务互相拖累

• 现象：同时训练时，一个任务指标上升，另一个任务指标下降。
• 解决方案：
  1. 检查共享主干：两个任务共享Backbone，可能特征存在冲突。可以尝试在Backbone的较高层（如C4或C5阶段）就将特征图分叉，让两个任务拥有相对独立的低层特征提取路径。
  2. 调整损失权重：使用前面提到的“不确定性加权”方法进行自动化调优。
  3. 课程学习：先训练一个任务较多轮次，再引入第二个任务进行微调。

5.4 模型速度严重下降

• 原因：分割头增加了大量计算，特别是上采样和特征融合操作。
• 优化方案：
  1. 轻量化分割头：减少分割头中的通道数。例如，将FPN融合层的通道数从256减至128或64。
  2. 降低分割输出分辨率：模型内部计算可以用1/8或1/16的下采样率，只在最终输出前上采样到1/4。推理时，分割图可以通过插值快速放大到原图尺寸。
  3. 知识蒸馏：训练一个大型的教师网络，然后用它来蒸馏一个小型的学生网络，在精度和速度间取得平衡。

这个项目就像搭积木，把YOLOv5、旋转框、语义分割这几个成熟的“积木块”以正确的方式组合起来。最大的收获不是最终模型的效果，而是在这个缝合过程中，你对数据流、损失函数、多任务平衡、调试技巧的深刻理解。我自己的经验是，先从最简单的部分跑通（比如只改旋转框），再加分割，每一步都做好验证和可视化，这样出了问题也能快速定位。WoodScape数据集的复杂性迫使你去思考更底层的问题，这比在标准数据集上刷分数要有价值得多。

本文还有配套的精品资源，点击获取