95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
import torch
|
|
from torch import nn
|
|
from models.convnext import ConvNeXt, CNBlock, CNBlockConfig
|
|
from typing import Any, Callable, List, Optional, Sequence
|
|
|
|
|
|
class VariationalEncoder(nn.Module):
|
|
|
|
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7):
|
|
super(VariationalEncoder, self).__init__()
|
|
|
|
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size)
|
|
|
|
self.mean_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
|
|
self.var_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
|
|
|
|
self.flatten_mean = nn.Flatten()
|
|
self.flatten_var = nn.Flatten()
|
|
|
|
|
|
def call(self, inputs, training):
|
|
x = inputs
|
|
|
|
x = self.encoder(x, training=training)
|
|
mean = self.mean_branch(x)
|
|
log_var = self.var_branch(x)
|
|
|
|
mean = self.flatten_mean(mean)
|
|
log_var = self.flatten_var(log_var)
|
|
|
|
return mean, log_var
|
|
|
|
|
|
class AEEncoder(nn.Module):
|
|
|
|
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', device='cuda'):
|
|
super(AEEncoder, self).__init__()
|
|
|
|
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device)
|
|
if activation == 'tanh':
|
|
self.act = nn.Tanh()
|
|
elif activation == 'sigmoid':
|
|
self.act = nn.Sigmoid()
|
|
else:
|
|
raise ValueError('Activation function not supported')
|
|
|
|
|
|
def forward(self, inputs):
|
|
if len(inputs) == 2:
|
|
x, _ = inputs
|
|
else:
|
|
x = inputs
|
|
|
|
x = self.encoder(x)
|
|
x = self.act(x)
|
|
|
|
if len(inputs) == 2:
|
|
return x, inputs[1]
|
|
else:
|
|
return x
|
|
|
|
|
|
class Decoder(nn.Module):
|
|
|
|
def __init__(self, block_setting: CNBlockConfig, kernel_size=7, is_label=False, nr_of_classes=None, device='cuda'):
|
|
super(Decoder, self).__init__()
|
|
|
|
self.is_label = is_label
|
|
|
|
self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_classes, device=device, reversed=True)
|
|
|
|
self.softmax = nn.Softmax(dim=1)
|
|
|
|
def forward(self, inputs):
|
|
if len(inputs) == 2:
|
|
x, mask = inputs
|
|
else:
|
|
x = inputs
|
|
|
|
x = self.decoder(x)
|
|
|
|
if self.is_label:
|
|
x = x * mask[:,:,:,None]
|
|
|
|
background = torch.ones((mask.shape[0],512,256,1), device=mask.get_device()) * (1 - mask[:,:,:,None])
|
|
x = torch.concat((background, x), axis=-1)
|
|
x = self.softmax(x)
|
|
|
|
if len(inputs) == 2:
|
|
if len(mask.shape) == 2:
|
|
x = x * mask[None,None,:,:]
|
|
else:
|
|
x = x * mask[:,None,:,:]
|
|
|
|
return x |