爆改YOLOv8|利用图像分割网络UNetV2改进yolov8主干-即插即用

1，本文介绍

U-Net v2 通过引入创新的跳跃连接设计来提升医学图像分割的精度。这一版本专注于更有效地融合不同层级的特征，包括高级特征中的语义信息和低级特征中的细节信息。通过这种优化，U-Net v2 能够在低级特征中注入丰富的语义，同时细化高级特征，从而实现更精准的对象边界描绘和小结构提取。

其主要技术创新包括：

• 多级特征提取：使用深度神经网络编码器从输入图像中提取不同层次的特征。
• 语义与细节融合(Semantics and Detail Infusion, SDI)模块：通过哈达玛积操作，将高级特征中的语义信息与低级特征中的细节信息融合到各层级的特征图中。
• 改进的跳跃连接：这些新型跳跃连接增强了各层特征的语义和细节表现，从而在解码器阶段实现更高精度的分割。

关于UNetV2的详细介绍可以看论文：https://arxiv.org/abs/2311.17791

本文将讲解如何将UNetV2融合进yolov8

话不多说，上代码！

2， 将UNetV2融合进yolov8

2.1 步骤一

找到如下的目录'ultralytics/nn/modules'，然后在这个目录下创建一个UNetV2.py文件，文件名字可以根据你自己的习惯起，然后将UNetV2的核心代码复制进去

import os.path
import warnings
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
import math
 
__all__ = ['pvt_v2_b0', 'pvt_v2_b1', 'pvt_v2_b2', 'pvt_v2_b3', 'pvt_v2_b4', 'pvt_v2_b5']
 
class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
 
        self.fc1 = nn.Conv2d(in_planes, in_planes // 16, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // 16, in_planes, 1, bias=False)
 
        self.sigmoid = nn.Sigmoid()
 
    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)
 
 
class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
 
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
 
        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()
 
    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv1(x)
        return self.sigmoid(x)
 
 
class BasicConv2d(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1):
        super(BasicConv2d, self).__init__()
 
        self.conv = nn.Conv2d(in_planes, out_planes,
                              kernel_size=kernel_size, stride=stride,
                              padding=padding, dilation=dilation, bias=False)
        self.bn = nn.BatchNorm2d(out_planes)
        self.relu = nn.ReLU(inplace=True)
 
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x
 
 
class Encoder(nn.Module):
    def __init__(self, pretrain_path):
        super().__init__()
        self.backbone = pvt_v2_b2()
 
        if pretrain_path is None:
            warnings.warn('please provide the pretrained pvt model. Not using pretrained model.')
        elif not os.path.isfile(pretrain_path):
            warnings.warn(f'path: {pretrain_path} does not exists. Not using pretrained model.')
        else:
            print(f"using pretrained file: {pretrain_path}")
            save_model = torch.load(pretrain_path)
            model_dict = self.backbone.state_dict()
            state_dict = {k: v for k, v in save_model.items() if k in model_dict.keys()}
            model_dict.update(state_dict)
 
            self.backbone.load_state_dict(model_dict)
 
    def forward(self, x):
        f1, f2, f3, f4 = self.backbone(x)  # (x: 3, 352, 352)
        return f1, f2, f3, f4
 
 
class SDI(nn.Module):
    def __init__(self, channel):
        super().__init__()
 
        self.convs = nn.ModuleList(
            [nn.Conv2d(channel, channel, kernel_size=3, stride=1, padding=1) for _ in range(4)])
 
    def forward(self, xs, anchor):
        ans = torch.ones_like(anchor)
        target_size = anchor.shape[-1]
 
        for i, x in enumerate(xs):
            if x.shape[-1] > target_size:
                x = F.adaptive_avg_pool2d(x, (target_size, target_size))
            elif x.shape[-1] < target_size:
                x = F.interpolate(x, size=(target_size, target_size),
                                      mode='bilinear', align_corners=True)
 
            ans = ans * self.convs[i](x)
 
        return ans
 
 
class UNetV2(nn.Module):
    """
    use SpatialAtt + ChannelAtt
    """
    def __init__(self, channel=3, n_classes=1, deep_supervision=True, pretrained_path=None):
        super().__init__()
        self.deep_supervision = deep_supervision
 
        self.encoder = Encoder(pretrained_path)
 
        self.ca_1 = ChannelAttention(64)
        self.sa_1 = SpatialAttention()
 
        self.ca_2 = ChannelAttention(128)
        self.sa_2 = SpatialAttention()
 
        self.ca_3 = ChannelAttention(320)
        self.sa_3 = SpatialAttention()
 
        self.ca_4 = ChannelAttention(512)
        self.sa_4 = SpatialAttention()
 
        self.Translayer_1 = BasicConv2d(64, channel, 1)
        self.Translayer_2 = BasicConv2d(128, channel, 1)
        self.Translayer_3 = BasicConv2d(320, channel, 1)
        self.Translayer_4 = BasicConv2d(512, channel, 1)
 
        self.sdi_1 = SDI(channel)
        self.sdi_2 = SDI(channel)
        self.sdi_3 = SDI(channel)
        self.sdi_4 = SDI(channel)
 
        self.seg_outs = nn.ModuleList([
            nn.Conv2d(channel, n_classes, 1, 1) for _ in range(4)])
 
        self.deconv2 = nn.ConvTranspose2d(channel, channel, kernel_size=4, stride=2, padding=1,
                                          bias=False)
        self.deconv3 = nn.ConvTranspose2d(channel, channel, kernel_size=4, stride=2,
                                          padding=1, bias=False)
        self.deconv4 = nn.ConvTranspose2d(channel, channel, kernel_size=4,