edge-enhanced model

model_edge_enhance.py

@@ -0,0 +1,306 @@
+
+# Model Architecture
+# Author: Landy Xu, created on Nov. 12, 2022
+# Last modified by Simon on Nov. 13
+# Version 2: add attention to shallow feature, change first conv to 1x1 kernal
+'''
+Change log: 
+
+- Landy: create feature extractor and DILRAN
+- Simon: revise some writing style of module configs (e.g., replace = True),
+refine the FE module, add recon module
+- Simon: create full model pipeline
+- Simon: add leaky relu to recon module
+'''
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+
+class ConvLeakyRelu2d(nn.Module):
+    # convolution
+    # leaky relu
+    def __init__(self, in_channels, out_channels, kernel_size=3, padding=1, stride=1, dilation=1, groups=1):
+        super(ConvLeakyRelu2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation, groups=groups)
+        # self.bn   = nn.BatchNorm2d(out_channels)
+    def forward(self,x):
+        # print(x.size())
+        return F.leaky_relu(self.conv(x), negative_slope=0.2)
+
+class Sobelxy(nn.Module):
+    def __init__(self,channels, kernel_size=3, padding=1, stride=1, dilation=1, groups=1):
+        super(Sobelxy, self).__init__()
+        sobel_filter = np.array([[1, 0, -1],
+                                 [2, 0, -2],
+                                 [1, 0, -1]])
+        self.convx=nn.Conv2d(channels, channels, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation, groups=channels,bias=False)
+        self.convx.weight.data.copy_(torch.from_numpy(sobel_filter))
+        self.convy=nn.Conv2d(channels, channels, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation, groups=channels,bias=False)
+        self.convy.weight.data.copy_(torch.from_numpy(sobel_filter.T))
+    def forward(self, x):
+        sobelx = self.convx(x)
+        sobely = self.convy(x)
+        x=torch.abs(sobelx) + torch.abs(sobely)
+        return x
+
+class Conv1(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, padding=0, stride=1, dilation=1, groups=1):
+        super(Conv1, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation, groups=groups)
+    def forward(self,x):
+        return self.conv(x)
+
+class DenseBlock(nn.Module):
+    def __init__(self,channels):
+        super(DenseBlock, self).__init__()
+        # self.conv_in = nn.Conv2d(1, 64, kernel_size=1)
+        self.conv1 = ConvLeakyRelu2d(channels, channels)
+        self.conv2 = ConvLeakyRelu2d(2*channels, 2*channels)
+        # self.conv3 = ConvLeakyRelu2d(3*channels, channels)
+    def forward(self,x):
+        # x = self.conv_in(x)
+        x = torch.cat((x,self.conv1(x)),dim=1)
+        x = self.conv2(x)
+        # x = torch.cat((x, self.conv3(x)), dim=1)
+        return x
+
+class Edge_Enhancer(nn.Module):
+    def __init__(self,in_channels,out_channels):
+        super(Edge_Enhancer, self).__init__()
+        self.dense =DenseBlock(in_channels)
+        self.convdown=Conv1(2*in_channels,out_channels)
+        self.sobelconv=Sobelxy(in_channels)
+        self.convup = Conv1(in_channels,out_channels)
+        self.dropout = nn.Dropout(0.2)
+    def forward(self,x):
+        x1=self.dropout(self.dense(x))
+        x1=self.convdown(x1)
+        x2=self.sobelconv(x)
+        x2=self.convup(x2)
+        return self.dropout(x1+x2)
+
+
+class DILRAN(nn.Module):
+    def __init__(self):
+        super(DILRAN, self).__init__()
+        # TODO: confirm convolution
+        self.conv = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.down = nn.AvgPool2d(2, 2)
+        self.lu = nn.ReLU(replace = True)
+
+    def forward(self, x):
+        prev = self.conv(x) + self.conv(self.conv(x)) + self.conv(self.conv(self.conv(x)))
+        return torch.mul(self.lu(self.up(self.down(x))), prev) + x
+
+
+class FeatureExtractor(nn.Module):
+    def __init__(self, level):
+        super(FeatureExtractor, self).__init__()
+        # TODO: confirm dilated convolution
+        self.conv = nn.Conv2d(1, 64, (1, 1), (1, 1), (0, 0), dilation = 2)
+        self.network = DILRAN()
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.down = nn.AvgPool2d(2, 2)
+        self.lu = nn.ReLU(replace = True)
+
+    def forward(self, x):
+        n1 = self.network(self.conv(x[0]))
+        n2 = self.network(self.conv(x[1]))
+        n3 = self.network(self.conv(x[2]))
+        return torch.cat((n1, n2, n3), 0)
+
+
+class DILRAN_V1(nn.Module):
+    '''
+    V1: concat the output of three (conv-d,DILRAN) paths channel wise and add the low level feature to the concat output
+    temporary, will edit if necessary
+    '''
+    def __init__(self, cat_first = False, use_leaky = False):
+        super(DILRAN_V1, self).__init__()
+        # cat_first, whether to perform channel-wise concat before DILRAN
+        # convolution in DILRAN, in channel is the channel from the previous block
+        if not cat_first:
+            self.conv_d = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding="same")
+            self.bnorm = nn.BatchNorm2d(num_features=64)
+        else:
+            self.conv_d = nn.Conv2d(in_channels=64*3, out_channels=64*3, kernel_size=3, stride=1, padding="same")
+            self.bnorm = nn.BatchNorm2d(num_features=64*3)
+
+        if not use_leaky:
+            self.relu = nn.ReLU()
+        else:
+            self.lrelu = nn.LeakyReLU(0.2, inplace=True)
+
+        self.down = nn.AvgPool2d(2, 2)
+        self.up = nn.Upsample(scale_factor=2, mode="nearest")
+
+
+    def forward(self, x):
+        # pooling -> upsample -> ReLU block
+        pur_path = self.relu(self.up(self.down(x)))
+        # 3*3, 5*5, 7*7 multiscale addition block
+        conv_path = self.conv_d(x) + self.conv_d(self.conv_d(x)) + self.conv_d(self.conv_d(self.conv_d(x)))
+        # attention
+        attn = torch.mul(pur_path, conv_path)
+        # residual + attention
+        resid_x = x + attn
+        return resid_x
+
+
+class FE_V1(nn.Module):
+    '''
+    feature extractor block (temporary, will edit if necessary)
+    '''
+    def __init__(self):
+        super(FE_V1, self).__init__()
+
+        # multiscale dilation conv2d
+        self.convd1 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, dilation=1, padding="same")
+        self.convd2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, dilation=3, padding="same")
+        self.convd3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, dilation=5, padding="same")
+
+        self.reduce = nn.Conv2d(in_channels=64*3, out_channels=64, kernel_size=1, stride=1, padding="same")
+        self.relu = nn.ReLU()
+
+        self.bnorm1 = nn.BatchNorm2d(num_features=64)
+
+        self.dilran = DILRAN_V1()
+
+
+    def forward(self, x):
+
+        # dilated convolution
+        dilf1 = self.convd1(x)
+        dilf2 = self.convd2(x)
+        dilf3 = self.convd3(x)
+
+        diltotal = torch.cat((dilf1, dilf2, dilf3), dim = 1)
+        diltotal = self.reduce(diltotal)
+        diltotal = self.bnorm1(diltotal)
+
+        # single DILRAN
+        out = self.dilran(diltotal)
+        out = self.bnorm1(out)
+        #out = self.relu(out)
+        return out
+
+        # DILRAN
+        # dilran_o1 = self.dilran(dilf1)
+        # # batchnorm
+        # dilran_o1 = self.bnorm1(dilran_o1)
+        # dilran_o2 = self.dilran(dilf2)
+        # # batchnorm
+        # dilran_o2 = self.bnorm1(dilran_o2)
+        # dilran_o3 = self.dilran(dilf3)
+        # # batchnorm
+        # dilran_o3 = self.bnorm1(dilran_o3)
+        # # element-wise addition
+        # cat_o = dilran_o1 + dilran_o2 + dilran_o3
+
+        # return cat_o
+
+class MSFuNet(nn.Module):
+    '''
+    the whole network (from input image -> feature maps to be used in fusion strategy)
+    temporary, will edit if necessary
+    '''
+    def __init__(self):
+        super(MSFuNet, self).__init__()
+
+        self.conv_id = nn.Sequential(nn.Conv2d(in_channels=64*3, out_channels=64, kernel_size=1, stride=1, padding="same"))
+                                    #nn.BatchNorm2d(num_features = 64))
+                                    #nn.ReLU(inplace=True))
+        self.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=1, stride=1, padding="same")
+
+        self.conv2 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding="same"),
+                                    nn.BatchNorm2d(num_features=64),
+                                    nn.ReLU(),
+                                    nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding="same"),
+                                    nn.BatchNorm2d(num_features=64))
+
+        self.relu = nn.ReLU()
+        self.down = nn.AvgPool2d(2, 2)
+        self.bnorm = nn.BatchNorm2d(num_features=64)
+        self.up = nn.Upsample(scale_factor=2, mode="nearest")
+
+        self.fe = FE_V1()
+        self.edge_enhance = Edge_Enhancer(64,64)
+
+    def forward(self, x):
+        # x: input image
+        resid = self.conv1(x)
+        temp0 = self.conv1(x) # shallow feature, 64 x (1x1)
+        pur_orig = self.relu(self.up(self.down(x)))
+        attn = torch.mul(pur_orig, temp0)
+        x = x + attn
+        # feature returned from feature extractor
+        deep_fe = self.fe(x)
+        pur_x = self.relu(self.up(self.down(x)))
+        attn2 = torch.mul(pur_x, deep_fe)
+        add = attn2 + x
+
+        # addded for edge enhance
+        edge_x = self.edge_enhance(resid)
+        add = add + edge_x
+        return add
+        #x = x + cat_feature
+        # short cut connection 
+        # expand_x = self.conv_id(x)
+        # add = expand_x + cat_feature
+
+        #add = self.conv2(add)
+        # add = self.conv2(resid) # should get shape [b, 64, 256, 256]
+        # return add
+
+
+class Recon(nn.Module):
+    '''
+    reconstruction module (temporary, will edit if necessary)
+    '''
+    def __init__(self):
+        super(Recon, self).__init__()
+
+        # version 1
+        # self.recon_conv = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding="same"),
+        #                                 nn.LeakyReLU(0.2, inplace=True),
+        #                                 nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding="same"),
+        #                                 nn.LeakyReLU(0.2, inplace=True),
+        #                                 nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, stride=1, padding="same"),
+        #                                 nn.LeakyReLU(0.2, inplace=True),
+        #                                 nn.Conv2d(in_channels=16, out_channels=1, kernel_size=3, stride=1, padding="same"),
+        #                                 nn.LeakyReLU(0.2, inplace=True))
+
+        # version 2
+        self.recon_conv = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding="same"),
+                                        nn.LeakyReLU(0.2, inplace=True),    
+                                        #nn.ReLU(),
+                                        nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding="same"),
+                                        #nn.ReLU(),
+                                        nn.LeakyReLU(0.2, inplace=True),
+                                        nn.Conv2d(in_channels=32, out_channels=1, kernel_size=3, stride=1, padding="same"))
+                                        #nn.ReLU())
+    def forward(self, x):
+        x = self.recon_conv(x)
+        return x # should get shape [b, 1, 256, 256]
+
+
+class fullModel(nn.Module):
+    '''
+    Feature extractor + reconstruction
+    a full model pipeline
+    '''
+    def __init__(self):
+        super(fullModel, self).__init__()
+
+        self.fe = MSFuNet()
+        self.recon = Recon()
+
+    def forward(self, x):
+        deep_fe = self.fe(x)
+        recon_img = self.recon(deep_fe)
+        return recon_img