This commit is contained in:
2023-06-06 16:48:48 +00:00
parent e3ca2129fd
commit b103b09414
+5 -5
View File
@@ -99,8 +99,8 @@ class CVAE(nn.Module):
z_body = self.body_enc(uv_body_posed)
z_static = self.static_enc(uv_template)
if self.use_normals:
z_curv = self.norm_enc(uv_curv)
z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
z_norm = self.norm_enc(uv_curv)
z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
else:
z_cond = torch.concat([z_static, z_body], dim=1)
@@ -108,13 +108,13 @@ class CVAE(nn.Module):
def forward(self, inputs):
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
uv_unposed, uv_template, uv_body_posed, uv_norm, uv_mask = inputs
x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_norm], dim=1)
mean, logvar = self.encoder(x)
z: torch.Tensor = self.reparameterize(mean, logvar)
z = z.view((mean.shape[0], *self.latent_shape))
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_norm)
z = torch.concat([z, z_cond], dim=1)
x = self.decoder(z)