From f3c1e3308e923c8e19073ef41212f115582a6b38 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 31 May 2023 21:40:18 +0000 Subject: [PATCH] direct training with normals --- configs/direct_3.py | 82 +++++++++++++++++++++++++++++++++++++++++ configs/direct_4.py | 82 +++++++++++++++++++++++++++++++++++++++++ run.py | 2 +- scripts/direct_train.py | 11 +++++- 4 files changed, 174 insertions(+), 3 deletions(-) create mode 100644 configs/direct_3.py create mode 100644 configs/direct_4.py diff --git a/configs/direct_3.py b/configs/direct_3.py new file mode 100644 index 0000000..37ac1fd --- /dev/null +++ b/configs/direct_3.py @@ -0,0 +1,82 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'direct_3' +config['notes'] = 'with normals input' +config['tags'] = ['torch', 'direct'] +config['script'] = 'scripts/direct_train.py' + +config['pretrained_model'] = '' +config['use_curvature'] = False +config['use_normals'] = 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/valid' +config['test_dir'] = '/data-net/dataset/test_ready' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +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_4.py b/configs/direct_4.py new file mode 100644 index 0000000..2fa83a7 --- /dev/null +++ b/configs/direct_4.py @@ -0,0 +1,82 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'direct_5' +config['notes'] = 'with normals input but only single garment sequence' +config['tags'] = ['torch', 'direct'] +config['script'] = 'scripts/direct_train.py' + +config['pretrained_model'] = '' +config['use_curvature'] = False +config['use_normals'] = 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/run.py b/run.py index 13272a1..22c264b 100644 --- a/run.py +++ b/run.py @@ -1,6 +1,6 @@ import os import sys -from configs.direct_1 import config +from configs.direct_3 import config import argparse # parser = argparse.ArgumentParser(description='') diff --git a/scripts/direct_train.py b/scripts/direct_train.py index fe8c2b1..1bdf526 100644 --- a/scripts/direct_train.py +++ b/scripts/direct_train.py @@ -76,6 +76,8 @@ def train_step(batch): # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp): if config['use_curvature']: input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1) + elif config['use_normals']: + 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']]) @@ -121,6 +123,8 @@ def valid_step(batch): with torch.inference_mode(): if config['use_curvature']: input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1) + elif config['use_normals']: + 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']]) @@ -154,6 +158,8 @@ def plot_everything(batch: dict, type: str): with torch.inference_mode(): if config['use_curvature']: input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1) + elif config['use_normals']: + 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']]).detach() @@ -164,8 +170,9 @@ def plot_everything(batch: dict, type: str): orig_img = torch.permute(batch['uv_unposed'], [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 pred_img = torch.permute(pred, [0,2,3,1]) - img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2) + img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals_img, orig_img, pred_img), axis=2) imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)] F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") @@ -273,7 +280,7 @@ if __name__ == '__main__': if not args['debug']: # model = torch.compile(model) - wandb.watch(model, log='all', log_freq=1000, log_graph=True) + wandb.watch(model, log='all', log_freq=10000, log_graph=True) train()