Spaces:
Runtime error
Runtime error
| """Pytorch implementation of AESc model architecture""" | |
| __author__ = "Jonas Rabensteiner" | |
| import torch | |
| from torch import nn | |
| from torchvision.transforms.functional import crop | |
| # create model | |
| def resize_layer(conv, deconv): | |
| """If needed resize the feature map of the deconvolution to the size of the corresponding feature map from the encoder part. | |
| Args: | |
| conv (tensor): feature map of the encoder part | |
| deconv (tensor): corresponding feature map of the decoder part | |
| Returns: | |
| tensor: resized decoder feature map | |
| """ | |
| height = deconv.shape[2] | |
| width= deconv.shape[3] | |
| if deconv.shape[2] > conv.shape[2]: | |
| deconv = crop(deconv, top=0, left=0, height=height-1, width=width) | |
| #nn.Cropping2D(cropping=((0, 1), (0, 0)))(deconv) | |
| if deconv.shape[3] > conv.shape[3]: | |
| deconv = crop(deconv, top=0, left=0, height=height, width=width-1) | |
| #nn.Cropping2D(cropping=((0, 0), (0, 1)))(deconv) | |
| return deconv | |
| class conv_block(nn.Module): | |
| """Layer "block" of 2D (de)convolution, batch normalization and activation""" | |
| def __init__(self, in_c, out_c, kernel_size=5, stride=True, activation=nn.LeakyReLU()): | |
| super().__init__() | |
| if not stride: | |
| self.conv = nn.Conv2d(in_c, out_c, kernel_size=kernel_size, padding="same") | |
| else: | |
| self.conv = nn.Conv2d(in_c, out_c, kernel_size=kernel_size, stride=2, padding=2) | |
| self.bn = nn.BatchNorm2d(out_c) | |
| self.activation = activation | |
| def forward(self, inputs): | |
| x = self.conv(inputs) | |
| x = self.bn(x) | |
| x = self.activation(x) | |
| return x | |
| class encoder_block(nn.Module): | |
| """Layer "block" of conv_block and dropout, which corresponds to one "step" of the encoder""" | |
| def __init__(self, in_c, out_c, kernel_size=5, activation=nn.LeakyReLU(), dropout_rate=0.0): | |
| super().__init__() | |
| self.conv = conv_block(in_c, out_c, kernel_size, activation) | |
| self.dropout = nn.Dropout(dropout_rate) | |
| def forward(self, inputs): | |
| x = self.conv(inputs) | |
| x = self.dropout(x) | |
| return x | |
| class decoder_block(nn.Module): | |
| """Layer "block" of upsampling, conv_block, dropout and skip connections which corresponds to one "step" of the decoder""" | |
| def __init__(self, in_c, out_c, kernel_size=5, activation=nn.LeakyReLU(),dropout_rate=0.0, skip_connections=True): | |
| super().__init__() | |
| self.up = nn.UpsamplingNearest2d(scale_factor=2) | |
| self.conv = conv_block(in_c, out_c, kernel_size, stride=False) | |
| self.dropout = nn.Dropout(dropout_rate) | |
| self.skip_connections = skip_connections | |
| def forward(self, inputs, skip): | |
| #print(inputs.shape) | |
| x = self.conv(inputs) | |
| #print("skip shape", skip.shape) | |
| x = resize_layer(skip, x) | |
| #print("x resized shape", x.shape) | |
| if self.skip_connections: | |
| x = skip + x | |
| x = self.up(x) | |
| x = self.dropout(x) | |
| #print(x.shape) | |
| return x | |
| class AESc(nn.Module): | |
| """Autoencoder with skip connections (AESc) according to the original paper by Anne-Sophie Collin.""" | |
| def __init__(self, cmap = "rgb", kernel_size=5, activation=nn.LeakyReLU(), dropout_rate=0.0): | |
| """Instantiates the model layers | |
| Args: | |
| cmap (str, optional): color map to use for the model. Either "rgb" or "gray". Defaults to "rgb". | |
| kernel_size (int, optional): kernel size. Defaults to 5. | |
| activation (nn.Module, optional): activation function. Defaults to nn.LeakyReLU(). | |
| dropout_rate (float, optional): dropout rate. Defaults to 0.0. | |
| """ | |
| super().__init__() | |
| #Encoder | |
| if cmap == "gray": | |
| self.e1 = encoder_block(1, 16, kernel_size, activation, dropout_rate) | |
| else: | |
| self.e1 = encoder_block(3, 16, kernel_size, activation, dropout_rate) | |
| self.e2 = encoder_block(16, 32, kernel_size, activation, dropout_rate) | |
| self.e3 = encoder_block(32, 64, kernel_size, activation, dropout_rate) | |
| self.e4 = encoder_block(64, 128, kernel_size, activation, dropout_rate) | |
| self.e5 = encoder_block(128, 256, kernel_size, activation, dropout_rate) | |
| self.e6 = encoder_block(256, 512, kernel_size, activation, dropout_rate) | |
| #Decoder | |
| self.d1 = nn.UpsamplingNearest2d(scale_factor=2) | |
| self.d2 = decoder_block(512, 256, kernel_size, activation, dropout_rate) | |
| self.d3 = decoder_block(256, 128, kernel_size, activation, dropout_rate) | |
| self.d4 = decoder_block(128, 64, kernel_size, activation, dropout_rate) | |
| self.d5 = decoder_block(64, 32, kernel_size, activation, dropout_rate) | |
| self.d6 = decoder_block(32, 16, kernel_size, activation, dropout_rate) | |
| #Output | |
| if cmap == "gray": | |
| self.outputs = conv_block(16, 1, stride=False, activation=nn.Identity()) | |
| else: | |
| self.outputs = conv_block(16, 3, stride=False, activation=nn.Identity()) | |
| self.sigmoid = nn.Sigmoid() | |
| def forward(self, inputs): | |
| #Encoder | |
| p1 = self.e1(inputs) | |
| p2 = self.e2(p1) | |
| p3 = self.e3(p2) | |
| p4 = self.e4(p3) | |
| p5 = self.e5(p4) | |
| p6 = self.e6(p5) | |
| #Decoder | |
| d1 = self.d1(p6) | |
| d2 = self.d2(d1, p5) | |
| d3 = self.d3(d2, p4) | |
| d4 = self.d4(d3, p3) | |
| d5 = self.d5(d4, p2) | |
| d6 = self.d6(d5, p1) | |
| #Output | |
| outputs = self.outputs(d6) | |
| outputs = self.sigmoid(outputs) | |
| outputs = resize_layer(inputs, outputs) | |
| return outputs | |