|
|
|
|
| from torch import nn |
| import torch.nn.functional as F |
| import torch |
|
|
|
|
| class ResBlock2d(nn.Module): |
| def __init__(self, in_features, kernel_size, padding): |
| super(ResBlock2d, self).__init__() |
| self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, |
| padding=padding) |
| self.conv2 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, |
| padding=padding) |
| |
| self.norm1 = nn.Conv2d( |
| in_channels=in_features, out_channels=in_features, kernel_size=1) |
| self.norm2 = nn.Conv2d( |
| in_channels=in_features, out_channels=in_features, kernel_size=1) |
|
|
| def forward(self, x): |
| out = self.norm1(x) |
| out = F.relu(out, inplace=True) |
| out = self.conv1(out) |
| out = self.norm2(out) |
| out = F.relu(out, inplace=True) |
| out = self.conv2(out) |
| out += x |
| return out |
|
|
|
|
| class RGBADecoderNet(nn.Module): |
| def __init__(self, c=64, out_planes=4, num_bottleneck_blocks=1): |
| super(RGBADecoderNet, self).__init__() |
| self.conv_rgba = nn.Sequential(nn.Conv2d(c, out_planes, kernel_size=3, stride=1, |
| padding=1, dilation=1, bias=True)) |
| self.bottleneck = torch.nn.Sequential() |
| for i in range(num_bottleneck_blocks): |
| self.bottleneck.add_module( |
| 'r' + str(i), ResBlock2d(c, kernel_size=(3, 3), padding=(1, 1))) |
|
|
| def forward(self, features_weighted_mask_atfeaturesscale_list=[]): |
| return torch.sigmoid(self.conv_rgba(self.bottleneck(features_weighted_mask_atfeaturesscale_list.pop(0)))) |
|
|
