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

update cvae_low version

Browse Source
This commit is contained in:
luni
2022-09-26 12:54:03 +00:00
parent fdf36b27e7
commit 83294bc00c
3 changed files with 53 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae.py
+5 -3
View File
@@ -87,21 +87,23 @@ def train_step(batch):
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
e_loss = edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])
c_loss = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
# ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
reconstr_loss = l2_loss
reconstr_loss = l2_loss + e_loss + c_loss
reconstr_loss = tf.reduce_mean(reconstr_loss)
kl = kl_loss(mean, logvar)
loss = reconstr_loss + balancing_coeff * kl
ms_ssim = tf.reduce_mean(ms_ssim)
l1_loss = tf.reduce_mean(l1_loss)
l2_loss = tf.reduce_mean(l2_loss)
e_loss = tf.reduce_mean(e_loss)
c_loss = tf.reduce_mean(c_loss)
mean = tf.reduce_mean(mean)
logvar = tf.reduce_mean(logvar)