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

propagate training parameter

Browse Source
This commit is contained in:
luni
2022-10-12 21:05:49 +00:00
parent fb0fa16ed0
commit 4bf72f8b55
1 changed files with 11 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_model.py
+11 -11
View File
@@ -49,22 +49,22 @@ class CVAE(tf.keras.Model):
self.avg_garment = None
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
@tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None):
if noise is None:
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
else:
eps = noise
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
x = self.decoder(z, training=False)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if uv_mask is not None:
x = x * uv_mask[:,:,:,tf.newaxis]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
@@ -78,13 +78,13 @@ class CVAE(tf.keras.Model):
return eps * tf.exp(logvar * .5) + mean
def encode_conds(self, uv_template, uv_body_posed, uv_curv):
z_body = self.body_enc(uv_body_posed)
z_static = self.static_enc(uv_template)
z_curv = self.curv_enc(uv_curv)
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training=None):
z_body = self.body_enc(uv_body_posed, training=training)
z_static = self.static_enc(uv_template, training=training)
z_curv = self.curv_enc(uv_curv, training=training)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
z_cond = self.cond_encoder(z_cond, training=training)
return z_cond
@@ -101,7 +101,7 @@ class CVAE(tf.keras.Model):
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x, training=training)
z = self.reparameterize(mean, logvar)
z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high)
z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high, training=training)
z = tf.concat([z, z_cond], axis=-1)
x = self.decoder(z, training=training)