small updates

scripts/generate_results.py

@@ -7,17 +7,19 @@ import argparse
 parser = argparse.ArgumentParser(description='')
 parser.add_argument('--name', type=str)
 parser.add_argument('--config', type=str)
+parser.add_argument('--iteration', type=str, default='')
 parser.add_argument('--debug', action='store_true')
+parser.add_argument('-g', '--gpu', type=str)
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64 import config
+    from configs.iccvw.pyr_64_128_cape_f_2 import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+os.environ["CUDA_VISIBLE_DEVICES"] = args['gpu']
 if args['debug']:
     os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
@@ -70,7 +72,10 @@ def step_pyramid(batch):
             shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
 
             mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
-            mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+            if not config['is_cape']:
+                mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+            else:
+                mesh_pred_posed = mesh_pred_shaped
         else:
             mesh_pred_posed = mesh_pred
 
@@ -102,15 +107,32 @@ def step(batch):
 
 def generate():
     for batch in tqdm(test_dataloader):
-        if is_pyramid:
-            mesh_pred = step_pyramid(batch)
-        else: 
-            mesh_pred = step(batch)
+        with torch.inference_mode():
+            if is_pyramid:
+                mesh_pred = step_pyramid(batch)
+            else: 
+                mesh_pred = step(batch)
 
         ids = batch['id']
         frame_nrs = batch['frame_nr']
         garments = batch['garment']
 
+        # import plotly.graph_objects as go
+        # mesh_pred = mesh_pred.cpu().numpy()
+        # mask = batch['mesh_mask'].cpu().numpy()
+        # i = 3
+        # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,mask[i],0], y=mesh_pred[i,mask[i],2], z=mesh_pred[i,mask[i],1], mode='markers', marker=dict(size=1))])
+        # fig.update_layout(scene_aspectmode='cube')
+        # fig.show()
+
+        # import plotly.graph_objects as go
+        # mesh_pred = mesh_pred.cpu().numpy()
+        # i = 3
+        # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,:,0], y=mesh_pred[i,:,2], z=mesh_pred[i,:,1], mode='markers', marker=dict(size=1))])
+        # fig.update_layout(scene_aspectmode='cube')
+        # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+        # fig.show()
+
         for i in range(len(ids)):
             np.savez_compressed(f'{out_folder}/{ids[i]}_{frame_nrs[i]}_{garments[i]}', mesh_pred[i].cpu().numpy())
 
@@ -121,11 +143,17 @@ if __name__ == '__main__':
 
     smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
-    dataset_resolutions = config['pyramid']['resolutions'].copy()
+    dataset_resolutions = []
+    if 'pyramid' in config:
+        dataset_resolutions = config['pyramid']['resolutions'].copy()
     if 512 not in dataset_resolutions:
         dataset_resolutions.append(512)
 
+    # remove all values from config['signature] that start with mesh_
+    config['signature'] = set([x for x in config['signature'] if not x.startswith('mesh_')])
+
     dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file=config['test_frames'], use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    # dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
     test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
 
     is_pyramid = config['name'].startswith('pyr')
@@ -135,7 +163,8 @@ if __name__ == '__main__':
     else:
         model: CVAE = CVAE(config)
 
-    it = load_model(model, 'checkpoints/iccv/torch/' + config['name'])
+    it = load_model(model, 'checkpoints/iccvw/' + config['name'], iteration=args['iteration'])
+    model.eval()
     # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
 
     name = config['name'] + '_' + str(it)