From 73ef4fe4d11118a36ea1bc00378990f2094e850b Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 29 Oct 2022 20:15:31 +0000 Subject: [PATCH] add noise option to cvae model --- models/cvae_model.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/models/cvae_model.py b/models/cvae_model.py index cfff553..7770d5b 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -38,7 +38,7 @@ class CVAE(tf.keras.Model): self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'), tf.keras.layers.Dense(256, activation='relu'), tf.keras.layers.Dense(128, activation='relu'), - tf.keras.layers.Dense(config['latent_dim']),]) + tf.keras.layers.Dense(config['latent_dim'], activation='tanh')]) self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size']) self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size']) @@ -48,6 +48,9 @@ class CVAE(tf.keras.Model): else: self.avg_garment = None + self.add_noise = config['cvae']['add_noise'] + self.noise_amount = config['cvae']['noise_amount'] + @tf.function def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None): @@ -103,6 +106,10 @@ class CVAE(tf.keras.Model): z = self.reparameterize(mean, logvar) z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high, training=training) + if self.add_noise and training == True: + z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount) + + z = tf.concat([z, z_cond], axis=-1) x = self.decoder(z, training=training)