cape_m pyramid working

configs/pyr_64_cape_m.py

@@ -3,7 +3,7 @@ from utils.convnext_utils import CNBlockConfig
 config = dict()
 
 config['gpu'] = '5'
-config['name'] = 'pyr_torch_64_cape_m'
+config['name'] = 'pyr_torch_64_cape_m_2'
 config['notes'] = ''
 config['tags'] = ['pyr', 'iccv', 'torch']
 config['script'] = 'scripts/cvae_pyramid.py'
@@ -44,7 +44,7 @@ if config['is_cape']:
   config['smpl_res'] = 3638
 
 config['lr'] = 0.0001
-config['starting_balancing_coeff'] = 0.001
+config['starting_balancing_coeff'] = 0.0001
 config['balancing_activation_limit'] = 0.015
 config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
 config['avg_garment_path'] = ''

scripts/cvae_pyramid.py

@@ -13,7 +13,8 @@ args = vars(parser.parse_args())
 config_module = importlib.import_module(args['config'])
 config = config_module.config
 if args['debug']:
-    from configs.pyr_64_cape_m import config
+    from configs.pyr_32 import config
+    # from configs.pyr_64_cape_m import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -259,9 +260,9 @@ def train_step(batch):
     kl = kl_loss(mean, logvar)
 
     if config['use_extra_losses']:
-        reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
+        reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + l2_loss + l2_mesh_loss)
         if not config['is_cape']:
-            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] 
+            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (c_loss) * config['reg_loss_coeff'] 
     else:
         reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
 
@@ -344,9 +345,9 @@ def valid_step(batch):
         kl = kl_loss(mean, logvar)
 
         if config['use_extra_losses']:
-            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
+            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + l2_loss + l2_mesh_loss)
             if not config['is_cape']:
-                reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] 
+                reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (c_loss) * config['reg_loss_coeff'] 
         else:
             reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
 
@@ -440,10 +441,14 @@ def plot_everything(batch: dict, type: str):
         levels_input = group_inputs_for_levels(batch, config)
         output_levels = model(levels_input)
 
-    mesh_pred_levels = [uv2mesh(output_levels[i][0]).detach().cpu().numpy() for i in range(len(output_levels))]
+    mesh_pred_levels = [uv2mesh(output_levels[i][0]).detach().cpu() for i in range(len(output_levels))]
-    mesh_orig_posed = batch['mesh_posed'].cpu().numpy()
+    mesh_orig_posed = batch['mesh_posed'].cpu()
     body = batch['mesh_body_posed'].cpu().numpy()
     
+    # mesh_pred = torch.cat((mesh_pred[:, :, 0:1], -mesh_pred[:, :, 2:3], mesh_pred[:, :, 1:2]), dim=2)
+    # mesh_pred_levels = [torch.cat((mesh_pred_levels[i][:, :, 0:1], -mesh_pred_levels[i][:, :, 2:3], mesh_pred_levels[i][:, :, 1:2]), dim=2).numpy() for i in range(len(mesh_pred_levels))]
+    # mesh_orig_posed = torch.cat((mesh_orig_posed[:, :, 0:1], -mesh_orig_posed[:, :, 2:3], mesh_orig_posed[:, :, 1:2]), dim=2).numpy()
+
     orig_img = (torch.permute(uv_orig, [0,2,3,1]) + 1) / 2
     uv_static_img = (torch.permute(batch['uv_static'], [0,2,3,1]) + 1) / 2
     uv_body_posed_img = (torch.permute(batch['uv_body_posed'], [0,2,3,1]) + 1) / 2
@@ -477,6 +482,14 @@ def plot_everything(batch: dict, type: str):
         for i in range(mask.shape[0]):
             mesh_unposed_shaped_full[i, mask[i], :] = mesh_unposed_shaped[i]
         mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_unposed_shaped_full, device=config['device']))[0].cpu().numpy()
+    
+        # rotate gt 180 degrees along z axis
+        # gt = [np.matmul(gt[i], np.array([[np.cos(np.pi), -np.sin(np.pi), 0], [np.sin(np.pi), np.cos(np.pi), 0], [0, 0, 1]])) for i in range(mask.shape[0])]
+
+        # gt_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32)
+        # for i in range(mask.shape[0]):
+        #     gt_full[i, mask[i], :] = gt[i]
+        # mesh_orig_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(np.stack(gt_full, axis=0), device=config['device'], dtype=torch.float32))[0].cpu().numpy()
     else:
         mesh_posed = mesh_pred_levels[-1]
         gt = [mesh_orig[i][mask[i], :] for i in range(mesh_posed.shape[0])]
@@ -819,9 +832,9 @@ if __name__ == '__main__':
         # assert len(config['pyramid']['resolutions']) == 1
         step = load_model(model, config['pretrained_model'], optimizer=optimizer, scaler=scaler, iteration=config['pretrained_model_iteration'])
 
-    if not args['debug']:
+    # if not args['debug']:
-        model = torch.compile(model)
+    #     model = torch.compile(model)
-        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+    #     wandb.watch(model, log='all', log_freq=1000, log_graph=True)
 
 
     if len(config['pyramid']['resolutions']) == 1:
@@ -852,7 +865,7 @@ if __name__ == '__main__':
     # load_model(model.levels[-1], 'checkpoints/cvae_19')
 
     if not args['debug']:
-        model = torch.compile(model)
+        # model = torch.compile(model)
         wandb.watch(model, log='all', log_freq=1000, log_graph=True)
 
     train()