YOLOv8-multi-task

一次性实时通用多任务识别

本仓库(Yolov8多任务)是论文《一次性实时通用多任务识别》的官方PyTorch实现。

一次性实时通用多任务识别

作者: 王家园、吴乔明 ^:email:、张宁

(^:email:) 通讯作者

IEEE车辆技术汇刊

A-YOLOM示意图

贡献

• 我们开发了一个轻量级模型,能将三个任务整合到一个统一的模型中。这对需要实时处理的多任务特别有益。
• 我们专门为分割架构的颈部设计了一个新颖的自适应拼接模块。该模块可以自适应地拼接特征,无需手动设计,进一步增强了模型的通用性。
• 我们设计了一个轻量级、简单且通用的分割头。我们对同类型的任务头使用统一的损失函数,这意味着我们不需要为特定任务定制设计。它仅由一系列卷积层构建而成。
• 我们基于公开可用的自动驾驶数据集进行了广泛的实验,结果表明我们的模型在性能上优于现有工作,特别是在推理时间和可视化方面。此外,我们还使用真实道路数据集进行了进一步的实验,结果也表明我们的模型明显优于最先进的方法。

结果

参数和速度

交通目标检测结果

可行驶区域分割结果

车道线检测结果:

消融研究 1: 自适应拼接模块:

消融研究 2: 不同多任务模型和分割结构的结果:

YOLOv8(多任务)和YOLOM(n)仅显示两个分割头的总参数。它们实际上有三个头,我们忽略了检测头的参数,因为这是针对分割结构的消融研究。

注意:

• 我们参考的作品包括 Multinet（论文，代码），DLT-Net（论文），Faster R-CNN（论文，代码），YOLOv5s（代码），PSPNet（论文，代码），ENet（论文，代码），SCNN（论文，代码），SAD-ENet（论文，代码），YOLOP（论文，代码），HybridNets（论文，代码），YOLOv8（代码）。感谢他们出色的工作。

可视化

实际道路

要求

本代码库使用Python==3.7.16和PyTorch==1.13.1开发。

你可以使用一张1080Ti GPU，批量大小设为16。这样就足够了，只是训练时间会更长。我们推荐使用4090或更强大的GPU，这样会更快。

我们强烈建议你创建一个纯净的环境，并按照我们的说明来构建你的环境。否则，你可能会遇到一些问题，因为YOLOv8有许多机制会自动检测你的环境包。然后它会改变一些变量值，进而影响代码运行。

cd YOLOv8-multi-task
pip install -e .

数据准备和预训练模型

下载

• 从这里下载图像。

• 预训练模型：A-YOLOM # 包括两个版本，尺度"n"和"s"。

• 从这里下载检测标注。

• 从这里下载可驾驶区域分割标注。

• 从这里下载车道线分割标注。

我们建议数据集目录结构如下：

# id代表对应关系
├─数据集根目录
│ ├─images
│ │ ├─train2017
│ │ ├─val2017
│ ├─detection-object
│ │ ├─labels
│ │ │ ├─train2017
│ │ │ ├─val2017
│ ├─seg-drivable-10
│ │ ├─labels
│ │ │ ├─train2017
│ │ │ ├─val2017
│ ├─seg-lane-11
│ │ ├─labels
│ │ │ ├─train2017
│ │ │ ├─val2017

在./ultralytics/datasets/bdd-multi.yaml中更新你的数据集路径。

训练

你可以在./ultralytics/yolo/cfg/default.yaml中设置训练配置。

python train.py

你可以在train.py中更改设置

# 设置

sys.path.insert(0, "/home/jiayuan/ultralytics-main/ultralytics")
# 你应该将路径更改为你本地的"ultralytics"文件路径
model = YOLO('/home/jiayuan/ultralytics-main/ultralytics/models/v8/yolov8-bdd-v4-one-dropout-individual.yaml', task='multi')
# 你需要更改模型路径为你的路径。
# 模型文件保存在"./ultralytics/models/v8"下
model.train(data='/home/jiayuan/ultralytics-main/ultralytics/datasets/bdd-multi-toy.yaml', batch=4, epochs=300, imgsz=(640,640), device=[4], name='v4_640', val=True, task='multi',classes=[2,3,4,9,10,11],combine_class=[2,3,4,9],single_cls=True)

• data：请将"data"路径更改为你的路径。你可以在"./ultralytics/datasets"下找到它。

• device：如果你有多个GPU，请列出你的GPU编号，例如[0,1,2,3,4,5,6,7,8]

• name：你的项目名称，结果和训练好的模型将保存在"./ultralytics/runs/multi/你的项目名称"下

• task：如果你想使用多任务模型，请在这里保持"multi"

• classes：你可以更改此项来控制训练中的分类，10和11表示可驾驶区域和车道线分割。你可以在"./ultralytics/datasets/bdd-multi.yaml"下创建或更改数据集映射

• combine_class：表示模型将"classes"合并为一个类，例如我们的项目将"汽车"、"公交车"、"卡车"和"火车"合并为"车辆"。

• single_cls：这将把整个检测类别组合成一个类别。例如，如果你的数据集中有7个类别，当你使用"single_cls"时，它会自动将它们组合成一个类别。当你设置single_cls=False或从model.train()中删除single_cls时，请按照下面的注意事项更改dataset.yaml和model.yaml中的"tnc"，dataset.yaml中的"nc_list"，以及检测头的输出。

评估

你可以在./ultralytics/yolo/cfg/default.yaml中设置评估配置

python val.py

你可以在val.py中更改设置

# 设置

sys.path.insert(0, "/home/jiayuan/yolom/ultralytics")
# 与训练相同，你应该将路径更改为你的路径。

model = YOLO('/home/jiayuan/ultralytics-main/ultralytics/runs/best.pt')
# 请将此路径更改为你训练好的模型。你可以使用我们提供的预训练模型或你在"./ultralytics/runs/multi/Your Project Name/weight/best.pt"下的模型
metrics = model.val(data='/home/jiayuan/ultralytics-main/ultralytics/datasets/bdd-multi.yaml',device=[3],task='multi',name='val',iou=0.6,conf=0.001, imgsz=(640,640),classes=[2,3,4,9,10,11],combine_class=[2,3,4,9],single_cls=True)

• data：请将"data"路径更改为你的路径。你可以在"./ultralytics/datasets"下找到它
• device：如果你有多个GPU，请列出你的GPU编号，例如[0,1,2,3,4,5,6,7,8]。我们不建议你在验证时使用多GPU，因为通常一个GPU就足够了。
• speed：如果你想计算FPS，你应该设置"speed=True"。这个FPS计算方法参考自HybridNets(代码)
• single_cls：应该保持与训练时相同的布尔值。

预测

python predict.py

你可以在predict.py中更改设置

# 设置

sys.path.insert(0, "/home/jiayuan/ultralytics-main/ultralytics")
number = 3 #输入你工作中的任务数量，如果你有1个检测和3个分割任务，这里应该是4。
model = YOLO('/home/jiayuan/ultralytics-main/ultralytics/runs/best.pt')  
model.predict(source='/data/jiayuan/dash_camara_dataset/daytime', imgsz=(384,672), device=[3],name='v4_daytime', save=True, conf=0.25, iou=0.45, show_labels=False)
# 预测结果将保存在"runs"文件夹下

注意：如果你想使用我们提供的预训练模型，请确保你的输入图像大小为(720,1280)，并保持"imgsz=(384,672)"以达到最佳性能，你可以更改"imgsz"的值，但结果可能会不同，因为它与训练尺寸不同。

• source：你的输入或想要预测的图像文件夹。
• show_labels=False：关闭标签的显示。请记住，当你使用"single cell=True"的预训练模型时，标签默认会显示第一个类别名称。
• boxes=False：关闭分割任务的边界框。

注意

• 这个代码很容易扩展到任何多分割和检测任务，只需修改模型yaml和数据集yaml文件信息，并按照我们的标签格式创建你的数据集，请记住，你应该在检测任务名称中保留"det"，在分割任务名称中保留"seg"。然后代码就可以工作了。无需修改基本代码，我们已经在基本代码中完成了必要的工作。

• 请记住，当你更改检测任务的类别数量时，请更改dataset.yaml和model.yaml中的"tnc"。"tcn"表示总类别数，包括检测和分割。例如，如果你有7个检测类别，1个分割和另一个分割，"tnc"应该设置为9。

  • "nc_list"也需要更新，它应该与你的"labels_list"顺序匹配。例如，如果你的"labels_list"中有detection-object、seg-drivable、seg-lane，那么"nc_list"应该是[7,1,1]。这意味着你在detection-object中有7个类别，在drivable分割中有1个类别，在lane分割中有1个类别。

  • 你还需要更改检测头输出数量，这在model.yaml中，例如" - [[15, 18, 21], 1, Detect, [int number for detection class]] # 36 Detect(P3, P4, P5)"，请将"int number for detection class"更改为你的检测任务中的类别数量，按照上面的例子，这里应该是7。

• 如果你想更改一些基本代码来实现你的想法，请搜索"###### Jiayuan"或"######Jiayuan"，我们基于YOLOv8(代码)更改了这些部分，以在单个模型中实现多任务。

引用

如果你发现我们的论文和代码对你的研究有用，请考虑给予星标:star:和引用:pencil:：

@ARTICLE{wang2024you,
  author={Wang, Jiayuan and Wu, Q. M. Jonathan and Zhang, Ning},
  journal={IEEE Transactions on Vehicular Technology}, 
  title={You Only Look at Once for Real-Time and Generic Multi-Task}, 
  year={2024},
  pages={1-13},
  keywords={Multi-task learning;panoptic driving perception;object detection;drivable area segmentation;lane line segmentation},
  doi={10.1109/TVT.2024.3394350}}