223 lines
6.6 KiB
Python
223 lines
6.6 KiB
Python
import tensorflow as tf
|
|
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
|
|
from models.resnet import resnet_18, make_basic_block_layer
|
|
|
|
|
|
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):
|
|
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):
|
|
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):
|
|
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, features):
|
|
super(VariationalEncoderNew, self).__init__()
|
|
|
|
self.encoder = resnet_18(features)
|
|
|
|
self.mean_branch = tf.keras.Sequential([
|
|
tf.keras.layers.Conv2D(filters=features,
|
|
kernel_size=(3, 3),
|
|
strides=1,
|
|
padding="same"),
|
|
tf.keras.layers.BatchNormalization(),
|
|
tf.keras.layers.Conv2D(filters=features,
|
|
kernel_size=(3, 3),
|
|
strides=1,
|
|
padding="same")
|
|
])
|
|
|
|
self.var_branch = tf.keras.Sequential([
|
|
tf.keras.layers.Conv2D(filters=features,
|
|
kernel_size=(3, 3),
|
|
strides=1,
|
|
padding="same"),
|
|
tf.keras.layers.BatchNormalization(),
|
|
tf.keras.layers.Conv2D(filters=features,
|
|
kernel_size=(3, 3),
|
|
strides=1,
|
|
padding="same")
|
|
])
|
|
|
|
self.flatten_mean = tf.keras.layers.Flatten()
|
|
self.flatten_var = tf.keras.layers.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 AEEncoderNew(tf.keras.Model):
|
|
|
|
def __init__(self, features, activation='tanh'):
|
|
super(AEEncoderNew, self).__init__()
|
|
|
|
self.encoder = resnet_18(features)
|
|
self.act = tf.keras.layers.Activation(activation)
|
|
|
|
|
|
def call(self, inputs, training):
|
|
if len(inputs) == 2:
|
|
x, _ = inputs
|
|
else:
|
|
x = inputs
|
|
|
|
x = self.encoder(x, training=training)
|
|
x = self.act(x)
|
|
|
|
if len(inputs) == 2:
|
|
return x, inputs[1]
|
|
else:
|
|
return x
|
|
|
|
|
|
class DecoderNew(tf.keras.Model):
|
|
|
|
def __init__(self, features, is_label=False, is_cvae=False):
|
|
super(DecoderNew, self).__init__()
|
|
|
|
self.is_label = is_label
|
|
|
|
if not is_label:
|
|
self.decoder = resnet_18(features, reversed=True, is_cvae=is_cvae)
|
|
else:
|
|
self.decoder = resnet_18(features, reversed=True, labels=6, is_cvae=is_cvae)
|
|
|
|
|
|
def call(self, inputs, training):
|
|
if len(inputs) == 2:
|
|
x, mask = inputs
|
|
else:
|
|
x = inputs
|
|
|
|
x = self.decoder(x, training=training)
|
|
|
|
if self.is_label:
|
|
x = x * mask[:,:,:,tf.newaxis]
|
|
|
|
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)
|
|
|
|
if len(inputs) == 2:
|
|
if len(mask.shape) == 2:
|
|
x = x * mask[tf.newaxis,:,:,tf.newaxis]
|
|
else:
|
|
x = x * mask[:,:,:,tf.newaxis]
|
|
|
|
return x |