| 105 | return self.sigmoid(x) |
| 106 | |
| 107 | class CASCADE_Cat(nn.Module): |
| 108 | def __init__(self, channels=[512,320,128,64]): |
| 109 | super(CASCADE_Cat,self).__init__() |
| 110 | |
| 111 | self.Conv_1x1 = nn.Conv2d(channels[0],channels[0],kernel_size=1,stride=1,padding=0) |
| 112 | self.ConvBlock4 = conv_block(ch_in=channels[0], ch_out=channels[0]) |
| 113 | |
| 114 | self.Up3 = up_conv(ch_in=channels[0],ch_out=channels[1]) |
| 115 | self.AG3 = Attention_block(F_g=channels[1],F_l=channels[1],F_int=channels[2]) |
| 116 | self.ConvBlock3 = conv_block(ch_in=2*channels[1], ch_out=channels[1]) |
| 117 | |
| 118 | self.Up2 = up_conv(ch_in=channels[1],ch_out=channels[2]) |
| 119 | self.AG2 = Attention_block(F_g=channels[2],F_l=channels[2],F_int=channels[3]) |
| 120 | self.ConvBlock2 = conv_block(ch_in=2*channels[2], ch_out=channels[2]) |
| 121 | |
| 122 | self.Up1 = up_conv(ch_in=channels[2],ch_out=channels[3]) |
| 123 | self.AG1 = Attention_block(F_g=channels[3],F_l=channels[3],F_int=int(channels[3]/2)) |
| 124 | self.ConvBlock1 = conv_block(ch_in=2*channels[3], ch_out=channels[3]) |
| 125 | |
| 126 | self.CA4 = ChannelAttention(channels[0]) |
| 127 | self.CA3 = ChannelAttention(2*channels[1]) |
| 128 | self.CA2 = ChannelAttention(2*channels[2]) |
| 129 | self.CA1 = ChannelAttention(2*channels[3]) |
| 130 | |
| 131 | self.SA = SpatialAttention() |
| 132 | |
| 133 | def forward(self,x, skips): |
| 134 | |
| 135 | d4 = self.Conv_1x1(x) |
| 136 | |
| 137 | # CAM4 |
| 138 | d4 = self.CA4(d4)*d4 |
| 139 | d4 = self.SA(d4)*d4 |
| 140 | d4 = self.ConvBlock4(d4) |
| 141 | |
| 142 | # upconv3 |
| 143 | d3 = self.Up3(d4) |
| 144 | |
| 145 | # AG3 |
| 146 | x3 = self.AG3(g=d3,x=skips[0]) |
| 147 | |
| 148 | # Concat 3 |
| 149 | d3 = torch.cat((x3,d3),dim=1) |
| 150 | |
| 151 | # CAM3 |
| 152 | d3 = self.CA3(d3)*d3 |
| 153 | d3 = self.SA(d3)*d3 |
| 154 | d3 = self.ConvBlock3(d3) |
| 155 | |
| 156 | # upconv2 |
| 157 | d2 = self.Up2(d3) |
| 158 | |
| 159 | # AG2 |
| 160 | x2 = self.AG2(g=d2,x=skips[1]) |
| 161 | |
| 162 | # Concat 2 |
| 163 | d2 = torch.cat((x2,d2),dim=1) |
| 164 | |