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

compare models

Browse Source
This commit is contained in:
luni
2023-02-21 16:06:25 +00:00
parent ebbc52d0fd
commit 2100d364f9
4 changed files with 202 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
evaluate_ae.py
+11 -3
View File
@@ -39,7 +39,7 @@ else:
from losses import collision_ratio, kl_loss #, concept_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, UV2MeshLabel, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAENew as CVAE
@@ -48,6 +48,7 @@ from custom_utils.supersmpl_io import writeOBJ
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
uv2mesh_label = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
@tf.function
def step(batch):
@@ -56,9 +57,10 @@ def step(batch):
else:
uv_orig = batch['uv_posed']
pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
pred, _, _, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
mesh_weigths = uv2mesh_label(weights)
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
@@ -68,7 +70,13 @@ def 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)
else:
mesh_pred_posed = mesh_pred
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped, mesh_weigths)
# mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])