EEM

import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
import matplotlib.pyplot as plt

#多方向sobel 这个不是多方向的
def get_sobel(in_chan, out_chan):
    filter_x = np.array([
        [1, 0, -1],
        [2, 0, -2],
        [1, 0, -1],
    ]).astype(np.float32)
    filter_y = np.array([
        [1, 2, 1],
        [0, 0, 0],
        [-1, -2, -1],
    ]).astype(np.float32)

    filter_x = filter_x.reshape((1, 1, 3, 3))
    filter_x = np.repeat(filter_x, in_chan, axis=1)
    filter_x = np.repeat(filter_x, out_chan, axis=0)

    filter_y = filter_y.reshape((1, 1, 3, 3))
    filter_y = np.repeat(filter_y, in_chan, axis=1)
    filter_y = np.repeat(filter_y, out_chan, axis=0)

    filter_x = torch.from_numpy(filter_x)
    filter_y = torch.from_numpy(filter_y)
    filter_x = nn.Parameter(filter_x, requires_grad=False)
    filter_y = nn.Parameter(filter_y, requires_grad=False)
    conv_x = nn.Conv2d(in_chan, out_chan, kernel_size=3, stride=1, padding=1, bias=False)
    conv_x.weight = filter_x
    conv_y = nn.Conv2d(in_chan, out_chan, kernel_size=3, stride=1, padding=1, bias=False)
    conv_y.weight = filter_y
    sobel_x = nn.Sequential(conv_x, nn.BatchNorm2d(out_chan))
    sobel_y = nn.Sequential(conv_y, nn.BatchNorm2d(out_chan))

    return sobel_x, sobel_y

def run_sobel(conv_x, conv_y, input):
    g_x = conv_x(input)
    g_y = conv_y(input)
    g = torch.sqrt(torch.pow(g_x, 2) + torch.pow(g_y, 2))
    return torch.sigmoid(g) * input

class Downsample_block(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Downsample_block, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        y = F.relu(self.bn2(self.conv2(x)))
        x = F.max_pool2d(y, 2, stride=2)

        return x, y

#针对E1
class Edge1(nn.Module):   
    def __init__(self, in_channels):
        super(Edge1, self).__init__()

        self.sobel_x1, self.sobel_y1 = get_sobel(in_channels, 1)
        self.conv11 = nn.Conv2d(in_channels, in_channels, kernel_size=1, padding=1)
        self.bn = nn.BatchNorm2d(in_channels)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        y = run_sobel(self.sobel_x1, self.sobel_y1, x)
        out = F.relu(self.bn(y))
        out = self.conv11(out)
        sigmoid_out = self.sigmoid(out)
        EM = E1 = sigmoid_out
        
        return EM

#针对E2
class Edge2(nn.Module):   
    def __init__(self, in_channels):
        super(Edge2, self).__init__()

        self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(in_channels, out_channels)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        avg_out = self.global_avg_pool(x)
        out = self.fc(avg_out)
        out = self.relu(out)
        out = self.softmax(out)
        out = out
        E2S = out * x
        
        return E2S

#上采样
class Upsample(nn.Module):
    def __init__(self, in_channels, out_channels, scale_factor=2):
        super(Upsample, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
        self.upsample = nn.Upsample(scale_factor=scale_factor, mode='bilinear', align_corners=False)

    def forward(self, x):
        x = self.conv(x)
        x = self.upsample(x)
        return x


class EdgeFeatureExtraction(nn.Module):
def __init__(self):
        super(EdgeFeatureExtraction, self).__init__()

        self.upsample = Upsample(in_channels, in_channels)

    def forward(self, EM, E2S):
        ES = self.upsample(E2S)
        E = EM + E2S
        
        return E


class Depthwise_Conv(nn.Module):
    def __init__(self,in_ch,out_ch,groups):
        super(Depthwise_Conv,self).__init__()
        self.conv=nn.Conv2d(in_channels=in_ch,
                            out_channels=out_ch,
                            kernel_size=3,
                            stride=1,
                            padding=0,
                            groups=groups,
                            bias=False)
    def forward(self,input):
        return self.conv(input) 

class DepthWiseConv(nn.Module):
    def __init__(self,in_channel,out_channel):
        super(DepthWiseConv, self).__init__()
        # 逐通道卷积
        self.depth_conv = nn.Conv2d(in_channels=in_channel,
                                    out_channels=in_channel,
                                    kernel_size=3,
                                    stride=1,
                                    padding=1,
                                    groups=in_channel)
        # groups是一个数，当groups=in_channel时,表示做逐通道卷积
        #逐点卷积
        self.point_conv = nn.Conv2d(in_channels=in_channel,
                                    out_channels=out_channel,
                                    kernel_size=1,
                                    stride=1,
                                    padding=0,
                                    groups=1)
    
   def forward(self,input):
        out = self.depth_conv(input)
        out_conv33DW = self.point_conv(out)
        return out_conv33DW
        
class EdgeEnhancement(nn.Module):  
    def __init__(self，ch_in):
        super(EdgeEnhancement, self).__init__()

        self.conv33 = nn.Conv2d(2, 1, kernel_size=3, padding=kernel_size//2, bias=False)  # 3x3 卷积
        self.conv13 = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 3), padding=(0, 1), bias=False)
        self.conv31 = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 1), padding=(1, 0), bias=False)

        self.conv11 = nn.Conv2d(2, 1, kernel_size=1, padding=kernel_size//2，bias=False)  # 1x1 卷积

    def forward(self, E，out_conv33DW):
        out_conv33 = self.conv33(E)
        out_conv13 = self.conv13(E)
        out_conv31 = self.conv31(E)
        
        out1 = out_conv13 + out_conv31
        out2 = out + out_conv33 + out_conv33DW
        
        out = self.conv11(out2)

        return out