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

文章目录

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


前言

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

一、数据准备

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

二、数据输入处修改代码

对标签label增加分类分支
基于YoloV5 multi task网络增加分类分支_第2张图片
基于YoloV5 multi task网络增加分类分支_第3张图片
基于YoloV5 multi task网络增加分类分支_第4张图片
基于YoloV5 multi task网络增加分类分支_第5张图片

三、网络修改

基于YoloV5 multi task网络增加分类分支_第6张图片
基于YoloV5 multi task网络增加分类分支_第7张图片

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

基于YoloV5 multi task网络增加分类分支_第8张图片
基于YoloV5 multi task网络增加分类分支_第9张图片

五、最终可视化结果

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

总结

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

你可能感兴趣的:(深度学习,计算机视觉,人工智能)