luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cape_m pyramid working

Browse Source
This commit is contained in:
luni
2023-05-22 23:20:30 +02:00
parent 7d17b1a222
commit d6002f30f6
2 changed files with 26 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/cvae_pyramid.py
+24 -11
View File
@@ -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()