Create CMFNet.py

model/CMFNet.py

@@ -0,0 +1,191 @@
+import torch
+import torch.nn as nn
+from model.block import SAB, CAB, PAB, conv, SAM, conv3x3, conv_down
+
+##########################################################################
+## U-Net
+bn = 2  # block number-1
+
+class Encoder(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, block):
+        super(Encoder, self).__init__()
+        if block == 'CAB':
+            self.encoder_level1 = [CAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.encoder_level2 = [CAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.encoder_level3 = [CAB(n_feat + (scale_unetfeats * 2), kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+        elif block == 'PAB':
+            self.encoder_level1 = [PAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.encoder_level2 = [PAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.encoder_level3 = [PAB(n_feat + (scale_unetfeats * 2), kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+        elif block == 'SAB':
+            self.encoder_level1 = [SAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.encoder_level2 = [SAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.encoder_level3 = [SAB(n_feat + (scale_unetfeats * 2), kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+        self.encoder_level1 = nn.Sequential(*self.encoder_level1)
+        self.encoder_level2 = nn.Sequential(*self.encoder_level2)
+        self.encoder_level3 = nn.Sequential(*self.encoder_level3)
+        self.down12 = DownSample(n_feat, scale_unetfeats)
+        self.down23 = DownSample(n_feat + scale_unetfeats, scale_unetfeats)
+
+    def forward(self, x):
+        enc1 = self.encoder_level1(x)
+        x = self.down12(enc1)
+        enc2 = self.encoder_level2(x)
+        x = self.down23(enc2)
+        enc3 = self.encoder_level3(x)
+        return [enc1, enc2, enc3]
+
+class Decoder(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, block):
+        super(Decoder, self).__init__()
+        if block == 'CAB':
+            self.decoder_level1 = [CAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.decoder_level2 = [CAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.decoder_level3 = [CAB(n_feat + (scale_unetfeats * 2), kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+        elif block == 'PAB':
+            self.decoder_level1 = [PAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.decoder_level2 = [PAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.decoder_level3 = [PAB(n_feat + (scale_unetfeats * 2), kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+        elif block == 'SAB':
+            self.decoder_level1 = [SAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.decoder_level2 = [SAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+            self.decoder_level3 = [SAB(n_feat + (scale_unetfeats * 2), kernel_size, reduction, bias=bias, act=act) for _ in range(bn)]
+        self.decoder_level1 = nn.Sequential(*self.decoder_level1)
+        self.decoder_level2 = nn.Sequential(*self.decoder_level2)
+        self.decoder_level3 = nn.Sequential(*self.decoder_level3)
+        if block == 'CAB':
+            self.skip_attn1 = CAB(n_feat, kernel_size, reduction, bias=bias, act=act)
+            self.skip_attn2 = CAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act)
+        if block == 'PAB':
+            self.skip_attn1 = PAB(n_feat, kernel_size, reduction, bias=bias, act=act)
+            self.skip_attn2 = PAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act)
+        if block == 'SAB':
+            self.skip_attn1 = SAB(n_feat, kernel_size, reduction, bias=bias, act=act)
+            self.skip_attn2 = SAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act)
+        self.up21 = SkipUpSample(n_feat, scale_unetfeats)
+        self.up32 = SkipUpSample(n_feat + scale_unetfeats, scale_unetfeats)
+
+    def forward(self, outs):
+        enc1, enc2, enc3 = outs
+        dec3 = self.decoder_level3(enc3)
+        x = self.up32(dec3, self.skip_attn2(enc2))
+        dec2 = self.decoder_level2(x)
+        x = self.up21(dec2, self.skip_attn1(enc1))
+        dec1 = self.decoder_level1(x)
+        return [dec1, dec2, dec3]
+
+##########################################################################
+##---------- Resizing Modules ----------
+class DownSample(nn.Module):
+    def __init__(self, in_channels, s_factor):
+        super(DownSample, self).__init__()
+        self.down = nn.Sequential(nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False),
+                                  nn.Conv2d(in_channels, in_channels + s_factor, 1, stride=1, padding=0, bias=False))
+
+    def forward(self, x):
+        x = self.down(x)
+        return x
+
+class UpSample(nn.Module):
+    def __init__(self, in_channels, s_factor):
+        super(UpSample, self).__init__()
+        self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+                                nn.Conv2d(in_channels + s_factor, in_channels, 1, stride=1, padding=0, bias=False))
+
+    def forward(self, x):
+        x = self.up(x)
+        return x
+
+class SkipUpSample(nn.Module):
+    def __init__(self, in_channels, s_factor):
+        super(SkipUpSample, self).__init__()
+        self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+                                nn.Conv2d(in_channels + s_factor, in_channels, 1, stride=1, padding=0, bias=False))
+
+    def forward(self, x, y):
+        x = self.up(x)
+        x = x + y
+        return x
+
+##########################################################################
+# Mixed Residual Module
+class Mix(nn.Module):
+    def __init__(self, m=1):
+        super(Mix, self).__init__()
+        w = nn.Parameter(torch.FloatTensor([m]), requires_grad=True)
+        w = nn.Parameter(w, requires_grad=True)
+        self.w = w
+        self.mix_block = nn.Sigmoid()
+
+    def forward(self, fea1, fea2, feat3):
+        factor = self.mix_block(self.w)
+        other = (1 - factor)/2
+        output = fea1 * other.expand_as(fea1) + fea2 * factor.expand_as(fea2) + feat3 * other.expand_as(feat3)
+        return output, factor
+
+##########################################################################
+# Architecture
+class CMFNet(nn.Module):
+    def __init__(self, in_c=3, out_c=3, n_feat=96, scale_unetfeats=48, kernel_size=3, reduction=4, bias=False):
+        super(CMFNet, self).__init__()
+
+        p_act = nn.PReLU()
+        self.shallow_feat1 = nn.Sequential(conv(in_c, n_feat // 2, kernel_size, bias=bias), p_act,
+                                           conv(n_feat // 2, n_feat, kernel_size, bias=bias))
+        self.shallow_feat2 = nn.Sequential(conv(in_c, n_feat // 2, kernel_size, bias=bias), p_act,
+                                           conv(n_feat // 2, n_feat, kernel_size, bias=bias))
+        self.shallow_feat3 = nn.Sequential(conv(in_c, n_feat // 2, kernel_size, bias=bias), p_act,
+                                           conv(n_feat // 2, n_feat, kernel_size, bias=bias))
+
+        self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, p_act, bias, scale_unetfeats, 'CAB')
+        self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, p_act, bias, scale_unetfeats, 'CAB')
+
+        self.stage2_encoder = Encoder(n_feat, kernel_size, reduction, p_act, bias, scale_unetfeats, 'PAB')
+        self.stage2_decoder = Decoder(n_feat, kernel_size, reduction, p_act, bias, scale_unetfeats, 'PAB')
+
+        self.stage3_encoder = Encoder(n_feat, kernel_size, reduction, p_act, bias, scale_unetfeats, 'SAB')
+        self.stage3_decoder = Decoder(n_feat, kernel_size, reduction, p_act, bias, scale_unetfeats, 'SAB')
+
+        self.sam1o = SAM(n_feat, kernel_size=3, bias=bias)
+        self.sam2o = SAM(n_feat, kernel_size=3, bias=bias)
+        self.sam3o = SAM(n_feat, kernel_size=3, bias=bias)
+
+        self.mix = Mix(1)
+        self.add123 = conv(out_c, out_c, kernel_size, bias=bias)
+        self.concat123 = conv(n_feat*3, n_feat, kernel_size, bias=bias)
+        self.tail = conv(n_feat, out_c, kernel_size, bias=bias)
+
+
+    def forward(self, x):
+        ## Compute Shallow Features
+        shallow1 = self.shallow_feat1(x)
+        shallow2 = self.shallow_feat2(x)
+        shallow3 = self.shallow_feat3(x)
+
+        ## Enter the UNet-CAB
+        x1 = self.stage1_encoder(shallow1)
+        x1_D = self.stage1_decoder(x1)
+        ## Apply SAM
+        x1_out, x1_img = self.sam1o(x1_D[0], x)
+
+        ## Enter the UNet-PAB
+        x2 = self.stage2_encoder(shallow2)
+        x2_D = self.stage2_decoder(x2)
+        ## Apply SAM
+        x2_out, x2_img = self.sam2o(x2_D[0], x)
+
+        ## Enter the UNet-SAB
+        x3 = self.stage3_encoder(shallow3)
+        x3_D = self.stage3_decoder(x3)
+        ## Apply SAM
+        x3_out, x3_img = self.sam3o(x3_D[0], x)
+
+        ## Aggregate SAM features of Stage 1, Stage 2 and Stage 3
+        mix_r = self.mix(x1_img, x2_img, x3_img)
+        mixed_img = self.add123(mix_r[0])
+
+        ## Concat SAM features of Stage 1, Stage 2 and Stage 3
+        concat_feat = self.concat123(torch.cat([x1_out, x2_out, x3_out], 1))
+        x_final = self.tail(concat_feat)
+
+        return x_final + mixed_img