luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

upscael AE output if needed (body ae)

Browse Source
This commit is contained in:
luni
2022-09-27 21:25:50 +00:00
parent 64264364be
commit 1da83294ee
2 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/base.py
+6 -1
View File
@@ -72,11 +72,12 @@ class AEEncoder(tf.keras.Model):
class Decoder(tf.keras.Model): class Decoder(tf.keras.Model):
def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False): def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
super(Decoder, self).__init__() super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size)) self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
self.reshape = tf.keras.layers.Reshape(input_size) self.reshape = tf.keras.layers.Reshape(input_size)
self.upscale = upscale
if not is_label: if not is_label:
self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size) self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
@@ -94,6 +95,10 @@ class Decoder(tf.keras.Model):
x = self.reshape(x) x = self.reshape(x)
x = self.decoder(x, training=training) 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(inputs) == 2:
if len(mask.shape) == 2: if len(mask.shape) == 2:
x = x * mask[tf.newaxis,:,:,tf.newaxis] x = x * mask[tf.newaxis,:,:,tf.newaxis]
models/cond_encoder.py
+1 -1
View File
@@ -8,7 +8,7 @@ class AutoEncoder(tf.keras.Model):
super(AutoEncoder, self).__init__() super(AutoEncoder, self).__init__()
self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size']) self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label']) self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
def call(self, inputs, training=None): def call(self, inputs, training=None):