<filename>EF-SAI-main/codes/Networks/submodules.py
import math
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
import functools



# Channel Attention from CBAM
class ChannelAttentionv2(nn.Module):
    def __init__(self, in_planes, in_planes2, ratio=16, kernel_size=3):
        super(ChannelAttentionv2, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.conv1 = nn.Sequential(nn.Conv2d(in_planes2*2, 1, kernel_size, padding=kernel_size//2, bias=False), nn.ReLU())
        self.conv2 = nn.Sequential(nn.Conv2d(in_planes2, 1, kernel_size, padding=kernel_size//2, bias=False), nn.ReLU())
        self.conv3 = nn.Sequential(nn.Conv2d(in_planes2, 1, kernel_size, padding=kernel_size//2, bias=False), nn.ReLU())
        #self.conv4 = nn.Sequential(nn.Conv2d(in_planes+1, 1, kernel_size, padding=kernel_size//2, bias=False), nn.ReLU())
        #self.conv5 = nn.Sequential(nn.Conv2d(in_planes+1, 1, kernel_size, padding=kernel_size // 2, bias=False),nn.ReLU())
        #self.conv6 = nn.Sequential(nn.Conv2d(in_planes+1, 1, kernel_size, padding=kernel_size // 2, bias=False),nn.ReLU())
        in_planes3 = in_planes*3+3
        self.fc = nn.Sequential(nn.Conv2d(in_planes3, in_planes3 // ratio, 1, bias=False),
                                nn.ReLU(),
                                nn.Conv2d(in_planes3 // ratio, 3, 1, bias=False))
        self.sigmoid = nn.Sigmoid()

    def forward(self, x1, x2, y1, y2, z1, z2):
        x = self.conv1(x2)
        x = torch.cat([x1, x], dim=1)
        # x = self.conv4(x)
        y = self.conv2(y2)
        y = torch.cat([y1, y], dim=1)
        # y = self.conv5(y)
        z = self.conv3(z2)
        z = torch.cat([z1, z], dim=1)
        # z = self.conv6(z)
        x = torch.cat([x, y, z], dim=1)
        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)


#Generator
class Identity(nn.Module):
    def forward(self, x):
        return x

def get_norm_layer(norm_type='instance'):
    """Return a normalization layer

    Parameters:
        norm_type (str) -- the name of the normalization layer: batch | instance | none

    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
    """
    if norm_type == 'batch':
        momentum = 0.01
        norm_layer = functools.partial(nn.BatchNorm2d,momentum = momentum,  affine=True, track_running_stats=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
    elif norm_type == 'none':
        def norm_layer(x): return Identity()
    else:
        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
    return norm_layer
class ResnetGenerator(nn.Module):
    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.

    We adapt Torch code and idea from <NAME>'s neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
    """

    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
        """Construct a Resnet-based generator

        Parameters:
            input_nc (int)      -- the number of channels in input images
            output_nc (int)     -- the number of channels in output images
            ngf (int)           -- the number of filters in the last conv layer
            norm_layer          -- normalization layer
            use_dropout (bool)  -- if use dropout layers
            n_blocks (int)      -- the number of ResNet blocks
            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
        """
        assert(n_blocks >= 0)
        super(ResnetGenerator, self).__init__()
        if type(norm_layer) == functools.partial:
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d

        model = [nn.ReflectionPad2d(3),
                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
                 norm_layer(ngf),
                 nn.ReLU(True)]

        n_downsampling = 2
        for i in range(n_downsampling):  # add downsampling layers
            mult = 2 ** i
            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
                      norm_layer(ngf * mult * 2),
                      nn.ReLU(True)]

        mult = 2 ** n_downsampling
        for i in range(n_blocks):       # add ResNet blocks

            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]

        for i in range(n_downsampling):  # add upsampling layers
            mult = 2 ** (n_downsampling - i)
            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
                                         kernel_size=3, stride=2,
                                         padding=1, output_padding=1,
                                         bias=use_bias),
                      norm_layer(int(ngf * mult / 2)),
                      nn.ReLU(True)]
        model += [nn.ReflectionPad2d(3)]
        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
        model += [nn.Tanh()]

        self.model = nn.Sequential(*model)

    def forward(self, input):
        """Standard forward"""
        return self.model(input)
class ResnetBlock(nn.Module):
    """Define a Resnet block"""

    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        """Initialize the Resnet block

        A resnet block is a conv block with skip connections
        We construct a conv block with build_conv_block function,
        and implement skip connections in <forward> function.
        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
        """
        super(ResnetBlock, self).__init__()
        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)

    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        """Construct a convolutional block.

        Parameters:
            dim (int)           -- the number of channels in the conv layer.
            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
            norm_layer          -- normalization layer
            use_dropout (bool)  -- if use dropout layers.
            use_bias (bool)     -- if the conv layer uses bias or not

        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
        """
        conv_block = []
        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)

        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
        if use_dropout:
            conv_block += [nn.Dropout(0.5)]

        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]

        return nn.Sequential(*conv_block)

    def forward(self, x):
        """Forward function (with skip connections)"""
        out = x + self.conv_block(x)  # add skip connections
        return out


class UnetGenerator(nn.Module):
    """Create a Unet-based generator"""

    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
        """Construct a Unet generator
        Parameters:
            input_nc (int)  -- the number of channels in input images
            output_nc (int) -- the number of channels in output images
            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
                                image of size 128x128 will become of size 1x1 # at the bottleneck
            ngf (int)       -- the number of filters in the last conv layer
            norm_layer      -- normalization layer

        We construct the U-Net from the innermost layer to the outermost layer.
        It is a recursive process.
        """
        super(UnetGenerator, self).__init__()
        # construct unet structure
        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer
        for i in range(num_downs - 5):          # add intermediate layers with ngf * 8 filters
            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
        # gradually reduce the number of filters from ngf * 8 to ngf
        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer

    def forward(self, input):
        """Standard forward"""
        return self.model(input)


class UnetSkipConnectionBlock(nn.Module):
    """Defines the Unet submodule with skip connection.
        X -------------------identity----------------------
        |-- downsampling -- |submodule| -- upsampling --|
    """

    def __init__(self, outer_nc, inner_nc, input_nc=None,
                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
        """Construct a Unet submodule with skip connections.

        Parameters:
            outer_nc (int) -- the number of filters in the outer conv layer
            inner_nc (int) -- the number of filters in the inner conv layer
            input_nc (int) -- the number of channels in input images/features
            submodule (UnetSkipConnectionBlock) -- previously defined submodules
            outermost (bool)    -- if this module is the outermost module
            innermost (bool)    -- if this module is the innermost module
            norm_layer          -- normalization layer
            use_dropout (bool)  -- if use dropout layers.
        """
        super(UnetSkipConnectionBlock, self).__init__()
        self.outermost = outermost
        if type(norm_layer) == functools.partial:
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d
        if input_nc is None:
            input_nc = outer_nc
        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
                             stride=2, padding=1, bias=use_bias)
        downrelu = nn.LeakyReLU(0.2, True)
        downnorm = norm_layer(inner_nc)
        uprelu = nn.ReLU(True)
        upnorm = norm_layer(outer_nc)

        if outermost:
            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
                                        kernel_size=4, stride=2,
                                        padding=1)
            down = [downconv]
            up = [uprelu, upconv, nn.Tanh()]
            model = down + [submodule] + up
        elif innermost:
            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
                                        kernel_size=4, stride=2,
                                        padding=1, bias=use_bias)
            down = [downrelu, downconv]
            up = [uprelu, upconv, upnorm]
            model = down + up
        else:
            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
                                        kernel_size=4, stride=2,
                                        padding=1, bias=use_bias)
            down = [downrelu, downconv, downnorm]
            up = [uprelu, upconv, upnorm]

            if use_dropout:
                model = down + [submodule] + up + [nn.Dropout(0.5)]
            else:
                model = down + [submodule] + up

        self.model = nn.Sequential(*model)

    def forward(self, x):
        if self.outermost:
            return self.model(x)
        else:   # add skip connections
            return torch.cat([x, self.model(x)], 1)
def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
    Parameters:
        net (network)      -- the network to be initialized
        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        gain (float)       -- scaling factor for normal, xavier and orthogonal.
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Return an initialized network.
    """
    if len(gpu_ids) > 0:
        assert(torch.cuda.is_available())
        net.to(gpu_ids[0])
        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
    init_weights(net, init_type, init_gain=init_gain)
    return net
def init_weights(net, init_type='normal', init_gain=0.02):
    """Initialize network weights.

    Parameters:
        net (network)   -- network to be initialized
        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.

    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
    work better for some applications. Feel free to try yourself.
    """
    def init_func(m):  # define the initialization function
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, 0.0, init_gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=init_gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
            init.normal_(m.weight.data, 1.0, init_gain)
            init.constant_(m.bias.data, 0.0)

    print('initialize network with %s' % init_type)
    net.apply(init_func)  # apply the initialization function <init_func>
def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
    """Create a generator

    Parameters:
        input_nc (int) -- the number of channels in input images
        output_nc (int) -- the number of channels in output images
        ngf (int) -- the number of filters in the last conv layer
        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
        norm (str) -- the name of normalization layers used in the network: batch | instance | none
        use_dropout (bool) -- if use dropout layers.
        init_type (str)    -- the name of our initialization method.
        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Returns a generator

    Our current implementation provides two types of generators:
        U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
        The original U-Net paper: https://arxiv.org/abs/1505.04597

        Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
        Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
        We adapt Torch code from <NAME>'s neural style transfer project (https://github.com/jcjohnson/fast-neural-style).


    The generator has been initialized by <init_net>. It uses RELU for non-linearity.
    """
    net = None
    norm_layer = get_norm_layer(norm_type=norm)

    if netG == 'resnet_9blocks':
        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
    elif netG == 'resnet_6blocks':
        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
    elif netG == 'unet_128':
        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
    elif netG == 'unet_256':
        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
    else:
        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
    return init_net(net, init_type, init_gain, gpu_ids)









#Fusion Swin Transformer
class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


def window_partition(x, window_size):
    """
    Args:
        x: (B, H, W, C)
        window_size (int): window size

    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    return windows


def window_reverse(windows, window_size, H, W):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image

    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x

class PatchEmbed(nn.Module):
    r""" Image to Patch Embedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None, firstembed=False):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]
        self.in_chans = in_chans
        self.embed_dim = embed_dim
        self.firstembed = firstembed
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None
        if(firstembed):

            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)  # 输入嵌入投影
    def forward(self, x):

        if(self.firstembed):
            #x = self.proj(x).flatten(2).transpose(1, 2)  # shape = (B, P_h*P_w, C)
            x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
        else:
            x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
        if self.norm is not None:
            x = self.norm(x)
        return x

    def flops(self):
        flops = 0
        H, W = self.img_size
        if self.norm is not None:
            flops += H * W * self.embed_dim
        return flops


class PatchUnEmbed(nn.Module):
    r""" Image to Patch Unembedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None, firstembed = False):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim
        self.firstembed = firstembed
        if(firstembed):
            self.proj = nn.Conv2d(embed_dim, in_chans, kernel_size=patch_size, stride=patch_size)  # 输入嵌入投影
    def forward(self, x, x_size):
        B, HW, C = x.shape
        if(self.firstembed):
            #self.proj(x).transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
            x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
        else:
            x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
        return x

    def flops(self):
        flops = 0
        return flops
class WindowAttention_eef(nn.Module):
    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.

    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
    """

    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):

        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wh, Ww
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        # define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH


        # get pair-wise relative position index for each token inside the window
        coords_h = torch.arange(self.window_size[0])
        coords_w = torch.arange(self.window_size[1])
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
        self.register_buffer("relative_position_index", relative_position_index)

        #self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.qkv_eoa = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.qkv_eaa = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.qkv_f = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)

        # self.proj = nn.Linear(dim, dim)
        # self.proj_drop = nn.Dropout(proj_drop)

        self.proj_eoa = nn.Linear(dim, dim)
        self.proj_drop_eoa = nn.Dropout(proj_drop)

        self.proj_eaa = nn.Linear(dim, dim)
        self.proj_drop_eaa = nn.Dropout(proj_drop)

        self.proj_f = nn.Linear(dim, dim)
        self.proj_drop_f = nn.Dropout(proj_drop)

        trunc_normal_(self.relative_position_bias_table, std=.02)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x,y,z, mask=None): #x:eoa y:frame z:eaa
        """
        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        B_, N, C = x.shape
        #print(x.size())
        # qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        # q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)

        qkv_eoa = self.qkv_eoa(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q_eoa, k_eoa, v_eoa = qkv_eoa[0], qkv_eoa[1], qkv_eoa[2]  # make torchscript happy (cannot use tensor as tuple)
        #print(q_eoa.size())
        #print(k_eoa.size())
        qkv_f = self.qkv_f(y).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q_f, k_f, v_f = qkv_f[0], qkv_f[1], qkv_f[2]  # make torchscript happy (cannot use tensor as tuple)

        qkv_eaa = self.qkv_eaa(z).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q_eaa, k_eaa, v_eaa = qkv_eaa[0], qkv_eaa[1], qkv_eaa[2]  # make torchscript happy (cannot use tensor as tuple)

        q_eoa = q_eoa * self.scale
        q_f = q_f * self.scale
        q_eaa = q_eaa * self.scale

        attn_eoa = (q_eoa @ k_eoa.transpose(-2, -1))
        attn_f = (q_f @ k_f.transpose(-2, -1))
        attn_eaa = (q_eaa @ k_eaa.transpose(-2, -1))


        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww

        attn_eoa = attn_eoa + relative_position_bias.unsqueeze(0)
        attn_f = attn_f + relative_position_bias.unsqueeze(0)
        attn_eaa = attn_eaa + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nW = mask.shape[0]
            attn_eoa = attn_eoa.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn_eoa = attn_eoa.view(-1, self.num_heads, N, N)
            attn_eoa = self.softmax(attn_eoa)

            attn_f = attn_f.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn_f = attn_f.view(-1, self.num_heads, N, N)
            attn_f = self.softmax(attn_f)

            attn_eaa = attn_eaa.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn_eaa = attn_eaa.view(-1, self.num_heads, N, N)
            attn_eaa = self.softmax(attn_eaa)
        else:
            attn_eoa = self.softmax(attn_eoa)
            attn_eaa = self.softmax(attn_eaa)
            attn_f = self.softmax(attn_f)

        #attn = self.attn_drop(attn)

        attn_all = attn_eoa + attn_eaa + attn_f

        #print(attn_all.size())
        out_x = (attn_all @ v_eoa).transpose(1, 2).reshape(B_, N, C)
        out_y = (attn_all @ v_f).transpose(1, 2).reshape(B_, N, C)
        out_z = (attn_all @ v_eaa).transpose(1, 2).reshape(B_, N, C)

        out_x = self.proj_eaa(out_x)
        out_x = self.proj_drop_eaa(out_x)

        out_y = self.proj_f(out_y)
        out_y = self.proj_drop_f(out_y)

        out_z = self.proj_eoa(out_z)
        out_z = self.proj_drop_eoa(out_z)
        return out_x, out_y, out_z
class FusionSwinTransformerBlock(nn.Module):
    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_size:
            # if window size is larger than input resolution, we don't partition windows
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"

        if self.shift_size > 0:
            attn_mask = self.calculate_mask(self.input_resolution)
        else:
            attn_mask = None

        self.register_buffer("attn_mask", attn_mask)

        ##########
        self.norm1_eoa = norm_layer(dim)
        self.norm1_eaa = norm_layer(dim)
        self.norm1_f = norm_layer(dim)
        self.attn = WindowAttention_eef(
            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

        self.norm2_eoa = norm_layer(dim)
        self.mlp_eoa = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
        self.norm2_eaa = norm_layer(dim)
        self.mlp_eaa = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
        self.norm2_f = norm_layer(dim)
        self.mlp_f = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)

    def calculate_mask(self, x_size):
        # calculate attention mask for SW-MSA
        H, W = x_size
        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
        h_slices = (slice(0, -self.window_size),
                    slice(-self.window_size, -self.shift_size),
                    slice(-self.shift_size, None))
        w_slices = (slice(0, -self.window_size),
                    slice(-self.window_size, -self.shift_size),
                    slice(-self.shift_size, None))
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask[:, h, w, :] = cnt
                cnt += 1

        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))

        return attn_mask

    def forward(self, x, y , z , x_size):
        H, W = x_size
        B, L, C = x.shape
        # assert L == H * W, "input feature has wrong size"

        x_shortcut = x
        y_shortcut = y
        z_shortcut = z

        x = self.norm1_eoa(x)
        x = x.view(B, H, W, C)
        y = self.norm1_f(y)
        y = y.view(B, H, W, C)
        z = self.norm1_eaa(z)
        z = z.view(B, H, W, C)

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
            shifted_y = torch.roll(y, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
            shifted_z = torch.roll(z, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
        else:
            shifted_x = x
            shifted_y = y
            shifted_z = z

        # partition windows
        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
        y_windows = window_partition(shifted_y, self.window_size)  # nW*B, window_size, window_size, C
        y_windows = y_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
        z_windows = window_partition(shifted_z, self.window_size)  # nW*B, window_size, window_size, C
        z_windows = z_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C

        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
        if self.input_resolution == x_size:
            x_attn_windows, y_attn_windows, z_attn_windows = self.attn(x_windows, y_windows, z_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
        else:
            x_attn_windows, y_attn_windows, z_attn_windows = self.attn(x_windows, y_windows, z_windows, mask=self.calculate_mask(x_size).to(x.device))

        # merge windows
        x_attn_windows = x_attn_windows.view(-1, self.window_size, self.window_size, C)
        y_attn_windows = y_attn_windows.view(-1, self.window_size, self.window_size, C)
        z_attn_windows = z_attn_windows.view(-1, self.window_size, self.window_size, C)

        shifted_x = window_reverse(x_attn_windows, self.window_size, H, W)  # B H' W' C
        shifted_y = window_reverse(y_attn_windows, self.window_size, H, W)  # B H' W' C
        shifted_z = window_reverse(z_attn_windows, self.window_size, H, W)  # B H' W' C


        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
            y = torch.roll(shifted_y, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
            z = torch.roll(shifted_z, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
        else:
            x = shifted_x
            y = shifted_y
            z = shifted_z

        x = x.view(B, H * W, C)
        y = y.view(B, H * W, C)
        z = z.view(B, H * W, C)

        # FFN
        x = x_shortcut + self.drop_path(x)
        y = y_shortcut + self.drop_path(y)
        z = z_shortcut + self.drop_path(z)
        #print(x.size())
        x = x + self.drop_path(self.mlp_eoa(self.norm2_eoa(x)))
        y = y + self.drop_path(self.mlp_f(self.norm2_f(y)))
        z = z + self.drop_path(self.mlp_eaa(self.norm2_eaa(z)))

        return x, y, z
<filename>EF-SAI-main/codes/Networks/submodules.py
import math
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
import functools



# Channel Attention from CBAM
class ChannelAttentionv2(nn.Module):
    def __init__(self, in_planes, in_planes2, ratio=16, kernel_size=3):
        super(ChannelAttentionv2, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.conv1 = nn.Sequential(nn.Conv2d(in_planes2*2, 1, kernel_size, padding=kernel_size//2, bias=False), nn.ReLU())
        self.conv2 = nn.Sequential(nn.Conv2d(in_planes2, 1, kernel_size, padding=kernel_size//2, bias=False), nn.ReLU())
        self.conv3 = nn.Sequential(nn.Conv2d(in_planes2, 1, kernel_size, padding=kernel_size//2, bias=False), nn.ReLU())
        #self.conv4 = nn.Sequential(nn.Conv2d(in_planes+1, 1, kernel_size, padding=kernel_size//2, bias=False), nn.ReLU())
        #self.conv5 = nn.Sequential(nn.Conv2d(in_planes+1, 1, kernel_size, padding=kernel_size // 2, bias=False),nn.ReLU())
        #self.conv6 = nn.Sequential(nn.Conv2d(in_planes+1, 1, kernel_size, padding=kernel_size // 2, bias=False),nn.ReLU())
        in_planes3 = in_planes*3+3
        self.fc = nn.Sequential(nn.Conv2d(in_planes3, in_planes3 // ratio, 1, bias=False),
                                nn.ReLU(),
                                nn.Conv2d(in_planes3 // ratio, 3, 1, bias=False))
        self.sigmoid = nn.Sigmoid()

    def forward(self, x1, x2, y1, y2, z1, z2):
        x = self.conv1(x2)
        x = torch.cat([x1, x], dim=1)
        # x = self.conv4(x)
        y = self.conv2(y2)
        y = torch.cat([y1, y], dim=1)
        # y = self.conv5(y)
        z = self.conv3(z2)
        z = torch.cat([z1, z], dim=1)
        # z = self.conv6(z)
        x = torch.cat([x, y, z], dim=1)
        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)


#Generator
class Identity(nn.Module):
    def forward(self, x):
        return x

def get_norm_layer(norm_type='instance'):
    """Return a normalization layer

    Parameters:
        norm_type (str) -- the name of the normalization layer: batch | instance | none

    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
    """
    if norm_type == 'batch':
        momentum = 0.01
        norm_layer = functools.partial(nn.BatchNorm2d,momentum = momentum,  affine=True, track_running_stats=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
    elif norm_type == 'none':
        def norm_layer(x): return Identity()
    else:
        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
    return norm_layer
class ResnetGenerator(nn.Module):
    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.

    We adapt Torch code and idea from <NAME>'s neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
    """

    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
        """Construct a Resnet-based generator

        Parameters:
            input_nc (int)      -- the number of channels in input images
            output_nc (int)     -- the number of channels in output images
            ngf (int)           -- the number of filters in the last conv layer
            norm_layer          -- normalization layer
            use_dropout (bool)  -- if use dropout layers
            n_blocks (int)      -- the number of ResNet blocks
            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
        """
        assert(n_blocks >= 0)
        super(ResnetGenerator, self).__init__()
        if type(norm_layer) == functools.partial:
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d

        model = [nn.ReflectionPad2d(3),
                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
                 norm_layer(ngf),
                 nn.ReLU(True)]

        n_downsampling = 2
        for i in range(n_downsampling):  # add downsampling layers
            mult = 2 ** i
            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
                      norm_layer(ngf * mult * 2),
                      nn.ReLU(True)]

        mult = 2 ** n_downsampling
        for i in range(n_blocks):       # add ResNet blocks

            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]

        for i in range(n_downsampling):  # add upsampling layers
            mult = 2 ** (n_downsampling - i)
            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
                                         kernel_size=3, stride=2,
                                         padding=1, output_padding=1,
                                         bias=use_bias),
                      norm_layer(int(ngf * mult / 2)),
                      nn.ReLU(True)]
        model += [nn.ReflectionPad2d(3)]
        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
        model += [nn.Tanh()]

        self.model = nn.Sequential(*model)

    def forward(self, input):
        """Standard forward"""
        return self.model(input)
class ResnetBlock(nn.Module):
    """Define a Resnet block"""

    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        """Initialize the Resnet block

        A resnet block is a conv block with skip connections
        We construct a conv block with build_conv_block function,
        and implement skip connections in <forward> function.
        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
        """
        super(ResnetBlock, self).__init__()
        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)

    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        """Construct a convolutional block.

        Parameters:
            dim (int)           -- the number of channels in the conv layer.
            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
            norm_layer          -- normalization layer
            use_dropout (bool)  -- if use dropout layers.
            use_bias (bool)     -- if the conv layer uses bias or not

        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
        """
        conv_block = []
        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)

        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
        if use_dropout:
            conv_block += [nn.Dropout(0.5)]

        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]

        return nn.Sequential(*conv_block)

    def forward(self, x):
        """Forward function (with skip connections)"""
        out = x + self.conv_block(x)  # add skip connections
        return out


class UnetGenerator(nn.Module):
    """Create a Unet-based generator"""

    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
        """Construct a Unet generator
        Parameters:
            input_nc (int)  -- the number of channels in input images
            output_nc (int) -- the number of channels in output images
            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
                                image of size 128x128 will become of size 1x1 # at the bottleneck
            ngf (int)       -- the number of filters in the last conv layer
            norm_layer      -- normalization layer

        We construct the U-Net from the innermost layer to the outermost layer.
        It is a recursive process.
        """
        super(UnetGenerator, self).__init__()
        # construct unet structure
        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer
        for i in range(num_downs - 5):          # add intermediate layers with ngf * 8 filters
            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
        # gradually reduce the number of filters from ngf * 8 to ngf
        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer

    def forward(self, input):
        """Standard forward"""
        return self.model(input)


class UnetSkipConnectionBlock(nn.Module):
    """Defines the Unet submodule with skip connection.
        X -------------------identity----------------------
        |-- downsampling -- |submodule| -- upsampling --|
    """

    def __init__(self, outer_nc, inner_nc, input_nc=None,
                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
        """Construct a Unet submodule with skip connections.

        Parameters:
            outer_nc (int) -- the number of filters in the outer conv layer
            inner_nc (int) -- the number of filters in the inner conv layer
            input_nc (int) -- the number of channels in input images/features
            submodule (UnetSkipConnectionBlock) -- previously defined submodules
            outermost (bool)    -- if this module is the outermost module
            innermost (bool)    -- if this module is the innermost module
            norm_layer          -- normalization layer
            use_dropout (bool)  -- if use dropout layers.
        """
        super(UnetSkipConnectionBlock, self).__init__()
        self.outermost = outermost
        if type(norm_layer) == functools.partial:
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d
        if input_nc is None:
            input_nc = outer_nc
        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
                             stride=2, padding=1, bias=use_bias)
        downrelu = nn.LeakyReLU(0.2, True)
        downnorm = norm_layer(inner_nc)
        uprelu = nn.ReLU(True)
        upnorm = norm_layer(outer_nc)

        if outermost:
            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
                                        kernel_size=4, stride=2,
                                        padding=1)
            down = [downconv]
            up = [uprelu, upconv, nn.Tanh()]
            model = down + [submodule] + up
        elif innermost:
            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
                                        kernel_size=4, stride=2,
                                        padding=1, bias=use_bias)
            down = [downrelu, downconv]
            up = [uprelu, upconv, upnorm]
            model = down + up
        else:
            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
                                        kernel_size=4, stride=2,
                                        padding=1, bias=use_bias)
            down = [downrelu, downconv, downnorm]
            up = [uprelu, upconv, upnorm]

            if use_dropout:
                model = down + [submodule] + up + [nn.Dropout(0.5)]
            else:
                model = down + [submodule] + up

        self.model = nn.Sequential(*model)

    def forward(self, x):
        if self.outermost:
            return self.model(x)
        else:   # add skip connections
            return torch.cat([x, self.model(x)], 1)
def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
    Parameters:
        net (network)      -- the network to be initialized
        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        gain (float)       -- scaling factor for normal, xavier and orthogonal.
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Return an initialized network.
    """
    if len(gpu_ids) > 0:
        assert(torch.cuda.is_available())
        net.to(gpu_ids[0])
        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
    init_weights(net, init_type, init_gain=init_gain)
    return net
def init_weights(net, init_type='normal', init_gain=0.02):
    """Initialize network weights.

    Parameters:
        net (network)   -- network to be initialized
        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.

    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
    work better for some applications. Feel free to try yourself.
    """
    def init_func(m):  # define the initialization function
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, 0.0, init_gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=init_gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
            init.normal_(m.weight.data, 1.0, init_gain)
            init.constant_(m.bias.data, 0.0)

    print('initialize network with %s' % init_type)
    net.apply(init_func)  # apply the initialization function <init_func>
def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
    """Create a generator

    Parameters:
        input_nc (int) -- the number of channels in input images
        output_nc (int) -- the number of channels in output images
        ngf (int) -- the number of filters in the last conv layer
        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
        norm (str) -- the name of normalization layers used in the network: batch | instance | none
        use_dropout (bool) -- if use dropout layers.
        init_type (str)    -- the name of our initialization method.
        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Returns a generator

    Our current implementation provides two types of generators:
        U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
        The original U-Net paper: https://arxiv.org/abs/1505.04597

        Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
        Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
        We adapt Torch code from <NAME>'s neural style transfer project (https://github.com/jcjohnson/fast-neural-style).


    The generator has been initialized by <init_net>. It uses RELU for non-linearity.
    """
    net = None
    norm_layer = get_norm_layer(norm_type=norm)

    if netG == 'resnet_9blocks':
        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
    elif netG == 'resnet_6blocks':
        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
    elif netG == 'unet_128':
        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
    elif netG == 'unet_256':
        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
    else:
        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
    return init_net(net, init_type, init_gain, gpu_ids)









#Fusion Swin Transformer
class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


def window_partition(x, window_size):
    """
    Args:
        x: (B, H, W, C)
        window_size (int): window size

    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    return windows


def window_reverse(windows, window_size, H, W):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image

    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x

class PatchEmbed(nn.Module):
    r""" Image to Patch Embedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None, firstembed=False):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]
        self.in_chans = in_chans
        self.embed_dim = embed_dim
        self.firstembed = firstembed
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None
        if(firstembed):

            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)  # 输入嵌入投影
    def forward(self, x):

        if(self.firstembed):
            #x = self.proj(x).flatten(2).transpose(1, 2)  # shape = (B, P_h*P_w, C)
            x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
        else:
            x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
        if self.norm is not None:
            x = self.norm(x)
        return x

    def flops(self):
        flops = 0
        H, W = self.img_size
        if self.norm is not None:
            flops += H * W * self.embed_dim
        return flops


class PatchUnEmbed(nn.Module):
    r""" Image to Patch Unembedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None, firstembed = False):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim
        self.firstembed = firstembed
        if(firstembed):
            self.proj = nn.Conv2d(embed_dim, in_chans, kernel_size=patch_size, stride=patch_size)  # 输入嵌入投影
    def forward(self, x, x_size):
        B, HW, C = x.shape
        if(self.firstembed):
            #self.proj(x).transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
            x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
        else:
            x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
        return x

    def flops(self):
        flops = 0
        return flops
class WindowAttention_eef(nn.Module):
    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.

    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
    """

    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):

        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wh, Ww
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        # define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH


        # get pair-wise relative position index for each token inside the window
        coords_h = torch.arange(self.window_size[0])
        coords_w = torch.arange(self.window_size[1])
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
        self.register_buffer("relative_position_index", relative_position_index)

        #self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.qkv_eoa = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.qkv_eaa = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.qkv_f = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)

        # self.proj = nn.Linear(dim, dim)
        # self.proj_drop = nn.Dropout(proj_drop)

        self.proj_eoa = nn.Linear(dim, dim)
        self.proj_drop_eoa = nn.Dropout(proj_drop)

        self.proj_eaa = nn.Linear(dim, dim)
        self.proj_drop_eaa = nn.Dropout(proj_drop)

        self.proj_f = nn.Linear(dim, dim)
        self.proj_drop_f = nn.Dropout(proj_drop)

        trunc_normal_(self.relative_position_bias_table, std=.02)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x,y,z, mask=None): #x:eoa y:frame z:eaa
        """
        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        B_, N, C = x.shape
        #print(x.size())
        # qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        # q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)

        qkv_eoa = self.qkv_eoa(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q_eoa, k_eoa, v_eoa = qkv_eoa[0], qkv_eoa[1], qkv_eoa[2]  # make torchscript happy (cannot use tensor as tuple)
        #print(q_eoa.size())
        #print(k_eoa.size())
        qkv_f = self.qkv_f(y).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q_f, k_f, v_f = qkv_f[0], qkv_f[1], qkv_f[2]  # make torchscript happy (cannot use tensor as tuple)

        qkv_eaa = self.qkv_eaa(z).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q_eaa, k_eaa, v_eaa = qkv_eaa[0], qkv_eaa[1], qkv_eaa[2]  # make torchscript happy (cannot use tensor as tuple)

        q_eoa = q_eoa * self.scale
        q_f = q_f * self.scale
        q_eaa = q_eaa * self.scale

        attn_eoa = (q_eoa @ k_eoa.transpose(-2, -1))
        attn_f = (q_f @ k_f.transpose(-2, -1))
        attn_eaa = (q_eaa @ k_eaa.transpose(-2, -1))


        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww

        attn_eoa = attn_eoa + relative_position_bias.unsqueeze(0)
        attn_f = attn_f + relative_position_bias.unsqueeze(0)
        attn_eaa = attn_eaa + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nW = mask.shape[0]
            attn_eoa = attn_eoa.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn_eoa = attn_eoa.view(-1, self.num_heads, N, N)
            attn_eoa = self.softmax(attn_eoa)

            attn_f = attn_f.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn_f = attn_f.view(-1, self.num_heads, N, N)
            attn_f = self.softmax(attn_f)

            attn_eaa = attn_eaa.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn_eaa = attn_eaa.view(-1, self.num_heads, N, N)
            attn_eaa = self.softmax(attn_eaa)
        else:
            attn_eoa = self.softmax(attn_eoa)
            attn_eaa = self.softmax(attn_eaa)
            attn_f = self.softmax(attn_f)

        #attn = self.attn_drop(attn)

        attn_all = attn_eoa + attn_eaa + attn_f

        #print(attn_all.size())
        out_x = (attn_all @ v_eoa).transpose(1, 2).reshape(B_, N, C)
        out_y = (attn_all @ v_f).transpose(1, 2).reshape(B_, N, C)
        out_z = (attn_all @ v_eaa).transpose(1, 2).reshape(B_, N, C)

        out_x = self.proj_eaa(out_x)
        out_x = self.proj_drop_eaa(out_x)

        out_y = self.proj_f(out_y)
        out_y = self.proj_drop_f(out_y)

        out_z = self.proj_eoa(out_z)
        out_z = self.proj_drop_eoa(out_z)
        return out_x, out_y, out_z
class FusionSwinTransformerBlock(nn.Module):
    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_size:
            # if window size is larger than input resolution, we don't partition windows
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"

        if self.shift_size > 0:
            attn_mask = self.calculate_mask(self.input_resolution)
        else:
            attn_mask = None

        self.register_buffer("attn_mask", attn_mask)

        ##########
        self.norm1_eoa = norm_layer(dim)
        self.norm1_eaa = norm_layer(dim)
        self.norm1_f = norm_layer(dim)
        self.attn = WindowAttention_eef(
            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

        self.norm2_eoa = norm_layer(dim)
        self.mlp_eoa = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
        self.norm2_eaa = norm_layer(dim)
        self.mlp_eaa = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
        self.norm2_f = norm_layer(dim)
        self.mlp_f = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)

    def calculate_mask(self, x_size):
        # calculate attention mask for SW-MSA
        H, W = x_size
        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
        h_slices = (slice(0, -self.window_size),
                    slice(-self.window_size, -self.shift_size),
                    slice(-self.shift_size, None))
        w_slices = (slice(0, -self.window_size),
                    slice(-self.window_size, -self.shift_size),
                    slice(-self.shift_size, None))
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask[:, h, w, :] = cnt
                cnt += 1

        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))

        return attn_mask

    def forward(self, x, y , z , x_size):
        H, W = x_size
        B, L, C = x.shape
        # assert L == H * W, "input feature has wrong size"

        x_shortcut = x
        y_shortcut = y
        z_shortcut = z

        x = self.norm1_eoa(x)
        x = x.view(B, H, W, C)
        y = self.norm1_f(y)
        y = y.view(B, H, W, C)
        z = self.norm1_eaa(z)
        z = z.view(B, H, W, C)

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
            shifted_y = torch.roll(y, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
            shifted_z = torch.roll(z, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
        else:
            shifted_x = x
            shifted_y = y
            shifted_z = z

        # partition windows
        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
        y_windows = window_partition(shifted_y, self.window_size)  # nW*B, window_size, window_size, C
        y_windows = y_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
        z_windows = window_partition(shifted_z, self.window_size)  # nW*B, window_size, window_size, C
        z_windows = z_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C

        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
        if self.input_resolution == x_size:
            x_attn_windows, y_attn_windows, z_attn_windows = self.attn(x_windows, y_windows, z_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
        else:
            x_attn_windows, y_attn_windows, z_attn_windows = self.attn(x_windows, y_windows, z_windows, mask=self.calculate_mask(x_size).to(x.device))

        # merge windows
        x_attn_windows = x_attn_windows.view(-1, self.window_size, self.window_size, C)
        y_attn_windows = y_attn_windows.view(-1, self.window_size, self.window_size, C)
        z_attn_windows = z_attn_windows.view(-1, self.window_size, self.window_size, C)

        shifted_x = window_reverse(x_attn_windows, self.window_size, H, W)  # B H' W' C
        shifted_y = window_reverse(y_attn_windows, self.window_size, H, W)  # B H' W' C
        shifted_z = window_reverse(z_attn_windows, self.window_size, H, W)  # B H' W' C


        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
            y = torch.roll(shifted_y, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
            z = torch.roll(shifted_z, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
        else:
            x = shifted_x
            y = shifted_y
            z = shifted_z

        x = x.view(B, H * W, C)
        y = y.view(B, H * W, C)
        z = z.view(B, H * W, C)

        # FFN
        x = x_shortcut + self.drop_path(x)
        y = y_shortcut + self.drop_path(y)
        z = z_shortcut + self.drop_path(z)
        #print(x.size())
        x = x + self.drop_path(self.mlp_eoa(self.norm2_eoa(x)))
        y = y + self.drop_path(self.mlp_f(self.norm2_f(y)))
        z = z + self.drop_path(self.mlp_eaa(self.norm2_eaa(z)))

        return x, y, z
# Channel Attention from CBAM #self.conv4 = nn.Sequential(nn.Conv2d(in_planes+1, 1, kernel_size, padding=kernel_size//2, bias=False), nn.ReLU()) #self.conv5 = nn.Sequential(nn.Conv2d(in_planes+1, 1, kernel_size, padding=kernel_size // 2, bias=False),nn.ReLU()) #self.conv6 = nn.Sequential(nn.Conv2d(in_planes+1, 1, kernel_size, padding=kernel_size // 2, bias=False),nn.ReLU()) # x = self.conv4(x) # y = self.conv5(y) # z = self.conv6(z) #Generator Return a normalization layer      Parameters:         norm_type (str) -- the name of the normalization layer: batch | instance | none      For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).     For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics. Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.      We adapt Torch code and idea from <NAME>'s neural style transfer project(https://github.com/jcjohnson/fast-neural-style) Construct a Resnet-based generator          Parameters:             input_nc (int)      -- the number of channels in input images             output_nc (int)     -- the number of channels in output images             ngf (int)           -- the number of filters in the last conv layer             norm_layer          -- normalization layer             use_dropout (bool)  -- if use dropout layers             n_blocks (int)      -- the number of ResNet blocks             padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero # add downsampling layers # add ResNet blocks # add upsampling layers Standard forward Define a Resnet block Initialize the Resnet block          A resnet block is a conv block with skip connections         We construct a conv block with build_conv_block function,         and implement skip connections in <forward> function.         Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf Construct a convolutional block.          Parameters:             dim (int)           -- the number of channels in the conv layer.             padding_type (str)  -- the name of padding layer: reflect | replicate | zero             norm_layer          -- normalization layer             use_dropout (bool)  -- if use dropout layers.             use_bias (bool)     -- if the conv layer uses bias or not          Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU)) Forward function (with skip connections) # add skip connections Create a Unet-based generator Construct a Unet generator         Parameters:             input_nc (int)  -- the number of channels in input images             output_nc (int) -- the number of channels in output images             num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,                                 image of size 128x128 will become of size 1x1 # at the bottleneck             ngf (int)       -- the number of filters in the last conv layer             norm_layer      -- normalization layer          We construct the U-Net from the innermost layer to the outermost layer.         It is a recursive process. # construct unet structure # add the innermost layer # add intermediate layers with ngf * 8 filters # gradually reduce the number of filters from ngf * 8 to ngf # add the outermost layer Standard forward Defines the Unet submodule with skip connection.         X -------------------identity----------------------         |-- downsampling -- |submodule| -- upsampling --| Construct a Unet submodule with skip connections.          Parameters:             outer_nc (int) -- the number of filters in the outer conv layer             inner_nc (int) -- the number of filters in the inner conv layer             input_nc (int) -- the number of channels in input images/features             submodule (UnetSkipConnectionBlock) -- previously defined submodules             outermost (bool)    -- if this module is the outermost module             innermost (bool)    -- if this module is the innermost module             norm_layer          -- normalization layer             use_dropout (bool)  -- if use dropout layers. # add skip connections Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights     Parameters:         net (network)      -- the network to be initialized         init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal         gain (float)       -- scaling factor for normal, xavier and orthogonal.         gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2      Return an initialized network. # multi-GPUs Initialize network weights.      Parameters:         net (network)   -- network to be initialized         init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal         init_gain (float)    -- scaling factor for normal, xavier and orthogonal.      We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might     work better for some applications. Feel free to try yourself. # define the initialization function # BatchNorm Layer's weight is not a matrix; only normal distribution applies. # apply the initialization function <init_func> Create a generator      Parameters:         input_nc (int) -- the number of channels in input images         output_nc (int) -- the number of channels in output images         ngf (int) -- the number of filters in the last conv layer         netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128         norm (str) -- the name of normalization layers used in the network: batch | instance | none         use_dropout (bool) -- if use dropout layers.         init_type (str)    -- the name of our initialization method.         init_gain (float)  -- scaling factor for normal, xavier and orthogonal.         gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2      Returns a generator      Our current implementation provides two types of generators:         U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)         The original U-Net paper: https://arxiv.org/abs/1505.04597          Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)         Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.         We adapt Torch code from <NAME>'s neural style transfer project (https://github.com/jcjohnson/fast-neural-style).       The generator has been initialized by <init_net>. It uses RELU for non-linearity. #Fusion Swin Transformer Args:         x: (B, H, W, C)         window_size (int): window size      Returns:         windows: (num_windows*B, window_size, window_size, C) Args:         windows: (num_windows*B, window_size, window_size, C)         window_size (int): Window size         H (int): Height of image         W (int): Width of image      Returns:         x: (B, H, W, C) Image to Patch Embedding      Args:         img_size (int): Image size.  Default: 224.         patch_size (int): Patch token size. Default: 4.         in_chans (int): Number of input image channels. Default: 3.         embed_dim (int): Number of linear projection output channels. Default: 96.         norm_layer (nn.Module, optional): Normalization layer. Default: None # 输入嵌入投影 #x = self.proj(x).flatten(2).transpose(1, 2)  # shape = (B, P_h*P_w, C) # B Ph*Pw C # B Ph*Pw C Image to Patch Unembedding      Args:         img_size (int): Image size.  Default: 224.         patch_size (int): Patch token size. Default: 4.         in_chans (int): Number of input image channels. Default: 3.         embed_dim (int): Number of linear projection output channels. Default: 96.         norm_layer (nn.Module, optional): Normalization layer. Default: None # 输入嵌入投影 #self.proj(x).transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C # B Ph*Pw C # B Ph*Pw C Window based multi-head self attention (W-MSA) module with relative position bias.     It supports both of shifted and non-shifted window.      Args:         dim (int): Number of input channels.         window_size (tuple[int]): The height and width of the window.         num_heads (int): Number of attention heads.         qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True         qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set         attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0         proj_drop (float, optional): Dropout ratio of output. Default: 0.0 # Wh, Ww # define a parameter table of relative position bias # 2*Wh-1 * 2*Ww-1, nH # get pair-wise relative position index for each token inside the window # 2, Wh, Ww # 2, Wh*Ww # 2, Wh*Ww, Wh*Ww # Wh*Ww, Wh*Ww, 2 # shift to start from 0 # Wh*Ww, Wh*Ww #self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) # self.proj = nn.Linear(dim, dim) # self.proj_drop = nn.Dropout(proj_drop) #x:eoa y:frame z:eaa Args:             x: input features with shape of (num_windows*B, N, C)             mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None #print(x.size()) # qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) # q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple) # make torchscript happy (cannot use tensor as tuple) #print(q_eoa.size()) #print(k_eoa.size()) # make torchscript happy (cannot use tensor as tuple) # make torchscript happy (cannot use tensor as tuple) # Wh*Ww,Wh*Ww,nH # nH, Wh*Ww, Wh*Ww #attn = self.attn_drop(attn) #print(attn_all.size()) # if window size is larger than input resolution, we don't partition windows ########## # calculate attention mask for SW-MSA # 1 H W 1 # nW, window_size, window_size, 1 # assert L == H * W, "input feature has wrong size" # cyclic shift # partition windows # nW*B, window_size, window_size, C # nW*B, window_size*window_size, C # nW*B, window_size, window_size, C # nW*B, window_size*window_size, C # nW*B, window_size, window_size, C # nW*B, window_size*window_size, C # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size # nW*B, window_size*window_size, C # merge windows # B H' W' C # B H' W' C # B H' W' C # reverse cyclic shift # FFN #print(x.size())