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

add sampling to cvae norm

Browse Source
This commit is contained in:
luni
2023-07-12 15:22:58 +02:00
parent 5a2138b45b
commit 80b1c668b1
1 changed files with 6 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_norm_model.py
+6 -17
View File
@@ -36,27 +36,16 @@ class CVAENorm(nn.Module):
self.latent_shape = (config['cvae']['latent_dim'], 8, 4)
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
def sample(self, uv_template, uv_body_posed, uv_mask, noise=None):
if noise is None:
eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
eps = torch.randn(uv_template.shape[0], self.latent_shape[0], self.latent_shape[1], self.latent_shape[2]).to(uv_template.device)
else:
eps = noise
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
if add_noise:
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z, training=False)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if self.use_offsets_over_body:
x = x + uv_body_posed
x = x * uv_mask[:,:,:,tf.newaxis]
z_cond = self.encode_conds(uv_template, uv_body_posed)
z = torch.cat([eps, z_cond], dim=1)
x = self.decoder(z)
x = x * uv_mask[:,None,:,:]
return x