| 378 | is used to predict and regress only strand data, with no scalp |
| 379 | ''' |
| 380 | class StrandCodec(nn.Module): |
| 381 | def __init__(self, do_vae=True, scale_init=30.0, decode_type="dir", decode_random_verts=False, nr_verts_per_strand=256, nr_values_to_decode=255, dim_per_value_decoded=3): |
| 382 | super(StrandCodec, self).__init__() |
| 383 | |
| 384 | self.do_vae = do_vae |
| 385 | self.decode_type = decode_type |
| 386 | self.decode_random_verts = decode_random_verts |
| 387 | |
| 388 | # encode |
| 389 | |
| 390 | self.encoder = StrandEncoder1dCNNWN(self.do_vae, out_channels=64) #Uses LinearWN |
| 391 | |
| 392 | # decoder |
| 393 | self.decoder = StrandGeneratorSiren(in_channels=64, modulation_hidden_dim=128, siren_hidden_dim=128, |
| 394 | scale_init=scale_init, decode_type=decode_type, decode_random_verts=decode_random_verts, |
| 395 | nr_verts_per_strand=nr_verts_per_strand, nr_values_to_decode=nr_values_to_decode, |
| 396 | dim_per_value_decoded=dim_per_value_decoded, |
| 397 | ) |
| 398 | |
| 399 | |
| 400 | |
| 401 | def save(self, root_folder, experiment_name, hyperparams, iter_nr, info=None): |
| 402 | name=str(iter_nr) |
| 403 | if info is not None: |
| 404 | name+="_"+info |
| 405 | models_path = os.path.join(root_folder, experiment_name, name, "models") |
| 406 | if not os.path.exists(models_path): |
| 407 | os.makedirs(models_path, exist_ok=True) |
| 408 | torch.save(self.state_dict(), os.path.join(models_path, "strand_codec.pt")) |
| 409 | |
| 410 | hyperparams_params_path=os.path.join(models_path, "hyperparams.json") |
| 411 | with open(hyperparams_params_path, 'w', encoding='utf-8') as f: |
| 412 | json.dump(vars(hyperparams), f, ensure_ascii=False, indent=4) |
| 413 | |
| 414 | |
| 415 | def forward(self, gt_dict, hyperparams, normalization_dict): |
| 416 | if hyperparams.normalize_input: |
| 417 | gt_dict = normalize_gt_data(gt_dict, normalization_dict) |
| 418 | encoded_dict = self.encoder(gt_dict) |
| 419 | pred_dict=self.decoder(encoded_dict["z"], hyperparams, normalization_dict) |
| 420 | return pred_dict, encoded_dict |