large refactor
This commit is contained in:
+130
-151
@@ -1,223 +1,202 @@
|
||||
import tensorflow as tf
|
||||
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
|
||||
|
||||
|
||||
|
||||
|
||||
class Bottleneck(tf.keras.Model):
|
||||
|
||||
def __init__(self, dim, latent_dim):
|
||||
super(Bottleneck, self).__init__()
|
||||
|
||||
self.bottleneck = tf.keras.Sequential(
|
||||
[Block(dim=dim, drop_path=0,
|
||||
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)
|
||||
x = self.flatten(x)
|
||||
x = self.dense(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Decoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, bottleneck_size, latent_dim):
|
||||
super(Decoder, self).__init__()
|
||||
|
||||
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
|
||||
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
|
||||
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
x = inputs
|
||||
|
||||
x = self.dense(x)
|
||||
x = self.reshape(x)
|
||||
x = self.decoder(x)
|
||||
|
||||
return x
|
||||
from models.base import AEEncoder, Bottleneck, Decoder
|
||||
|
||||
|
||||
class AutoEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self, size):
|
||||
super(AutoEncoder, self).__init__()
|
||||
|
||||
if size == 512:
|
||||
bottleneck_size = 64
|
||||
latent_dim = 128
|
||||
kernel_size = 7
|
||||
encoder_depths = [3, 3, 9, 9, 3]
|
||||
encoder_dims = [32, 64, 128, 256, bottleneck_size]
|
||||
decoder_input = (8, 4, bottleneck_size)
|
||||
decoder_depths = [3, 3, 9, 3, 3]
|
||||
decoder_dims = [bottleneck_size, 256, 128, 64, 16]
|
||||
elif size == 128:
|
||||
raise NotImplementedError
|
||||
bottleneck_size = 512
|
||||
latent_dim = 128
|
||||
kernel_size = 5
|
||||
encoder_depths =[16, 16, 9, 9]
|
||||
encoder_dims = [64, 128, 256, bottleneck_size]
|
||||
decoder_input = (4, 2, bottleneck_size)
|
||||
decoder_depths = [16, 16, 9, 9]
|
||||
decoder_dims = [bottleneck_size, 256, 128, 64]
|
||||
elif size == 32:
|
||||
raise NotImplementedError
|
||||
bottleneck_size = 512
|
||||
latent_dim = 128
|
||||
kernel_size = 3
|
||||
encoder_depths = [27, 27]
|
||||
encoder_dims = [256, bottleneck_size]
|
||||
decoder_input = (4, 2, bottleneck_size)
|
||||
decoder_depths = [27, 27]
|
||||
decoder_dims = [bottleneck_size, 256]
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
bottleneck_size = 512
|
||||
latent_dim = 128
|
||||
|
||||
self.encoder = ConvNeXt(depths=[3, 3, 9, 9, 3, 3], dims=[32, 64, 128, 256, 256, bottleneck_size], include_top=False)
|
||||
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
|
||||
self.decoder = Decoder(bottleneck_size, latent_dim)
|
||||
self.encoder : AEEncoder = AEEncoder(latent_dim, encoder_depths, encoder_dims, bottleneck_size, kernel_size)
|
||||
self.decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size=kernel_size)
|
||||
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
x = inputs
|
||||
def call(self, inputs, training=None, use_mask=True):
|
||||
if use_mask:
|
||||
x, mask = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = tf.transpose(x, [0,2,3,1])
|
||||
x = self.encoder(x, training=training)
|
||||
# x = self.pooling(x)
|
||||
# x = self.reshape(x)
|
||||
# x = self.flatten(x)
|
||||
|
||||
x = self.bottleneck(x, training=training)
|
||||
|
||||
x = self.decoder(x, training=training)
|
||||
x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
x = tf.transpose(x, [0,3,1,2])
|
||||
|
||||
# x = x * uv_mask[:,:, tf.newaxis]
|
||||
# x = x * uv_mask_body[tf.newaxis,tf.newaxis,:,:]
|
||||
if use_mask:
|
||||
x = x * mask[tf.newaxis,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class Bottleneck32(tf.keras.Model):
|
||||
# class Bottleneck32(tf.keras.Model):
|
||||
|
||||
def __init__(self, dim, latent_dim):
|
||||
super(Bottleneck32, self).__init__()
|
||||
# def __init__(self, dim, latent_dim):
|
||||
# super(Bottleneck32, self).__init__()
|
||||
|
||||
self.bottleneck = tf.keras.Sequential(
|
||||
[Block(dim=dim, drop_path=0,
|
||||
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)
|
||||
# self.bottleneck = tf.keras.Sequential(
|
||||
# [Block(dim=dim, drop_path=0,
|
||||
# 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
|
||||
# def call(self, inputs, training=None):
|
||||
# x = inputs
|
||||
|
||||
x = self.bottleneck(x)
|
||||
x = self.flatten(x)
|
||||
x = self.dense(x)
|
||||
# x = self.bottleneck(x)
|
||||
# x = self.flatten(x)
|
||||
# x = self.dense(x)
|
||||
|
||||
return x
|
||||
# return x
|
||||
|
||||
|
||||
class Decoder32(tf.keras.Model):
|
||||
# class Decoder32(tf.keras.Model):
|
||||
|
||||
def __init__(self, bottleneck_size, latent_dim):
|
||||
super(Decoder32, self).__init__()
|
||||
# def __init__(self, bottleneck_size, latent_dim):
|
||||
# super(Decoder32, self).__init__()
|
||||
|
||||
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
|
||||
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
|
||||
self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
|
||||
# self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
|
||||
# self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
|
||||
# self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
x = inputs
|
||||
# def call(self, inputs, training=None):
|
||||
# x = inputs
|
||||
|
||||
x = self.dense(x)
|
||||
x = self.reshape(x)
|
||||
x = self.decoder(x)
|
||||
# x = self.dense(x)
|
||||
# x = self.reshape(x)
|
||||
# x = self.decoder(x)
|
||||
|
||||
return x
|
||||
# return x
|
||||
|
||||
|
||||
class AutoEncoder32(tf.keras.Model):
|
||||
# class AutoEncoder32(tf.keras.Model):
|
||||
|
||||
def __init__(self):
|
||||
super(AutoEncoder32, self).__init__()
|
||||
# def __init__(self):
|
||||
# super(AutoEncoder32, self).__init__()
|
||||
|
||||
|
||||
bottleneck_size = 512
|
||||
latent_dim = 128
|
||||
# bottleneck_size = 512
|
||||
# latent_dim = 128
|
||||
|
||||
self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
|
||||
self.bottleneck = Bottleneck32(bottleneck_size, latent_dim)
|
||||
self.decoder = Decoder32(bottleneck_size, latent_dim)
|
||||
# self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
|
||||
# self.bottleneck = Bottleneck32(bottleneck_size, latent_dim)
|
||||
# self.decoder = Decoder32(bottleneck_size, latent_dim)
|
||||
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
x = inputs
|
||||
# def call(self, inputs, training=None):
|
||||
# x = inputs
|
||||
|
||||
x = tf.transpose(x, [0,2,3,1])
|
||||
x = self.encoder(x, training=training)
|
||||
# x = tf.transpose(x, [0,2,3,1])
|
||||
# x = self.encoder(x, training=training)
|
||||
|
||||
x = self.bottleneck(x, training=training)
|
||||
# x = self.bottleneck(x, training=training)
|
||||
|
||||
x = self.decoder(x, training=training)
|
||||
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
x = tf.transpose(x, [0,3,1,2])
|
||||
# x = self.decoder(x, training=training)
|
||||
# # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# x = tf.transpose(x, [0,3,1,2])
|
||||
|
||||
return x
|
||||
# return x
|
||||
|
||||
|
||||
|
||||
|
||||
class Bottleneck128(tf.keras.Model):
|
||||
# class Bottleneck128(tf.keras.Model):
|
||||
|
||||
def __init__(self, dim, latent_dim):
|
||||
super(Bottleneck128, self).__init__()
|
||||
# def __init__(self, dim, latent_dim):
|
||||
# super(Bottleneck128, self).__init__()
|
||||
|
||||
self.bottleneck = tf.keras.Sequential(
|
||||
[Block(dim=dim, drop_path=0,
|
||||
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)
|
||||
# self.bottleneck = tf.keras.Sequential(
|
||||
# [Block(dim=dim, drop_path=0,
|
||||
# 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
|
||||
# def call(self, inputs, training=None):
|
||||
# x = inputs
|
||||
|
||||
x = self.bottleneck(x)
|
||||
x = self.flatten(x)
|
||||
x = self.dense(x)
|
||||
# x = self.bottleneck(x)
|
||||
# x = self.flatten(x)
|
||||
# x = self.dense(x)
|
||||
|
||||
return x
|
||||
# return x
|
||||
|
||||
|
||||
class Decoder128(tf.keras.Model):
|
||||
# class Decoder128(tf.keras.Model):
|
||||
|
||||
def __init__(self, bottleneck_size, latent_dim):
|
||||
super(Decoder128, self).__init__()
|
||||
# def __init__(self, bottleneck_size, latent_dim):
|
||||
# super(Decoder128, self).__init__()
|
||||
|
||||
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
|
||||
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
|
||||
self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
|
||||
# self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
|
||||
# self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
|
||||
# self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
x = inputs
|
||||
# def call(self, inputs, training=None):
|
||||
# x = inputs
|
||||
|
||||
x = self.dense(x)
|
||||
x = self.reshape(x)
|
||||
x = self.decoder(x)
|
||||
# x = self.dense(x)
|
||||
# x = self.reshape(x)
|
||||
# x = self.decoder(x)
|
||||
|
||||
return x
|
||||
# return x
|
||||
|
||||
|
||||
class AutoEncoder128(tf.keras.Model):
|
||||
# class AutoEncoder128(tf.keras.Model):
|
||||
|
||||
def __init__(self):
|
||||
super(AutoEncoder128, self).__init__()
|
||||
# def __init__(self):
|
||||
# super(AutoEncoder128, self).__init__()
|
||||
|
||||
|
||||
bottleneck_size = 512
|
||||
latent_dim = 128
|
||||
# bottleneck_size = 512
|
||||
# latent_dim = 128
|
||||
|
||||
self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
|
||||
self.bottleneck = Bottleneck128(bottleneck_size, latent_dim)
|
||||
self.decoder = Decoder128(bottleneck_size, latent_dim)
|
||||
# self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
|
||||
# self.bottleneck = Bottleneck128(bottleneck_size, latent_dim)
|
||||
# self.decoder = Decoder128(bottleneck_size, latent_dim)
|
||||
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
x = inputs
|
||||
# def call(self, inputs, training=None):
|
||||
# x = inputs
|
||||
|
||||
x = tf.transpose(x, [0,2,3,1])
|
||||
x = self.encoder(x, training=training)
|
||||
# x = tf.transpose(x, [0,2,3,1])
|
||||
# x = self.encoder(x, training=training)
|
||||
|
||||
x = self.bottleneck(x, training=training)
|
||||
# x = self.bottleneck(x, training=training)
|
||||
|
||||
x = self.decoder(x, training=training)
|
||||
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
x = tf.transpose(x, [0,3,1,2])
|
||||
# x = self.decoder(x, training=training)
|
||||
# # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# x = tf.transpose(x, [0,3,1,2])
|
||||
|
||||
return x
|
||||
# return x
|
||||
Reference in New Issue
Block a user