diff --git a/configs/iccvw/abl_all.py b/configs/iccvw/abl_all.py index b9d9639..8008e19 100644 --- a/configs/iccvw/abl_all.py +++ b/configs/iccvw/abl_all.py @@ -3,15 +3,15 @@ from utils.convnext_utils import CNBlockConfig config = dict() config['gpu'] = '2' -config['name'] = 'abl_iccvw_all_1' -config['notes'] = '' +config['name'] = 'abl_iccvw_all_2' +config['notes'] = 'fixed collision loss' config['tags'] = ['iccv', 'abl_all', 'torch'] config['script'] = 'scripts/cvae.py' -config['pretrained_model'] = '' +config['pretrained_model'] = 'checkpoints/iccvw/abl_iccvw_all_1' config['use_skinning'] = True # and deshaping -config['use_normals'] = True +config['use_normals'] = True config['use_extra_losses'] = True # edge, collision, 3d mesh config['use_norm_loss'] = True @@ -34,14 +34,16 @@ config['device'] = 'cuda' config['train_dir'] = '/data-local/127/cloth3d_processed' config['valid_dir'] = '/data-local/127/cloth3d_processed' -config['test_dir'] = '/data-net/dataset/test_ready' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' if config['is_skirt']: config['train_frames'] = 'frames_train_skirt.txt' config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' else: config['train_frames'] = 'frames_train.txt' config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' diff --git a/configs/iccvw/abl_all_3.py b/configs/iccvw/abl_all_3.py new file mode 100644 index 0000000..b5fb5f7 --- /dev/null +++ b/configs/iccvw/abl_all_3.py @@ -0,0 +1,196 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'abl_iccvw_all_3' +config['notes'] = 'fixed collision loss, train from zero' +config['tags'] = ['iccv', 'abl_all', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = '' + +config['use_skinning'] = True # and deshaping +config['use_normals'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True + +ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False) +assert ablation_count == 1 or ablation_count == 0 + +config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv +if config['use_skinning']: + assert (config['avg_garment_path'] != '') + +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision + +config['batch_size'] = 4 +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +############### +config['lr'] = 0.0001# 0.0001 +############### +config['starting_balancing_coeff'] = 0.0001 +config['balancing_activation_limit'] = 0.01 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') + +config['upscale_enhancer'] = '' +config['freeze_encoders'] = True +config['use_one_hot'] = False +config['bottleneck_scale'] = 2 + +config['cvae'] = dict() +config['cvae']['device'] = config['device'] +config['cvae']['is_label'] = False +config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 +config['cvae']['kernel_size'] = 7 +config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale'] +config['cvae']['latent_dim'] = config['cvae']['bottleneck_size'] + +block_setting = [ + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 256, 3), + CNBlockConfig(256, config['cvae']['bottleneck_size'], 9), + CNBlockConfig(config['cvae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3), + CNBlockConfig(256, 128, 9), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 3, 3), +] + +config['cvae']['encoder_block_settings'] = block_setting +config['cvae']['decoder_block_settings'] = block_setting_reverse + + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/abl_all_4.py b/configs/iccvw/abl_all_4.py new file mode 100644 index 0000000..75b6d81 --- /dev/null +++ b/configs/iccvw/abl_all_4.py @@ -0,0 +1,196 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'abl_iccvw_all_4' +config['notes'] = 'fixed collision loss, train from zero' +config['tags'] = ['iccv', 'abl_all', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = 'checkpoints/iccvw/abl_iccvw_all_3' + +config['use_skinning'] = True # and deshaping +config['use_normals'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True + +ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False) +assert ablation_count == 1 or ablation_count == 0 + +config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv +if config['use_skinning']: + assert (config['avg_garment_path'] != '') + +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision + +config['batch_size'] = 4 +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +############### +config['lr'] = 0.0001# 0.0001 +############### +config['starting_balancing_coeff'] = 0.0001 +config['balancing_activation_limit'] = 0.01 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') + +config['upscale_enhancer'] = '' +config['freeze_encoders'] = True +config['use_one_hot'] = False +config['bottleneck_scale'] = 2 + +config['cvae'] = dict() +config['cvae']['device'] = config['device'] +config['cvae']['is_label'] = False +config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 +config['cvae']['kernel_size'] = 7 +config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale'] +config['cvae']['latent_dim'] = config['cvae']['bottleneck_size'] + +block_setting = [ + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 256, 3), + CNBlockConfig(256, config['cvae']['bottleneck_size'], 9), + CNBlockConfig(config['cvae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3), + CNBlockConfig(256, 128, 9), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 3, 3), +] + +config['cvae']['encoder_block_settings'] = block_setting +config['cvae']['decoder_block_settings'] = block_setting_reverse + + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/abl_all_5.py b/configs/iccvw/abl_all_5.py new file mode 100644 index 0000000..3afe5e4 --- /dev/null +++ b/configs/iccvw/abl_all_5.py @@ -0,0 +1,196 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'abl_iccvw_all_5' +config['notes'] = 'fixed collision loss, train from zero, fixed latest checkpoint loading' +config['tags'] = ['iccv', 'abl_all', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = 'checkpoints/iccvw/abl_iccvw_all_3' + +config['use_skinning'] = True # and deshaping +config['use_normals'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True + +ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False) +assert ablation_count == 1 or ablation_count == 0 + +config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv +if config['use_skinning']: + assert (config['avg_garment_path'] != '') + +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision + +config['batch_size'] = 4 +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +############### +config['lr'] = 0.0001# 0.0001 +############### +config['starting_balancing_coeff'] = 0.0001 +config['balancing_activation_limit'] = 0.01 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') + +config['upscale_enhancer'] = '' +config['freeze_encoders'] = True +config['use_one_hot'] = False +config['bottleneck_scale'] = 2 + +config['cvae'] = dict() +config['cvae']['device'] = config['device'] +config['cvae']['is_label'] = False +config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 +config['cvae']['kernel_size'] = 7 +config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale'] +config['cvae']['latent_dim'] = config['cvae']['bottleneck_size'] + +block_setting = [ + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 256, 3), + CNBlockConfig(256, config['cvae']['bottleneck_size'], 9), + CNBlockConfig(config['cvae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3), + CNBlockConfig(256, 128, 9), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 3, 3), +] + +config['cvae']['encoder_block_settings'] = block_setting +config['cvae']['decoder_block_settings'] = block_setting_reverse + + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/abl_all_skirt.py b/configs/iccvw/abl_all_skirt.py new file mode 100644 index 0000000..d82a784 --- /dev/null +++ b/configs/iccvw/abl_all_skirt.py @@ -0,0 +1,196 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '2' +config['name'] = 'abl_iccvw_all_skirt' +config['notes'] = 'fixed collision loss, train from zero' +config['tags'] = ['iccv', 'abl_all', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = '' + +config['use_skinning'] = True # and deshaping +config['use_normals'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True + +ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False) +assert ablation_count == 1 or ablation_count == 0 + +config['avg_garment_path'] = 'custom_data/avg/train_skirt_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv +if config['use_skinning']: + assert (config['avg_garment_path'] != '') + +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision + +config['batch_size'] = 4 +config['resolution'] = 512 +config['is_skirt'] = True +config['dataset_in_memory'] = False +config['is_cape'] = False +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_1_skirt' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1_skirt' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1_skirt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +############### +config['lr'] = 0.0001# 0.0001 +############### +config['starting_balancing_coeff'] = 0.0001 +config['balancing_activation_limit'] = 0.01 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') + +config['upscale_enhancer'] = '' +config['freeze_encoders'] = True +config['use_one_hot'] = False +config['bottleneck_scale'] = 2 + +config['cvae'] = dict() +config['cvae']['device'] = config['device'] +config['cvae']['is_label'] = False +config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 +config['cvae']['kernel_size'] = 7 +config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale'] +config['cvae']['latent_dim'] = config['cvae']['bottleneck_size'] + +block_setting = [ + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 256, 3), + CNBlockConfig(256, config['cvae']['bottleneck_size'], 9), + CNBlockConfig(config['cvae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3), + CNBlockConfig(256, 128, 9), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 3, 3), +] + +config['cvae']['encoder_block_settings'] = block_setting +config['cvae']['decoder_block_settings'] = block_setting_reverse + + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/pyr_32.py b/configs/iccvw/pyr_32.py new file mode 100644 index 0000000..c5e7f06 --- /dev/null +++ b/configs/iccvw/pyr_32.py @@ -0,0 +1,302 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '1' +config['name'] = 'pyr_iccvw_32' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = False +config['device'] = 'cuda' + + +config['use_skinning'] = True # and deshaping +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True +config['use_normals'] = True + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [32] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.012 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'mesh_unposed', + 'uv_unposed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + config['signature'].add('uv_unposed' + res_modifier) + config['signature'].add('uv_normals' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' +config['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(level['bottleneck_size'], None, 16), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 3, 16), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['is_label'] = False +# config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +# config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/pyr_64.py b/configs/iccvw/pyr_64.py new file mode 100644 index 0000000..fa6b943 --- /dev/null +++ b/configs/iccvw/pyr_64.py @@ -0,0 +1,302 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '3' +config['name'] = 'pyr_iccvw_64' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = False +config['device'] = 'cuda' + + +config['use_skinning'] = True # and deshaping +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True +config['use_normals'] = True + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.012 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'mesh_unposed', + 'uv_unposed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + config['signature'].add('uv_unposed' + res_modifier) + config['signature'].add('uv_normals' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' +config['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(level['bottleneck_size'], None, 16), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 3, 16), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['is_label'] = False +# config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +# config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/pyr_64_2.py b/configs/iccvw/pyr_64_2.py new file mode 100644 index 0000000..cf7affb --- /dev/null +++ b/configs/iccvw/pyr_64_2.py @@ -0,0 +1,302 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'pyr_iccvw_64_2' +config['notes'] = 'continue from previous' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = False +config['device'] = 'cuda' + + +config['use_skinning'] = True # and deshaping +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True +config['use_normals'] = True + +config['pretrained_model'] = 'checkpoints/iccvw/pyr_iccvw_64' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.012 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'mesh_unposed', + 'uv_unposed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + config['signature'].add('uv_unposed' + res_modifier) + config['signature'].add('uv_normals' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' +config['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(level['bottleneck_size'], None, 16), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 3, 16), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['is_label'] = False +# config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +# config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/pyr_64_cape_f.py b/configs/iccvw/pyr_64_cape_f.py new file mode 100644 index 0000000..6c58a1f --- /dev/null +++ b/configs/iccvw/pyr_64_cape_f.py @@ -0,0 +1,295 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '2' +config['name'] = 'pyr_iccvw_64_cape_f' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = True +config['device'] = 'cuda' + + +config['use_skinning'] = True # and deshaping +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True +config['use_normals'] = True + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127_2/cape_processed/train/female' +config['valid_dir'] = '/data-local/127_2/cape_processed/test/female' +config['test_dir'] = '/data-local/127_2/cape_processed/test/female' +config['train_frames'] = None +config['valid_frames'] = None + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.012 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'mesh_unposed', + 'uv_unposed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + config['signature'].add('uv_unposed' + res_modifier) + config['signature'].add('uv_normals' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_cape_f' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_cape_f' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_cape_f' +config['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(level['bottleneck_size'], None, 16), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 3, 16), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['is_label'] = False +# config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +# config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/pyr_64_cape_m.py b/configs/iccvw/pyr_64_cape_m.py new file mode 100644 index 0000000..e639ff8 --- /dev/null +++ b/configs/iccvw/pyr_64_cape_m.py @@ -0,0 +1,295 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'pyr_iccvw_64_cape_m' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = True +config['device'] = 'cuda' + + +config['use_skinning'] = True # and deshaping +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True +config['use_normals'] = True + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127_2/cape_processed/train/male' +config['valid_dir'] = '/data-local/127_2/cape_processed/test/male' +config['test_dir'] = '/data-local/127_2/cape_processed/test/male' +config['train_frames'] = None +config['valid_frames'] = None + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.012 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'mesh_unposed', + 'uv_unposed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + config['signature'].add('uv_unposed' + res_modifier) + config['signature'].add('uv_normals' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_cape_m' +config['body_model'] = 'checkpoints/iccvw/body_torch_cape_m' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_cape_m' +config['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(level['bottleneck_size'], None, 16), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 3, 16), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['is_label'] = False +# config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +# config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/pyr_64_norm.py b/configs/iccvw/pyr_64_norm.py new file mode 100644 index 0000000..a64fafa --- /dev/null +++ b/configs/iccvw/pyr_64_norm.py @@ -0,0 +1,288 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'pyr_iccvw_64_norm' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = False +config['device'] = 'cuda' + + +config['use_skinning'] = True # and deshaping +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True +config['use_normals'] = False + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.012 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + config['signature'].add('uv_unposed' + res_modifier) + config['signature'].add('uv_normals' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' +config['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2 - 40*2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2 - 40*2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2 - 40*2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2 - 40*2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + raise NotImplementedError + else: + raise NotImplementedError + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['is_label'] = False +# config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +# config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/pyr_64_reg.py b/configs/iccvw/pyr_64_reg.py new file mode 100644 index 0000000..65b00d0 --- /dev/null +++ b/configs/iccvw/pyr_64_reg.py @@ -0,0 +1,306 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '1' +config['name'] = 'pyr_iccvw_64_reg' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = False +config['device'] = 'cuda' + + +config['use_skinning'] = True # and deshaping +config['use_extra_losses'] = False # edge, collision, 3d mesh +config['use_norm_loss'] = True +config['use_normals'] = True + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.012 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + config['signature'].add('uv_unposed' + res_modifier) + config['signature'].add('uv_normals' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' +config['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(level['bottleneck_size'], None, 16), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 3, 16), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['is_label'] = False +# config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +# config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/pyr_64_skin.py b/configs/iccvw/pyr_64_skin.py new file mode 100644 index 0000000..8d2e6a3 --- /dev/null +++ b/configs/iccvw/pyr_64_skin.py @@ -0,0 +1,309 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'pyr_iccvw_64_skin' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = False +config['device'] = 'cuda' + + +config['use_skinning'] = False # and deshaping +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True +config['use_normals'] = True + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.012 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + if config['use_skinning']: + config['signature'].add('uv_unposed' + res_modifier) + else: + config['signature'].add('uv_posed' + res_modifier) + config['signature'].add('uv_normals' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' +config['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(level['bottleneck_size'], None, 16), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 3, 16), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['is_label'] = False +# config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +# config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/pyr_64_skirt.py b/configs/iccvw/pyr_64_skirt.py new file mode 100644 index 0000000..6c1549e --- /dev/null +++ b/configs/iccvw/pyr_64_skirt.py @@ -0,0 +1,302 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '1' +config['name'] = 'pyr_iccvw_64_skirt' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = True +config['is_cape'] = False +config['device'] = 'cuda' + + +config['use_skinning'] = True # and deshaping +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_norm_loss'] = True +config['use_normals'] = True + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-local/127_2/cloth3d_processed_test' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' + config['test_frames'] = 'frames_test_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + config['test_frames'] = 'frames_test.txt' + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.012 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'mesh_unposed', + 'uv_unposed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + config['signature'].add('uv_unposed' + res_modifier) + config['signature'].add('uv_normals' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_1_skirt' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1_skirt' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1_skirt' +config['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256_skirt' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128_skirt' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64_skirt' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32_skirt' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 + + block_setting = [ + CNBlockConfig(level['bottleneck_size'], None, 16), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 3, 16), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['is_label'] = False +# config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +# config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/data/dataset.py b/data/dataset.py index 7ebc6d5..4da75f7 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -140,7 +140,7 @@ class DatasetCloth3D(torch.utils.data.Dataset): self.resolutions = resolutions self.transforms = tv.transforms.ToTensor() - self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=True) + self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=False) self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device) @@ -208,7 +208,7 @@ class DatasetCloth3D(torch.utils.data.Dataset): item['uv_unposed' + res_string] = uv_unposed_no_shape if 'uv_body_posed' + res_string in self.signature: - uv_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed'] + uv_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string[1:], f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed'] item['uv_body_posed' + res_string] = uv_body_posed_no_orient if 'uv_posed' + res_string in self.signature: @@ -244,6 +244,9 @@ if __name__ == '__main__': import plotly.express as ex from plotly.subplots import make_subplots from custom_utils.custom_utils import _maskFaces + from losses.losses import collision_loss, normal_loss + from smpl.smpl_torch import SMPL + from torch.profiler import profile, record_function, ProfilerActivity F = np.load('smpl/smpl_faces.npy') F_skirt = np.load('smpl/smpl_faces_skirt.npy') @@ -267,61 +270,98 @@ if __name__ == '__main__': dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt') train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) + smpl = SMPL(orig_smpl=False, path='smpl', device='cuda') + # create a profiler object + profiler = torch.profiler.profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=torch.profiler.schedule(wait=2, warmup=1, active=3, repeat=2), + on_trace_ready=torch.profiler.tensorboard_trace_handler("log_dir"), + with_stack=True, + profile_memory=True, + ) + + + step = 0 for batch in train_dataloader: batch - b = 0 - mesh_mask = batch['mesh_mask'][b].cpu().numpy() - faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F - _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask) + print(batch['id'][0]) + with profiler: + m = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], batch['mesh_unposed'])[0] + profiler.step() + step += 1 + if step > 25: + break - mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask] - mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask] - mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask] - mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy() - - fig = make_subplots(rows=2, cols=2, - specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}], - [{'type': 'mesh3d'}, {'type': 'mesh3d'}]]) - - fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1) - fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2) - fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=1) - fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], opacity=1), row=2, col=2) - - fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False), - yaxis=dict(range=[-1, 1], autorange=False), - zaxis=dict(range=[-1, 1], autorange=False), - aspectmode='cube'), - scene2=dict(xaxis=dict(range=[-1, 1], autorange=False), - yaxis=dict(range=[-1, 1], autorange=False), - zaxis=dict(range=[-1, 1], autorange=False), - aspectmode='cube'), - scene3=dict(xaxis=dict(range=[-1, 1], autorange=False), - yaxis=dict(range=[-1, 1], autorange=False), - zaxis=dict(range=[-1, 1], autorange=False), - aspectmode='cube'), - scene4=dict(xaxis=dict(range=[-1, 1], autorange=False), - yaxis=dict(range=[-1, 1], autorange=False), - zaxis=dict(range=[-1, 1], autorange=False), - aspectmode='cube')) - - fig.show() + profiler.stop() + print(profiler.key_averages().table(sort_by="cpu_time_total", row_limit=25)) + print(profiler.key_averages().table(sort_by="cuda_time_total", row_limit=25)) + print(profiler.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=10)) - uv_mask = batch['uv_mask'][b].cpu().numpy() - uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0) - uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2 - uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2 - uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2 - uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2 - uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3 + # b = 0 + # mesh_mask = batch['mesh_mask'][b].cpu().numpy() + # faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F + # _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask) - img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2) - img = np.transpose(img, (1, 2, 0)) - ex.imshow(img).show() + # c = collision_loss(batch['mesh_posed'], batch['mesh_body_posed'], batch['mesh_mask']) + # fig = go.Figure(data=[go.Mesh3d(x=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,0], y=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,1], z=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,2], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), + # go.Mesh3d(x=batch['mesh_body_posed'].cpu().numpy()[0,:,0], y=batch['mesh_body_posed'].cpu().numpy()[0,:,1], z=batch['mesh_body_posed'].cpu().numpy()[0,:,2], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1), + # ]) + # fig.update_layout(scene_aspectmode='cube') + # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) + # fig.show() - pass + + + + # mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask] + # mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask] + # mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask] + # mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy() + + # fig = make_subplots(rows=2, cols=2, + # specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}], + # [{'type': 'mesh3d'}, {'type': 'mesh3d'}]]) + + # fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1) + # fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2) + # fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=1) + # fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], opacity=1), row=2, col=2) + + # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False), + # yaxis=dict(range=[-1, 1], autorange=False), + # zaxis=dict(range=[-1, 1], autorange=False), + # aspectmode='cube'), + # scene2=dict(xaxis=dict(range=[-1, 1], autorange=False), + # yaxis=dict(range=[-1, 1], autorange=False), + # zaxis=dict(range=[-1, 1], autorange=False), + # aspectmode='cube'), + # scene3=dict(xaxis=dict(range=[-1, 1], autorange=False), + # yaxis=dict(range=[-1, 1], autorange=False), + # zaxis=dict(range=[-1, 1], autorange=False), + # aspectmode='cube'), + # scene4=dict(xaxis=dict(range=[-1, 1], autorange=False), + # yaxis=dict(range=[-1, 1], autorange=False), + # zaxis=dict(range=[-1, 1], autorange=False), + # aspectmode='cube')) + + # fig.show() + + + # uv_mask = batch['uv_mask'][b].cpu().numpy() + # uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0) + # uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2 + # uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2 + # uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2 + # uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2 + # uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3 + + # img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2) + # img = np.transpose(img, (1, 2, 0)) + # ex.imshow(img).show() + + # pass pass diff --git a/data/dataset_cape.py b/data/dataset_cape.py index 4c4c370..ef6c9d6 100644 --- a/data/dataset_cape.py +++ b/data/dataset_cape.py @@ -82,7 +82,7 @@ class StaticCache(): res_string = '_' + str(res) if 'uv_static' + res_string in self.signature: - uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static + uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static.astype(np.float32) if 'uv_mask' + res_string in self.signature: if self.is_cape: data['uv_mask' + res_string] = self.uv_masks[res] @@ -158,6 +158,7 @@ class DatasetCape(torch.utils.data.Dataset): item['id'] = frame.id if 'frame_nr' in self.signature: item['frame_nr'] = frame.frame_nr + item['garment'] = 'cape' if 'pose_shape' in self.signature: pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32) pose_shape[:3] = 0 @@ -183,27 +184,18 @@ class DatasetCape(torch.utils.data.Dataset): if res != 512: res_string = '_' + str(res) - if 'uv_curvature' + res_string in self.signature: - uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz')); - if res_string != '': - uv_curvature = uv_curvature['uv_curvature'] - else: - uv_curvature = uv_curvature['curvature_uv'] - uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device) - item['uv_curvature' + res_string] = uv_curvature - if 'uv_unposed' + res_string in self.signature: uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape'] - item['uv_unposed' + res_string] = uv_unposed_no_shape + item['uv_unposed' + res_string] = uv_unposed_no_shape.astype(np.float32) if 'uv_body_posed' + res_string in self.signature: uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz')) uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient'] - item['uv_body_posed' + res_string] = uv_posed_body_no_orient + item['uv_body_posed' + res_string] = uv_posed_body_no_orient.astype(np.float32) if 'uv_posed' + res_string in self.signature: uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient'] - item['uv_posed' + res_string] = uv_posed + item['uv_posed' + res_string] = uv_posed.astype(np.float32) if 'uv_normals' + res_string in self.signature: uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals'] diff --git a/data/downscale_cape.py b/data/downscale_cape.py new file mode 100644 index 0000000..87cac14 --- /dev/null +++ b/data/downscale_cape.py @@ -0,0 +1,207 @@ +import argparse +import os +import tqdm +import cv2 +import scipy.interpolate as spi +import numpy as np +import plotly.graph_objects as go +import plotly.express as ex +import torch +from torchvision.transforms import Resize +from torchvision.transforms.functional import InterpolationMode + +os.environ["CUDA_VISIBLE_DEVICES"] = "0" + + +def get_img_idxs(): + img_idxs = dict() + img_coords = dict() + uv_masks = dict() + uv_mask = np.load('custom_data/garment_uv_mask_full.npy') + coords = np.transpose(np.stack(np.nonzero((1 - uv_mask))), (1,0)) + img_idx = np.transpose(np.stack(np.nonzero((uv_mask))), (1,0)) + img_idxs[512] = img_idx + img_coords[512] = coords + uv_masks[512] = uv_mask.astype(np.int32) + for res in [32, 64, 128, 256]: + down = resizers_nearest[res](torch.from_numpy(uv_mask[None,:,:]).to(device)).cpu().numpy() + down = np.squeeze(down.astype(bool)) + coords = np.transpose(np.stack(np.nonzero((1 - down))), (1,0)) + img_idx = np.transpose(np.stack(np.nonzero((down))), (1,0)) + img_idxs[res] = img_idx + img_coords[res] = coords + uv_masks[res] = down.astype(np.int32) + + + uv_masks_body = dict() + uv_mask_body = np.load('custom_utils/mask_256x512.npy') + uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2) + uv_masks_body[512] = uv_mask_body + for res in [32, 64, 128, 256]: + down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:]).to(device)).cpu().numpy() + down = np.squeeze(down.astype(bool)) + uv_masks_body[res] = down + + return img_idxs, img_coords, uv_masks, uv_masks_body + + +def fill(img, mask): + coords = img_coords[img.shape[0]] + img_idx = img_idxs[img.shape[0]] + # coords = np.transpose(np.stack(np.nonzero((1 - uv_masks[img.shape[0]]))), (1,0)) + # img_idx = np.transpose(np.stack(np.nonzero((uv_masks[img.shape[0]]))), (1,0)) + img_data = img[img_idx[:,0], img_idx[:,1], :] + + interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) + img_2 = interpolator(coords) + res = np.zeros((512,256,3), dtype=np.float32) + res[coords[:,0], coords[:,1], :] = img_2 + res[img_idx[:,0], img_idx[:,1], :] = img_data + + return res + + +def fill_label(img, mask): + coords = img_coords[img.shape[0]] + img_idx = img_idxs[img.shape[0]] + img_data = img[img_idx[:,0], img_idx[:,1]] + + interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) + img_2 = interpolator(coords) + res = np.zeros((512,256), dtype=np.float32) + res[coords[:,0], coords[:,1]] = img_2 + res[img_idx[:,0], img_idx[:,1]] = img_data + + return res + + +def generate_downscaled_uv_map(folder, id, size): + img = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); img = img['uv_mask'].astype(np.float32) + down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) + down = np.squeeze(down.numpy().astype(bool)) + np.savez_compressed(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz'), uv_mask = down) + + +class UVMaskStore(): + def __init__(self, folder, size): + self.size = size + + def __getitem__(self, id) -> dict: + return uv_masks[self.size] + + +def downscale_uv_maps_posed_body(folder, type, sizes): + for size in sizes: + os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) + + frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient'] + img = fill(img, uv_masks_body[512]) + for size in sizes: + mask = uv_masks_body[size] + down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() + down = down * mask[None,:,:] + np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_body_no_orient = down) + + +def downscale_uv_normals(folder, type, sizes): + for size in sizes: + os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) + + frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, frame)); img = img['uv_normals'] + img = fill(img, uv_masks_body[512]) + for size in sizes: + mask = uv_masks_body[size] + down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() + down = down * mask[None,:,:] + np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_normals = down) + + + +def downscale_uv_maps_unposed(folder, type, sizes): + for size in sizes: + os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) + + frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape'] + mask_full = uv_mask_store[512][None] + img = fill(img, mask_full) + for size in sizes: + mask = uv_mask_store[size][None] + down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() + down = down * mask[None,:,:] + np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down) + + +def downscale_uv_maps_static(folder, type, sizes): + for size in sizes: + os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) + + frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, frame)); img = img['uv_static'] + mask_full = uv_mask_store[512][None] + img = fill(img, mask_full) + + for size in sizes: + down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() + mask = uv_mask_store[size][None] + down = down * mask[None,:,:] + np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_static = down) + + +def downscale_uv_maps_mask(folder, type, size): + os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) + + frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32) + down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) + down = np.squeeze(down.numpy().astype(bool)) + np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down) + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/male') + parser.add_argument("-s", "--start", help = "start of range", default=0) + parser.add_argument("-e", "--end", help = "end of range", default=100000) + parser.add_argument("-n", "--name", help = "end of range", default='lhunor') + parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true') + args = parser.parse_args() + + device = 'cuda' + + sizes = [32, 64, 128, 256] + folder = args.folder + + resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes } + resizers_nearest = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT).to(device) for res in sizes } + + img_idxs, img_coords, uv_masks, uv_masks_body = get_img_idxs() + + uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes } + uv_mask_store[512] = UVMaskStore(folder, size=512) + + # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32) + # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64) + # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128) + # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256) + # downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes) + + downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes) + downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes) + downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes) + + + print() \ No newline at end of file diff --git a/data/preprocessing.py b/data/preprocessing.py index dff4038..ea09635 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -36,94 +36,6 @@ from smpl.smpl_torch import SMPL from custom_utils.custom_utils import _maskFaces - -def clean_dataset(folder): - ids = os.listdir(folder) - - for id in tqdm(ids): - - if not os.path.isdir(os.path.join(folder, id)): - continue - - files = os.listdir(os.path.join(folder, id)) - - for file in files: - if file.startswith('glmarks') or file.find('nricp') != -1: - os.remove(os.path.join(folder, id, file)) - - -def generate_frames_file(folder, file_name): - ids = os.listdir(folder) - frames = [] - - for id in tqdm(ids): - - if not os.path.isdir(os.path.join(folder, id)): - continue - - garments = [f.replace('.obj','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('.obj')] - info = loadInfo(os.path.join(folder, id, 'info.mat')) - - for garment in garments: - for i in range(info['poses'].shape[1]): - frames.append(Frame(id, i, garment)) - - with open(os.path.join(folder, file_name), 'w') as f: - for i, frame in enumerate(frames): - f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n') - - -def split_data(folder, ratio): - with open(os.path.join(folder, 'frames_all.txt'), "r") as f: - frames = f.readlines() - - ids = sorted(list(set([frame.split()[0] for frame in frames]))) - train_last_id = round(len(ids) * ratio) - - train_frames = [] - valid_frames = [] - - for frame in tqdm(frames): - id = frame.split()[0] - if id in ids[:train_last_id]: - train_frames.append(frame) - else: - valid_frames.append(frame) - - with open(os.path.join(folder, f'frames_train_{ratio}.txt'), 'w') as f: - f.writelines(train_frames) - - with open(os.path.join(folder, f'frames_valid_{round(1-ratio,1)}.txt'), 'w') as f: - f.writelines(valid_frames) - - -def genereated_registered_and_mask(folder, is_skirt = False): - ids = os.listdir(folder) - - for id in tqdm(ids): - garment_dir = os.path.join(folder, id) - if not os.path.isdir(garment_dir): - continue - garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and not f.endswith('.mat') and not f.endswith('.obj') and not f.startswith('Combined')] - - assert len(garments) == 1 - - if os.path.exists(os.path.join(garment_dir, garments[0] + '_registered.pc16')): - continue - - garment = garments[0] - - verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] - map = sio.loadmat(os.path.join(garment_dir, garment + '_map.mat'))['map'].reshape((-1)) - verts = verts[:, map, :] # ---------------> This is the garment after registration - raw_mask = sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].reshape((-1)) - mask = raw_mask > 0 - # boundaries = raw_mask == 2 - - writePC2(os.path.join(garment_dir, garment + '_registered.pc16'), verts, True) - np.save(os.path.join(garment_dir, garment + '_mask.npy'), mask) - - def create_extras_file(folder): # create a new pytables file root = '/'.join(folder.split('/')[:-1]) @@ -261,14 +173,15 @@ def create_mesh_body_posed_no_orient(folder, out_folder): continue info = sio.loadmat(os.path.join(folder, id, 'info.mat')) - - pose, shape, gender = info['poses'], info['shape'], info['gender'] + # if pose.shape[0] == 72: pose = np.transpose(pose, [1, 0]).astype(np.float32) pose[:, :3] = 0 beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32) + smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0]) + T = smpl(torch.tensor(pose, device=smpl.device), torch.tensor(np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), device=smpl.device))[0] T = T.cpu().numpy() @@ -456,6 +369,9 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1) pose = torch.tensor(pose, dtype=torch.float32, device='cuda') shape = torch.tensor(shape, dtype=torch.float32, device='cuda') + + smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0]) + for garment in garments: if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): @@ -627,6 +543,49 @@ def create_frames_file(folder, out_folder, split=0.9): f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') +def create_frames_file_test(folder, out_folder): + ids = os.listdir(folder) + + frames_non_skirt = [] + frames_skirt = [] + ids_non_skirt = set() + ids_skirt = set() + + for id in tqdm(ids): + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + info = sio.loadmat(os.path.join(folder, id, 'info.mat')) + garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + frames = info['poses'].shape[1] + + assert info['poses'].shape[0] == 72 + + for garment in garments: + for i in range(frames): + if garment == 'Skirt' or garment == 'Dress': + frames_skirt.append(Frame(id, i, garment)) + ids_skirt.add(id) + else: + frames_non_skirt.append(Frame(id, i, garment)) + ids_non_skirt.add(id) + + ids_skirt = list(ids_skirt) + ids_non_skirt = list(ids_non_skirt) + + with open(os.path.join(out_folder, 'frames_test.txt'), 'w') as f: + for frame in frames_non_skirt: + if frame.id in ids_non_skirt: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, 'frames_test_skirt.txt'), 'w') as f: + for frame in frames_skirt: + if frame.id in ids_skirt: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + + def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1): ids = os.listdir(folder) @@ -674,25 +633,26 @@ def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1): def check_irregular_poses(folder, folder_reg, out_folder): - frame: Frame = Frame('01512', 50, 'Dress') - mesh_posed = readPC2Frame(os.path.join(out_folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) - # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy') + for i in range(100, 200, 5): + frame: Frame = Frame('01512', i, 'Dress') + mesh_posed = readPC2Frame(os.path.join(out_folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy') - mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy') - mesh_body_posed = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V'] - mesh_body_posed = mesh_body_posed[int(mesh_body_posed.shape[0]/2)+frame.frame_nr] + mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy') + mesh_body_posed = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V'] + mesh_body_posed = mesh_body_posed[::2][frame.frame_nr] - F = np.load('smpl/smpl_faces_skirt.npy') - _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) + F = np.load('smpl/smpl_faces_skirt.npy') + _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) - fig = go.Figure(data=[go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)]) - fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1)) - fig.update_layout(scene = dict( - xaxis = dict(range=[-1,1],), - yaxis = dict(range=[-1,1],), - zaxis = dict(range=[-1,1],),),) - fig.update_layout(scene_aspectmode='cube') - fig.show() + fig = go.Figure(data=[go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)]) + fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1)) + fig.update_layout(scene = dict( + xaxis = dict(range=[-1,1],), + yaxis = dict(range=[-1,1],), + zaxis = dict(range=[-1,1],),),) + fig.update_layout(scene_aspectmode='cube') + fig.show() pass @@ -734,13 +694,118 @@ def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'): np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment) +def remove_irregular_poses(folder, folder_reg, out_folder): + ids = sorted(os.listdir(folder)) + + for id in tqdm(ids): + + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + info = loadInfo(os.path.join(folder, id, 'info.mat')) + + for garment in garments: + + verts_count = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'].shape[0] + if info['poses'].shape[1] != verts_count: + shutil.rmtree(os.path.join(folder, id)) + if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')): + os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')) + # os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.npy')) + + if len(garments) == 2: + if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.npy')) + break + # os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garments[1]}.npy')) + + +def double_check_irregular_poses(folder, folder_reg, out_folder): + # ids = sorted(os.listdir(folder)) + + # for id in tqdm(ids): + + # garment_dir = os.path.join(folder, id) + # if not os.path.isdir(garment_dir): + # continue + + # garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + # info = loadInfo(os.path.join(folder, id, 'info.mat')) + + # for garment in garments: + # body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0] + # verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0] + + # assert body_count == verts_count + + frames = sorted(os.listdir(out_folder + '/mesh_posed_no_orient')) + frames = [(id.split('_')[0], id.split('_')[1].split('.')[0]) for id in frames] + + for frame in tqdm(frames): + id = frame[0] + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + garment = frame[1] + body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0] + verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0] + + assert body_count == verts_count + + +def get_ids(): + ids = [] + with open(os.path.join('/data-local/127/cloth3d_processed', 'frames_train.txt'), 'r') as f: + for line in f: + ids.append(line.split()[0]) + + with open(os.path.join('/data-local/127/cloth3d_processed', 'frames_train_skirt.txt'), 'r') as f: + for line in f: + ids.append(line.split()[0]) + + ids = sorted(list(set(ids))) + + with open(os.path.join('/data-local/127/cloth3d_processed', 'frames_train_ids.txt'), 'w') as f: + for id in ids: + f.write(f'{id}\n') + + ids = [] + with open(os.path.join('/data-local/127_2/cloth3d_processed_test', 'frames_test.txt'), 'r') as f: + for line in f: + ids.append(line.split()[0]) + + with open(os.path.join('/data-local/127_2/cloth3d_processed_test', 'frames_test_skirt.txt'), 'r') as f: + for line in f: + ids.append(line.split()[0]) + + ids = sorted(list(set(ids))) + + with open(os.path.join('/data-local/127_2/cloth3d_processed_test', 'frames_test_ids.txt'), 'w') as f: + for id in ids: + f.write(f'{id}\n') + + pass + if __name__ == '__main__': # folder = '/data-net/dataset/cloth3d/test' # folder_reg = '/data-local/127/cloth3d_registration/test' - # out_folder = '/data-local/127/cloth3d_processed_test' + # out_folder = '/data-local/127_2/cloth3d_processed_test' folder = '/data-net/dataset/cloth3d/train' folder_reg = '/data-local/127/cloth3d_registration/train' @@ -754,6 +819,8 @@ if __name__ == '__main__': # create_extras_file(folder) # check_irregular_poses(folder, folder_reg, out_folder) + # remove_irregular_poses(folder, folder_reg, out_folder) + # double_check_irregular_poses(folder, folder_reg, out_folder) # create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder) # verify_registration_ans_mask(folder, out_folder) @@ -762,9 +829,11 @@ if __name__ == '__main__': # create_mesh_body_posed_no_orient(folder, out_folder) # create_mesh_normals(folder, out_folder) # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end) - calculate_average(folder, out_folder) + # calculate_average(folder, out_folder) # create_frames_file(folder, out_folder) + get_ids() + # create_frames_file_test(folder, out_folder) # create_frames_file_partial(folder, out_folder, part=0.5) # folder = '/data-local/127_2/cape_processed/test/female' diff --git a/losses/losses.py b/losses/losses.py index 5c59dfc..06fa74a 100644 --- a/losses/losses.py +++ b/losses/losses.py @@ -4,9 +4,9 @@ import numpy as np # import tensorflow_graphics as tfg import scipy as sp import scipy.io as sio +from scipy.spatial import KDTree from pytorch3d.structures import Meshes -from scipy.spatial import cKDTree from torch_cluster import knn # from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals # from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template @@ -33,15 +33,26 @@ def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor): batch_size = V.shape[0] vertices_count = V.shape[1] batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(V.shape[1]) - batch_y = torch.arange(batch_size, device=B.get_device()).repeat_interleave(B.shape[1]) V = V.reshape(-1, V.shape[-1]) B = B.reshape(-1, B.shape[-1]) - idx = knn(V, B.to(V.dtype), k=1, batch_x=batch_x, batch_y=batch_y)[1] + idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_x)[1] idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count) return idx +# def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor): +# # create an empty tensor to store the indices of the nearest neighbors +# batch_size = V.shape[0] +# vertices_count = V.shape[1] +# idx = torch.zeros((batch_size, vertices_count), dtype=torch.long, device=V.device) + +# # loop over the batch dimension and find the nearest neighbor for each point in B +# for i in range(batch_size): +# tree = KDTree(B[i].cpu().numpy()) +# _, nearest = tree.query(V[i].cpu().numpy(), k=1) +# idx[i] = torch.from_numpy(nearest) + +# return idx -# @tf.function def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004): if B.shape[1] == 14475: F_norm_local = F_norm diff --git a/models/pyramid_level.py b/models/pyramid_level.py index 62fcea1..b8abeb1 100644 --- a/models/pyramid_level.py +++ b/models/pyramid_level.py @@ -100,7 +100,7 @@ from models.upscaler import Upscaler class PyramidLevel(nn.Module): - def __init__(self, config): + def __init__(self, config, use_normals=True): super(PyramidLevel, self).__init__() self.size = config['resolution'] @@ -114,6 +114,7 @@ class PyramidLevel(nn.Module): load_model(self.upscale_enhancer, config['upscale_enhancer']) for param in self.upscale_enhancer.parameters(): param.requires_grad = False + self.use_normals = use_normals def enhance(self, low, mask): x = self.upscale_enhancer([low, mask]) @@ -152,9 +153,13 @@ class PyramidLevel(nn.Module): def forward(self, inputs, z_cond): - uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs + uv_unposed, uv_template, uv_body_posed, uv_normals, uv_mask, uv_mask_high = inputs - x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1) + if self.use_normals: + x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_normals], dim=1) + else: + x = torch.concat([uv_unposed, uv_template, uv_body_posed], dim=1) + mean, logvar = self.encoder(x) z = self.reparameterize(mean, logvar) z = torch.reshape(z, (mean.shape[0], -1, 8, 4)) diff --git a/models/pyramid_model.py b/models/pyramid_model.py index 0a041e8..447aab5 100644 --- a/models/pyramid_model.py +++ b/models/pyramid_model.py @@ -3,133 +3,12 @@ import time import random import torch from torch import nn -import torch.nn.functional as F -# import numpy as np -# import scipy as sp -# import plotly.graph_objects as go -# import plotly.express as ex from utils.checkpoints import load_model from models.base import AEEncoder from models.pyramid_level import PyramidLevel - - -# class Pyramid(nn.Module): - -# def __init__(self, config): -# super(Pyramid, self).__init__() - -# self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], kernel_size=config['cond_ae_static']['kernel_size'], activation=config['cond_ae_static']['activation']) -# if config['pyramid']['static_model'] != '': -# load_model(self.static_enc, config['pyramid']['static_model'] + '_enc') - -# self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], kernel_size=config['cond_ae_body']['kernel_size'], activation=config['cond_ae_body']['activation']) -# if config['pyramid']['body_model'] != '': -# load_model(self.body_enc, config['pyramid']['body_model'] + '_enc') - -# self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], kernel_size=config['cond_ae_curv']['kernel_size'], activation=config['cond_ae_curv']['activation']) -# if config['pyramid']['curvature_model'] != '': -# load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc') - -# if config['freeze_encoders']: -# for param in self.static_enc.parameters(): param.requires_grad = False -# for param in self.body_enc.parameters(): param.requires_grad = False -# for param in self.curv_enc.parameters(): param.requires_grad = False - - -# self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae_static': config['cond_ae_static'], 'cond_ae_body': config['cond_ae_body'], 'cond_ae_curv': config['cond_ae_curv'], } }) for i in range(len(config['pyramid']['levels']))] - -# # if config['is_cape']: -# # if len(config['pyramid']['resolutions']) == 2: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 3: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 4: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # else: - -# # if len(config['pyramid']['resolutions']) == 2: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 3: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 4: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) -# # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 5: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)]) -# # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) -# # self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) - -# # self.add_noise = config['pyramid']['add_noise'] -# # self.noise_amount = config['pyramid']['noise_amount'] -# self.use_offsets = config['use_offsets'] - -# # self.upscale_enhancer = None -# # if config['pyramid']['final_upscale_enhancer'] != '': -# # k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5)) -# # self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size) -# # load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer']) -# # self.upscale_enhancer.trainable = False - -# def enhance(self, low, mask): -# x = self.upscale_enhancer([low, mask]) -# return x - -# def encode_conds(self, uv_template, uv_body_posed, uv_curv): -# z_static = self.static_enc(uv_template) -# z_body = self.body_enc(uv_body_posed) -# z_curv = self.curv_enc(uv_curv) - -# z_cond = torch.concat([z_static, z_body, z_curv], dim=-1) - -# return z_cond - -# def call(self, inputs): -# levels_input, conds = inputs -# uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2] - -# z_cond = self.encode_conds(uv_static, uv_body, uv_curv) - -# # if self.add_noise == True and training == True: -# # z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount) - -# out_levels = [] -# pred, mean, logvar = self.levels[0](levels_input[0], z_cond) -# if self.use_offsets: -# pred = pred + levels_input[0][2] # uv_body_posed_32 -# if 1 != len(self.levels): -# pred = self.levels[0].enhance(pred, levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res -# pred = F.interpolate(pred, size=(self.levels[1].size, self.levels[1].size//2), mode='bilinear') -# out_levels.append([pred, mean, logvar]) -# pred = pred * levels_input[1][-2][:,None,:,:] -# else: -# pred = self.enhance(pred, levels_input[0][-2]) -# pred = F.interpolate(pred, size=(512, 256), mode='bilinear') -# pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask -# out_levels.append([pred, mean, logvar]) - -# for i in range(1, len(self.levels)): -# pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond) -# pred = pred + pred_level -# if i != len(self.levels) - 1: -# pred = self.levels[i].enhance(pred, levels_input[i][-2]) -# out_levels.append([pred, mean, logvar]) -# pred = F.interpolate(pred, size=(self.levels[i+1].size, self.levels[i+1].size//2), mode='bilinear') -# pred = pred * levels_input[i+1][-2][:,None,:,:] -# else: -# pred = self.enhance(pred, levels_input[i][-1]) -# pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask -# pred = F.interpolate(pred, size=(512, 256), mode='bilinear') -# out_levels.append([pred, mean, logvar]) - -# return out_levels - +from torchvision.transforms import Resize +from torchvision.transforms.functional import InterpolationMode class Pyramid(nn.Module): @@ -145,65 +24,41 @@ class Pyramid(nn.Module): if config['body_model'] != '': load_model(self.body_enc, config['body_model'] + '_enc') - self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], kernel_size=config['cond_ae_curv']['kernel_size'], activation=config['cond_ae_curv']['activation'], is_label=True, nr_of_labels=7) - if config['curvature_model'] != '': - load_model(self.curv_enc, config['curvature_model'] + '_enc') + if config['use_normals']: + self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device']) + if config['normals_model'] != '': + load_model(self.norm_enc, config['normals_model'] + '_enc') if config['freeze_encoders']: for param in self.static_enc.parameters(): param.requires_grad = False for param in self.body_enc.parameters(): param.requires_grad = False - for param in self.curv_enc.parameters(): param.requires_grad = False + if config['use_normals']: + for param in self.norm_enc.parameters(): param.requires_grad = False - self.levels = nn.ModuleList([PyramidLevel({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))]) + self.levels = nn.ModuleList([PyramidLevel({ **config['pyramid']['levels'][i] }, use_normals=config['use_normals']) for i in range(len(config['pyramid']['levels']))]) + self.resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BILINEAR).to(config['device']) for res in config['pyramid']['resolutions'] } + self.resizers[512] = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device']) - # if len(config['pyramid']['resolutions']) == 2: - # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) - # if len(config['pyramid']['resolutions']) == 3: - # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) - # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) - # if len(config['pyramid']['resolutions']) == 4: - # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)]) - # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) - # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) - # if len(config['pyramid']['resolutions']) == 5: - # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)]) - # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)]) - # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) - # self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) + self.use_normals = config['use_normals'] - # self.add_noise = config['pyramid']['add_noise'] - # self.noise_amount = config['pyramid']['noise_amount'] - # self.use_offsets = config['use_offsets'] - - # self.upscale_enhancer = None - # if config['pyramid']['final_upscale_enhancer'] != '': - # k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5)) - # self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size) - # load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer']) - # self.upscale_enhancer.trainable = False - - # def enhance(self, low, mask, training): - # x = self.upscale_enhancer([low, mask], training=training) - # return x - - def encode_conds(self, uv_template, uv_body_posed, uv_curv): + def encode_conds(self, uv_template, uv_body_posed, uv_norm): z_static = self.static_enc(uv_template) z_body = self.body_enc(uv_body_posed) - z_curv = self.curv_enc(uv_curv) - z_cond = torch.concat([z_static, z_body, z_curv], dim=1) + if self.use_normals: + z_norm = self.norm_enc(uv_norm) + z_cond = torch.concat([z_static, z_body, z_norm], dim=1) + else: + z_cond = torch.concat([z_static, z_body], dim=1) return z_cond def forward(self, inputs): levels_input, conds = inputs - uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2] + uv_static, uv_body, uv_norm = conds[0], conds[1], conds[2] - z_cond = self.encode_conds(uv_static, uv_body, uv_curv) - - # if self.add_noise == True and training == True: - # z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount) + z_cond = self.encode_conds(uv_static, uv_body, uv_norm) out_levels = [] pred, mean, logvar = self.levels[0](levels_input[0], z_cond) @@ -211,35 +66,28 @@ class Pyramid(nn.Module): # pred = pred + levels_input[0][2] # uv_body_posed_32 if 1 != len(self.levels): pred = self.levels[0].enhance(pred, levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res - pred = F.interpolate(pred, size=(self.levels[1].size, self.levels[1].size//2), mode='bilinear') + pred = self.resizers[self.levels[1].size](pred) out_levels.append([pred, mean, logvar]) pred = pred * levels_input[1][-2][:,None,:,:] else: - # pred = self.enhance(pred, levels_input[0][-2], training=training) pred = self.levels[0].enhance(pred, levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res - pred = F.interpolate(pred, size=(512, 256), mode='bilinear') + pred = self.resizers[512](pred) pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask out_levels.append([pred, mean, logvar]) for i in range(1, len(self.levels)): input = levels_input[i] - # input[0] = input[0] - pred - # input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3]) - # input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None] - # input[0] = input[0] * input[-2][:,None,:,:] # uv_mask_res + input[0] = input[0] - pred pred_level, mean, logvar = self.levels[i](input, z_cond) - # pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res pred = pred + pred_level if i != len(self.levels) - 1: pred = self.levels[i].enhance(pred, levels_input[i][-2]) out_levels.append([pred, mean, logvar]) - pred = F.interpolate(pred, size=(self.levels[i+1].size, self.levels[i+1].size//2), mode='bilinear') + pred = self.resizers[self.levels[i+1].size](pred) pred = pred * levels_input[i+1][-2][:,None,:,:] else: - # if self.upscale_enhancer is not None: - # pred = self.enhance(pred, levels_input[-1][-2], training=training) pred = self.levels[-1].enhance(pred, levels_input[-1][-2]) - pred = F.interpolate(pred, size=(512, 256), mode='bilinear') + pred = self.resizers[512](pred) pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask out_levels.append([pred, mean, logvar]) diff --git a/scripts/cvae.py b/scripts/cvae.py index 2ee71e9..948e6bf 100644 --- a/scripts/cvae.py +++ b/scripts/cvae.py @@ -10,9 +10,8 @@ parser.add_argument('--config', type=str) parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) - if args['debug']: - from configs.iccvw.abl_all import config + from configs.iccvw.abl_all_4 import config else: config_module = importlib.import_module(args['config']) config = config_module.config @@ -40,6 +39,7 @@ else: import torch import torch.nn as nn import torch.optim as optim +from torch.profiler import profile, record_function, ProfilerActivity from data.dataset import Dataset from models.cond_encoder import AutoEncoder from utils.checkpoints import CheckpointSaver, load_model @@ -82,11 +82,13 @@ def kl_loss(mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: def train_step(batch): + # with profiler: if config['use_skinning']: mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed'] else: mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed'] + # with record_function("forward"): model.train() pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) mesh_pred = uv2mesh(pred) @@ -115,20 +117,26 @@ def train_step(batch): if config['use_extra_losses']: reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff'] - loss = reconstr_loss + g_balancing_coeff * kl + loss = g_balancing_coeff * kl + reconstr_loss + # with record_function("backward"): optimizer.zero_grad() loss.backward() optimizer.step() + # with record_function("eval"): model.eval() with torch.inference_mode(): mean = torch.mean(mean) logvar = torch.mean(logvar) ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 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']) - - return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric, + + # profiler.step() + + + return { 'train/loss': loss, + 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric, 'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar, 'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss, @@ -253,8 +261,8 @@ def plot_everything(batch: dict, type: str): collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))] collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))] - figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) } - figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) } + figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'])) for i in range(4) } + figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'])) for i in range(4) } mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0)) logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0)) @@ -466,15 +474,16 @@ def train(): valid_accumulator.reset() print(step) - if train_loss['train/euclid_error'] < config['balancing_activation_limit']: - g_balancing_coeff += 1 / 10000 + # if train_loss['train/euclid_error'] < config['balancing_activation_limit']: + # g_balancing_coeff += 1 / 10000 - if step % 25000 == 0: + if step % 25000 == 0 and not args['debug']: ckpt_saver(model, step, optimizer, None) step += 1 + if __name__ == '__main__': if not args['debug']: init_wandb() @@ -500,10 +509,22 @@ if __name__ == '__main__': step = 0 if config['pretrained_model'] != '': - step = load_model(model, config['pretrained_model']) + step = load_model(model, config['pretrained_model'], optimizer) if not args['debug']: # model = torch.compile(model) wandb.watch(model, log='all', log_freq=10000, log_graph=True) - train() \ No newline at end of file + # create a profiler object + # profiler = torch.profiler.profile( + # activities=[ProfilerActivity.CPU], + # schedule=torch.profiler.schedule(wait=2, warmup=1, active=3, repeat=2), + # on_trace_ready=torch.profiler.tensorboard_trace_handler("log_dir"), + # ) + +# with profiler: +# with record_function("model_inference"): + train() + + # profiler.stop() + # print(profiler.key_averages().table(sort_by="cpu_time_total", row_limit=25)) \ No newline at end of file diff --git a/scripts/cvae_pyramid.py b/scripts/cvae_pyramid.py index 32c3fcd..9779384 100644 --- a/scripts/cvae_pyramid.py +++ b/scripts/cvae_pyramid.py @@ -10,11 +10,12 @@ parser.add_argument('--config', type=str) parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) -config_module = importlib.import_module(args['config']) -config = config_module.config if args['debug']: - from configs.pyr_32 import config - # from configs.pyr_64_cape_m import config + from configs.iccvw.pyr_64_reg import config +else: + config_module = importlib.import_module(args['config']) + config = config_module.config + os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: @@ -42,6 +43,7 @@ import torch.nn.functional as F from data.dataset import Dataset from utils.checkpoints import CheckpointSaver, load_model from utils.utils import InfiniteDataLoader as DataLoader +from utils.utils import text_to_html from utils.postprocess import fix_collision from losses.losses import collision_loss, normal_loss, kl_loss from pytorch_msssim import MS_SSIM @@ -52,7 +54,8 @@ from models.pyramid_model import Pyramid def init_wandb(): - run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) + run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/cvae_pyramid.py') artifact.add_file('losses/losses.py') @@ -99,8 +102,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col fig.add_trace(go.Mesh3d( x=mesh_posed[:,0], - y=mesh_posed[:,1], - z=mesh_posed[:,2], + y=mesh_posed[:,2], + z=mesh_posed[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2], @@ -235,47 +238,40 @@ def train_step(batch): res_modifier = '_' + str(res) if res != 512 else '' pred, mean, logvar = out_levels[i] - mesh_pred = uv2mesh(pred) - if config['use_skinning']: - shape = batch['pose_shape'][:, 72:] - gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) - 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_pred_posed = mesh_pred - - # is the output already masked? - # ms_ssim = torch.mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) - # ms_ssim = torch.mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3)) - - 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(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) - l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-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']) kl = kl_loss(mean, logvar) if config['use_extra_losses']: + mesh_pred = uv2mesh(pred) + + if config['use_skinning']: + shape = batch['pose_shape'][:, 72:] + gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) + 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_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']) + c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) + n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) + reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (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'] + if config['use_norm_loss'] and step > 2000: + reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] else: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss) - - if config['use_norm_loss'] and step > 2000: - reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] kl_total += kl - loss = reconstr_loss_total + g_balancing_coeff * kl_total - - # scaler.scale(loss).backward() - # scaler.step(optimizer) - # scaler.update() - # optimizer.zero_grad() + loss: torch.Tensor = reconstr_loss_total + g_balancing_coeff * kl_total optimizer.zero_grad() loss.backward() @@ -286,13 +282,33 @@ def train_step(batch): mean = torch.mean(mean) logvar = torch.mean(logvar) + if not config['use_extra_losses']: + mesh_pred = uv2mesh(pred) + + if config['use_skinning']: + shape = batch['pose_shape'][:, 72:] + gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) + 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_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']) + c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) + l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-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_static'] + 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']) output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl, 'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, 'train/normal_loss': n_loss, - 'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss, 'train/ms_ssim' + res_modifier: ms_ssim, } } # + 'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss, 'train/ms_ssim' + res_modifier: ms_ssim, + 'train/collision_ratio': c_ratio } } # return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output } @@ -332,28 +348,24 @@ def valid_step(batch): else: mesh_pred_posed = mesh_pred - # is the output already masked? - # ms_ssim = torch.mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) - # ms_ssim = torch.mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3)) - 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(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-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']) + 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) if not config['is_cape']: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] + if config['use_norm_loss'] and step > 2000: + reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] else: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss) - if config['use_norm_loss'] and step > 2000: - reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] - kl_total += kl loss = reconstr_loss_total + g_balancing_coeff * kl_total @@ -452,12 +464,12 @@ def plot_everything(batch: dict, type: str): 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 - uv_curvature = (torch.tile(torch.argmax(torch.permute(batch['uv_curvature'], [0,2,3,1]), dim=-1)[:,:,:,None], [1,1,1,3]) / 6) - pred_img = torch.concatenate([torch.permute(F.interpolate(output_levels[i][0], size=(512, 256), mode='bilinear'), [0,2,3,1]) for i in range(len(output_levels))], dim=3) + 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 = (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_curvature, orig_img, pred_img), axis=2) + img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals, orig_img, pred_img), axis=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") @@ -521,7 +533,14 @@ def plot_everything(batch: dict, type: str): collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))] figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) } - return { **figs, **figs_coll, **{ type + '/images': imgs, } } + batch_id = batch['id'] + batch_frame_nr = batch['frame_nr'] + batch_garment = batch['garment'] + batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))] + batch_info = text_to_html(batch_info) + batch_info = wandb.Html(batch_info) + + return { **figs, **figs_coll, **{ type + '/images': imgs, }, **{ f'{type}/batch_info': batch_info }, } def sample(folder): @@ -791,9 +810,9 @@ def train(): print(step) if train_loss['train/euclid_error'] < config['balancing_activation_limit']: - g_balancing_coeff += 1 / 10000 + g_balancing_coeff += 1 / 50000 - if step % 10000 == 0: + if step % 25000 == 0 and not args['debug']: ckpt_saver(model, step, optimizer, None) step += 1 @@ -821,7 +840,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/iccv/torch/' + config['name'], top_n=5) + ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=5) scaler = torch.cuda.amp.GradScaler(enabled=use_amp) g_balancing_coeff = config['starting_balancing_coeff'] @@ -840,12 +859,14 @@ if __name__ == '__main__': if len(config['pyramid']['resolutions']) == 1: for param in model.static_enc.parameters(): param.requires_grad = False for param in model.body_enc.parameters(): param.requires_grad = False - for param in model.curv_enc.parameters(): param.requires_grad = False + if config['use_normals']: + for param in model.norm_enc.parameters(): param.requires_grad = False else: - load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else '')) + load_model(model, config['pyramid']['pretrained_pyramid'], iteration=config['pyramid']['pretrained_pyramid_iteration']) for param in model.static_enc.parameters(): param.requires_grad = False for param in model.body_enc.parameters(): param.requires_grad = False - for param in model.curv_enc.parameters(): param.requires_grad = False + if config['use_normals']: + for param in model.norm_enc.parameters(): param.requires_grad = False # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer']) # model.levels[0].upscale_enhancer.trainable = False @@ -866,6 +887,6 @@ 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() \ No newline at end of file diff --git a/scripts/generate_results.py b/scripts/generate_results.py new file mode 100644 index 0000000..a57d483 --- /dev/null +++ b/scripts/generate_results.py @@ -0,0 +1,148 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +import importlib +import argparse +parser = argparse.ArgumentParser(description='') +parser.add_argument('--name', type=str) +parser.add_argument('--config', type=str) +parser.add_argument('--debug', action='store_true') +args = vars(parser.parse_args()) + +if args['debug']: + from configs.iccvw.pyr_64 import config +else: + config_module = importlib.import_module(args['config']) + config = config_module.config + +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] +if args['debug']: + os.environ["CUDA_LAUNCH_BLOCKING"] = '1' + +import time +import argparse +import wandb +import tqdm +import random +import numpy as np +import plotly.graph_objects as go +import plotly.express as ex +from PIL import Image + + +import torch +import torch.nn as nn +import torch.optim as optim +import torch.nn.functional as F +from tqdm import tqdm +from data.dataset import Dataset +from utils.checkpoints import CheckpointSaver, load_model +from torch.utils.data import DataLoader +from utils.utils import text_to_html +from utils.postprocess import fix_collision +from losses.losses import collision_loss, normal_loss, kl_loss +from pytorch_msssim import MS_SSIM +from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels +from smpl.smpl_torch import SMPL +from plotly.subplots import make_subplots +from models.pyramid_model import Pyramid +from models.cvae_model import CVAE + + +uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + + +def step_pyramid(batch): + model.eval() + with torch.inference_mode(): + levels_input, conds = group_inputs_for_levels(batch, config) + + out_levels = model([levels_input, conds]) + pred, _, _ = out_levels[-1] + + mesh_pred = uv2mesh(pred) + + if config['use_skinning']: + shape = batch['pose_shape'][:, 72:] + gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) + 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_pred_posed = mesh_pred + + return mesh_pred_posed + + +def step(batch): + if config['use_skinning']: + uv_orig = batch['uv_unposed'] + else: + uv_orig = batch['uv_posed'] + + model.eval() + pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) + mesh_pred = uv2mesh(pred) + + 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])]) + + 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_pred_posed = mesh_pred + + return mesh_pred_posed + + +def generate(): + for batch in tqdm(test_dataloader): + if is_pyramid: + mesh_pred = step_pyramid(batch) + else: + mesh_pred = step(batch) + + ids = batch['id'] + frame_nrs = batch['frame_nr'] + garments = batch['garment'] + + for i in range(len(ids)): + np.savez_compressed(f'{out_folder}/{ids[i]}_{frame_nrs[i]}_{garments[i]}', mesh_pred[i].cpu().numpy()) + + + + +if __name__ == '__main__': + + smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device']) + skirt_modifier = '_skirt' if config['is_skirt'] else '' + dataset_resolutions = config['pyramid']['resolutions'].copy() + if 512 not in dataset_resolutions: + dataset_resolutions.append(512) + + dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file=config['test_frames'], use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape']) + test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, ) + + is_pyramid = config['name'].startswith('pyr') + + if is_pyramid: + model: Pyramid = Pyramid(config) + else: + model: CVAE = CVAE(config) + + it = load_model(model, 'checkpoints/iccv/torch/' + config['name']) + # it = load_model(model, 'checkpoints/iccvw/' + config['name']) + + name = config['name'] + '_' + str(it) + + folder = '/data-local/127_2/results' + out_folder = f'{folder}/{name}' + + os.makedirs(out_folder, exist_ok=True) + + generate() \ No newline at end of file diff --git a/smpl/smpl_torch.py b/smpl/smpl_torch.py index 1b90f78..b18dfb7 100644 --- a/smpl/smpl_torch.py +++ b/smpl/smpl_torch.py @@ -10,7 +10,7 @@ from time import time from torch.nn import Module class SMPL(Module): - def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False): + def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False, batch_size=4): super(SMPL, self).__init__() self.J_regressor = [] @@ -56,6 +56,14 @@ class SMPL(Module): self.shape = (1,27578,3) self.device = device + self.batch_size = batch_size + + self.with_zeros_ones = torch.tensor( + [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32 + ).expand(self.batch_size,-1,-1).to(self.device) + self.results_zeros = torch.zeros((self.batch_size, 24, 1), dtype=torch.float32).to(self.device) + self.J_ones = torch.ones((self.batch_size, 24, 1), dtype=torch.float32, device=self.device) + self.rest_shape_h_ones_full = torch.ones((self.batch_size, self.weights_all.shape[1], 1), dtype=torch.float32).to(self.device) if orig_smpl: self.to_keep = self.to_keep[self.to_keep < 6890] @@ -67,6 +75,8 @@ class SMPL(Module): self.weights = self.weights_all[:,self.to_keep,:] self.shape = (1, 6890, 3) + self.rest_shape_h_ones = torch.ones((self.batch_size, self.to_keep.shape[0], 1), dtype=torch.float32).to(self.device) + @staticmethod def rodrigues(r): """ @@ -91,7 +101,7 @@ class SMPL(Module): m = torch.stack( (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1) - m = torch.reshape(m, (-1, 3, 3)) + m = m.view((-1, 3, 3)) i_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) \ + torch.zeros((theta_dim, 3, 3), dtype=torch.float32)).to(r.device) A = r_hat.permute(0, 2, 1) @@ -112,8 +122,8 @@ class SMPL(Module): Tensor after appending of shape [4,4] """ ones = torch.tensor( - [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32 - ).expand(x.shape[0],-1,-1).to(x.device) + [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32 + ).expand(x.shape[0],-1,-1).to(x.device) ret = torch.cat((x, ones), dim=1) return ret @@ -158,8 +168,8 @@ class SMPL(Module): # v_shaped = torch.tensordot(betas, self.shapedirs, dims=([1], [2])) + self.v_template v_shaped = [] - v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + torch.reshape(self.v_template[0], shape=self.shape)) # female - v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male + v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + self.v_template[0].view(self.shape)) # female + v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + self.v_template[1].view(self.shape)) # male v_shaped = torch.stack(v_shaped, axis=-1) # J = torch.matmul(self.J_regressor, v_shaped) @@ -167,7 +177,7 @@ class SMPL(Module): J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])]) else: J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])]) - J0 = torch.reshape(J[:,0,:], (-1, 1, 3)) + J0 = J[:,0,:].view((-1, 1, 3)) J = J - J0 R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3) @@ -177,12 +187,12 @@ class SMPL(Module): R_cube = R_cube_big[:, 1:, :, :] I_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) + \ torch.zeros((batch_num, R_cube.shape[1], 3, 3), dtype=torch.float32)).to(self.device) - lrotmin = (R_cube - I_cube).reshape(batch_num, -1, 1).squeeze(dim=2) + lrotmin = (R_cube - I_cube).view(batch_num, -1, 1).squeeze(dim=2) v_posed = torch.stack([v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[torch.greater(g, 0.5).to(torch.long)], dims=([0], [2])) for i, g in enumerate(betas[:,-1])]) results = [] results.append( - self.with_zeros(torch.cat((R_cube_big[:, 0], torch.reshape(J[:, 0, :], (-1, 3, 1))), dim=2)) + self.with_zeros(torch.cat((R_cube_big[:, 0], J[:, 0, :].view((-1, 3, 1))), dim=2)) ) for i in range(1, 24): results.append( @@ -190,7 +200,7 @@ class SMPL(Module): results[self.parent[i]], self.with_zeros( torch.cat( - (R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))), + (R_cube_big[:, i], (J[:, i, :] - J[:, self.parent[i], :]).view((-1, 3, 1))), dim=2 ) ) @@ -202,25 +212,25 @@ class SMPL(Module): self.pack( torch.matmul( stacked, - torch.reshape( - torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2), + torch.cat((J, self.results_zeros), dim=2).view( (batch_num, 24, 4, 1) ) ) ) - joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0] + joints = torch.matmul(results, torch.transpose(torch.concat([J, self.J_ones], axis=-1)[:,:,None,:],2,3))[:,:,:3,0] weights = self.weights if verts is not None else self.weights_all + rest_shape_h_ones_local = self.rest_shape_h_ones if verts is not None else self.rest_shape_h_ones_full T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) rest_shape_h = torch.cat( - (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2 + (v_posed, rest_shape_h_ones_local), dim=2 ) - v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1))) - v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3] + v = torch.matmul(T, rest_shape_h.view((batch_num, -1, 4, 1))) + v = v.view((batch_num, -1, 4))[:, :, :3] result = v - root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3)) + root_joint = joints[:,0,:].view((-1, 1, 3)) result = result - root_joint joints = joints - root_joint diff --git a/utils/checkpoints.py b/utils/checkpoints.py index b564e40..8d23549 100644 --- a/utils/checkpoints.py +++ b/utils/checkpoints.py @@ -34,11 +34,13 @@ class CheckpointSaver: def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=None, scaler: torch.nn.Module=None, iteration=''): # list of checkpoint files in the folder checkpoints checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')] - checkpoint = torch.load(checkpoints[-1]) - iteration = checkpoints[-1].split('_')[-1].split('.')[0] + iterations = np.asarray([int(f.split('_')[-1].split('.')[0]) for f in checkpoints]) + last_iteration = np.argmax(iterations) + checkpoint = torch.load(checkpoints[last_iteration]) + iteration = iterations[last_iteration] model.load_state_dict(checkpoint['model']) if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict()) if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler']) - print(f'Loaded model from {checkpoints[-1]} at iteration {iteration}') + print(f'Loaded model from {checkpoints[last_iteration]} at iteration {iteration}') return int(iteration) if iteration != '' else 0 diff --git a/utils/generate_figures.py b/utils/generate_figures.py index b08abbf..bd22685 100644 --- a/utils/generate_figures.py +++ b/utils/generate_figures.py @@ -3,33 +3,19 @@ import sys sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') import numpy as np +import cmapy as cm + +from utils import remove_long_edges from custom_utils.custom_utils import save_stl, write_obj from data.cloth3d_io import readPC2Frame from data.utils import Frame from custom_utils.custom_utils import _maskFaces import plotly.graph_objects as go +import plotly.express as ex from plotly.subplots import make_subplots -def remove_long_edges(vertices, faces): - # Compute edge lengths - edges = np.vstack((faces[:,[0,1]], faces[:,[1,2]], faces[:,[2,0]])) - edge_lengths = np.linalg.norm(vertices[edges[:,0]] - vertices[edges[:,1]], axis=1) - - # Find edges longer than 0.1 - long_edges = edges[edge_lengths > 0.1] - - # Find vertices that are part of long edges - vertices_to_remove = np.unique(long_edges) - - # Remove vertices and update faces - mask = np.ones(vertices.shape[0], dtype=bool) - mask[vertices_to_remove] = False - new_vertices = vertices[mask] - - - return new_vertices, mask def plot_unposing_example(): @@ -87,7 +73,7 @@ def plot_normals_example(): # mesh_unposed = readPC2Frame(os.path.join(folder_test, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) # mesh_unposed_shaped = readPC2Frame(os.path.join('temp', f'{frame.id}_{frame.garment}_unposed_shaped.pc16'), frame.frame_nr, True)[mesh_mask] mesh_normals = readPC2Frame(os.path.join(folder_test, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) - uv_normals = np.load(os.path.join(folder_test, 'uv_normals', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')) + uv_normals = np.load(os.path.join(folder_test, 'uv_normals', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_normals'] F = np.load('smpl/smpl_faces.npy') _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) @@ -101,12 +87,13 @@ def plot_normals_example(): new_mask[mesh_mask] = v_mask _, faces = _maskFaces(np.zeros((14475, 3)), F, new_mask) - fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}]]) + v_colors = (mesh_normals + 1.4) / 2.8 + v_colors = v_colors[v_mask] + # v_colors = (v_colors - v_colors.min()) / (v_colors.max() - v_colors.min()) - fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), row=1, col=1) - fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=1), row=1, col=2) - fig.add_trace(go.Mesh3d(x=mesh_unposed_shaped[:,0], y=mesh_unposed_shaped[:,2], z=mesh_unposed_shaped[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightgreen', opacity=1), row=1, col=3) - fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightyellow', opacity=1), row=1, col=4) + fig = go.Figure() + + fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], vertexcolor=v_colors, opacity=1),) fig.update_layout(scene=dict( xaxis=dict(range=[-1, 1]), @@ -114,11 +101,29 @@ def plot_normals_example(): zaxis=dict(range=[-1, 1]), ), scene_aspectmode='cube') + helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), + yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), + zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),) + fig.update_layout(scene = helper) + + # fig.show() + + uv_normals = (uv_normals + 1.4) / 2.8 + uv_normals = np.concatenate((uv_normals[:256], uv_normals[256:]), axis=1) + # flip uv_normals along y axis + uv_normals = uv_normals[:,::-1,:] + ex.imshow(uv_normals).show() + + + + + pass + if __name__ == '__main__': folder_train = '/data-local/127/cloth3d_processed' - folder_test = '/data-local/127/cloth3d_processed_test' + folder_test = '/data-local/127_2/cloth3d_processed_test' # plot_unposing_example() plot_normals_example()