From ef05cf57d48f503038c9f8ae811da3925dd05e86 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 29 Oct 2022 20:01:23 +0000 Subject: [PATCH] restore orig body enc --- auto_encoder_body.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/auto_encoder_body.py b/auto_encoder_body.py index 5473289..a2a15bb 100644 --- a/auto_encoder_body.py +++ b/auto_encoder_body.py @@ -45,7 +45,7 @@ def init_wandb(): artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) -uv2mesh = UV2Mesh(res=256, is_smpl_orig=config['is_cape']) +uv2mesh = UV2Mesh(res=256, is_smpl_orig=config['is_cape'], is_skirt=config['is_skirt']) @tf.function @@ -58,7 +58,7 @@ def train_step(batch): e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 3638 if config['is_cape'] else 14475)))) l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1)) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 2, batch['uv_body_posed' + resolution_modifier] + 2, max_val=4, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) loss = 1e3 * (1 - ms_ssim) @@ -78,7 +78,7 @@ def valid_step(batch): e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 3638 if config['is_cape'] else 14475)))) l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1)) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 2, batch['uv_body_posed' + resolution_modifier] + 2, max_val=4, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) loss = 1e3 * (1 - ms_ssim) @@ -172,6 +172,9 @@ if __name__ == '__main__': model = AutoEncoder(config['cond_ae']) optimizer = tf.optimizers.Adam(learning_rate=config['lr']) + # load_model(model.encoder, 'checkpoints/body_1_skirt_2_enc') + # load_model(model.decoder, 'checkpoints/body_1_skirt_2_dec') + ckpt_enc = tf.train.Checkpoint(model=model.encoder) manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20) ckpt_dec = tf.train.Checkpoint(model=model.decoder)