commit all

evaluate_ae.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_skin_skirt import config
+from configs.config_abl_all_post import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -42,11 +42,12 @@ 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 dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.cvae_model import CVAE
+from models.cvae_model import CVAENew as CVAE
 from custom_utils.supersmpl_io import writeOBJ
+# from unif_metric import chamfer_and_score
 
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
 
 @tf.function
 def step(batch):
@@ -80,11 +81,10 @@ if __name__ == '__main__':
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
 
-    smpl = SMPL(use_one_hot=config['use_one_hot'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     folder = config['test_dir']
-    smpl = SMPL(use_one_hot=False)
+    smpl = SMPL(use_one_hot=False, orig_smpl=config['is_cape'])
     model = CVAE(config)
 
     # str_split = config['name'].split('_')
@@ -92,18 +92,22 @@ if __name__ == '__main__':
     # checkpoint_name = '_'.join(str_split[:test_idx])
     checkpoint_name = config['name']
 
-    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+    if config['is_cape']:
+        F = np.load("custom_utils/smpl_faces_orig.npy").astype(np.int32)
+    else:
+        F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+
 
     ids = os.listdir(os.path.join(folder, 'pose_shape'))
     ids = [id.split('_')[0] for id in ids]
     ids = set(ids)
 
-    iteration = '70000'
+    iteration = ''
 
-    load_model(model, 'checkpoints/' + checkpoint_name, iteration=iteration)
+    load_model(model, 'checkpoints/post/' + checkpoint_name, iteration=iteration)
     # load_model(model, 'checkpoints/' + 'abl_all_120k')
 
-    output_folder = '/home/workspace/data/results/' + checkpoint_name + '-' + iteration
+    output_folder = '/home/workspace/data/results/post/' + checkpoint_name + '-' + iteration
 
 
     os.makedirs(output_folder, exist_ok=True)
@@ -120,12 +124,14 @@ if __name__ == '__main__':
                                        resolutions=[config['resolution']], 
                                        in_memory=False, 
                                        do_repeat=False,
-                                       use_one_hot=config['use_one_hot'],
-                                       use_no_orient=config['remove_global_orientation'])
+                                       use_no_orient=config['remove_global_orientation'],
+                                       frames_file=config['frames_file_test'] if 'frames_file_test' in config.keys() else None, 
+                                       is_cape=config['is_cape'])
     static_dataset_test = StaticDataset(folder=config['test_dir'], 
                                         signature=config['signature_static'], 
                                         resolutions=[config['resolution']], 
-                                        in_memory=True)
+                                        in_memory=True,
+                                        is_cape=config['is_cape'])
 
     # faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
     # ids = os.listdir(faces_folder)