YOLOv5+BiSeNet——同时进行目标检测和语义分割

前言

在Gayhub上看到个项目，有人在YOLOv5的基础上，新增了一个分割头，把BiSeNet语义分割算法加入到了目标检测中，使其能够同时进行目标检测和语义分割。
项目地址：https://github.com/TomMao23/multiyolov5

效果预览

先看我使用原作者提供的模型，复刻出来的效果：
(本来想放视频的，不过传了两次CSDN都莫名其妙消失了，那就放动图了)

模型架构

目标检测模型采用的是YOLOv5，具体原理在我之前的博文【目标检测】从YOLOv1到YOLOX(理论梳理)里已经详细解读过。
语义分割模型采用的是部分BiSeNet结构，因为我不是这个方向的，具体原理不做细究，放张BiSeNet的结构图[1]：

核心代码

原作者目标检测使用的Coco数据集，语义分割使用的是Cityscapes数据集。
模型主要是在YOLOv5-5.0版本上进行修改的，基准模型采用的是YOLOv5m，语义分割的实现主要是在模型输出的Head部分添加了一个头：
yolov5m_city_seg.yaml

# parameters
nc: 10  # number of classes
n_segcls: 19 # 分割类别数
depth_multiple: 0.67  # model depth multiple
width_multiple: 0.75  # layer channel multiple

# anchors
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, C3, [1024, False]],  # 9
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4  # PANet是add, yolov5是concat
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
                  #[类别/输出通道, C3的n, C3的c2, C3的shortcut(以base为例，其他头含义可能不同)] yolo.py解析代码, []内第一项必须是输出通道数
   #[[4, 19], 1, SegMaskLab, [n_segcls, 3, 256, False]],  # 语义分割头通道配置256,[]内n为３
   [[16, 19, 22], 1, SegMaskPSP, [n_segcls, 3, 256, False]],  # 语义分割头通道配置256
   #[[16, 19, 22], 1, SegMaskBiSe, [n_segcls, 3, 256, False]],  # 语义分割头通道配置无效
   #[[16], 1, SegMaskBase, [n_segcls, 3, 512, False]],  # 语义分割头通道配置512

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)  必须在最后一层, 原代码很多默认了Detect是最后, 并没有全改
  ]

代码中，在最后的输出部分，作者添加了3个和Detect平行的分割头，其中，SegMaskLab、SegMaskPSP、SegMaskBiSe、SegMaskBase分别是不同的独立结构，是作者实验所用。
在yolo.py中，可以看到它们详细的结构：

class SegMaskPSP(nn.Module):  # PSP头，多了RFB2和FFM，同样砍了通道数，没找到合适的位置加辅助损失，因此放弃辅助损失
    def __init__(self, n_segcls=19, n=1, c_hid=256, shortcut=False, ch=()):  # n是C3的, (接口保留了,没有使用)c_hid是隐藏层输出通道数（注意配置文件s*0.5,m*0.75,l*1）
        super(SegMaskPSP, self).__init__()
        self.c_in8 = ch[0]  # 16  # 用16,19,22宁可在融合处加深耗费一些时间，检测会涨点分割也很好。严格的消融实验证明用17,20,23分割可能还会微涨，但检测会掉３个点以上，所有头如此
        self.c_in16 = ch[1]  # 19
        self.c_in32 = ch[2]  # 22
        # self.c_aux = ch[0]  # 辅助损失  找不到合适地方放辅助，放弃
        self.c_out = n_segcls
        # 注意配置文件通道写256,此时s模型c_hid＝128
        self.out = nn.Sequential(  # 实验表明引入较浅非线性不太强的层做分割会退化成检测的辅助(分割会相对低如72退到70,71，检测会明显升高)，PP前应加入非线性强一点的层并适当扩大感受野
                                RFB2(c_hid*3, c_hid, d=[2,3], map_reduce=6),  # 3*128//6=64　RFB2和RFB无关，仅仅是历史遗留命名(训完与训练模型效果不错就没有改名重训了)
                                PyramidPooling(c_hid, k=[1, 2, 3, 6]),  # 按原文1,2,3,6，PSP加全局更好，但是ASPP加了全局后出现边界破碎
                                FFM(c_hid*2, c_hid, k=3, is_cat=False),  # FFM改用k=3, 相应的砍掉部分通道降低计算量(原则就是差距大的融合哪怕砍通道第一层也最好用3*3卷积，FFM融合效果又比一般卷积好，除base头外其他头都遵循这种融合方式)
                                nn.Conv2d(c_hid, self.c_out, kernel_size=1, padding=0),
                                nn.Upsample(scale_factor=8, mode='bilinear', align_corners=True),
                               )
        self.m8 = nn.Sequential(
                                Conv(self.c_in8, c_hid, k=1),
        )
        self.m32 = nn.Sequential(
                                Conv(self.c_in32, c_hid, k=1),
                                nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True),
        )
        self.m16 = nn.Sequential(
                                Conv(self.c_in16, c_hid, k=1),
                                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
        )
        # self.aux = nn.Sequential(
        #                        Conv(self.c_aux, 256, 3),  
        #                        nn.Dropout(0.1, False), 
        #                        nn.Conv2d(256, self.c_out, kernel_size=1),
        #                        nn.Upsample(scale_factor=8, mode='bilinear', align_corners=True),
        # )
    def forward(self, x):
        # 这个头三层融合输入做过消融实验，单独16:72.6三层融合:73.5,建议所有用1/8的头都采用三层融合，在Lab的实验显示三层融合的1/16输入也有增长
        feat = torch.cat([self.m8(x[0]), self.m16(x[1]), self.m32(x[2])], 1)
        # return self.out(feat) if not self.training else [self.out(feat), self.aux(x[0])]
        return self.out(feat)

下面是模型检测(detect.py)中的主要改动，在模型输出部分使用seg来获取语义分割结果，再利用提前定义好的颜色图Cityscapes_COLORMAP分别给分割部分上色。

seg = F.interpolate(seg, (im0.shape[0], im0.shape[1]), mode='bilinear', align_corners=True)[0]
mask = label2image(seg.max(axis=0)[1].cpu().numpy(), Cityscapes_COLORMAP)[:, :, ::-1]
dst = cv2.addWeighted(mask, 0.4, im0, 0.6, 0)

代码备份

其它改动还很多，可以去原作者的仓库阅读。
这里将其代码进行备份，包含作者提供的模型权重：
https://pan.baidu.com/s/1JtqCtlJwk5efkiTQqmNpVA?pwd=36bk

References

[1]https://blog.csdn.net/qq_40073354/article/details/120725919