import torch import torch.nn as nn import torch.nn.functional as F from torchvision import transforms import numpy as np from .SSN import Conv, Conv2DMod, Decoder, Up from .attention import AttentionBlock from .blocks import ResBlock, Res_Type, get_activation class Attention_Encoder(nn.Module): def __init__(self, in_channels=3, mid_act='gelu', dropout=0.0, num_heads=8, resnet=True): super(Attention_Encoder, self).__init__() self.in_conv = Conv(in_channels, 32-in_channels, stride=1, activation=mid_act, resnet=resnet) self.down_32_64 = Conv(32, 64, stride=2, activation=mid_act, resnet=resnet) self.down_64_64_1 = Conv(64, 64, activation=mid_act, resnet=resnet) self.down_64_128 = Conv(64, 128, stride=2, activation=mid_act, resnet=resnet) self.down_128_128_1 = Conv(128, 128, activation=mid_act, resnet=resnet) self.down_128_256 = Conv(128, 256, stride=2, activation=mid_act, resnet=resnet) self.down_256_256_1 = Conv(256, 256, activation=mid_act, resnet=resnet) self.down_256_256_1_attn = AttentionBlock(256, num_heads) self.down_256_512 = Conv(256, 512, stride=2, activation=mid_act, resnet=resnet) self.down_512_512_1 = Conv(512, 512, activation=mid_act, resnet=resnet) self.down_512_512_1_attn = AttentionBlock(512, num_heads) self.down_512_512_2 = Conv(512, 512, activation=mid_act, resnet=resnet) self.down_512_512_2_attn = AttentionBlock(512, num_heads) self.down_512_512_3 = Conv(512, 512, activation=mid_act, resnet=resnet) self.down_512_512_3_attn = AttentionBlock(512, num_heads) def forward(self, x): x1 = self.in_conv(x) # 32 x 256 x 256 x1 = torch.cat((x, x1), dim=1) x2 = self.down_32_64(x1) x3 = self.down_64_64_1(x2) x4 = self.down_64_128(x3) x5 = self.down_128_128_1(x4) x6 = self.down_128_256(x5) x7 = self.down_256_256_1(x6) x7 = self.down_256_256_1_attn(x7) x8 = self.down_256_512(x7) x9 = self.down_512_512_1(x8) x9 = self.down_512_512_1_attn(x9) x10 = self.down_512_512_2(x9) x10 = self.down_512_512_2_attn(x10) x11 = self.down_512_512_3(x10) x11 = self.down_512_512_3_attn(x11) return x11, x10, x9, x8, x7, x6, x5, x4, x3, x2, x1 class Attention_Decoder(nn.Module): def __init__(self, out_channels=3, mid_act='gelu', out_act='sigmoid', resnet = True, num_heads=8): super(Attention_Decoder, self).__init__() input_channel = 512 fea_dim = 100 self.to_style1 = nn.Linear(in_features=fea_dim, out_features=input_channel) self.up_16_16_1 = Conv(input_channel, 256, activation=mid_act, style=False, resnet=resnet) self.up_16_16_1_attn = AttentionBlock(256, num_heads=num_heads) self.up_16_16_2 = Conv(768, 512, activation=mid_act, resnet=resnet) self.up_16_16_2_attn = AttentionBlock(512, num_heads=num_heads) self.up_16_16_3 = Conv(1024, 512, activation=mid_act, resnet=resnet) self.up_16_16_3_attn = AttentionBlock(512, num_heads=num_heads) self.up_16_32 = Up(1024, 256, activation=mid_act, resnet=resnet) self.to_style2 = nn.Linear(in_features=fea_dim, out_features=512) self.up_32_32_1 = Conv(512, 256, activation=mid_act, style=False, resnet=resnet) self.up_32_32_1_attn = AttentionBlock(256, num_heads=num_heads) self.up_32_64 = Up(512, 128, activation=mid_act, resnet=resnet) self.to_style3 = nn.Linear(in_features=fea_dim, out_features=256) self.up_64_64_1 = Conv(256, 128, activation=mid_act, style=False, resnet=resnet) self.up_64_128 = Up(256, 64, activation=mid_act, resnet=resnet) self.to_style4 = nn.Linear(in_features=fea_dim, out_features=128) self.up_128_128_1 = Conv(128, 64, activation=mid_act, style=False, resnet=resnet) self.up_128_256 = Up(128, 32, activation=mid_act, resnet=resnet) self.out_conv = Conv(64, out_channels, activation=out_act) self.out_act = get_activation(out_act) def forward(self, x): x11, x10, x9, x8, x7, x6, x5, x4, x3, x2, x1 = x y = self.up_16_16_1(x11) # 256 x 16 x 16 y = self.up_16_16_1_attn(y) y = torch.cat((x10, y), dim=1) # 768 x 16 x 16 y = self.up_16_16_2(y, y) # 512 x 16 x 16 y = self.up_16_16_2_attn(y) y = torch.cat((x9, y), dim=1) # 1024 x 16 x 16 y = self.up_16_16_3(y, y) # 512 x 16 x 16 y = self.up_16_16_3_attn(y) y = torch.cat((x8, y), dim=1) # 1024 x 16 x 16 y = self.up_16_32(y, y) # 256 x 32 x 32 y = torch.cat((x7, y), dim=1) y = self.up_32_32_1(y) # 256 x 32 x 32 y = self.up_32_32_1_attn(y) y = torch.cat((x6, y), dim=1) y = self.up_32_64(y, y) y = torch.cat((x5, y), dim=1) y = self.up_64_64_1(y) # 128 x 64 x 64 y = torch.cat((x4, y), dim=1) y = self.up_64_128(y, y) y = torch.cat((x3, y), dim=1) y = self.up_128_128_1(y) # 64 x 128 x 128 y = torch.cat((x2, y), dim=1) y = self.up_128_256(y, y) # 32 x 256 x 256 y = torch.cat((x1, y), dim=1) y = self.out_conv(y, y) # 3 x 256 x 256 y = self.out_act(y) return y class Attention_Unet(nn.Module): def __init__(self, in_channels, out_channels, num_heads=8, resnet=True, mid_act='gelu', out_act='gelu'): super(Attention_Unet, self).__init__() self.encoder = Attention_Encoder(in_channels, mid_act, num_heads, resnet) self.decoder = Attention_Decoder(out_channels, mid_act, out_act, resnet) def forward(self, x): latent = self.encoder(x) pred = self.decoder(latent) return pred if __name__ == '__main__': test_input = torch.randn(5, 1, 256, 256) style = torch.randn(5, 100) model = SSN_v1(1, 1, mid_act='gelu', out_act='gelu', resnet=True) test_out = model(test_input, style) print('Ouptut shape: ', test_out.shape)