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

barely working decoder

Browse Source
This commit is contained in:
luni
2023-02-24 10:34:53 +00:00
parent c1113379ec
commit fb4e7dcc65
6 changed files with 92 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae.py
+12 -13
View File
@@ -88,7 +88,7 @@ def train_step(batch):
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
mesh_pred = uv2mesh(pred)
mesh_pred = pred
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
@@ -105,18 +105,18 @@ def train_step(batch):
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
# l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
# conc_loss = concept_loss(pred, batch['uv_unposed'])
kl = kl_loss(mean, logvar)
reconstr_loss = l2_loss
reconstr_loss = l2_mesh_loss
if config['use_extra_losses']:
if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff']
reconstr_loss = reconstr_loss + (e_loss) * config['reg_loss_coeff']
else:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
if config['use_norm_loss'] and step > 2000:
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
@@ -132,7 +132,7 @@ def train_step(batch):
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
# 'train/ms_ssim': ms_ssim,
}
@@ -146,7 +146,7 @@ def valid_step(batch):
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
mesh_pred = pred
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
@@ -163,15 +163,15 @@ def valid_step(batch):
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
# l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
# conc_loss = concept_loss(pred, batch['uv_unposed'])
kl = kl_loss(mean, logvar)
reconstr_loss = l2_loss
reconstr_loss = l2_mesh_loss
if config['use_extra_losses']:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
if config['use_norm_loss']:
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
@@ -256,13 +256,12 @@ def plot_everything(batch: dict, type: str):
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
mesh_pred = pred.numpy()
orig_img = uv_orig
uv_static_img = batch['uv_static']
uv_body_posed_img = batch['uv_body_posed']
pred_img = pred
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img), axis=2)
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")