Spaces:
Running
Running
| import torch.nn as nn | |
| import torch | |
| class ResidualConv(nn.Module): | |
| def __init__(self, input_dim, output_dim, stride, padding): | |
| super(ResidualConv, self).__init__() | |
| self.conv_block = nn.Sequential( | |
| nn.BatchNorm2d(input_dim), | |
| nn.ReLU(), | |
| nn.Conv2d( | |
| input_dim, output_dim, kernel_size=3, stride=stride, padding=padding | |
| ), | |
| nn.BatchNorm2d(output_dim), | |
| nn.ReLU(), | |
| nn.Conv2d(output_dim, output_dim, kernel_size=3, padding=1), | |
| ) | |
| self.conv_skip = nn.Sequential( | |
| nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=stride, padding=1), | |
| nn.BatchNorm2d(output_dim), | |
| ) | |
| def forward(self, x): | |
| return self.conv_block(x) + self.conv_skip(x) | |
| class Upsample(nn.Module): | |
| def __init__(self, input_dim, output_dim, kernel, stride): | |
| super(Upsample, self).__init__() | |
| self.upsample = nn.ConvTranspose2d( | |
| input_dim, output_dim, kernel_size=kernel, stride=stride | |
| ) | |
| def forward(self, x): | |
| return self.upsample(x) | |
| class Squeeze_Excite_Block(nn.Module): | |
| def __init__(self, channel, reduction=16): | |
| super(Squeeze_Excite_Block, self).__init__() | |
| self.avg_pool = nn.AdaptiveAvgPool2d(1) | |
| self.fc = nn.Sequential( | |
| nn.Linear(channel, channel // reduction, bias=False), | |
| nn.ReLU(inplace=True), | |
| nn.Linear(channel // reduction, channel, bias=False), | |
| nn.Sigmoid(), | |
| ) | |
| def forward(self, x): | |
| b, c, _, _ = x.size() | |
| y = self.avg_pool(x).view(b, c) | |
| y = self.fc(y).view(b, c, 1, 1) | |
| return x * y.expand_as(x) | |
| class ASPP(nn.Module): | |
| def __init__(self, in_dims, out_dims, rate=[6, 12, 18]): | |
| super(ASPP, self).__init__() | |
| self.aspp_block1 = nn.Sequential( | |
| nn.Conv2d( | |
| in_dims, out_dims, 3, stride=1, padding=rate[0], dilation=rate[0] | |
| ), | |
| nn.ReLU(inplace=True), | |
| nn.BatchNorm2d(out_dims), | |
| ) | |
| self.aspp_block2 = nn.Sequential( | |
| nn.Conv2d( | |
| in_dims, out_dims, 3, stride=1, padding=rate[1], dilation=rate[1] | |
| ), | |
| nn.ReLU(inplace=True), | |
| nn.BatchNorm2d(out_dims), | |
| ) | |
| self.aspp_block3 = nn.Sequential( | |
| nn.Conv2d( | |
| in_dims, out_dims, 3, stride=1, padding=rate[2], dilation=rate[2] | |
| ), | |
| nn.ReLU(inplace=True), | |
| nn.BatchNorm2d(out_dims), | |
| ) | |
| self.output = nn.Conv2d(len(rate) * out_dims, out_dims, 1) | |
| self._init_weights() | |
| def forward(self, x): | |
| x1 = self.aspp_block1(x) | |
| x2 = self.aspp_block2(x) | |
| x3 = self.aspp_block3(x) | |
| out = torch.cat([x1, x2, x3], dim=1) | |
| return self.output(out) | |
| def _init_weights(self): | |
| for m in self.modules(): | |
| if isinstance(m, nn.Conv2d): | |
| nn.init.kaiming_normal_(m.weight) | |
| elif isinstance(m, nn.BatchNorm2d): | |
| m.weight.data.fill_(1) | |
| m.bias.data.zero_() | |
| class Upsample_(nn.Module): | |
| def __init__(self, scale=2): | |
| super(Upsample_, self).__init__() | |
| self.upsample = nn.Upsample(mode="bilinear", scale_factor=scale) | |
| def forward(self, x): | |
| return self.upsample(x) | |
| class AttentionBlock(nn.Module): | |
| def __init__(self, input_encoder, input_decoder, output_dim): | |
| super(AttentionBlock, self).__init__() | |
| self.conv_encoder = nn.Sequential( | |
| nn.BatchNorm2d(input_encoder), | |
| nn.ReLU(), | |
| nn.Conv2d(input_encoder, output_dim, 3, padding=1), | |
| nn.MaxPool2d(2, 2), | |
| ) | |
| self.conv_decoder = nn.Sequential( | |
| nn.BatchNorm2d(input_decoder), | |
| nn.ReLU(), | |
| nn.Conv2d(input_decoder, output_dim, 3, padding=1), | |
| ) | |
| self.conv_attn = nn.Sequential( | |
| nn.BatchNorm2d(output_dim), | |
| nn.ReLU(), | |
| nn.Conv2d(output_dim, 1, 1), | |
| ) | |
| def forward(self, x1, x2): | |
| out = self.conv_encoder(x1) + self.conv_decoder(x2) | |
| out = self.conv_attn(out) | |
| return out * x2 |