diff --git a/configs/pyr_64_cape_m.py b/configs/pyr_64_cape_m.py index acf949a..cee4022 100644 --- a/configs/pyr_64_cape_m.py +++ b/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'] = '' diff --git a/scripts/cvae_pyramid.py b/scripts/cvae_pyramid.py index 56bd3c9..52bfd41 100644 --- a/scripts/cvae_pyramid.py +++ b/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_orig_posed = batch['mesh_posed'].cpu().numpy() + mesh_pred_levels = [uv2mesh(output_levels[i][0]).detach().cpu() for i in range(len(output_levels))] + 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']: - model = torch.compile(model) - wandb.watch(model, log='all', log_freq=1000, log_graph=True) + # if not args['debug']: + # model = torch.compile(model) + # 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() \ No newline at end of file