电脑基础 · 2023年3月23日

基于YoloV5 multi task网络增加分类分支

文章目录

  • 前言
  • 一、数据准备
  • 二、数据输入处修改代码
  • 三、网络修改
  • 四、模型输出的和标签接上去,计算loss
  • 五、最终可视化结果
  • 总结

前言

YoloV5做目标检测任务,但是当我们需要对图像中的目标进行检测,还需要根据图像整体对该图像做分类时,我们可以在Yolo-Detect检测头前加一个分类分支,这样可以在共享Backbone情况下实现做检测任务,又做分类任务,下面具体介绍增加分类分支的操作方法

一、数据准备

如下图,yolov5的标签 .txt 文件, 只需在第一行新增类别标签(0,1,2…),第二行开始是目标框bbox标签
基于YoloV5 multi task网络增加分类分支

二、数据输入处修改代码

对标签label增加分类分支
基于YoloV5 multi task网络增加分类分支
基于YoloV5 multi task网络增加分类分支
基于YoloV5 multi task网络增加分类分支
基于YoloV5 multi task网络增加分类分支

三、网络修改

基于YoloV5 multi task网络增加分类分支
基于YoloV5 multi task网络增加分类分支

四、模型输出的和标签接上去,计算loss

基于YoloV5 multi task网络增加分类分支
基于YoloV5 multi task网络增加分类分支

五、最终可视化结果

原版yolo输出3个检测头,增加后多输出一个分类分支
基于YoloV5 multi task网络增加分类分支

总结

简单记录一下,因为yolo版本也在不断更新,说一没有详细说明,简单对主要修改处标记了,如有类似修改需求。大致参照思路即可。不能全抄。