add noise option to cvae model

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)