small updates

scripts/cvae_pyramid.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_reg import config
+    from configs.iccvw.pyr_64_128_256_512_cape_f import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -217,6 +217,9 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
 
 
 def train_step(batch):
+    if config['is_cape']:
+        batch['mesh_fix'] = {'mesh_mask': batch['mesh_mask']}
+
     model.train()
     # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
     levels_input, conds = group_inputs_for_levels(batch, config)
@@ -313,6 +316,9 @@ def train_step(batch):
 
 
 def valid_step(batch):
+    if config['is_cape']:
+        batch['mesh_fix'] = {'mesh_mask': batch['mesh_mask']}
+
     model.eval()
     with torch.inference_mode():
         levels_input, conds = group_inputs_for_levels(batch, config)
@@ -384,6 +390,9 @@ def valid_step(batch):
 
 
 def plot_everything(batch: dict, type: str):
+    if config['is_cape']:
+        batch['mesh_fix'] = {'mesh_mask': batch['mesh_mask']}
+
     if config['use_skinning']:
         mesh_orig, uv_orig = batch['mesh_unposed'].cpu().numpy(), batch['uv_unposed']
     else:
@@ -394,7 +403,7 @@ 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() for i in range(len(output_levels))]
+    mesh_pred_levels = [uv2mesh(output_levels[i][0]).cpu() for i in range(len(output_levels))]
     mesh_orig_posed = batch['mesh_posed'].cpu()
     body = batch['mesh_body_posed'].cpu().numpy()
     
@@ -406,16 +415,16 @@ def plot_everything(batch: dict, type: str):
     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
     uv_normals = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.4) / 2.8
-    pred_img = torch.concatenate([torch.permute(model.resizers[512](output_levels[i][0]), [0,2,3,1]) for i in range(len(output_levels))], dim=3)
+    pred_img = torch.concatenate([torch.permute(model.resizers[512](output_levels[i][0]), [0,2,3,1]) for i in range(len(output_levels))], dim=2)
     # pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
     pred_img = (pred_img - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None]) / (pred_img.flatten().reshape(pred_img.shape[0], -1).max(1)[0][:,None,None,None] - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None])
 
-    img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals, orig_img, pred_img), axis=2)
+    img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals, orig_img, pred_img), dim=2)
     imgs = [wandb.Image(img[i].cpu().numpy()) for i in range(4)]
 
     _F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
     mask = batch['mesh_fix']['mesh_mask'].cpu().numpy().astype(bool)
-    if 'mesh_fix' in batch:
+    if 'mesh_fix' in batch and not config['is_cape']:
         faces = [np.load( (dataset_train if type == 'train' else dataset_valid).folder + "/mesh_fixed/" + batch['id'][i] + '_' + batch['garment'][i] + ".npz" )['faces'] for i in range(len(batch['id']))]
     else:
         _, faces = maskFacesBatch(np.ones_like(mesh_orig), _F, mask)
@@ -553,7 +562,7 @@ if __name__ == '__main__':
 
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
-    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=20)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
     scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
 
     g_balancing_coeff = config['starting_balancing_coeff']
@@ -561,13 +570,10 @@ if __name__ == '__main__':
 
     step = 0
     if config['pretrained_model'] != '':
-        # 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']:
-    #     model = torch.compile(model)
-    #     wandb.watch(model, log='all', log_freq=1000, log_graph=True)
-
+        if 'ignore_saved_optimizer' in config.keys()  and config['ignore_saved_optimizer']:
+            step = load_model(model, config['pretrained_model'])
+        else:
+            step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration'])
 
     if len(config['pyramid']['resolutions']) == 1:
         for param in model.static_enc.parameters(): param.requires_grad = False