108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
import tensorflow as tf
|
|
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
|
|
|
|
class Bottleneck(tf.keras.Model):
|
|
|
|
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=None):
|
|
x = inputs
|
|
|
|
x = self.bottleneck(x, training=training)
|
|
x = self.flatten(x)
|
|
x = self.dense(x)
|
|
|
|
return x
|
|
|
|
|
|
class VariationalEncoder(tf.keras.Model):
|
|
|
|
def __init__(self, latent_dim, depths, dims, kernel_size=7):
|
|
super(VariationalEncoder, self).__init__()
|
|
|
|
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, 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=None):
|
|
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=None):
|
|
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=None):
|
|
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 |