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

restore orig body enc

Browse Source
This commit is contained in:
luni
2022-10-29 20:01:23 +00:00
parent 8ef3b1045d
commit ef05cf57d4
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+6 -3
View File
@@ -45,7 +45,7 @@ def init_wandb():
artifact.add_file('convnext/models/convnext_tf.py') artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact) 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 @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)))) 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)) 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) 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)))) 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)) 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) loss = 1e3 * (1 - ms_ssim)
@@ -172,6 +172,9 @@ if __name__ == '__main__':
model = AutoEncoder(config['cond_ae']) model = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr']) 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) ckpt_enc = tf.train.Checkpoint(model=model.encoder)
manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20) manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20)
ckpt_dec = tf.train.Checkpoint(model=model.decoder) ckpt_dec = tf.train.Checkpoint(model=model.decoder)