From e23d0973981f37c1bf753fdc9ee93385d2066585 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 1 Jun 2023 12:55:42 +0000 Subject: [PATCH] direct training experiments --- configs/direct_6.py | 83 ++++++++++++++++++++++++++++++++++ configs/direct_7.py | 87 ++++++++++++++++++++++++++++++++++++ run.py | 2 +- scripts/direct_train.py | 98 +++++++++++++++++++++++++---------------- 4 files changed, 231 insertions(+), 39 deletions(-) create mode 100644 configs/direct_6.py create mode 100644 configs/direct_7.py diff --git a/configs/direct_6.py b/configs/direct_6.py new file mode 100644 index 0000000..15c9693 --- /dev/null +++ b/configs/direct_6.py @@ -0,0 +1,83 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'direct_6' +config['notes'] = 'with normals input but only single garment sequence with msssim and l2 mesh loss only' +config['tags'] = ['torch', 'direct'] +config['script'] = 'scripts/direct_train.py' + +config['pretrained_model'] = '' +config['use_curvature'] = False +config['use_normals'] = True +config['use_skinning'] = True + +# assert that only use curvature or normals +assert not (config['use_curvature'] and config['use_normals']) + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/train' +config['valid_dir'] = '/data-local/127/train' +config['test_dir'] = '/data-net/dataset/test_ready' +config['train_frames'] = 'frames_00761.txt' +config['valid_frames'] = 'frames_00761.txt' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = True +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' +config['input_features'] = 9 + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['is_label'] = False +config['kernel_size'] = 7 +config['bottleneck_size'] = 64 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 128, 3), + CNBlockConfig(128, 256, 9), + CNBlockConfig(256, config['bottleneck_size'], 9), + CNBlockConfig(config['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['bottleneck_size'], 256, 3), + CNBlockConfig(256, 128, 9), + CNBlockConfig(128, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['encoder_block_settings'] = block_setting +config['decoder_block_settings'] = block_setting_reverse +config['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['reg_loss_coeff'] = 0.1 + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_normals', + 'mesh_mask', + 'mesh_body_posed', + 'mesh_posed', + 'uv_unposed', + 'mesh_unposed', + 'uv_static', + 'uv_mask', + 'uv_body_posed', + } \ No newline at end of file diff --git a/configs/direct_7.py b/configs/direct_7.py new file mode 100644 index 0000000..d0f2404 --- /dev/null +++ b/configs/direct_7.py @@ -0,0 +1,87 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '6' +config['name'] = 'direct_7' +config['notes'] = 'with normals input but only single garment sequence with msssim and l2 mesh loss only and without skinning' +config['tags'] = ['torch', 'direct'] +config['script'] = 'scripts/direct_train.py' + +config['pretrained_model'] = '' +config['use_curvature'] = False +config['use_normals'] = True +config['use_skinning'] = False + +# assert that only use curvature or normals +assert not (config['use_curvature'] and config['use_normals']) + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/train' +config['valid_dir'] = '/data-local/127/train' +config['test_dir'] = '/data-net/dataset/test_ready' +config['train_frames'] = 'frames_00761.txt' +config['valid_frames'] = 'frames_00761.txt' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = True +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' +config['input_features'] = 9 + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['is_label'] = False +config['kernel_size'] = 7 +config['bottleneck_size'] = 64 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 128, 3), + CNBlockConfig(128, 256, 9), + CNBlockConfig(256, config['bottleneck_size'], 9), + CNBlockConfig(config['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['bottleneck_size'], 256, 3), + CNBlockConfig(256, 128, 9), + CNBlockConfig(128, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['encoder_block_settings'] = block_setting +config['decoder_block_settings'] = block_setting_reverse +config['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['reg_loss_coeff'] = 0.1 + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_normals', + 'mesh_mask', + 'mesh_body_posed', + 'mesh_posed', + 'uv_static', + 'uv_mask', + 'uv_body_posed', + } + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') diff --git a/run.py b/run.py index 22c264b..f1902c0 100644 --- a/run.py +++ b/run.py @@ -1,6 +1,6 @@ import os import sys -from configs.direct_3 import config +from configs.direct_6 import config import argparse # parser = argparse.ArgumentParser(description='') diff --git a/scripts/direct_train.py b/scripts/direct_train.py index 1bdf526..4e43694 100644 --- a/scripts/direct_train.py +++ b/scripts/direct_train.py @@ -13,7 +13,7 @@ args = vars(parser.parse_args()) config_module = importlib.import_module(args['config']) config = config_module.config if args['debug']: - from configs.direct_1 import config + from configs.direct_7 import config # from configs.ae_static_orig_2 import config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] @@ -52,7 +52,7 @@ from pytorch_msssim import MS_SSIM def init_wandb(): - run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) + run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127_2/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/direct_train.py') artifact.add_file('losses/losses.py') @@ -83,21 +83,23 @@ def train_step(batch): pred = model([input, batch['uv_mask']]) mesh_pred = uv2mesh(pred) - mesh_orig = batch['mesh_unposed'] - shape = batch['pose_shape'][:, 72:] - gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long) - 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] - - e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) - c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) - # l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2])) - l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1)) - n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + if config['use_skinning']: + mesh_orig = batch['mesh_unposed'] + uv_orig = batch['uv_unposed'] + shape = batch['pose_shape'][:, 72:] + gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long) + 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] + else: + mesh_orig = batch['mesh_posed'] + uv_orig = batch['uv_posed'] + mesh_pred_posed = mesh_pred - loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff'] + l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1)) + + ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1)) + loss = l2_mesh_loss + (1 - ms_ssim) * 1e3 optimizer.zero_grad() loss.backward() @@ -109,9 +111,13 @@ def train_step(batch): model.eval() with torch.inference_mode(): - ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1)) metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) - + e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) + c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) + # l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2])) + n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/euclid_error': metric, 'train/collision_ratio': c_ratio, 'train/ms_ssim': ms_ssim, 'train/collision_loss': c_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss } @@ -127,14 +133,21 @@ def valid_step(batch): input = torch.cat([batch['uv_normals'], batch['uv_body_posed'], batch['uv_static']], axis=1) else: input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1) + pred = model([input, batch['uv_mask']]) mesh_pred = uv2mesh(pred) - mesh_orig = batch['mesh_unposed'] - shape = batch['pose_shape'][:, 72:] - gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long) - 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 config['use_skinning']: + mesh_orig = batch['mesh_unposed'] + uv_orig = batch['uv_unposed'] + shape = batch['pose_shape'][:, 72:] + gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long) + 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] + else: + mesh_orig = batch['mesh_posed'] + uv_orig = batch['uv_posed'] + mesh_pred_posed = mesh_pred e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) @@ -143,8 +156,8 @@ def valid_step(batch): l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1)) n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) - ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1)) - loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff'] + ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1)) + loss = l2_mesh_loss + (1 - ms_ssim) * 1e3 metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) @@ -165,9 +178,14 @@ def plot_everything(batch: dict, type: str): pred = model([input, batch['uv_mask']]).detach() mesh_pred = uv2mesh(pred).cpu().numpy() - mesh_orig = batch['mesh_unposed'].cpu().numpy() + if config['use_skinning']: + mesh_orig = batch['mesh_unposed'].cpu().numpy() + uv_orig = batch['uv_unposed'] + else: + mesh_orig = batch['mesh_posed'].cpu().numpy() + uv_orig = batch['uv_posed'] - orig_img = torch.permute(batch['uv_unposed'], [0,2,3,1]) + orig_img = torch.permute(uv_orig, [0,2,3,1]) uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1]) uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1]) uv_normals_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2 @@ -178,20 +196,24 @@ def plot_everything(batch: dict, type: str): F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") mask = batch['mesh_mask'].cpu().numpy().astype(bool) - shape = batch['pose_shape'][:, 72:] - gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long) - 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])]).cpu().numpy() + if config['use_skinning']: + shape = batch['pose_shape'][:, 72:] + gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long) + 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])]).cpu().numpy() - garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask) - garment_unposed = [garment_unposed[i] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])] - gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(pred.shape[0])] + garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask) + garment_unposed = [garment_unposed[i] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])] + gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(pred.shape[0])] - for i in range(mesh_orig.shape[0]): - mesh_pred[i, mask[i], :] += shape_offsets[i, mask[i], :] - mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_pred, device=config['device']))[0].cpu().numpy() + for i in range(mesh_orig.shape[0]): + mesh_pred[i, mask[i], :] += shape_offsets[i, mask[i], :] + mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_pred, device=config['device']))[0].cpu().numpy() + else: + garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask) + mesh_posed = mesh_pred + gt = [mesh_orig[i][mask[i], :] for i in range(pred.shape[0])] gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])] - mesh_posed_fixed = fix_collision(torch.tensor(mesh_posed, device=config['device']), batch['mesh_body_posed'], torch.tensor(mask, device=config['device'])).cpu().numpy() body = batch['mesh_body_posed'].cpu().numpy()