replace curvature with normals

This commit is contained in:
2023-06-01 08:41:53 +00:00
parent a1470cfd6e
commit 534fbfcde1
4 changed files with 214 additions and 20 deletions
+9 -9
View File
@@ -27,16 +27,16 @@ class CVAE(nn.Module):
if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc')
if config['use_curvature']:
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'], device=config['device'], is_label=True, nr_of_labels=7)
if config['curvature_model'] != '':
load_model(self.curv_enc, config['curvature_model'] + '_enc')
if config['use_normals']:
self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device'])
if config['normals_model'] != '':
load_model(self.norm_enc, config['normals_model'] + '_enc')
if config['freeze_encoders']:
for param in self.static_enc.parameters(): param.requires_grad = False
for param in self.body_enc.parameters(): param.requires_grad = False
if config['use_curvature']:
for param in self.curv_enc.parameters(): param.requires_grad = False
if config['use_normals']:
for param in self.norm_enc.parameters(): param.requires_grad = False
self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])
@@ -48,7 +48,7 @@ class CVAE(nn.Module):
else:
self.avg_garment = None
self.use_curvature = config['use_curvature']
self.use_normals = config['use_normals']
self.use_offsets_over_body = not config['use_skinning']
if self.use_offsets_over_body:
@@ -98,8 +98,8 @@ class CVAE(nn.Module):
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)
if self.use_curvature:
z_curv = self.curv_enc(uv_curv)
if self.use_normals:
z_curv = self.norm_enc(uv_curv)
z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
else:
z_cond = torch.concat([z_static, z_body], dim=1)