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

add noise option to cvae model

Browse Source
This commit is contained in:
luni
2022-10-29 20:15:31 +00:00
parent 36a118a22e
commit 73ef4fe4d1
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_model.py
+8 -1
View File
@@ -38,7 +38,7 @@ class CVAE(tf.keras.Model):
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'), self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'), tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(128, 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.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']) 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: else:
self.avg_garment = None self.avg_garment = None
self.add_noise = config['cvae']['add_noise']
self.noise_amount = config['cvae']['noise_amount']
@tf.function @tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None): 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 = self.reparameterize(mean, logvar)
z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high, training=training) 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) z = tf.concat([z, z_cond], axis=-1)
x = self.decoder(z, training=training) x = self.decoder(z, training=training)