Upload 2 files

model/CyueNet_models.py

@@ -0,0 +1,696 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import einops
+
+from timm.models.layers import trunc_normal_
+
+from einops import rearrange
+import math
+
+from model.MobileNetV2 import mobilenet_v2
+from torch.nn import Parameter
+
+
+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)
+        x = self.relu(x)
+        return x
+
+
+class Reduction(nn.Module):
+    def __init__(self, in_channel, out_channel):
+        super(Reduction, self).__init__()
+        self.reduce = nn.Sequential(
+            BasicConv2d(in_channel, out_channel, 1),
+            BasicConv2d(out_channel, out_channel, 3, padding=1),
+            BasicConv2d(out_channel, out_channel, 3, padding=1)
+        )
+
+    def forward(self, x):
+        return self.reduce(x)
+
+
+class TopDownLayer(nn.Module):
+    def __init__(self, channel):
+        super(TopDownLayer, self).__init__()
+        self.conv = nn.Sequential(nn.Conv2d(channel, channel, 3, 1, 1, bias=False), nn.BatchNorm2d(channel))
+
+        self.relu = nn.ReLU()
+
+        self.channel_compress = nn.Sequential(
+            nn.Conv2d(channel * 2, channel, 1, bias=False),
+            nn.BatchNorm2d(channel),
+            nn.ReLU()
+        )
+
+    def forward(self, x, x2):
+        res1 = self.conv(x)
+        res1 = self.relu(res1)
+
+        res1 = F.interpolate(res1, x2.size()[2:], mode='bilinear', align_corners=True)
+
+        res_cat = torch.cat((res1, x2), dim=1)
+
+        resl = self.channel_compress(res_cat)
+
+        return resl
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, head=8, d_model=32, dropout=0.1):
+        super(MultiHeadAttention, self).__init__()
+        assert (d_model % head == 0)
+        self.d_k = d_model // head
+        self.head = head
+        self.d_model = d_model
+        self.linear_query = nn.Linear(d_model, d_model)
+        self.linear_key = nn.Linear(d_model, d_model)
+        self.linear_value = nn.Linear(d_model, d_model)
+
+        self.dropout = nn.Dropout(p=dropout)
+        self.attn = None
+        self.inb = nn.Linear(32, d_model)
+
+    def self_attention(self, query, key, value, mask=None):
+        d_k = query.shape[-1]
+        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
+        self_attn = F.softmax(scores, dim=-1)
+        # self.attn = self_attn if self.attn is None else self.attn + self_attn
+        if self.dropout is not None:
+            self_attn = self.dropout(self_attn)
+
+        return torch.matmul(self_attn, value), self_attn
+
+    def forward(self, query, key, value, mask=None):
+        n_batch = query.size(0)
+        query = query.flatten(start_dim=2).permute(0, 2, 1)
+
+        query = self.inb(query)
+
+        key = key.flatten(start_dim=2).permute(0, 2, 1)
+
+        key = self.inb(key)
+
+        value = value.flatten(start_dim=2).permute(0, 2, 1)
+
+        value = self.inb(value)
+
+        x, self.attn = self.self_attention(query, key, value, mask=mask)
+
+        x = x.permute(0, 2, 1)
+        embedding_dim = x.size(-1)
+
+        d_k = h = int(embedding_dim ** 0.5)
+
+        x = einops.rearrange(x, 'b n (d_k h) -> b n d_k h', d_k=d_k, h=h)
+
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self):
+        super(Upsample, self).__init__()
+
+    def forward(self, x, x2):
+        x = F.interpolate(x, size=x2.size()[2:], mode='bilinear', align_corners=True)
+
+        return x
+
+
+class MultiScaleAttention(nn.Module):
+    def __init__(self, channel):
+        super(MultiScaleAttention, self).__init__()
+        # SPatial attention for each branch
+        self.attention_branches = nn.ModuleList([SpatialAttention() for _ in range(5)])
+        self.upsample = Upsample()
+        self.conv_reduce = nn.Conv2d(channel * 6, channel, kernel_size=1)
+
+    def forward(self, x0, x1, x2, x3, x4, x5):
+        x0_att = self.attention_branches[0](x0) * x0
+        x1_att = self.attention_branches[0](x1) * x1
+        x2_att = self.attention_branches[0](x2) * x2
+        x3_att = self.attention_branches[0](x3) * x3
+        x4_att = self.attention_branches[0](x4) * x4
+        x5_att = self.attention_branches[0](x5) * x5
+
+        x1_att_up = self.upsample(x1_att, x0)
+        x2_att_up = self.upsample(x2_att, x0)
+        x3_att_up = self.upsample(x3_att, x0)
+        x4_att_up = self.upsample(x4_att, x0)
+        
+        x5_att_up = self.upsample(x5_att, x0)
+
+        x_cat = torch.cat((x0_att, x1_att_up, x2_att_up, x3_att_up, x4_att_up, x5_att_up), dim=1)
+
+        x_out = self.conv_reduce(x_cat)
+
+        return x_out
+
+
+class Basic2(nn.Module):
+    def __init__(self, in_channel, out_channel):
+        super(Basic2, self).__init__()
+        self.relu = nn.ReLU(True)
+        # join
+        self.channel_attention = ChannelAttention(out_channel)
+        self.channel_attention = SpatialAttention()
+        self.branch0 = nn.Sequential(
+            BasicConv2d(in_channel, out_channel, 1),
+        )
+        self.branch1 = nn.Sequential(
+            BasicConv2d(in_channel, out_channel, kernel_size=(1, 3), padding=(0, 1)),
+            BasicConv2d(out_channel, out_channel, kernel_size=(3, 1), padding=(1, 0)),
+            BasicConv2d(out_channel, out_channel, 3, padding=3, dilation=3)
+        )
+        self.branch2 = nn.Sequential(
+            BasicConv2d(in_channel, out_channel, kernel_size=(1, 5), padding=(0, 2)),
+            BasicConv2d(out_channel, out_channel, kernel_size=(5, 1), padding=(2, 0)),
+            BasicConv2d(out_channel, out_channel, 3, padding=5, dilation=5)
+        )
+        self.branch3 = nn.Sequential(
+            BasicConv2d(in_channel, out_channel, kernel_size=(1, 7), padding=(0, 3)),
+            BasicConv2d(out_channel, out_channel, kernel_size=(7, 1), padding=(3, 0)),
+            BasicConv2d(out_channel, out_channel, 3, padding=7, dilation=7)
+        )
+        self.branch4 = nn.Sequential(
+            BasicConv2d(in_channel, out_channel, kernel_size=(1, 9), padding=(0, 4)),
+            BasicConv2d(out_channel, out_channel, kernel_size=(9, 1), padding=(4, 0)),
+            BasicConv2d(out_channel, out_channel, 3, padding=9, dilation=9)
+        )
+        
+        self.branch5 = nn.Sequential(
+            BasicConv2d(in_channel, out_channel, kernel_size=(1, 11), padding=(0, 5)),
+            BasicConv2d(out_channel, out_channel, kernel_size=(11, 1), padding=(5, 0)),
+            BasicConv2d(out_channel, out_channel, 3, padding=11, dilation=11)
+        )
+
+        self.multi_scale_attention = MultiScaleAttention(out_channel)
+        self.conv_combine = BasicConv2d(in_channel, in_channel, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        x3 = self.branch3(x)
+        x4 = self.branch4(x)
+        x5 = self.branch5(x)
+
+        x_att = self.multi_scale_attention(x0, x1, x2, x3, x4, x5)
+
+        x_combind = self.conv_combine(x_att)
+
+        x = x_combind + x
+        return x
+
+
+class ChannelAttention(nn.Module):
+    def __init__(self, in_planes):
+        super(ChannelAttention, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc = nn.Sequential(nn.Conv2d(in_planes, in_planes // 2, 1, bias=False),
+                                nn.ReLU(),
+                                nn.Conv2d(in_planes // 2, in_planes, 1, bias=False))
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.fc(self.avg_pool(x))
+        max_out = self.fc(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__()
+
+        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, 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)
+        x1 = torch.cat([avg_out, max_out], dim=1)
+        x2 = self.conv1(x1)
+        return self.sigmoid(x2)
+
+
+class MModule(nn.Module):
+    def __init__(self, channel):
+        super(MModule, self).__init__()
+
+        self.basic = Basic2(channel, channel)
+        self.SA = SpatialAttention()
+        self.CA = ChannelAttention(channel)
+
+    def forward(self, x):
+        x_mix = self.basic(x)
+        x_mix = x_mix * self.CA(x_mix) + x_mix
+        x_mix1 = x_mix * self.SA(x_mix) + x_mix
+
+        x_mix1 = x_mix1 + x
+
+        return x_mix1
+
+
+class MNodule(nn.Module):
+    def __init__(self, channel):
+        super(MNodule, self).__init__()
+        self.atrconv1 = BasicConv2d(channel, channel, 3, padding=3, dilation=3)
+        self.atrconv2 = BasicConv2d(channel, channel, 3, padding=5, dilation=5)
+        self.atrconv3 = BasicConv2d(channel, channel, 3, padding=7, dilation=7)
+        self.branch1 = nn.Sequential(
+            BasicConv2d(channel, channel, 1),
+            BasicConv2d(channel, channel, kernel_size=(1, 3), padding=(0, 1)),
+            BasicConv2d(channel, channel, kernel_size=(3, 1), padding=(1, 0))
+        )
+        self.branch2 = nn.Sequential(
+            BasicConv2d(channel, channel, 1),
+            BasicConv2d(channel, channel, kernel_size=(1, 5), padding=(0, 2)),
+            BasicConv2d(channel, channel, kernel_size=(5, 1), padding=(2, 0))
+        )
+        self.branch3 = nn.Sequential(
+            BasicConv2d(channel, channel, 1),
+            BasicConv2d(channel, channel, kernel_size=(1, 7), padding=(0, 3)),
+            BasicConv2d(channel, channel, kernel_size=(7, 1), padding=(3, 0))
+        )
+
+        self.conv_cat1 = BasicConv2d(2 * channel, channel, 3, padding=1)
+        self.conv_cat2 = BasicConv2d(2 * channel, channel, 3, padding=1)
+        self.conv_cat3 = BasicConv2d(2 * channel, channel, 3, padding=1)
+        self.conv1_1 = BasicConv2d(channel, channel, 1)
+
+        self.SA = SpatialAttention()
+        self.CA = ChannelAttention(channel)
+
+        self.sal_conv = nn.Sequential(
+            BasicConv2d(channel, channel, 3, padding=1),
+            BasicConv2d(channel, channel, 3, padding=1)
+        )
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        x1 = self.branch1(x)
+        x_atr1 = self.atrconv1(x)
+        s_mfeb1 = self.conv_cat1(torch.cat((x1, x_atr1), 1)) + x
+        x2 = self.branch2(s_mfeb1)
+        x_atr2 = self.atrconv2(s_mfeb1)
+        s_mfeb2 = self.conv_cat2(torch.cat((x2, x_atr2), 1)) + s_mfeb1 + x
+        x3 = self.branch3(s_mfeb2)
+        x_atr3 = self.atrconv3(s_mfeb2)
+        s_mfeb3 = self.conv_cat3(torch.cat((x3, x_atr3), 1)) + s_mfeb1 + s_mfeb2 + x
+        x_m = self.conv1_1(s_mfeb3)
+
+        x_ca = self.CA(x_m) * x_m
+        x_e = self.CA(x_m) * x_m
+
+        x_mix = self.sal_conv((self.SA(x_ca)) * x_ca) + s_mfeb1 + s_mfeb2 + s_mfeb3 + x
+
+        return x_mix
+
+
+class TransBasicConv2d(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size=2, stride=2, padding=0, dilation=1, bias=False):
+        super(TransBasicConv2d, self).__init__()
+        self.Deconv = nn.ConvTranspose2d(in_planes, out_planes,
+                                         kernel_size=kernel_size, stride=stride,
+                                         padding=padding, dilation=dilation, bias=bias)
+        self.bn = nn.BatchNorm2d(out_planes)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.Deconv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class features(nn.Module):
+    def __init__(self, channel):
+        super(features, self).__init__()
+        self.conv1 = BasicConv2d(channel, channel, 1)
+        self.conv2 = BasicConv2d(channel, channel, 1)
+        self.conv3 = BasicConv2d(channel, channel, 1)
+        self.conv4 = BasicConv2d(channel, channel, 1)
+        self.conv5 = BasicConv2d(channel, channel, 1)
+
+    def forward(self, x1, x2, x3, x4, x5):
+        x1 = self.conv1(x1)
+        x2 = self.conv2(x2)
+        x3 = self.conv3(x3)
+        x4 = self.conv4(x4)
+        x5 = self.conv5(x5)
+
+        return x1, x2, x3, x4, x5
+
+
+class conv_upsamle(nn.Module):
+    def __init__(self, channel):
+        super(conv_upsamle, self).__init__()
+        self.conv = BasicConv2d(channel, channel, 3, padding=1)
+
+    def forward(self, x, target):
+        if x.size()[2:] != target.size()[2:]:
+            x = F.interpolate(x, size=target.size()[2:], mode='bilinear', align_corners=True)
+        x = self.conv(x)
+        return x
+
+
+class AP_MP(nn.Module):
+    def __init__(self, stride=2):
+        super(AP_MP, self).__init__()
+        self.sz = stride
+        self.gapLayer = nn.AvgPool2d(kernel_size=self.sz, stride=self.sz)
+        self.gmpLayer = nn.MaxPool2d(kernel_size=self.sz, stride=self.sz)
+
+    def forward(self, x1, x2):
+        B, C, H, W = x1.size()
+        apimg = self.gapLayer(x1)
+        mpimg = self.gmpLayer(x2)
+        byimg = torch.norm(abs(apimg - mpimg), p=2, dim=1, keepdim=True)
+        return byimg
+
+class MOM(nn.Module):
+    def __init__(self, channel):
+        super(MOM, self).__init__()
+        self.channel = channel
+
+        self.conv1 = BasicConv2d(channel, channel, 3, padding=1)
+        self.conv2 = BasicConv2d(channel, channel, 3, padding=1)
+
+        self.CA1 = ChannelAttention(self.channel)
+        self.CA2 = ChannelAttention(self.channel)
+        self.SA1 = SpatialAttention()
+        self.SA2 = SpatialAttention()
+
+        self.glbamp = AP_MP()
+        self.upsample2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.conv = BasicConv2d(channel * 2 , channel, kernel_size=1, stride=1)
+
+        self.upSA = SpatialAttention()
+
+    def forward(self, x1, x2):
+        x1 = self.conv1(x1)
+        x2 = self.conv2(x2)
+
+        x1 = x1 + x1 * self.CA1(x1)
+        x2 = x2 + x2 * self.CA2(x2)
+
+        nx1 = x1 + x1 * self.SA2(x2)
+        nx2 = x2 + x2 * self.SA1(x1)
+
+        res = self.conv(torch.cat([nx1, nx2], dim=1))
+
+        res = res + x1
+        edg = res
+        ske = res
+
+        return res, edg, ske
+
+
+class AFM(nn.Module):
+    def __init__(self, channel):
+        super(AFM, self).__init__()
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+        self.sigmoid = nn.Sigmoid()
+        self.conv1_1 = nn.Conv2d(channel, channel, kernel_size=1)
+
+        self.ca1 = ChannelAttention(channel)
+        self.ca2 = ChannelAttention(channel)
+
+        self.sa = SpatialAttention()
+
+        self.sal_conv = nn.Sequential(
+            BasicConv2d(channel, channel, 3, padding=1),
+            BasicConv2d(channel, channel, 3, padding=1)
+        )
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x1, x2):
+        x2 = self.sigmoid(self.max_pool(x2))
+        xb = x2 * x1
+        x = self.conv1_1(xb)
+        x_c = self.ca1(x) * x
+        x_d = self.ca2(x) * x
+
+        s_mea = self.sal_conv((self.sa(x_c)) * x_c) + x1 + x2 + xb
+        ske = s_mea
+        e_pred = s_mea
+        return s_mea, e_pred, ske
+    
+class DummyMOM(nn.Module):
+    def __init__(self, channel):
+        super(DummyMOM, self).__init__()
+        self.conv1 = nn.Identity()  # 保持输入输出一致
+        self.conv2 = nn.Identity()  # 保持输入输出一致
+        
+        # 调整为64个输入通道
+        self.conv = nn.Conv2d(64, 32, kernel_size=1)  # 1x1卷积调整通道数
+
+    def forward(self, x1, x2):
+        # 先做拼接，然后调整通道数为32
+        res = self.conv(torch.cat([x1, x2], dim=1))
+        edg = res
+        ske = res
+        return res, edg, ske
+    
+class YUEM(nn.Module):
+    def __init__(self, channel):
+        super(YUEM, self).__init__()
+        self.channel = channel
+        self.m1 = MModule(self.channel)
+        self.m2 = MNodule(self.channel)
+        self.mha = MultiHeadAttention(channel)
+
+    def forward(self, x1, x2):
+        x1 = self.m1(x1)
+        x21 = self.m2(x2)
+
+        res = self.mha(x1, x21, x2)
+
+        edg = res
+        ske = res
+        return res, edg, ske
+
+
+class MTG(nn.Module):
+    def __init__(self, channel):
+        super(MTG, self).__init__()
+        self.ccs = nn.ModuleList([nn.Sequential(
+            BasicConv2d(3 * channel, channel, kernel_size=3, padding=1),
+            BasicConv2d(channel, channel, kernel_size=3, padding=1)
+        ) for i in range(5)])
+
+    def forward(self, x_sal, x_edg, x_ske):
+
+        x_combined = torch.cat((x_sal, x_edg,x_ske), dim=1)
+
+        x_sal_n = self.ccs[0](x_combined)
+
+        return x_sal_n
+class MMS(nn.Module):
+    def __init__(self, pretrained=True, channel=32):
+        super(MMS, self).__init__()
+        self.backbone = mobilenet_v2(pretrained)
+
+        self.Translayer1 = Reduction(16, channel)
+        self.Translayer2 = Reduction(24, channel)
+        self.Translayer3 = Reduction(32, channel)
+        self.Translayer4 = Reduction(96, channel)
+        self.Translayer5 = Reduction(320, channel)
+
+        self.trans_conv1 = TransBasicConv2d(channel, channel, kernel_size=2, stride=2,
+                                            padding=0, dilation=1, bias=False)
+        self.trans_conv2 = TransBasicConv2d(channel, channel, kernel_size=2, stride=2,
+                                            padding=0, dilation=1, bias=False)
+        self.trans_conv3 = TransBasicConv2d(channel, channel, kernel_size=2, stride=2,
+                                            padding=0, dilation=1, bias=False)
+        self.trans_conv4 = TransBasicConv2d(channel, channel, kernel_size=2, stride=2,
+                                            padding=0, dilation=1, bias=False)
+
+        self.mom = MOM(channel)
+        # self.mom = DummyMOM(channel)
+        self.afm = AFM(channel)
+        # self.afm = DummyMOM(channel)
+        self.yuem = YUEM(channel)
+        # self.yuem = DummyMOM(channel)
+
+        self.sigmoid = nn.Sigmoid()
+
+        self.sal_features = features(channel)
+        self.edg_features = features(channel)
+        self.ske_features = features(channel)
+        self.MTG = MTG(channel)
+
+        self.ccs = nn.ModuleList([nn.Sequential(
+            BasicConv2d(3 * channel, channel, kernel_size=3, padding=1),
+            BasicConv2d(channel, channel, kernel_size=3, padding=1)
+        ) for i in range(5)])
+        self.cme = nn.ModuleList([nn.Sequential(
+            BasicConv2d(3 * channel, channel, kernel_size=3, padding=1),
+            BasicConv2d(channel, channel, kernel_size=3, padding=1)
+        ) for i in range(5)])
+        self.cms = nn.ModuleList([nn.Sequential(
+            BasicConv2d(3 * channel, channel, kernel_size=3, padding=1),
+            BasicConv2d(channel, channel, kernel_size=3, padding=1)
+        ) for i in range(5)])
+
+        self.conv_cats = nn.ModuleList([nn.Sequential(
+            BasicConv2d(2 * channel, channel, kernel_size=3, padding=1),
+            BasicConv2d(channel, channel, kernel_size=3, padding=1)
+        ) for i in range(12)])
+
+
+        self.cus = nn.ModuleList([conv_upsamle(channel) for i in range(12)])
+        self.prediction = nn.ModuleList([
+            nn.Sequential(
+                BasicConv2d(channel, channel, kernel_size=3, padding=1),
+                nn.Conv2d(channel, 1, kernel_size=1)
+            ) for i in range(3)
+        ])
+        
+        self.S1 = nn.Sequential(
+            BasicConv2d(channel, channel, 3, padding=1),
+            nn.Conv2d(channel, 1, 1)
+        )
+        self.S2 = nn.Sequential(
+            BasicConv2d(channel, channel, 3, padding=1),
+            nn.Conv2d(channel, 1, 1)
+        )
+        self.S3 = nn.Sequential(
+            BasicConv2d(channel, channel, 3, padding=1),
+            nn.Conv2d(channel, 1, 1)
+        )
+        self.S4 = nn.Sequential(
+            BasicConv2d(channel, channel, 3, padding=1),
+            nn.Conv2d(channel, 1, 1)
+        )
+        self.S5 = nn.Sequential(
+            BasicConv2d(channel, channel, 3, padding=1),
+            nn.Conv2d(channel, 1, 1)
+        )
+
+
+    def forward(self, x):
+        size = x.size()[2:]
+        conv1, conv2, conv3, conv4, conv5 = self.backbone(x)
+
+        conv1 = self.Translayer1(conv1)
+        conv2 = self.Translayer2(conv2)
+        conv3 = self.Translayer3(conv3)
+        conv4 = self.Translayer4(conv4)
+        conv5 = self.Translayer5(conv5)
+
+        rgc5, edg5, ske5 = self.afm(conv5, conv5)
+        rgc4, edg4, ske4 = self.yuem(conv4, self.trans_conv4(conv5))
+        rgc3, edg3, ske3 = self.yuem(conv3, self.trans_conv3(conv4))
+        rgc2, edg2, ske2 = self.mom(conv2, self.trans_conv2(conv3))
+        rgc1, edg1, ske1 = self.mom(conv1, self.trans_conv1(conv2))
+
+
+        x_sal1, x_sal2, x_sal3, x_sal4, x_sal5 = self.sal_features(rgc1, rgc2, rgc3, rgc4, rgc5)
+        x_edg1, x_edg2, x_edg3, x_edg4, x_edg5 = self.edg_features(edg1, edg2, edg3, edg4, edg5)
+        x_ske1, x_ske2, x_ske3, x_ske4, x_ske5 = self.ske_features(ske1, ske2, ske3, ske4, ske5)
+
+
+        x_sal5_n = self.ccs[0](torch.cat((x_sal5, x_edg5, x_sal5), 1)) + x_sal5
+        x_edg5_n = self.cme[0](torch.cat((x_sal5, x_edg5, x_sal5), 1)) + x_edg5
+        x_ske5_n = self.cms[0](torch.cat((x_sal5, x_edg5, x_ske5), 1)) + x_ske5
+
+
+        x_sal4 = self.conv_cats[0](torch.cat((x_sal4, self.cus[0](x_sal5_n, x_sal4)), 1))
+        x_edg4 = self.conv_cats[1](torch.cat((x_edg4, self.cus[1](x_edg5_n, x_edg4)), 1))
+        x_ske4 = self.conv_cats[2](torch.cat((x_ske4, self.cus[2](x_ske5_n, x_ske4)), 1))
+
+        x_sal4_n = self.MTG(x_sal4, x_edg4, x_ske4) + x_sal4
+        x_edg4_n = self.MTG(x_sal4, x_edg4, x_ske4) + x_edg4
+        x_ske4_n = self.MTG(x_sal4, x_edg4, x_ske4) + x_ske4
+
+
+        x_sal3 = self.conv_cats[3](torch.cat((x_sal3, self.cus[3](x_sal4_n, x_sal3)), 1))
+        x_edg3 = self.conv_cats[4](torch.cat((x_edg3, self.cus[4](x_edg4_n, x_edg3)), 1))
+        x_ske3 = self.conv_cats[5](torch.cat((x_ske3, self.cus[5](x_ske4_n, x_ske3)), 1))
+
+
+        x_sal3_n = self.MTG(x_sal3, x_edg3, x_ske3) + x_sal3
+        x_edg3_n = self.MTG(x_sal3, x_edg3, x_ske3) + x_edg3
+        x_ske3_n = self.MTG(x_sal3, x_edg3, x_ske3) + x_ske3
+
+
+
+        x_sal2 = self.conv_cats[6](torch.cat((x_sal2, self.cus[6](x_sal3_n, x_sal2)), 1))
+        x_edg2 = self.conv_cats[7](torch.cat((x_edg2, self.cus[7](x_edg3_n, x_edg2)), 1))
+        x_ske2 = self.conv_cats[8](torch.cat((x_ske2, self.cus[8](x_ske3_n, x_ske2)), 1))
+
+
+        x_sal2_n = self.MTG(x_sal2, x_edg2, x_ske2) + x_sal2
+        x_edg2_n = self.MTG(x_sal2, x_edg2, x_ske2) + x_edg2
+        x_ske2_n = self.MTG(x_sal2, x_edg2, x_ske2) + x_ske2
+
+        x_sal1 = self.conv_cats[9](torch.cat((x_sal1, self.cus[9](x_sal2_n, x_sal1)), 1))
+        x_edg1 = self.conv_cats[10](torch.cat((x_edg1, self.cus[10](x_edg2_n, x_edg1)), 1))
+        x_ske1 = self.conv_cats[11](torch.cat((x_ske1, self.cus[11](x_ske2_n, x_ske1)), 1))
+
+
+        x_sal1_n = self.MTG(x_sal1, x_edg1, x_ske1) + x_sal1
+        x_edg1_n = self.MTG(x_sal1, x_edg1, x_ske1) + x_edg1
+        x_ske1_n = self.MTG(x_sal1, x_edg1, x_ske1) + x_ske1
+
+        sal_out = self.prediction[0](x_sal1_n)
+        edg_out = self.prediction[1](x_edg1_n)
+        ske_out = self.prediction[2](x_ske1_n)
+
+        x_sal2_n = self.prediction[0](x_sal2_n)
+        x_edg2_n = self.prediction[1](x_edg2_n)
+        x_ske2_n = self.prediction[2](x_ske2_n)
+        x_sal3_n = self.prediction[0](x_sal3_n)
+        x_edg3_n = self.prediction[1](x_edg3_n)
+        x_ske3_n = self.prediction[2](x_ske3_n)
+
+        x_sal4_n = self.prediction[0](x_sal4_n)
+        x_edg4_n = self.prediction[1](x_edg4_n)
+        x_ske4_n = self.prediction[2](x_ske4_n)
+
+        x_sal5_n = self.prediction[0](x_sal5_n)
+        x_edg5_n = self.prediction[1](x_edg5_n)
+        x_ske5_n = self.prediction[2](x_ske5_n)
+        
+        sal_out = F.interpolate(sal_out, size=size, mode='bilinear', align_corners=True)
+        edg_out = F.interpolate(edg_out, size=size, mode='bilinear', align_corners=True)
+        ske_out = F.interpolate(ske_out, size=size, mode='bilinear', align_corners=True)
+        sal2 = F.interpolate(x_sal2_n, size=size, mode='bilinear', align_corners=True)
+        edg2 = F.interpolate(x_edg2_n, size=size, mode='bilinear', align_corners=True)
+        ske2 = F.interpolate(x_ske2_n, size=size, mode='bilinear', align_corners=True)
+        sal3 = F.interpolate(x_sal3_n, size=size, mode='bilinear', align_corners=True)
+        edg3 = F.interpolate(x_edg3_n, size=size, mode='bilinear', align_corners=True)
+        ske3 = F.interpolate(x_ske3_n, size=size, mode='bilinear', align_corners=True)
+        sal4 = F.interpolate(x_sal4_n, size=size, mode='bilinear', align_corners=True)
+        edg4 = F.interpolate(x_edg4_n, size=size, mode='bilinear', align_corners=True)
+        ske4 = F.interpolate(x_ske4_n, size=size, mode='bilinear', align_corners=True)
+        sal5 = F.interpolate(x_sal5_n, size=size, mode='bilinear', align_corners=True)
+        edg5 = F.interpolate(x_edg5_n, size=size, mode='bilinear', align_corners=True)
+        ske5 = F.interpolate(x_ske5_n, size=size, mode='bilinear', align_corners=True)
+
+
+        return x_sal1_n, sal_out, self.sigmoid(sal_out), edg_out, self.sigmoid(edg_out), sal2, edg2, self.sigmoid(
+            sal2), self.sigmoid(edg2), sal3, edg3, self.sigmoid(sal3), self.sigmoid(edg3), sal4, edg4, self.sigmoid(
+            sal4), self.sigmoid(edg4), sal5, edg5, self.sigmoid(sal5), self.sigmoid(edg5), ske_out, self.sigmoid(
+            ske_out), ske2, self.sigmoid(ske2), ske3, self.sigmoid(ske3), ske4, self.sigmoid(ske4), ske5, self.sigmoid(
+            ske5)
+        # return x_sal1_n, sal_out, self.sigmoid(sal_out), edg_out, self.sigmoid(edg_out), sal2, edg2, self.sigmoid(
+        #     sal2), self.sigmoid(edg2), sal3, edg3, self.sigmoid(sal3), self.sigmoid(edg3), sal4, edg4, self.sigmoid(
+        #     sal4), self.sigmoid(edg4), sal5, edg5, self.sigmoid(sal5), self.sigmoid(edg5)
+
+

model/MobileNetV2.py

@@ -0,0 +1,123 @@
+from torch import nn
+import torch
+
+try:
+    from torchvision.models.utils import load_state_dict_from_url  # torchvision 0.4+
+except ModuleNotFoundError:
+    try:
+        from torch.hub import load_state_dict_from_url  # torch 1.x
+    except ModuleNotFoundError:
+        from torch.utils.model_zoo import load_url as load_state_dict_from_url  # torch 0.4.1
+
+model_urls = {
+    'mobilenet_v2': 'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth',
+}
+
+
+class ConvBNReLU(nn.Sequential):
+    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, dilation=1):
+        padding = (kernel_size - 1) // 2
+        if dilation != 1:
+            padding = dilation
+        super(ConvBNReLU, self).__init__(
+            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, dilation=dilation,
+                      bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU6(inplace=True)
+        )
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio, dilation=1):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = int(round(inp * expand_ratio))
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        layers = []
+        if expand_ratio != 1:
+            # pw
+            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
+        layers.extend([
+            # dw
+            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, dilation=dilation),
+            # pw-linear
+            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(oup),
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, pretrained=None, num_classes=1000, width_mult=1.0):
+        super(MobileNetV2, self).__init__()
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        inverted_residual_setting = [
+            # t, c, n, s, d
+            [1, 16, 1, 1, 1], # conv1 112*112*16
+            [6, 24, 2, 2, 1], # conv2 56*56*24
+            [6, 32, 3, 2, 1], # conv3 28*28*32
+            [6, 64, 4, 2, 1],
+            [6, 96, 3, 1, 1], # conv4 14*14*96
+            [6, 160, 3, 2, 1],
+            [6, 320, 1, 1, 1], # conv5 7*7*320
+        ]
+
+        # building first layer
+        input_channel = int(input_channel * width_mult)
+        self.last_channel = int(last_channel * max(1.0, width_mult))
+        features = [ConvBNReLU(3, input_channel, stride=2)]
+        # building inverted residual blocks
+        for t, c, n, s, d in inverted_residual_setting:
+            output_channel = int(c * width_mult)
+            for i in range(n):
+                stride = s if i == 0 else 1
+                dilation = d if i == 0 else 1
+                features.append(block(input_channel, output_channel, stride, expand_ratio=t, dilation=d))
+                input_channel = output_channel
+        # building last several layers
+        features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1))
+        # make it nn.Sequential
+        self.features = nn.Sequential(*features)
+
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        res = []
+        for idx, m in enumerate(self.features):
+            x = m(x)
+            if idx in [1, 3, 6, 13, 17]:
+                res.append(x)
+        return res
+
+
+def mobilenet_v2(pretrained=True, progress=True, **kwargs):
+    model = MobileNetV2(**kwargs)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls['mobilenet_v2'],
+                                              progress=progress)
+        print("loading imagenet pretrained mobilenetv2")
+        model.load_state_dict(state_dict, strict=False)
+        print("loaded imagenet pretrained mobilenetv2")
+    return model