static encoder runs, but doesn't work yet
This commit is contained in:
+37
-149
@@ -1,134 +1,21 @@
|
||||
import tensorflow as tf
|
||||
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
|
||||
import torch
|
||||
from torch import nn
|
||||
from models.convnext import ConvNeXt, CNBlock, CNBlockConfig
|
||||
from typing import Any, Callable, List, Optional, Sequence
|
||||
|
||||
class Bottleneck(tf.keras.Model):
|
||||
|
||||
class VariationalEncoder(nn.Module):
|
||||
|
||||
def __init__(self, dim, latent_dim, kernel_size=7):
|
||||
super(Bottleneck, self).__init__()
|
||||
|
||||
self.bottleneck = tf.keras.Sequential(
|
||||
[Block(dim=dim, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
self.dense = tf.keras.layers.Dense(latent_dim)
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.bottleneck(x, training=training)
|
||||
x = self.flatten(x)
|
||||
x = self.dense(x, training=training)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class VariationalEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, latent_dim, depths, dims, kernel_size=7):
|
||||
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7):
|
||||
super(VariationalEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size)
|
||||
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
self.mean = tf.keras.layers.Dense(latent_dim)
|
||||
self.log_var = tf.keras.layers.Dense(latent_dim)
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.flatten(x)
|
||||
mean = self.mean(x)
|
||||
log_var = self.log_var(x)
|
||||
|
||||
return mean, log_var
|
||||
|
||||
|
||||
class AEEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7):
|
||||
super(AEEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, _ = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.bottleneck(x, training=training)
|
||||
|
||||
if len(inputs) == 2:
|
||||
return x, inputs[1]
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class Decoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
|
||||
super(Decoder, self).__init__()
|
||||
|
||||
self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
|
||||
self.reshape = tf.keras.layers.Reshape(input_size)
|
||||
self.upscale = upscale
|
||||
|
||||
if not is_label:
|
||||
self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
|
||||
else:
|
||||
self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, mask = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.dense(x)
|
||||
x = self.reshape(x)
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
if self.upscale:
|
||||
x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
|
||||
if len(inputs) == 2:
|
||||
if len(mask.shape) == 2:
|
||||
x = x * mask[tf.newaxis,:,:,tf.newaxis]
|
||||
else:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class VariationalEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7):
|
||||
super(VariationalEncoderNew, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
|
||||
self.mean_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.var_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.flatten_mean = tf.keras.layers.Flatten()
|
||||
self.flatten_var = tf.keras.layers.Flatten()
|
||||
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):
|
||||
@@ -144,23 +31,27 @@ class VariationalEncoderNew(tf.keras.Model):
|
||||
return mean, log_var
|
||||
|
||||
|
||||
|
||||
class AEEncoderNew(tf.keras.Model):
|
||||
class AEEncoder(nn.Module):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
|
||||
super(AEEncoderNew, self).__init__()
|
||||
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', device='cuda'):
|
||||
super(AEEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.act = tf.keras.layers.Activation(activation)
|
||||
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 call(self, inputs, training):
|
||||
def forward(self, inputs):
|
||||
if len(inputs) == 2:
|
||||
x, _ = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.encoder(x)
|
||||
x = self.act(x)
|
||||
|
||||
if len(inputs) == 2:
|
||||
@@ -169,39 +60,36 @@ class AEEncoderNew(tf.keras.Model):
|
||||
return x
|
||||
|
||||
|
||||
class DecoderNew(tf.keras.Model):
|
||||
class Decoder(nn.Module):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7, is_label=False):
|
||||
super(DecoderNew, self).__init__()
|
||||
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
|
||||
|
||||
if not is_label:
|
||||
self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
|
||||
else:
|
||||
self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
|
||||
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 call(self, inputs, training):
|
||||
def forward(self, inputs):
|
||||
if len(inputs) == 2:
|
||||
x, mask = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.decoder(x, training=training)
|
||||
x = self.decoder(x)
|
||||
|
||||
if self.is_label:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
x = x * mask[:,:,:,None]
|
||||
|
||||
background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis])
|
||||
x = tf.concat((background, x), axis=-1)
|
||||
|
||||
x = tf.math.softmax(x)
|
||||
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[tf.newaxis,:,:,tf.newaxis]
|
||||
x = x * mask[None,None,:,:]
|
||||
else:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
x = x * mask[:,None,:,:]
|
||||
|
||||
return x
|
||||
Reference in New Issue
Block a user