From afd482522db43ae3a7bb44a89b0b9ce937fa3845 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 24 Aug 2023 11:36:22 +0000 Subject: [PATCH] save all --- .vscode/launch.json | 2 +- configs/iccvw/pyr_64_128_256_512_e2e_var.py | 302 ++++++++++++++ .../iccvw/pyr_64_128_256_512_skirt_e2e_var.py | 302 ++++++++++++++ data/convert_to_uv.py | 6 +- data/convert_to_uv_cape.py | 67 ++- data/dataset.py | 12 +- data/downscale.py | 28 +- data/downscale_cape.py | 34 +- data/preprocessing.py | 194 ++++++++- data/utils.py | 24 +- docker/Dockerfile2 | 5 +- docker/Dockerfile_hood | 48 +++ docker/docker-compose.yml | 23 +- losses/losses.py | 58 ++- models/pyramid_level.py | 2 +- models/pyramid_model.py | 57 ++- run.py | 15 +- run_parallel.py | 13 +- scripts/calculate.py | 34 ++ scripts/calculate_area_diff.py | 47 +++ scripts/calculate_chamfer.py | 207 ++++++++-- scripts/calculate_error_all.py | 14 +- scripts/cvae.py | 52 ++- scripts/cvae_pyramid.py | 390 ++++++++++++++++-- scripts/cvae_pyramid_e2e.py | 9 +- scripts/generate_results.py | 2 +- scripts/generate_results_all.py | 4 +- scripts/get_error.py | 32 +- scripts/get_error_hood.py | 74 ++-- scripts/plot_ablation.py | 71 ++-- scripts/plot_ablation_pyramid.py | 32 +- scripts/plot_cloth3d.py | 148 +++++++ scripts/plot_hood.py | 131 ++++++ scripts/upscaler_train.py | 13 +- utils/postprocess.py | 2 + utils/utils.py | 31 +- 36 files changed, 2194 insertions(+), 291 deletions(-) create mode 100644 configs/iccvw/pyr_64_128_256_512_e2e_var.py create mode 100644 configs/iccvw/pyr_64_128_256_512_skirt_e2e_var.py create mode 100644 docker/Dockerfile_hood create mode 100644 scripts/calculate.py create mode 100644 scripts/calculate_area_diff.py create mode 100644 scripts/plot_cloth3d.py create mode 100644 scripts/plot_hood.py diff --git a/.vscode/launch.json b/.vscode/launch.json index d71bb99..e24572d 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -12,7 +12,7 @@ "console": "integratedTerminal", "justMyCode": true, // true false "cwd": "${workspaceFolder}", - "args": ["--config", "configs.pyr_32", "--debug"] + "args": ["--debug"] } ] } \ No newline at end of file diff --git a/configs/iccvw/pyr_64_128_256_512_e2e_var.py b/configs/iccvw/pyr_64_128_256_512_e2e_var.py new file mode 100644 index 0000000..90d41f9 --- /dev/null +++ b/configs/iccvw/pyr_64_128_256_512_e2e_var.py @@ -0,0 +1,302 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '6' +config['name'] = 'pyr_iccvw_64_128_256_512_e2e_var_2' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid_e2e.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_128_256_512_e2e_2' +config['pretrained_model_iteration'] = '250000' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64, 128, 256, 512] +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'] = 1 +config['balancing_activation_limit'] = 0.01 +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, 64, 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, 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 == 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_128_256_512_skirt_e2e_var.py b/configs/iccvw/pyr_64_128_256_512_skirt_e2e_var.py new file mode 100644 index 0000000..219d883 --- /dev/null +++ b/configs/iccvw/pyr_64_128_256_512_skirt_e2e_var.py @@ -0,0 +1,302 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'pyr_iccvw_64_128_256_512_skirt_e2e_var_3' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid_e2e.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'] = 'checkpoints/iccvw/pyr_iccvw_64_128_256_512_skirt_e2e' +config['pretrained_model_iteration'] = '200000' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64, 128, 256, 512] +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'] = 1 +config['balancing_activation_limit'] = 0.01 +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, 64, 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, 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 == 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/convert_to_uv.py b/data/convert_to_uv.py index 674cde0..59ceb0c 100644 --- a/data/convert_to_uv.py +++ b/data/convert_to_uv.py @@ -399,10 +399,11 @@ def get_shapedirs(): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127/cloth3d_processed') + parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cloth3d_processed_test') parser.add_argument("-s", "--start", help = "start of range", default=0, type=int) parser.add_argument("-e", "--end", help = "end of range", default=100, type=int) parser.add_argument("-n", "--name", help = "end of range", default='lhunor') + parser.add_argument('--debug', action='store_true') args = parser.parse_args() F = {'non-skirt': np.load("smpl/smpl_faces.npy"), 'skirt': np.load("smpl/smpl_faces_skirt.npy")} @@ -442,6 +443,7 @@ if __name__ == '__main__': ######################################################################################### folder = args.folder + # data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed'] # 'mesh_unposed', data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed', ids = sorted(os.listdir(folder + '/mesh_body_posed_no_orient'))[args.start:args.end] @@ -460,6 +462,8 @@ if __name__ == '__main__': os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True) os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True) + # ids = ['00438'] + for id in tqdm.tqdm(ids): garments = glob.glob(os.path.join(folder, 'mesh_mask', f'{id}_*')) diff --git a/data/convert_to_uv_cape.py b/data/convert_to_uv_cape.py index 8fe86ee..9325f83 100644 --- a/data/convert_to_uv_cape.py +++ b/data/convert_to_uv_cape.py @@ -401,7 +401,7 @@ def get_shapedirs(): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/test/female') + parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/female') parser.add_argument("-s", "--start", help = "start of range", default=0, type=int) parser.add_argument("-e", "--end", help = "end of range", default=100, type=int) parser.add_argument("-n", "--name", help = "end of range", default='lhunor') @@ -448,11 +448,11 @@ if __name__ == '__main__': folder = args.folder # data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed', - data_to_convert = ['mesh_static', 'mesh_normals'] # 'mesh_unposed', + data_to_convert = ['mesh_posed', 'mesh_static'] # 'mesh_unposed', - frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))#[args.start:args.end] + frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient')) frames = [[frame.split('.')[0], frame.split('.')[1]] for frame in frames] - ids = list(set([frame[0] for frame in frames])) + ids = sorted(list(set([frame[0] for frame in frames])))[args.start:args.end] frames_by_id = {id: [frame[1] for frame in frames if frame[0] == id] for id in ids} shapedirs = get_shapedirs() @@ -462,11 +462,11 @@ if __name__ == '__main__': os.makedirs(folder, exist_ok=True) os.makedirs(os.path.join(folder, 'uv_posed_no_orient'), exist_ok=True) - os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True) - os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True) - os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True) - os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True) - os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True) + # os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True) + # os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True) + # os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True) + # os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True) + # os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True) mask = np.ones((3638), dtype=bool) @@ -476,15 +476,14 @@ if __name__ == '__main__': # x = readPC2(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'] # mesh_posed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32) # mesh_posed[:,mask, :] = x - if 'mesh_unposed' in data_to_convert: - mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32) - mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V'] + # if 'mesh_unposed' in data_to_convert: + # mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32) + # mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V'] if 'mesh_static' in data_to_convert: mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static'] - if 'mesh_normals' in data_to_convert: - mesh_normals = np.stack([np.load(os.path.join(folder, 'mesh_normals', f'{id}.{int(i):06d}.npz'))['mesh_normals'] for i in frames_by_id[id]], axis=0) - if 'body_posed' in data_to_convert and k != 1: - body_posed = readPC2(os.path.join(folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'] + if 'mesh_posed' in data_to_convert: + mesh_posed = np.stack([np.load(os.path.join(folder, 'mesh_posed_no_orient', f'{id}.{int(i):06d}.npz'))['mesh_posed_no_orient'] for i in frames_by_id[id]], axis=0) + start = time.perf_counter() @@ -493,13 +492,13 @@ if __name__ == '__main__': uv_static = uv_static[0].astype(np.float32) if 'mesh_posed' in data_to_convert: - uv_posed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border) - if 'mesh_unposed' in data_to_convert: - uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border) - if 'mesh_normals' in data_to_convert: - uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border) - if 'body_posed' in data_to_convert and k != 1: - uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border) + uv_posed_no_orient = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border) + # if 'mesh_unposed' in data_to_convert: + # uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border) + # if 'mesh_normals' in data_to_convert: + # uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border) + # if 'body_posed' in data_to_convert and k != 1: + # uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border) print(time.perf_counter() - start) @@ -515,16 +514,16 @@ if __name__ == '__main__': # np.savez_compressed(os.path.join(folder, 'uv_mask', f'{id}_{garment}.npz'), uv_mask=uv_mask_local) if 'mesh_posed' in data_to_convert: - for i in range(len(uv_posed)): - np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient', f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i]) - if 'mesh_unposed' in data_to_convert: - for i in range(len(uv_unposed)): - np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i]) - if 'mesh_normals' in data_to_convert: - for i, frame_nr in enumerate(frames_by_id[id]): - np.savez_compressed(os.path.join(folder, 'uv_normals', f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i]) - if 'body_posed' in data_to_convert and k != 1: - for i in range(len(uv_body_posed)): - np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i]) + for k, i in enumerate(frames_by_id[id]): + np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient', f'{id}.{int(i):06d}.npz'), uv_posed_no_orient=uv_posed_no_orient[k]) + # if 'mesh_unposed' in data_to_convert: + # for i in range(len(uv_unposed)): + # np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i]) + # if 'mesh_normals' in data_to_convert: + # for i, frame_nr in enumerate(frames_by_id[id]): + # np.savez_compressed(os.path.join(folder, 'uv_normals', f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i]) + # if 'body_posed' in data_to_convert and k != 1: + # for i in range(len(uv_body_posed)): + # np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i]) # pass diff --git a/data/dataset.py b/data/dataset.py index 2c0ab84..cf94cf5 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -276,7 +276,7 @@ if __name__ == '__main__': 'uv_mask', } - dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt') + dataset = Dataset('/data-local/127_2/cloth3d_processed_test', signature, resolutions=[512], device=device, frames_file='frames_temp.txt') train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) smpl = SMPL(orig_smpl=False, path='smpl', device='cuda') @@ -293,6 +293,16 @@ if __name__ == '__main__': step = 0 for batch in train_dataloader: batch + + from custom_utils.custom_utils import save_stl + fix_data = np.load(os.path.join(dataset.folder, 'mesh_fixed', f'{batch["id"][0]}_{batch["garment"][0]}.npz')) + faces = fix_data['faces'] + mask = fix_data['mask'] + + save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_posed.stl', batch['mesh_posed'][0].cpu().numpy()[mask], faces) + save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_unposed.stl', batch['mesh_unposed'][0].cpu().numpy()[mask], faces) + + print(batch['id'][0]) with profiler: m = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], batch['mesh_unposed'])[0] diff --git a/data/downscale.py b/data/downscale.py index 1924cae..018595a 100644 --- a/data/downscale.py +++ b/data/downscale.py @@ -11,7 +11,7 @@ import plotly.express as ex from torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode -os.environ["CUDA_VISIBLE_DEVICES"] = "4" +os.environ["CUDA_VISIBLE_DEVICES"] = "3" sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from data.utils import Frame @@ -165,32 +165,19 @@ def downscale_uv_maps_static(folder, type, sizes): np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame[0]}_{frame[1]}.npz'), uv_static = down) -def downscale_uv_maps_mask(folder, type, size): - os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) - - frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)] - - for frame in tqdm.tqdm(frames): - id, _ = frame.split('.') - 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/cloth3d_processed') + parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cloth3d_processed_test') parser.add_argument("-s", "--start", help = "start of range", default=0) - parser.add_argument("-e", "--end", help = "end of range", default=25000) + parser.add_argument("-e", "--end", help = "end of range", default=300000) 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() # sizes = [32, 64, 128, 256] device = 'cuda' - sizes = [32, 64, 128, 256] + sizes = [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 } @@ -199,19 +186,14 @@ if __name__ == '__main__': 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.astype(bool) - for res in [32, 64, 128, 256]: + for res in [64, 128, 256]: down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:].astype(np.int32)).to(device)) down = np.squeeze(down.cpu().numpy().astype(bool)) uv_masks_body[res] = down - 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) diff --git a/data/downscale_cape.py b/data/downscale_cape.py index 87cac14..66a4d99 100644 --- a/data/downscale_cape.py +++ b/data/downscale_cape.py @@ -23,7 +23,7 @@ def get_img_idxs(): img_idxs[512] = img_idx img_coords[512] = coords uv_masks[512] = uv_mask.astype(np.int32) - for res in [32, 64, 128, 256]: + for res in [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)) @@ -37,7 +37,7 @@ def get_img_idxs(): 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]: + for res in [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 @@ -113,6 +113,8 @@ def downscale_uv_normals(folder, type, sizes): frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] for frame in tqdm.tqdm(frames): + if not frame.startswith('03223_shortshort_chicken_wings'): + continue img = np.load(os.path.join(folder, type, frame)); img = img['uv_normals'] img = fill(img, uv_masks_body[512]) for size in sizes: @@ -140,6 +142,23 @@ def downscale_uv_maps_unposed(folder, type, sizes): np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down) +def downscale_uv_maps_posed(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_no_orient'] + 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_posed_no_orient = 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) @@ -173,7 +192,7 @@ def downscale_uv_maps_mask(folder, type, size): 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("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/female') 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') @@ -182,7 +201,7 @@ if __name__ == '__main__': device = 'cuda' - sizes = [32, 64, 128, 256] + sizes = [64, 128, 256] folder = args.folder resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes } @@ -199,9 +218,10 @@ if __name__ == '__main__': # 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) + # downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes) + downscale_uv_maps_posed(folder=folder, type='uv_posed_no_orient', 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 5089977..3d7238a 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -12,6 +12,7 @@ args = parser.parse_args() os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu +os.environ["CUDA_LAUNCH_BLOCKING"] = '1' import math import shutil @@ -76,6 +77,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder): os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True) os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True) + ids = ['00021'] + for id in tqdm(ids): # id = '01467' @@ -87,8 +90,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder): info = loadInfo(os.path.join(folder, id, 'info.mat')) for garment in garments: - if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): - continue + # if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): + # continue verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] # if info['poses'].shape[1] == verts.shape[0]: @@ -488,6 +491,176 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1) pass +F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda') +F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda') +F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [1,1,1]), device='cuda') +F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [1,1,1]), device='cuda') + +from scipy.spatial import KDTree +from pytorch3d.structures import Meshes +from torch_cluster import knn + +def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor): + # flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function + batch_size = V.shape[0] + vertices_count = V.shape[1] + batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(vertices_count) + batch_y = torch.arange(batch_size, device=V.get_device()).repeat_interleave(14475) + V = V.reshape(-1, V.shape[-1]) + B = B.reshape(-1, B.shape[-1]) + idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_y, batch_y=batch_x)[1] + idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count) + return idx + + +def collision_loss(V: torch.Tensor, B: torch.Tensor, thr: float = .004): + F_norm_local = F_norm + mesh = Meshes(B, F_norm_local) + B_vn = torch.stack(mesh.verts_normals_list()) + + # compute correspondences + idx = _nearest_neighbour(V, B) + # compute loss + D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])]) + b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])]) + dot = torch.einsum('abc,abc->ab', D, b_vn) + thresholded = torch.minimum(dot - thr, torch.zeros_like(dot)) + loss = torch.sum(thresholded ** 2, axis=-1) + + return loss + + +def area_loss(pred, gt, faces): + """ + Calculates the area loss between predicted and ground truth meshes. + + Args: + pred (torch.Tensor): Predicted mesh vertices of shape (batch_size, num_vertices, 3). + gt (torch.Tensor): Ground truth mesh vertices of shape (batch_size, num_vertices, 3). + faces (torch.Tensor): Mesh faces of shape (num_faces, 3). + + Returns: + torch.Tensor: Area loss between predicted and ground truth meshes. + """ + # Calculate the area of each face for predicted and ground truth meshes + pred_faces = pred[:, faces, :] + gt_faces = gt[:, faces, :] + pred_edges1 = pred_faces[:, :, 1, :] - pred_faces[:, :, 0, :] + pred_edges2 = pred_faces[:, :, 2, :] - pred_faces[:, :, 0, :] + gt_edges1 = gt_faces[:, :, 1, :] - gt_faces[:, :, 0, :] + gt_edges2 = gt_faces[:, :, 2, :] - gt_faces[:, :, 0, :] + pred_cross = torch.cross(pred_edges1, pred_edges2, dim=2) + gt_cross = torch.cross(gt_edges1, gt_edges2, dim=2) + pred_area = 0.5 * torch.norm(pred_cross, dim=2) + gt_area = 0.5 * torch.norm(gt_cross, dim=2) + + # Calculate the area loss + area_loss = torch.mean(torch.abs(pred_area - gt_area)) + + return area_loss + + + +def create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=0, end=1): + ids = sorted(os.listdir(folder))[start:end] + os.makedirs(os.path.join(out_folder, 'mesh_static_unposed_no_shape'), exist_ok=True) + F = np.load('smpl/smpl_faces.npy') + # F_skirt = np.load('smpl/smpl_faces_skirt.npy') + + ids = ['00005'] + + 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')] + shape = np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:] + + shape = torch.tensor(shape, dtype=torch.float32, device='cuda') + + smpl = SMPL(path='smpl', device='cuda', batch_size=1) + pose = np.zeros((24, 3), np.float32) + pose[1,2] = .15 + pose[2,2] = -.15 + pose = torch.tensor(pose.flatten()[None,:], dtype=torch.float32, device='cuda') + + + for garment in garments: + if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): + continue + + mesh_static, faces, _, _ = readOBJ(os.path.join(folder, id, f'{garment}.obj')) + faces = quads2tris(faces) + mesh_static = np.matmul(mesh_static, np.array([[1,0,0],[0,0,-1],[0,1,0]])) + + unposed = mesh_static + + unposed = torch.tensor(unposed[None,:,:], dtype=torch.float32, device='cuda') + unposed.requires_grad = True + mesh_static = torch.tensor(mesh_static[None,:,:], dtype=torch.float32, device='cuda') + + optimizer = torch.optim.Adam([unposed], lr=0.0001) + + for i in range(10000): + optimizer.zero_grad() + + body = smpl(pose, shape)[0] + c_loss = collision_loss(unposed, body) + a_loss = area_loss(unposed, mesh_static, faces) + loss : torch.Tensor = c_loss + a_loss + + # Backward pass + loss.backward() + optimizer.step() + + pose = pose / 1.0005 + + # Print the loss + if i % 50 == 0: + pose_sum = np.abs(pose.detach().cpu().numpy()).sum().item() + loss_str = loss.detach().cpu().numpy()[0] + c_loss_str = c_loss.detach().cpu().numpy()[0] + a_loss_str = a_loss.detach().cpu().numpy() + print(f"it: {i} loss: {loss_str:.10f} c_loss: {c_loss_str:.10f}, a_loss: {a_loss_str:.8f}, pose: {pose_sum:.8f}") + + + unposed = unposed.detach().cpu().numpy() + body = body.detach().cpu().numpy() + + gender = int(shape[0].cpu().numpy()[-1] > 0) + shape_offsets = torch.tensordot(shape[:,:10], smpl.shapedirs[gender], dims=[[1], [2]]) + shape_offsets = shape_offsets.cpu().numpy()[:,smpl.to_keep,:] + + unshaped = (unposed - shape_offsets) + + writePC2(os.path.join(out_folder, 'mesh_static_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True) + + verts = unposed + # verts = unshaped + frame = 0 + fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)]) + # add body mesh + fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,2], z=body[frame,:,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() + + # verts = pred.detach().cpu().numpy()[:,mask,:] + # frame = 9 + # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', 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 + + def create_frames_file(folder, out_folder, split=0.9): ids = os.listdir(folder) @@ -879,13 +1052,13 @@ def generate_fixed_masks(folder, out_folder): if __name__ == '__main__': - # folder = '/data-net/dataset/cloth3d/test' - # folder_reg = '/data-local/127/cloth3d_registration/test' - # out_folder = '/data-local/127_2/cloth3d_processed_test' + folder = '/data-net/dataset/cloth3d/test' + folder_reg = '/data-local/127/cloth3d_registration/test' + out_folder = '/data-local/127_2/cloth3d_processed_test' - folder = '/data-net/dataset/cloth3d/train' - folder_reg = '/data-local/127/cloth3d_registration/train' - out_folder = '/data-local/127/cloth3d_processed' + # folder = '/data-net/dataset/cloth3d/train' + # folder_reg = '/data-local/127/cloth3d_registration/train' + # out_folder = '/data-local/127/cloth3d_processed' smpl = SMPL(path='smpl', device='cuda') @@ -898,13 +1071,14 @@ if __name__ == '__main__': # 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) + create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder) # verify_registration_ans_mask(folder, out_folder) # create_pose_shape(folder, out_folder) # create_mesh_static(folder, out_folder) # 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) + # create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=args.start, end=args.end) # calculate_average(folder, out_folder) # create_frames_file(folder, out_folder) @@ -915,4 +1089,4 @@ if __name__ == '__main__': # folder = '/data-local/127_2/cape_processed/test/female' # create_mesh_normals_cape(folder, None) - generate_fixed_masks(folder, out_folder) \ No newline at end of file + # generate_fixed_masks(folder, out_folder) \ No newline at end of file diff --git a/data/utils.py b/data/utils.py index 2dde4cd..cb1995a 100644 --- a/data/utils.py +++ b/data/utils.py @@ -1,10 +1,26 @@ import tables as tb class Frame: - def __init__(self, id: str, frame_nr: int, garment: str=''): - self.id: str = id - self.frame_nr: int = frame_nr - self.garment: str = garment + def __init__(self, id: str ='', frame_nr: int= 0, garment: str='', frame = None, frame_str: str = ''): + if frame is not None: + self.id: str = frame.id + self.frame_nr: int = frame.frame_nr + self.garment: str = frame.garment + assert id == '' and frame_nr == 0 and garment == '' and frame_str == '' + elif frame_str != '': + if ' ' in frame_str: + self.id: str = frame_str.split()[0] + self.frame_nr: int = int(frame_str.split()[1]) + self.garment: str = frame_str.split()[2] + else: + self.id: str = frame_str.split('_')[0] + self.frame_nr: int = int(frame_str.split('_')[1]) + self.garment: str = frame_str.split('_')[2] + assert id == '' and frame_nr == 0 and garment == '' and frame is None + else: + self.id: str = id + self.frame_nr: int = frame_nr + self.garment: str = garment def __hash__(self): return hash((self.id, self.frame_nr, self.garment)) diff --git a/docker/Dockerfile2 b/docker/Dockerfile2 index 13752ff..d3c866b 100644 --- a/docker/Dockerfile2 +++ b/docker/Dockerfile2 @@ -9,8 +9,11 @@ ARG USER_ID=11017 ARG GROUP_NAME=guest ARG GROUP_ID=11000 +# fix for tzdata asking for user input +ARG DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC + # Update the system and install required packages -RUN apt-get update && apt-get install -y sudo git apt-utils htop +RUN apt-get update && apt-get install -y sudo git apt-utils htop ffmpeg libsm6 libxext6 # Create a new group with the same GID as the host RUN groupadd --gid $GROUP_ID $GROUP_NAME diff --git a/docker/Dockerfile_hood b/docker/Dockerfile_hood new file mode 100644 index 0000000..5292bd7 --- /dev/null +++ b/docker/Dockerfile_hood @@ -0,0 +1,48 @@ +# Use an existing image as the base +# FROM nvcr.io/nvidia/pytorch:23.04-py3 +# FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel +FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel +# +# Add the current user to the image +ARG USER_NAME=lhunor +ARG USER_ID=11017 +ARG GROUP_NAME=guest +ARG GROUP_ID=11000 + +# Update the system and install required packages +RUN apt-get update && apt-get install -y sudo git apt-utils htop + +# Create a new group with the same GID as the host +RUN groupadd --gid $GROUP_ID $GROUP_NAME + +# Create a new user with the same UID and GID as the host +RUN useradd --uid $USER_ID --gid $GROUP_ID --shell /bin/bash --create-home $USER_NAME + +# don't require password with sudo, for convenience +# not the safest thing to do, but hopefully okay inside the container +RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers + +# Add the new user to the sudo group +RUN usermod -aG sudo $USER_NAME + +# install conda stuff still as root +RUN conda update -n base conda +RUN conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y +RUN conda install -c bottler nvidiacub -y +RUN conda install pytorch3d -c pytorch3d -y + +# Set the default user to the new user +USER $USER_NAME + +WORKDIR /workspace + +ENV PATH="$PATH:/home/$USER_NAME/.local/bin" +RUN pip install --upgrade pip + +# defining the volumes that will be mounted +VOLUME [ "/data-net" ] +VOLUME [ "/data-local" ] +VOLUME [ "/home-net" ] + +# Start the container with a bash shell +CMD ["/bin/bash"] \ No newline at end of file diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 76d69ab..8393c9c 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -104,4 +104,25 @@ services: resources: reservations: devices: - - capabilities: [gpu] \ No newline at end of file + - capabilities: [gpu] + + lhunor_hood: + image: lhunor_hood:pt2.0.1 + container_name: lhunor_hood_pt2.0.1 + ipc: host + tty: true + ulimits: + memlock: -1 + stack: 67108864 + volumes: + - '/home-net/lhunor:/home-net' + - '/data-net/guest/lhunor/data:/data-net' + - '/data-fast/126-data1/lhunor/data:/data-local/126' + - '/data-fast/127-data1/lhunor/data:/data-local/127' + - '/data-fast/127-data2/lhunor/data:/data-local/127_2' + - '/data-fast/128-data1/lhunor/data:/data-local/128' + deploy: + resources: + reservations: + devices: + - capabilities: [gpu] diff --git a/losses/losses.py b/losses/losses.py index 06fa74a..99c7d2c 100644 --- a/losses/losses.py +++ b/losses/losses.py @@ -32,7 +32,7 @@ def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor): # flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function 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_x = torch.arange(batch_size, device=V.device).repeat_interleave(V.shape[1]) V = V.reshape(-1, V.shape[-1]) B = B.reshape(-1, B.shape[-1]) idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_x)[1] @@ -58,7 +58,9 @@ def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: fl F_norm_local = F_norm elif B.shape[1] == 3638: F_norm_local = F_norm_orig - mesh = Meshes(B, F_norm_local) + if V.device != F_norm_local.device: + F_norm_local = F_norm_local.to(V.device) + mesh = Meshes(B, F_norm_local[:B.shape[0],:,:]) B_vn = torch.stack(mesh.verts_normals_list()) # compute correspondences @@ -74,6 +76,54 @@ def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: fl return loss, ratio +def _nearest_neighbour_error(V: torch.Tensor, B: torch.Tensor): + # flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function + batch_size = V.shape[0] + vertices_count = V.shape[1] + batch_x = torch.arange(batch_size, device=V.device).repeat_interleave(B.shape[1]) + batch_y = torch.arange(batch_size, device=V.device).repeat_interleave(V.shape[1]) + V = V.reshape(-1, V.shape[-1]) + B = B.reshape(-1, B.shape[-1]) + idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_y)[1] + idx = (idx - batch_y * 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 + +def collision_error(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = 0.0005): + if B.shape[1] == 14475: + F_norm_local = F_norm + elif B.shape[1] == 3638: + F_norm_local = F_norm_orig + if V.device != F_norm_local.device: + F_norm_local = F_norm_local.to(V.device) + mesh = Meshes(B, F_norm_local[:B.shape[0],:,:]) + B_vn = torch.stack(mesh.verts_normals_list()) + + # compute correspondences + idx = _nearest_neighbour_error(V, B) + # compute loss + D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])]) + b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])]) + dot = torch.einsum('abc,abc->ab', D, b_vn) + thresholded = torch.minimum(dot - thr, torch.zeros_like(dot)) + loss = torch.sum(thresholded ** 2, axis=-1) + ratio = torch.sum((thresholded < 0).to(torch.float32)) / V.shape[1] + + return loss, ratio + """ COLLISION FUNCTIONS END """ @@ -218,9 +268,9 @@ def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False): else: raise NotImplementedError - mesh = Meshes(pred, f_norm) + mesh = Meshes(pred, f_norm[:pred.shape[0], :, :]) pred_norm = torch.stack(mesh.verts_normals_list()) - mesh = Meshes(gt, f_norm) + mesh = Meshes(gt, f_norm[:pred.shape[0], :, :]) gt_norm = torch.stack(mesh.verts_normals_list()) return torch.mean(torch.mean((torch.sum(torch.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1)) diff --git a/models/pyramid_level.py b/models/pyramid_level.py index 482f3e1..33ae185 100644 --- a/models/pyramid_level.py +++ b/models/pyramid_level.py @@ -39,7 +39,7 @@ class PyramidLevel(nn.Module): def sample(self, z_cond, uv_mask, noise=None): if noise is None: - eps = torch.randn(z_cond.shape).to(uv_mask.device) + eps = torch.normal(mean=0.0, std=1.0, size=z_cond.shape).to(uv_mask.device) else: eps = noise diff --git a/models/pyramid_model.py b/models/pyramid_model.py index 6f81627..27f82e1 100644 --- a/models/pyramid_model.py +++ b/models/pyramid_model.py @@ -10,6 +10,29 @@ from models.pyramid_level import PyramidLevel from torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode +# class SampleParallel(nn.Module): +# def __init__(self, levels): +# super().__init__() +# self.levels = nn.ModuleList(levels) + +# def forward(self, z_cond, uv_masks, noise=None): +# level_outs = nn.parallel.data_parallel(self._forward_level, [(self.levels[i], z_cond, uv_masks[i], noise) for i in range(len(uv_masks))], dim=0, device_ids=[0, 0, 0, 0], output_device=0) +# # level_outs = nn.parallel.data_parallel(self.levels, [(z_cond, uv_masks[i], noise) for i in range(len(uv_masks))], device_ids=[0, 0, 0, 0], output_device=0) +# out = torch.zeros((uv_masks[0].shape[0], 3, 64, 32)).to(uv_masks[0].device) +# for i in range(len(self.levels)): +# level_out = level_outs[i] +# if i != len(self.levels) - 1: +# out = out + level_out +# out = self.levels[i].enhance(out, uv_masks[i]) +# out = self.resizers[self.levels[i+1].size](out) +# out = out * uv_masks[i+1][:,None,:,:] +# else: +# out = out + level_out +# out = out * uv_masks[0][:,None,:,:] +# return out + +# def _forward_level(self, level, z_cond, uv_mask, noise=None): +# return level(z_cond, uv_mask, noise=noise) class Pyramid(nn.Module): @@ -43,17 +66,24 @@ class Pyramid(nn.Module): self.use_normals = config['use_normals'] + + def sample(self, uv_template, uv_body_posed, uv_normals, uv_mask, noise=None): z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals) + # uv_masks = [uv_mask[64], uv_mask[128], uv_mask[256], uv_mask[512]] + # sampleP = SampleParallel(self.levels) + # out = sampleP(z_cond, uv_masks, noise=noise) + # return out + assert len(self.levels) == 4 out = torch.zeros((uv_template.shape[0], 3, 64, 32)).to(uv_template.device) for i in range(len(self.levels)): - level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise) + level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise[i] if noise is not None else None) if i != len(self.levels) - 1: - out = out + level_out + out = out + level_out out = self.levels[i].enhance(out, uv_mask[self.levels[i].size]) out = self.resizers[self.levels[i+1].size](out) out = out * uv_mask[self.levels[i+1].size][:,None,:,:] @@ -63,7 +93,30 @@ class Pyramid(nn.Module): out = out * uv_mask[512][:,None,:,:] return out + + + def sample_level(self, uv_template, uv_body_posed, uv_normals, uv_mask, prev, level, noise=None): + + z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals) + assert len(self.levels) == 4 + + i = level + out = prev + + level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise[i] if noise is not None else None) + if i != len(self.levels) - 1: + out = out + level_out + out = self.levels[i].enhance(out, uv_mask[self.levels[i].size]) + out = self.resizers[self.levels[i+1].size](out) + out = out * uv_mask[self.levels[i+1].size][:,None,:,:] + else: + out = out + level_out + out = out * uv_mask[512][:,None,:,:] + + return out + + 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) diff --git a/run.py b/run.py index bd248e4..e0c326d 100644 --- a/run.py +++ b/run.py @@ -1,6 +1,6 @@ import os import sys -from configs.cvae_norm_1 import config +from configs.iccvw.pyr_64_128_256_512_skirt_e2e_var import config import argparse # parser = argparse.ArgumentParser(description='') @@ -12,14 +12,15 @@ import argparse if __name__ == '__main__': import wandb # wandb.login(key='******************************************') - wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838') + # wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838') + # wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) name = config['name'] script = config['script'] config = [key for key in sys.modules.keys() if key.startswith('configs.')][-1] current_file = '/'.join(__file__.split('/')[-2:]) - print(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &') - if args['debug']: - os.system(f'python {script} --name lhunor_{name} --config {config}') - else: - os.system(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &') \ No newline at end of file + print(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/iccvw/{name}.out 2>&1 &') + # if args['debug']: + # os.system(f'python {script} --name lhunor_{name} --config {config}') + # else: + # os.system(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &') \ No newline at end of file diff --git a/run_parallel.py b/run_parallel.py index da12a64..d6aee25 100644 --- a/run_parallel.py +++ b/run_parallel.py @@ -12,19 +12,24 @@ if __name__ == '__main__': # ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:] skirt = 0 obj_folder = '/data-net/dataset/cloth3d/train' - folder = '/data-local/127/train' + folder = '/data-local/127_2/cape_processed/train/male' + # folder = '/data-local/127_2/cloth3d_processed_test' count = 0 - step = 100 + step = 35 # !!!!!!! change skirt parameter in script - while count <1200: + while count < 550: # os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python convert_to_uv_local.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &') - os.system(f'nohup python data/calculate_normals.py -f {folder} -f_obj {obj_folder} -s {count} -e {(count+step)} -n normals_train_calc_{count}_{(count+step)} > nohup/normals/normals_train_calc_{count}_{(count+step)}.out' + ' 2>&1 &') + # os.system(f'nohup python data/convert_to_uv_cape.py -f {folder} -s {count} -e {(count+step)} -n to_uv_cape_m_test_{count}_{(count+step)} > nohup/to_uv/cape_m_test_{count}_{(count+step)}.out' + ' 2>&1 &') + # os.system(f'nohup python data/convert_to_uv.py -f {folder} -s {count} -e {(count+step)} -n to_uv_{count}_{(count+step)} > nohup/to_uv/test_{count}_{(count+step)}.out' + ' 2>&1 &') + # os.system(f'nohup python data/downscale.py -f {folder} -s {count} -e {(count+step)} -n down_test_{count}_{(count+step)} > nohup/down/test_{count}_{(count+step)}.out' + ' 2>&1 &') + # os.system(f'nohup python data/downscale_cape.py -f {folder} -s {count} -e {(count+step)} -n down_train_cape_m_{count}_{(count+step)} > nohup/down/train_cape_m_{count}_{(count+step)}.out' + ' 2>&1 &') + os.system(f'nohup python scripts/calculate_chamfer.py -s {count} -e {(count+step)} -n hood_130k_{count}_{(count+step)} > nohup/hood/hood_130k_{count}_{(count+step)}.out' + ' 2>&1 &') # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_valid_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_valid_skirt_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') diff --git a/scripts/calculate.py b/scripts/calculate.py new file mode 100644 index 0000000..f0687aa --- /dev/null +++ b/scripts/calculate.py @@ -0,0 +1,34 @@ +import numpy as np + +error_1 = dict() +error_2 = dict() +#end2end with reg +# error_1[0] = {'Jumpsuit': 0.00888, 'Trousers': 0.01193, 'Tshirt': 0.01065, 'Top': 0.00725} +# error_1[1] = {'Jumpsuit': 0.00491, 'Trousers': 0.00704, 'Tshirt': 0.00531, 'Top': 0.00393} +# error_1[2] = {'Jumpsuit': 0.00379, 'Trousers': 0.00547, 'Tshirt': 0.00404, 'Top': 0.00309} +# error_1[3] = {'Jumpsuit': 0.00349, 'Trousers': 0.00502, 'Tshirt': 0.00382, 'Top': 0.00281} +# error_2[0] = {'Dress': 0.01009, 'Skirt': 0.01379} +# error_2[1] = {'Dress': 0.0055, 'Skirt': 0.0077} +# error_2[2] = {'Dress': 0.00418, 'Skirt': 0.00582} +# error_2[3] = {'Dress': 0.00359, 'Skirt': 0.00496} + +#end2end without reg +error_1[0] = {'Jumpsuit': 0.00927, 'Trousers': 0.01229, 'Tshirt': 0.01113, 'Top': 0.0076} +error_1[1] = {'Jumpsuit': 0.00503, 'Trousers': 0.00721, 'Tshirt': 0.0054, 'Top': 0.004} +error_1[2] = {'Jumpsuit': 0.00402, 'Trousers': 0.0059, 'Tshirt': 0.00421, 'Top': 0.00319} +error_1[3] = {'Jumpsuit': 0.00347, 'Trousers': 0.00528, 'Tshirt': 0.00365, 'Top': 0.00271} +error_2[0] = {'Dress': 0.01005, 'Skirt': 0.01395} +error_2[1] = {'Dress': 0.0056, 'Skirt': 0.00787} +error_2[2] = {'Dress': 0.00441, 'Skirt': 0.00634} +error_2[3] = {'Dress': 0.00379, 'Skirt': 0.0054} + +for i in range(4): + errors_mean = dict() + error = {**error_1[i], **error_2[i]} + for category in error.keys(): + errors_mean[category] = round(np.mean(error[category]), 5) + print(f'Level {i}: {errors_mean}') + + mean = round(np.mean(list(errors_mean.values())), 5) + + print(f'Final level {i}: {mean}') diff --git a/scripts/calculate_area_diff.py b/scripts/calculate_area_diff.py new file mode 100644 index 0000000..27921b5 --- /dev/null +++ b/scripts/calculate_area_diff.py @@ -0,0 +1,47 @@ +import trimesh +import numpy as np + +cases = [ + # ('06461_237_Trousers', 'temp/sampling_pyramid/06461_237_Trousers', [0,8,12,13]), + # ('01171_213_Dress', 'temp/sampling_pyramid/01171_213_Dress', [0,6,10,15]), + ('01852_70_Tshirt_lvl_0', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,10,1,5,13]), + ('01852_70_Tshirt_lvl_1', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,14,3,5,0 ]), + ('01852_70_Tshirt_lvl_2', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,12,0,8,6]), + ('01852_70_Tshirt_lvl_3', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,2,8,1,0]), +] + +for case in cases: + # Load the original mesh + folder = case[1] + original_mesh = trimesh.load(f'{folder}/{case[0]}_{case[2][0]}.stl') + + # Load the variation meshes + variation_meshes = [] + for i in case[2][1:]: + variation_meshes.append(trimesh.load(f'{folder}/{case[0]}_{i}.stl')) + + # Calculate area of original mesh + # original_area = original_mesh.area + + # Calculate area difference for each variation + area_diffs = [] + for variation_mesh in variation_meshes: + # calculate per vertex euclid distance + euclid = np.linalg.norm(original_mesh.vertices - variation_mesh.vertices, axis=1) * 1000 + euclid_diff = np.mean(euclid) + + + # area_diff = abs(variation_mesh.area - original_area) / original_area + area_diffs.append(euclid_diff) + + # Average area differences + # avg_area_diff = np.mean(area_diffs) + + variance = np.var(area_diffs) + + # Print results + print(f'Case: {case[0]}') + # print('Original area:', original_area) + print('Euclidean differences:', area_diffs) + print('Euclidean variance:', variance) + # print('Average Euclidean difference:', avg_area_diff) \ No newline at end of file diff --git a/scripts/calculate_chamfer.py b/scripts/calculate_chamfer.py index 8ead6c6..37f2ded 100644 --- a/scripts/calculate_chamfer.py +++ b/scripts/calculate_chamfer.py @@ -4,47 +4,40 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" -os.environ["CUDA_VISIBLE_DEVICES"] = '0' +os.environ["CUDA_VISIBLE_DEVICES"] = '6' 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 data.utils import Frame -from data.cloth3d_io import readPC2Frame, readPC2 -from data.preprocessing import deorient -from utils.checkpoints import CheckpointSaver, load_model -from torch.utils.data import DataLoader -from utils.utils import text_to_html, pickle_load -from utils.postprocess import fix_collision -from losses.losses import collision_loss, normal_loss, kl_loss -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 plotly.subplots import make_subplots -import pytorch3d.loss -from scipy.spatial.distance import cdist +from data.cloth3d_io import readPC2Frame, readPC2, readOBJ +from utils.utils import pickle_load from scipy.spatial import cKDTree from smpl.smpl_torch import SMPL from scipy.spatial.transform import Rotation as R +import trimesh as tm +from losses.losses import collision_loss, normal_loss, kl_loss, collision_error +parser = argparse.ArgumentParser() +parser.add_argument("-s", "--start", help = "start of range", default=0, type=int) +parser.add_argument("-e", "--end", help = "end of range", default=550, type=int) +parser.add_argument("-n", "--name", help = "end of range", default='lhunor') +parser.add_argument('--debug', action='store_true') +args = parser.parse_args() def chamfer(points1, points2, points2_tree: cKDTree): tree1 = cKDTree(points1) - dist1, _ = tree1.query(points2, workers=40) - dist2, _ = points2_tree.query(points1, workers=40) + dist1, _ = tree1.query(points2, workers=10) + dist2, _ = points2_tree.query(points1, workers=10) chamfer_dist = np.mean(dist1) + np.mean(dist2) return chamfer_dist @@ -60,34 +53,116 @@ def deorient(verts, pose, inverse=True): return rotated + +def get_collision_depth(m1, m2): + cm = tm.collision.CollisionManager() + cm.add_object('m1', m1) + cm.add_object('m2', m2) + data = cm.in_collision_internal(return_data=True)[1] + # c = np.asarray([d.index('m2') for d in data]) + # count = np.unique(c).shape[0] + depth_sum = np.sum([d.depth for d in data]) + return depth_sum + + +def get_area_diff(m1, m2): + return np.abs(m1.area - m2.area) + +device = 'cuda' +def get_collision_loss(pred, body, mask): + c_loss, c_ratio = collision_error(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device)) + return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy() + +def get_collision_loss_2(pred, body, mask): + _pred = np.zeros((14475, 3), dtype=np.float32) + _pred[mask] = pred + c_loss, c_ratio = collision_loss(torch.from_numpy(_pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005) + return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy() + def calculate(): errors_hood = [] - errors_ours = [] - garments = sorted(os.listdir(hood_folder)) + collisions_hood = [] + areas_hood = [] + collision_ratios_hood = [] + if calculate_ours: + errors_ours = [] + collisions_ours = [] + areas_ours = [] + collision_ratios_ours = [] + + garments = sorted(os.listdir(hood_folder))[args.start:args.end] garments = [g[:-4] for g in garments] + body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32) for garment in tqdm(garments): picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt' id, garment_type = garment.split('_') - pred_hood = pickle_load(f'{hood_folder}/{garment}.pkl')['verts'] + pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl') + pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces'] gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:] + template, template_faces, _, _ = readOBJ(os.path.join(cloth3d_folder, id, garment_type + '.obj')) mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask'] pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:] gt = deorient(gt, pose, True) + body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:] + fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz') + faces = fix_data['faces'] + m_template = tm.Trimesh(vertices=template, faces=template_faces, process=False, validate=False) for i in range(pred_hood.shape[0]): gt_tree = cKDTree(gt[i]) try: - pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3'] + if calculate_ours: + # pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3'] + pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'] except: print('missed one') continue + # m_body = tm.Trimesh(vertices=body[i], faces=body_faces, process=False, validate=False) + errors_hood.append((f'{id}_{i}_{garment_type}', chamfer(pred_hood[i], gt[i], gt_tree))) - errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask,:], gt[i], gt_tree))) + m_hood = tm.Trimesh(vertices=pred_hood[i], faces=pred_hood_faces, process=False, validate=False) + # collision_hood = get_collision_depth(m_body, m_hood) + # collisions_hood.append((f'{id}_{i}_{garment_type}', collision_hood)) + area_hood = get_area_diff(m_template, m_hood) + areas_hood.append((f'{id}_{i}_{garment_type}', area_hood)) + collision_ratio_hood = get_collision_loss(pred_hood[i], body[i], mask)[1] + collision_ratios_hood.append((f'{id}_{i}_{garment_type}', collision_ratio_hood)) + + if calculate_ours: + errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask], gt[i], gt_tree))) + m_ours = tm.Trimesh(vertices=pred_ours[mask], faces=faces, process=False, validate=False) + # collision_ours = get_collision_depth(m_body, m_ours) + # collisions_ours.append((f'{id}_{i}_{garment_type}', collision_ours)) + area_ours = get_area_diff(m_template, m_ours) + areas_ours.append((f'{id}_{i}_{garment_type}', area_ours)) + collision_ratio_ours = get_collision_loss(pred_ours[mask], body[i], mask)[1] + # collision_ratio_ours_2 = get_collision_loss_2(pred_ours[mask], body[i], mask)[1] + collision_ratios_ours.append((f'{id}_{i}_{garment_type}', collision_ratio_ours)) + + + # fig = go.Figure(data=[go.Mesh3d(x=pred_ours[mask][:,0], y=pred_ours[mask][:,2], z=pred_ours[mask][:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')]) + # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) + # 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', scene_aspectratio=dict(x=1, y=1, z=1)) + # fig.show() + + # fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')]) + # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) + # 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', scene_aspectratio=dict(x=1, y=1, z=1)) + # fig.show() + + # fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')]) + # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) + # 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', scene_aspectratio=dict(x=1, y=1, z=1)) + # fig.show() + # pass # fig = go.Figure() # fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1))) @@ -95,6 +170,7 @@ def calculate(): # 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 # fig = go.Figure() # fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1))) @@ -105,30 +181,83 @@ def calculate(): # save errors to a csv file - with open(f'{results_folder}/hood/{name}_hood.csv', 'w') as f: - for frame, error in errors_hood: - frame_str = ','.join(str(frame).split('_')) - f.write(f'{frame_str},{error}\n') + os.makedirs(f'{results_folder}/hood/{name}_hood', exist_ok=True) + os.makedirs(f'{results_folder}/hood/{name}_hood/chamfer', exist_ok=True) + os.makedirs(f'{results_folder}/hood/{name}_hood/collision', exist_ok=True) + os.makedirs(f'{results_folder}/hood/{name}_hood/area', exist_ok=True) + os.makedirs(f'{results_folder}/hood/{name}_hood/collision_ratio', exist_ok=True) - with open(f'{results_folder}/hood/{name}_ours.csv', 'w') as f: - for frame, error in errors_ours: - frame_str = ','.join(str(frame).split('_')) - f.write(f'{frame_str},{error}\n') + for frame, error in errors_hood: + np.save(f'{results_folder}/hood/{name}_hood/chamfer/{frame}.npy', error) + for frame, error in collisions_hood: + np.save(f'{results_folder}/hood/{name}_hood/collision/{frame}.npy', error) + for frame, error in areas_hood: + np.save(f'{results_folder}/hood/{name}_hood/area/{frame}.npy', error) + for frame, error in collision_ratios_hood: + np.save(f'{results_folder}/hood/{name}_hood/collision_ratio/{frame}.npy', error) + + os.makedirs(f'{results_folder}/hood/{name}_ours', exist_ok=True) + os.makedirs(f'{results_folder}/hood/{name}_ours/chamfer', exist_ok=True) + os.makedirs(f'{results_folder}/hood/{name}_ours/collision', exist_ok=True) + os.makedirs(f'{results_folder}/hood/{name}_ours/area', exist_ok=True) + os.makedirs(f'{results_folder}/hood/{name}_ours/collision_ratio', exist_ok=True) + + for frame, error in errors_ours: + np.save(f'{results_folder}/hood/{name}_ours/chamfer/{frame}.npy', error) + for frame, error in collisions_ours: + np.save(f'{results_folder}/hood/{name}_ours/collision/{frame}.npy', error) + for frame, error in areas_ours: + np.save(f'{results_folder}/hood/{name}_ours/area/{frame}.npy', error) + for frame, error in collision_ratios_ours: + np.save(f'{results_folder}/hood/{name}_ours/collision_ratio/{frame}.npy', error) + + # with open(f'{results_folder}/hood/{name}/chamfer_hood.csv', 'w') as f: + # for frame, error in errors_hood: + # frame_str = ','.join(str(frame).split('_')) + # f.write(f'{frame_str},{error}\n') + + # with open(f'{results_folder}/hood/{name}/collision_hood.csv', 'w') as f: + # for frame, error in collisions_hood: + # frame_str = ','.join(str(frame).split('_')) + # f.write(f'{frame_str},{error}\n') + + # with open(f'{results_folder}/hood/{name}/area_hood.csv', 'w') as f: + # for frame, error in areas_hood: + # frame_str = ','.join(str(frame).split('_')) + # f.write(f'{frame_str},{error}\n') + + # if calculate_ours: + # with open(f'{results_folder}/hood/{name}_chamfer_ours.csv', 'w') as f: + # for frame, error in errors_ours: + # frame_str = ','.join(str(frame).split('_')) + # f.write(f'{frame_str},{error}\n') + + # with open(f'{results_folder}/hood/{name}/collision_ours.csv', 'w') as f: + # for frame, error in collisions_ours: + # frame_str = ','.join(str(frame).split('_')) + # f.write(f'{frame_str},{error}\n') + + # with open(f'{results_folder}/hood/{name}/area_ours.csv', 'w') as f: + # for frame, error in areas_ours: + # frame_str = ','.join(str(frame).split('_')) + # f.write(f'{frame_str},{error}\n') if __name__ == '__main__': - name = 'hood_2' + name = 'hood_3_var' - hood_folder = '/data-local/127_2/results/hood/all_2' + hood_folder = '/data-local/127_2/results/hood/all_3' ours_folder = dict() - ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000' - ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000' + ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_var_3_450000' + ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_var_2_600000' cloth3d_folder = '/data-net/dataset/cloth3d/test' results_folder = '/data-local/127_2/results' - out_folder = f'{results_folder}/errors/{name}' + + calculate_ours = True + calculate() \ No newline at end of file diff --git a/scripts/calculate_error_all.py b/scripts/calculate_error_all.py index 60ce969..9beee1a 100644 --- a/scripts/calculate_error_all.py +++ b/scripts/calculate_error_all.py @@ -10,6 +10,7 @@ parser.add_argument('--config', type=str) parser.add_argument('--debug', action='store_true') parser.add_argument('--iteration', type=str, default='200000') parser.add_argument('-g', '--gpu', type=str, default='0') +parser.add_argument('-f', '--final', action='store_true') args = vars(parser.parse_args()) if args['debug']: @@ -68,7 +69,7 @@ def get_normal_loss(pred, gt, mask): return n_loss.cpu().numpy() def get_collision_loss(pred, body, mask): - c_loss, c_ratio = collision_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device)) + c_loss, c_ratio = collision_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005) return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy() @@ -115,10 +116,15 @@ def calculate(): errors_edge.append((str(frame), get_edge_loss(pred, gt, new_mask))) else: - for i in range(4): + levels = [0,1,2,3] if not args['final'] else [3] + for i in levels: _pred = pred.item()[f'lvl{i}'] c_loss, c_ratio = get_collision_loss(_pred, body, new_mask) + errors_collision.append((f'{frame}_lvl{i}', c_loss)) + errors_collision_ratio.append((f'{frame}_lvl{i}', c_ratio)) + errors_normal.append((f'{frame}_lvl{i}', get_normal_loss(_pred, gt, new_mask))) + errors_edge.append((f'{frame}_lvl{i}', get_edge_loss(_pred, gt, new_mask))) # faces = fix_data['faces'] # body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32) @@ -143,10 +149,6 @@ def calculate(): # fig.show() - errors_collision.append((f'{frame}_lvl{i}', c_loss)) - errors_collision_ratio.append((f'{frame}_lvl{i}', c_ratio)) - errors_normal.append((f'{frame}_lvl{i}', get_normal_loss(_pred, gt, new_mask))) - errors_edge.append((f'{frame}_lvl{i}', get_edge_loss(_pred, gt, new_mask))) # save errors to a csv file diff --git a/scripts/cvae.py b/scripts/cvae.py index 20d4140..4491792 100644 --- a/scripts/cvae.py +++ b/scripts/cvae.py @@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) if args['debug']: - from configs.iccvw.abl_all_5_var import config + from configs.iccvw.abl_all_5 import config else: config_module = importlib.import_module(args['config']) config = config_module.config @@ -45,7 +45,7 @@ from data.dataset import Dataset from models.cond_encoder import AutoEncoder from utils.checkpoints import CheckpointSaver, load_model from utils.utils import InfiniteDataLoader as DataLoader -from utils.utils import text_to_html +from utils.utils import text_to_html, save_stl, create_temp_frames from utils.postprocess import fix_collision from losses.losses import collision_loss, normal_loss from pytorch_msssim import MS_SSIM @@ -352,7 +352,52 @@ def sample(): break +def generate(): + from data.utils import Frame + load_model(model, f'checkpoints/iccvw/{config["name"]}') + frame = '01852_70_Tshirt' + # frame = '03849_27_Dress' # 03849_27_Dress 00054_1_Dress + frame = Frame(frame_str=frame) + create_temp_frames(config['train_dir'], frame) + + dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory']) + train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, ) + + batch = next(iter(train_dataloader)) + + 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.eval() + with torch.inference_mode(): + 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])]) + # shape_offsets = torch.stack([torch.matmul(shape[i, :10], smpl.shapedirs[gender[i]][:, 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 + + fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz') + mask = fix_data['mask'] + faces = fix_data['faces'] + + pred = mesh_pred_posed.cpu().numpy()[0, mask, :] + body = batch['mesh_body_posed'].cpu().numpy()[0] + body_faces = smpl.f + save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_mesh.stl', pred, faces) + save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_body.stl', body, body_faces) + pass def train(): @@ -410,6 +455,9 @@ if __name__ == '__main__': model = CVAE(config) + generate() + exit() + sample() exit() diff --git a/scripts/cvae_pyramid.py b/scripts/cvae_pyramid.py index bd3c471..909fd67 100644 --- a/scripts/cvae_pyramid.py +++ b/scripts/cvae_pyramid.py @@ -11,7 +11,8 @@ parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) if args['debug']: - from configs.iccvw.pyr_64_128_256_512 import config + # from configs.iccvw.pyr_64_128_256_512 import config + from configs.iccvw.pyr_64_128_256_512_e2e import config else: config_module = importlib.import_module(args['config']) config = config_module.config @@ -24,7 +25,7 @@ if args['debug']: import time import argparse import wandb -import tqdm +from tqdm import tqdm import random import numpy as np import plotly.graph_objects as go @@ -501,6 +502,32 @@ def speed_step(levels_input, conds): return model.sample([levels_input, conds], training=False) +from scipy.spatial import KDTree +import trimesh + +def connect_unconnected_vertices(mesh): + # Find unconnected vertices + unconnected_verts = np.where(np.array(mesh.vertex_degree) == 0)[0] + + # Create KD-tree from mesh vertices + tree = KDTree(mesh.vertices) + + new_faces = mesh.faces.copy() + + # Connect unconnected vertices to their two closest neighbors + for vert in unconnected_verts: + # Find the two closest vertices + _, closest_verts = tree.query(mesh.vertices[vert], k=3) + + # Create a new face with the unconnected vertex and its two closest neighbors + new_face = [vert, closest_verts[1], closest_verts[2]] + new_faces = np.vstack((new_faces, new_face)) + + mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=new_faces, process=False) + + return mesh + + def sample(): from data.utils import Frame from utils.utils import save_stl @@ -515,11 +542,10 @@ def sample(): load_model(model, f'checkpoints/iccvw/{config["name"]}') count = 0 - frames = ['01852_70_Tshirt'] + frames = ['06461_237_Trousers'] # 03528_132_Tshirt 01852_70_Tshirt 06461_237_Trousers 01893_173_Jumpsuit + frames = ['01171_213_Dress'] # 01039_118_Dress 01171_213_Dress frames = [Frame(frame_str=frame) for frame in frames] - - for frame in frames: create_temp_frames(config['train_dir'], frame) dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape']) @@ -532,6 +558,8 @@ def sample(): uv_normals = batch['uv_normals'] batch_size = uv_static.size(0) + body = batch['mesh_body_posed'][0].cpu().numpy() + body_faces = smpl.f for i in range(batch_size): @@ -541,8 +569,11 @@ def sample(): mask = fix_data['mask'] faces = fix_data['faces'] + os.makedirs(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True) + gt = batch['mesh_posed'][i].cpu().numpy()[mask] - save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces) + save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces) + save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces) uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1) uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1) @@ -551,6 +582,8 @@ def sample(): pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1) body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1) mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1) + + noise = [] for _ in range(4): with torch.inference_mode(): @@ -570,9 +603,12 @@ def sample(): for k in range(4): _mesh_posed_fixed = mesh_posed_fixed[k, mask, :] m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False) + m = connect_unconnected_vertices(m) + tm.smoothing.filter_humphrey(m, iterations=2) verts = m.vertices - save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces) + save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces) + save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}_body.stl', body, body_faces) counter += 1 pass @@ -582,6 +618,265 @@ def sample(): break + +def sample_pro(): + from data.utils import Frame + from utils.utils import save_stl + from pathlib import Path + import trimesh as tm + + def create_temp_frames(folder: str, frame: Frame): + with open(Path(folder) / 'frames_temp.txt', 'w') as f: + f.writelines([f'{frame}\n' for _ in range(4)]) + + model.eval() + load_model(model, f'checkpoints/iccvw/{config["name"]}') + count = 0 + + # frames = ['06461_237_Trousers'] # 03528_132_Tshirt 01852_70_Tshirt 06461_237_Trousers 01893_173_Jumpsuit + frames = ['01852_70_Tshirt'] # 01039_118_Dress 01171_213_Dress 01852_70_Tshirt + frames = [Frame(frame_str=frame) for frame in frames] + + for frame in frames: + create_temp_frames(config['train_dir'], frame) + dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape']) + dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, ) + batch = next(iter(dataloader)) + + uv_static = batch['uv_static'] + uv_body_posed = batch['uv_body_posed'] + uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] } + uv_normals = batch['uv_normals'] + + batch_size = uv_static.size(0) + body = batch['mesh_body_posed'][0].cpu().numpy() + body_faces = smpl.f + + + i=0 + frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i]) + fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz') + mask = fix_data['mask'] + faces = fix_data['faces'] + + os.makedirs(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True) + + gt = batch['mesh_posed'][i].cpu().numpy()[mask] + save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces) + save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces) + + uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1) + uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1) + uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)} + uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1) + pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1) + body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1) + mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1) + + out = torch.zeros((uv_static_tiled.shape[0], 3, 64, 32)).to(uv_static_tiled.device) + for l in range(4): + with torch.inference_mode(): + for b in range(4): + pred_level = model.sample_level(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled, out, l) + pred = pred_level + + mesh_pred = uv2mesh(pred) + + shape = 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(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0] + + mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, mask_tiled).cpu().numpy() + + for k in range(4): + _mesh_posed_fixed = mesh_posed_fixed[k, mask, :] + m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False) + m = connect_unconnected_vertices(m) + + tm.smoothing.filter_humphrey(m, iterations=2) + verts = m.vertices + save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_lvl_{l}_{b*4+k}.stl', verts, faces) + save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_lvl_{l}_{b*4+k}_body.stl', body, body_faces) + + #replace the contentent of out with its own first element tiled + if b == 3: + out = pred_level[3,:,:,:].repeat(4, 1, 1, 1) + + pass + + + # if count > 100: + # break + + + + +def pred_pro(): + from data.utils import Frame + from utils.utils import save_stl + from pathlib import Path + import trimesh as tm + from typing import List + from models.cvae_norm_model import CVAENorm + from data.downscale import fill + from torchvision.transforms import Resize + from torchvision.transforms.functional import InterpolationMode + + def create_temp_frames(folder: str, frames: List[Frame]): + with open(Path(folder) / 'frames_temp.txt', 'w') as f: + f.writelines([f'{frame}\n' for frame in frames]) + + model.eval() + load_model(model, f'checkpoints/iccvw/{config["name"]}') + count = 0 + + from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_skirt + model_skirt: Pyramid = Pyramid(config_skirt) + load_model(model_skirt, f'checkpoints/iccvw/{config_skirt["name"]}') + models = {'non-skirt': model, 'skirt': model_skirt} + from configs.iccvw.cvae_norm_5 import config as config_norm + model_norm = CVAENorm(config_norm) + model_norm.eval() + load_model(model_norm, f'checkpoints/iccvw/{config_norm["name"]}') + + uv2mesh_skirt = UV2Mesh(res=256, is_skirt=config_skirt['is_skirt'], is_smpl_orig=config['is_cape']) + uv2meshes = {'non-skirt': uv2mesh, 'skirt': uv2mesh_skirt} + + # '01320_30_Dress' + frames = ['01852_70_Tshirt', '04717_276_Trousers', '07532_152_Tshirt', '02792_282_Jumpsuit'] # 01039_118_Dress 01171_213_Dress + frames = [Frame(frame_str=frame) for frame in frames] + + create_temp_frames(config['train_dir'], frames) + dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape']) + dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False, ) + batch_orig = next(iter(dataloader)) + + batch_size = 4 + body_faces = smpl.f + + masks = [] + faces_orig = [] + for frame in frames: + fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz') + masks.append(fix_data['mask']) + faces_orig.append(fix_data['faces']) + masks = np.stack(masks, axis=0) + # faces = np.stack(faces, axis=0) + + resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to('cuda') for res in [64, 128, 256, 512] } + + mode = 'body' + + for i in range(batch_size): + + batch = batch_orig.copy() + uv_static = batch['uv_static'] + uv_body_posed = batch['uv_body_posed'] + uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] } + uv_normals = batch['uv_normals'] + + counter = 0 + frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i]) + # fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz') + # mask = fix_data['mask'] + # faces = fix_data['faces'] + type = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'non-skirt' + + os.makedirs(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True) + + # gt = batch['mesh_posed'][i].cpu().numpy()[mask] + # save_stl(Path('temp') / 'pred_pro' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces) + # save_stl(Path('temp') / 'pred_pro' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces) + + uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1) + uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1) + uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)} + uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1) + pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1) + body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1) + # masks_tiled = torch.from_numpy(masks[i][None,:]).to(body_posed.device).repeat(4, 1) + + if mode == 'static': + batch['uv_body_posed'] = uv_body_posed_tiled + batch['mesh_body_posed'] = body_posed + batch['pose_shape'] = pose_shape + masks_tiled = masks + for res in [64, 128, 256]: + batch[f'uv_body_posed_{res}'] = batch[f'uv_body_posed_{res}'][i][None,:,:,:].repeat(4, 1, 1, 1) + # orig_img = uv_body_posed_tiled.clone() + # downscaled = {64: [], 128: [], 256: []} + # for i in range(4): + # filled = fill(orig_img[i].cpu().numpy().transpose(1,2,0), uv_mask[512][i].cpu().numpy()) + # filled = torch.from_numpy(filled.transpose(2,0,1)).to(orig_img.device) + # for res in [64, 128, 256]: + # downscaled[res].append(resizers[res](filled) * uv_mask[res][i]) + # for res in [64, 128, 256]: + # batch[f'uv_body_posed_{res}'] = torch.stack(downscaled[res], dim=0) + faces = faces_orig + elif mode == 'body': + batch['uv_static'] = uv_static_tiled + for res in [64, 128, 256]: + batch[f'uv_static_{res}'] = batch[f'uv_static_{res}'][i][None,:,:,:].repeat(4, 1, 1, 1) + batch['mesh_mask'] = batch['mesh_mask'][i][None,:].repeat(4, 1) + for key in uv_mask_tiled.keys(): + batch[f'uv_mask_{key}'] = uv_mask_tiled[key] + batch['uv_mask'] = uv_mask_tiled[512] + masks_tiled = torch.from_numpy(masks[i][None,:]).to(body_posed.device).repeat(4, 1).cpu().numpy() + faces = [faces_orig[i], faces_orig[i], faces_orig[i], faces_orig[i]] + else: + raise Exception('mode not supported') + + body_posed = batch['mesh_body_posed'] + pose_shape = batch['pose_shape'] + + with torch.inference_mode(): + pred_norm = model_norm.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']) * 1.4 + batch['uv_normals'] = pred_norm + orig_img = pred_norm.clone() + downscaled = {64: [], 128: [], 256: []} + for i in range(4): + filled = fill(orig_img[i].cpu().numpy().transpose(1,2,0), uv_mask[512][i].cpu().numpy()) + filled = torch.from_numpy(filled.transpose(2,0,1)).to(orig_img.device) + for res in [64, 128, 256]: + downscaled[res].append(resizers[res](filled) * uv_mask[res][i]) + for res in [64, 128, 256]: + batch[f'uv_normals_{res}'] = torch.stack(downscaled[res], dim=0) + + levels_input, conds = group_inputs_for_levels(batch, config) + + with torch.inference_mode(): + out_levels = models[type]([levels_input, conds]) + pred, _, _ = out_levels[-1] + + mesh_pred = uv2meshes[type](pred) + + shape = 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(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0] + + # mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, torch.from_numpy(masks).to(body_posed.device)).cpu().numpy() + mesh_posed_fixed = mesh_pred_posed.cpu().numpy() + for k in range(4): + _mesh_posed_fixed = mesh_posed_fixed[k, masks_tiled[k], :] + m = tm.Trimesh(_mesh_posed_fixed, faces[k], process=False, validate=False) + m = connect_unconnected_vertices(m) + + tm.smoothing.filter_humphrey(m, iterations=2) + verts = m.vertices + save_stl(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces[k]) + save_stl(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}_body.stl', batch['mesh_body_posed'][k].cpu().numpy(), body_faces) + counter += 1 + + pass + + + def train(): global g_balancing_coeff global step @@ -620,9 +915,44 @@ def train(): step += 1 +def speed_test(): + from data.utils import Frame + from utils.utils import save_stl + from pathlib import Path + import trimesh as tm + + def create_temp_frames(folder: str, frame: Frame): + with open(Path(folder) / 'frames_temp.txt', 'w') as f: + f.writelines([f'{frame}\n' for _ in range(4)]) + + model.eval() + load_model(model, f'checkpoints/iccvw/{config["name"]}') + + frames = ['01852_70_Tshirt'] + frames = [Frame(frame_str=frame) for frame in frames] + + create_temp_frames(config['train_dir'], frames[0]) + dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape']) + dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, ) + batch = next(iter(dataloader)) + + uv_static = batch['uv_static'] + uv_body_posed = batch['uv_body_posed'] + uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] } + uv_normals = batch['uv_normals'] + + for i in tqdm(range(1000)): + with torch.inference_mode(): + if i ==5: + start = time.time() + model.sample(uv_static, uv_body_posed, uv_normals, uv_mask) + + print(f'per batch: {(time.time() - start) / 1000}') + + if __name__ == '__main__': - if not args['debug']: - init_wandb() + # if not args['debug']: + # init_wandb() use_amp = True smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device']) @@ -640,7 +970,14 @@ if __name__ == '__main__': model: Pyramid = Pyramid(config) - sample() + # count number of parameters + total_params = sum(p.numel() for p in model.parameters()) + + # speed_test() + # exit() + + # sample_pro() + pred_pro() exit() optimizer = optim.Adam(model.parameters(), lr=config['lr']) @@ -652,23 +989,23 @@ if __name__ == '__main__': reg_coeff = 1 step = 0 - if config['pretrained_model'] != '': - if 'ignore_saved_optimizer' in config.keys() and config['ignore_saved_optimizer']: - step = load_model(model, config['pretrained_model']) - else: - step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration']) + # if config['pretrained_model'] != '': + # if 'ignore_saved_optimizer' in config.keys() and config['ignore_saved_optimizer']: + # step = load_model(model, config['pretrained_model']) + # else: + # step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration']) - 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 - 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']) - for param in model.static_enc.parameters(): param.requires_grad = False - for param in model.body_enc.parameters(): param.requires_grad = False - if config['use_normals']: - for param in model.norm_enc.parameters(): param.requires_grad = False + # 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 + # 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']) + # for param in model.static_enc.parameters(): param.requires_grad = False + # for param in model.body_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 @@ -691,4 +1028,5 @@ if __name__ == '__main__': # model = torch.compile(model) wandb.watch(model, log='all', log_freq=10000, log_graph=True) + model.eval() train() \ No newline at end of file diff --git a/scripts/cvae_pyramid_e2e.py b/scripts/cvae_pyramid_e2e.py index 5b08e35..6a0230d 100644 --- a/scripts/cvae_pyramid_e2e.py +++ b/scripts/cvae_pyramid_e2e.py @@ -547,7 +547,7 @@ def train(): valid_accumulator.update(valid_loss) if step % 1000 == 0 or (step < 1000 and step % 100 == 0): - to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } }) + to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step, 'balancing_coeff': g_balancing_coeff } }) if step % 5000 == 0 or (step < 5000 and step % 1000 == 0): train_figs = plot_everything(batch, 'train') valid_figs = plot_everything(data_valid, 'valid') @@ -561,10 +561,9 @@ def train(): valid_accumulator.reset() print(step) - # if train_loss['train/euclid_error'] < config['balancing_activation_limit']: - # g_balancing_coeff += 1 / 50000 + g_balancing_coeff += 1 / 100000 - if step % 5000 == 0 and not args['debug']: + if step % 10000 == 0 and not args['debug']: ckpt_saver(model, step, optimizer, None) step += 1 @@ -593,7 +592,7 @@ if __name__ == '__main__': # optimizer = optim.Adam(model.parameters(), lr=config['lr']) optimizer = optim.Adam(model.parameters(), lr=config['lr']) _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3) - ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50) + ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=20) scaler = torch.cuda.amp.GradScaler(enabled=use_amp) resizer = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device']) diff --git a/scripts/generate_results.py b/scripts/generate_results.py index 30d7b90..ca00c17 100644 --- a/scripts/generate_results.py +++ b/scripts/generate_results.py @@ -14,7 +14,7 @@ parser.add_argument('-l', '--level', type=int, default=-1) args = vars(parser.parse_args()) if args['debug']: - from configs.iccvw.pyr_64_128_cape_f_2 import config_1 + from configs.iccvw.pyr_64_128_256_512_e2e_var import config_1 else: config_module = importlib.import_module(args['config']) config_1 = config_module.config diff --git a/scripts/generate_results_all.py b/scripts/generate_results_all.py index 853a3c6..1771bf1 100644 --- a/scripts/generate_results_all.py +++ b/scripts/generate_results_all.py @@ -10,6 +10,7 @@ parser.add_argument('--config', type=str) parser.add_argument('--iteration', type=str, default='') parser.add_argument('--debug', action='store_true') parser.add_argument('-g', '--gpu', type=str, default='1') +parser.add_argument('-fr', '--frames_file', type=str, default='') args = vars(parser.parse_args()) if args['debug']: @@ -117,7 +118,8 @@ if __name__ == '__main__': # remove all values from config['signature] that start with mesh_ config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')]) - dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape']) + frames_file = args['frames_file'] if args['frames_file'] != '' else config_1['test_frames'] + dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=frames_file, use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape']) # dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', 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, ) diff --git a/scripts/get_error.py b/scripts/get_error.py index 1a0bbfd..d111e74 100644 --- a/scripts/get_error.py +++ b/scripts/get_error.py @@ -8,13 +8,16 @@ parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) parser.add_argument('--config', type=str) parser.add_argument('--debug', action='store_true') -parser.add_argument('--iteration', type=str, default='190000') +parser.add_argument('--iteration', type=str, default='') parser.add_argument('-g', '--gpu', type=str) +parser.add_argument('-n', '--normal', action='store_true') +parser.add_argument('-e', '--edge', action='store_true') +parser.add_argument('-cr', '--collision_ratio', action='store_true') parser.add_argument('-a', '--all', action='store_true') args = vars(parser.parse_args()) if args['debug']: - from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1 + from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1 else: config_module = importlib.import_module(args['config']) config_1 = config_module.config @@ -32,7 +35,7 @@ from data.utils import Frame -def calculate(): +def calculate(type=''): all_levels = args['all'] if all_levels: errors = {0: dict(), 1: dict(), 2: dict(), 3: dict()} @@ -41,8 +44,7 @@ def calculate(): errors = dict() errors_mean = dict() - - with open(f'{results_folder}/errors/{name}.csv', 'r') as f: + with open(f'{results_folder}/errors/{name}{type}.csv', 'r') as f: lines = f.readlines() for line in lines: @@ -50,11 +52,11 @@ def calculate(): category = line[2] if all_levels: level = int(line[3][-1]) - if category not in errors: + if category not in errors[level].keys(): errors[level][category] = [] errors[level][category].append(float(line[4])) else: - if category not in errors: + if category not in errors.keys(): errors[category] = [] errors[category].append(float(line[3])) @@ -62,17 +64,18 @@ def calculate(): for i in range(4): for category in errors[i]: errors_mean[i][category] = round(np.mean(errors[i][category]), 5) - print(f'Level {i}: {errors_mean[i]}') + print(f'Level {i} error{type}: {errors_mean[i]}') mean = np.mean(list(errors_mean[i].values())) - print(f'Final error level {i}: {mean}') + print(f'Final error{type} level {i}: {mean}') + pass else: for category in errors: errors_mean[category] = round(np.mean(errors[category]), 5) print(errors_mean) - mean = np.mean(list(errors_mean.values())) - print(f'Final error: {mean}') + mean = round(np.mean(list(errors_mean.values())), 5) + print(f'Final error{type}: {mean}') @@ -98,5 +101,10 @@ if __name__ == '__main__': name = f'{name}_{args["iteration"]}' out_folder = f'{results_folder}/{name}' - + if args['normal']: + calculate('_normal') + if args['edge']: + calculate('_edge') + if args['collision_ratio']: + calculate('_collision_ratio') calculate() \ No newline at end of file diff --git a/scripts/get_error_hood.py b/scripts/get_error_hood.py index 9221939..bb5b375 100644 --- a/scripts/get_error_hood.py +++ b/scripts/get_error_hood.py @@ -2,30 +2,9 @@ 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') -parser.add_argument('--iteration', type=str, default='') -parser.add_argument('-g', '--gpu', type=str) -parser.add_argument('-n', '--normal', action='store_true') -parser.add_argument('-e', '--edge', action='store_true') -parser.add_argument('-cr', '--collision_ratio', action='store_true') -parser.add_argument('-a', '--all', action='store_true') -args = vars(parser.parse_args()) - -if args['debug']: - from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1 -else: - config_module = importlib.import_module(args['config']) - config_1 = config_module.config - - import time import argparse -import tqdm +from tqdm import tqdm import random import numpy as np import plotly.graph_objects as go @@ -35,50 +14,53 @@ from data.utils import Frame -def calculate(type=''): +def calculate(type='chamfer'): errors = dict() errors_mean = dict() - with open(f'{folder}/{hood_csv}', 'r') as f: - lines_hood = f.readlines() - with open(f'{folder}/{ours_csv}', 'r') as f: - lines_ours = f.readlines() + files_hood = sorted(os.listdir(f'{folder}/{run}_hood/{type}')) + errors_hood = [np.load(f'{folder}/{run}_hood/{type}/{file}') for file in tqdm(files_hood)] + files_ours = sorted(os.listdir(f'{folder}/hood_3_var_ours/{type}')) + errors_ours = [np.load(f'{folder}/hood_3_var_ours/{type}/{file}') for file in tqdm(files_ours)] - for line in lines_hood: - line = line.split(',') - category = line[2] + for error, frame in zip(errors_hood, files_hood): + frame_id, frame_nr, frame_garment = frame.split('_') + category = frame_garment if category not in errors: errors[category] = [] - errors[category].append(float(line[3])) + errors[category].append(error) for category in errors: errors_mean[category] = round(np.mean(errors[category]), 5) print(errors_mean) mean = round(np.mean(list(errors_mean.values())), 5) - print(f'Final HOOD error{type}: {mean}') + print(f'Final HOOD error_{type}: {mean}') - for line in lines_ours: - line = line.split(',') - category = line[2] + # errors = dict() + # errors_mean = dict() - if category not in errors: - errors[category] = [] - errors[category].append(float(line[3])) + # for error, frame in zip(errors_ours, files_ours): + # frame_id, frame_nr, frame_garment = frame.split('_') + # category = frame_garment - for category in errors: - errors_mean[category] = round(np.mean(errors[category]), 5) - print(errors_mean) + # if category not in errors: + # errors[category] = [] + # errors[category].append(error) - mean = round(np.mean(list(errors_mean.values())), 5) - print(f'Final OURS error{type}: {mean}') + # for category in errors: + # errors_mean[category] = round(np.mean(errors[category]), 5) + # print(errors_mean) + + # mean = round(np.mean(list(errors_mean.values())), 5) + # print(f'Final OURS error_{type}: {mean}') if __name__ == '__main__': folder = '/data-local/127_2/results/hood' - hood_csv = 'hood_2_hood.csv' - ours_csv = 'hood_2_ours.csv' + run = 'hood_3' - calculate() \ No newline at end of file + for type in ['chamfer', 'area', 'collision_ratio']: # ['chamfer', 'collision', 'area', 'collision_ratio']: + calculate(type) \ No newline at end of file diff --git a/scripts/plot_ablation.py b/scripts/plot_ablation.py index bb9fe27..0585c7f 100644 --- a/scripts/plot_ablation.py +++ b/scripts/plot_ablation.py @@ -25,7 +25,7 @@ import importlib from data.utils import Frame from data.dataset import read_frames_file -from utils.utils import save_stl +from utils.utils import save_stl, create_temp_frames from utils.checkpoints import CheckpointSaver, load_model from utils.postprocess import fix_collision from losses.losses import collision_loss, normal_loss, kl_loss @@ -70,25 +70,21 @@ def step_pyramid(model, batch, config): return mesh_pred_posed -def create_temp_frames(folder: str, frame: Frame): - with open(Path(folder) / 'frames_temp.txt', 'w') as f: - f.writelines([f'{frame}\n' for _ in range(4)]) - - if __name__ == '__main__': results_folder = Path('/data-local/127_2/results') - out_folder = Path('temp/ablation') + out_folder = Path('temp/supplementary') + os.makedirs(out_folder, exist_ok=True) smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device']) dataset_resolutions = [] dataset_resolutions = [64, 128, 256, 512] experiments = { - 'final': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) }, - 'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 195000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 200000) }, - 'skin': { 'nonskirt': ('pyr_64_128_256_512_skin', 225000), 'skirt': ('pyr_64_128_256_512_skin_skirt', 200000) }, - 'norm': { 'nonskirt': ('pyr_64_128_256_512_norm', 200000), 'skirt': ('pyr_64_128_256_512_norm_skirt', 200000) }, - 'reg': { 'nonskirt': ('pyr_64_128_256_512_reg', 200000), 'skirt': ('pyr_64_128_256_512_reg_skirt', 200000) }, + # 'inc': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) }, + 'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 250000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 200000) }, + # 'skin': { 'nonskirt': ('pyr_64_128_256_512_skin', 225000), 'skirt': ('pyr_64_128_256_512_skin_skirt', 200000) }, + # 'norm': { 'nonskirt': ('pyr_64_128_256_512_norm', 200000), 'skirt': ('pyr_64_128_256_512_norm_skirt', 200000) }, + # 'reg': { 'nonskirt': ('pyr_64_128_256_512_reg', 200000), 'skirt': ('pyr_64_128_256_512_reg_skirt', 200000) }, } configs = dict() @@ -111,8 +107,23 @@ if __name__ == '__main__': # frames = ['00054_1_Dress', '00005_1_Jumpsuit'] - frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt'] - frames = [Frame(frame_str=frame) for frame in frames] + # frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt'] + # frames = ['01852_70_Tshirt'] # , '03849_27_Dress', '00054_1_Dress' + # frames = [Frame(frame_str=frame) for frame in frames] + + frames_file = configs['e2e']['nonskirt']['train_frames'] + with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f: + frames = f.readlines() + frames_1 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames] + + frames_file = configs['e2e']['skirt']['train_frames'] + with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f: + frames = f.readlines() + frames_2 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames] + + frames = frames_1 + frames_2 + random.shuffle(frames) + body_faces = smpl.f for frame in frames: @@ -130,8 +141,8 @@ if __name__ == '__main__': body = batch['mesh_body_posed'][0].cpu().numpy() gt = batch['mesh_posed'][0].cpu().numpy() - save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces) - save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces) + # save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces) + # save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces) for experiment in experiments.keys(): @@ -139,20 +150,24 @@ if __name__ == '__main__': mesh_posed_fixed = fix_collision(pred, batch['mesh_body_posed'], batch['mesh_mask']).cpu().numpy() mesh_posed_fixed = mesh_posed_fixed[0, mask, :] - m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False) - tm.smoothing.filter_humphrey(m, iterations=2) - verts = m.vertices + try: + m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False) + tm.smoothing.filter_humphrey(m, iterations=2) + verts = m.vertices + except: + verts = mesh_posed_fixed + print('failed to apply smoothing') - save_stl(out_folder / f'{frame}_{experiment}_body.stl', body, body_faces) - save_stl(out_folder / f'{frame}_{experiment}_mesh.stl', verts, faces) + # save_stl(out_folder / f'{frame}_{experiment}_body.stl', body, body_faces) + # save_stl(out_folder / f'{frame}_{experiment}_mesh.stl', verts, faces) - # fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), - # go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1) - # ]) - # fig.update_layout(scene=dict(aspectmode='cube')) - # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), ) - # fig.update_layout(title_text=f'{frame} {experiment}') - # fig.show() + fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), + go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1) + ]) + fig.update_layout(scene=dict(aspectmode='cube')) + fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), ) + fig.update_layout(title_text=f'{frame} {experiment}') + fig.show() pass pass diff --git a/scripts/plot_ablation_pyramid.py b/scripts/plot_ablation_pyramid.py index 58168af..50704fb 100644 --- a/scripts/plot_ablation_pyramid.py +++ b/scripts/plot_ablation_pyramid.py @@ -24,7 +24,7 @@ from pathlib import Path from data.utils import Frame from data.dataset import read_frames_file -from utils.utils import save_stl +from utils.utils import save_stl, create_temp_frames from utils.checkpoints import CheckpointSaver, load_model from utils.postprocess import fix_collision from losses.losses import collision_loss, normal_loss, kl_loss @@ -38,7 +38,6 @@ from torch.utils.data import DataLoader import trimesh as tm from zipfile import ZipFile - from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1 from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_2 @@ -76,15 +75,13 @@ def step_pyramid(model, batch, config): return out -def create_temp_frames(folder: str, frame: Frame): - with open(Path(folder) / 'frames_temp.txt', 'w') as f: - f.writelines([f'{frame}\n' for _ in range(4)]) + if __name__ == '__main__': results_folder = Path('/data-local/127_2/results') out_folder = Path('temp/pyramid_ablation') - iteration = {'nonskirt': 190000, 'skirt': 200000 } + iteration = {'nonskirt': 250000, 'skirt': 200000 } smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device']) dataset_resolutions = [] @@ -99,8 +96,9 @@ if __name__ == '__main__': model['nonskirt'].eval() model['skirt'].eval() - frames = ['00054_1_Dress', '01070_13_Jumpsuit', '03849_27_Dress', '06057_152_Dress', + frames = ['01070_13_Jumpsuit', '00054_1_Dress', '03849_27_Dress', '06057_152_Dress', '03179_274_Dress', '04247_102_Tshirt', '01852_70_Tshirt'] + frames = ['00027_1_Dress'] frames_corrupt = ['00300_69_Jumpsuit', '07608_77_Jumpsuit', '06955_117_Dress'] frames = [Frame(frame_str=frame) for frame in frames] @@ -109,8 +107,8 @@ if __name__ == '__main__': config = { 'nonskirt': config_1, 'skirt': config_2 } for frame in tqdm(frames): - create_temp_frames(config_1['train_dir'], frame) - dataset = Dataset(config_1['train_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape']) + create_temp_frames(config_1['test_dir'], frame) + dataset = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape']) dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, ) picker = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'nonskirt' @@ -136,15 +134,15 @@ if __name__ == '__main__': tm.smoothing.filter_humphrey(m, iterations=2) verts = m.vertices - save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces) - save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces) + # save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces) + # save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces) - # fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), - # go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1) - # ]) - # fig.update_layout(scene=dict(aspectmode='cube')) - # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), ) - # fig.show() + fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), + go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1) + ]) + fig.update_layout(scene=dict(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 pass diff --git a/scripts/plot_cloth3d.py b/scripts/plot_cloth3d.py new file mode 100644 index 0000000..1a4c231 --- /dev/null +++ b/scripts/plot_cloth3d.py @@ -0,0 +1,148 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + + +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"] = '0' + + +import time +import argparse +import random +import tqdm +import numpy as np +import plotly.graph_objects as go +import plotly.express as ex +from utils.postprocess import fix_collision + + +import torch +from tqdm import tqdm +from data.utils import Frame +from data.cloth3d_io import readPC2Frame, readPC2, readOBJ +from utils.utils import pickle_load, save_stl +from scipy.spatial import cKDTree +from smpl.smpl_torch import SMPL +from scipy.spatial.transform import Rotation as R +import trimesh as tm +from plotly.subplots import make_subplots + + +smpl = SMPL(path='smpl', device='cpu') +def deorient(verts, pose, inverse=True): + assert pose.shape[1] == 72 + rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy() + rotated = np.zeros_like(verts) + for frame in range(verts.shape[0]): + r = R.from_matrix(rodr[frame, 0, :, :]) + rotated[frame] = r.apply(verts[frame], inverse=inverse) + + return rotated + + +def readBIN(fname): + if '.' not in fname: fname += '.bin' + elif '.' in fname and not fname.endswith('.bin'): + print("File name extension should be '.bin'") + return + + with open(fname, 'rb') as f: + F = np.frombuffer(f.read(), dtype=np.uint16).astype(np.int32) + return F.reshape((-1,3)) + + +def calculate(): + + garments = sorted(os.listdir('/data-local/127_2/cloth3d_processed_test/mesh_mask')) + garments = [g[:-4] for g in garments] + # random.shuffle(garments) + # garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)] + body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32) + + for garment in tqdm(garments): + # garment = garment_data[0] + + picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt' + id, garment_type = garment.split('_') + + if not os.path.exists(f'temp/cloth3d/{id}/outfit.pc16'): + continue + if not os.path.exists(os.path.join(cloth3d_folder, id, garment_type + '.pc16')): + continue + + pred_cloth3d = readPC2(f'temp/cloth3d/{id}/outfit.pc16', True)['V'] + pred_cloth3d_faces = readBIN(f'temp/cloth3d/{id}/outfit.bin') + gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'] + mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask'] + pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy') + gt = deorient(gt, pose, True) + # gt = deorient(gt, pose, True) + body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'] + fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz') + faces = fix_data['faces'] + + + i = np.random.randint(0, len(pred_cloth3d)) + # i = garment_data[1] + + try: + if calculate_ours: + pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3'] + except: + print('missed one') + continue + + mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy() + mesh_posed_fixed = mesh_posed_fixed[0, mask, :] + try: + m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False) + tm.smoothing.filter_humphrey(m, iterations=2) + pred_ours = m.vertices + except: + pred_ours = mesh_posed_fixed + print('failed to apply smoothing') + + # fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')]) + # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) + # 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', scene_aspectratio=dict(x=1, y=1, z=1)) + # fig.show() + + fig = go.Figure(data=[go.Mesh3d(x=pred_cloth3d[i][:,0], y=pred_cloth3d[i][:,1], z=pred_cloth3d[i][:,2], i=pred_cloth3d_faces[:,0], j=pred_cloth3d_faces[:,1], k=pred_cloth3d_faces[:,2], color='blue')]) + fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) + 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', scene_aspectratio=dict(x=1, y=1, z=1)) + fig.show() + + # fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')]) + # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) + # 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', scene_aspectratio=dict(x=1, y=1, z=1)) + # fig.show() + pass + + # save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces) + # save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces) + # save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces) + # save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces) + # save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces) + # save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces) + + + +if __name__ == '__main__': + + hood_folder = '/data-local/127_2/results/hood/all_2_105k' + ours_folder = dict() + ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000' + ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000' + cloth3d_folder = '/data-net/dataset/cloth3d/test' + + results_folder = '/data-local/127_2/results' + + calculate_ours = True + + + + calculate() \ No newline at end of file diff --git a/scripts/plot_hood.py b/scripts/plot_hood.py new file mode 100644 index 0000000..6597a81 --- /dev/null +++ b/scripts/plot_hood.py @@ -0,0 +1,131 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + + +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"] = '0' + + +import time +import argparse +import random +import tqdm +import numpy as np +import plotly.graph_objects as go +import plotly.express as ex +from utils.postprocess import fix_collision + + +import torch +from tqdm import tqdm +from data.utils import Frame +from data.cloth3d_io import readPC2Frame, readPC2, readOBJ +from utils.utils import pickle_load, save_stl +from scipy.spatial import cKDTree +from smpl.smpl_torch import SMPL +from scipy.spatial.transform import Rotation as R +import trimesh as tm +from plotly.subplots import make_subplots + + +smpl = SMPL(path='smpl', device='cpu') +def deorient(verts, pose, inverse=True): + assert pose.shape[1] == 72 + rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy() + rotated = np.zeros_like(verts) + for frame in range(verts.shape[0]): + r = R.from_matrix(rodr[frame, 0, :, :]) + rotated[frame] = r.apply(verts[frame], inverse=inverse) + + return rotated + + + +def calculate(): + + # garments = sorted(os.listdir(hood_folder)) + # garments = [g[:-4] for g in garments] + # random.shuffle(garments) + garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)] + body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32) + + for garment_data in tqdm(garments): + garment = garment_data[0] + picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt' + id, garment_type = garment.split('_') + + pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl') + pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces'] + gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:] + mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask'] + pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:] + gt = deorient(gt, pose, True) + body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:] + fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz') + faces = fix_data['faces'] + + + # i = np.random.randint(0, len(pred_hood)) + i = garment_data[1] + + try: + if calculate_ours: + pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3'] + except: + print('missed one') + continue + + mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy() + mesh_posed_fixed = mesh_posed_fixed[0, mask, :] + try: + m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False) + tm.smoothing.filter_humphrey(m, iterations=2) + pred_ours = m.vertices + except: + pred_ours = mesh_posed_fixed + print('failed to apply smoothing') + + # fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')]) + # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) + # 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', scene_aspectratio=dict(x=1, y=1, z=1)) + # fig.show() + + fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')]) + fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) + 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', scene_aspectratio=dict(x=1, y=1, z=1)) + fig.show() + + # fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')]) + # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) + # 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', scene_aspectratio=dict(x=1, y=1, z=1)) + # fig.show() + pass + + # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces) + # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces) + # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces) + # save_stl(f'temp/hood/ne2w/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces) + # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces) + # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces) + + + +if __name__ == '__main__': + + hood_folder = '/data-local/127_2/results/hood/all_2_105k' + ours_folder = dict() + ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000' + ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000' + cloth3d_folder = '/data-net/dataset/cloth3d/test' + + results_folder = '/data-local/127_2/results' + + calculate_ours = True + + + + calculate() \ No newline at end of file diff --git a/scripts/upscaler_train.py b/scripts/upscaler_train.py index f581ecd..9d7ab99 100644 --- a/scripts/upscaler_train.py +++ b/scripts/upscaler_train.py @@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" -os.environ["CUDA_VISIBLE_DEVICES"] = '2' +os.environ["CUDA_VISIBLE_DEVICES"] = '3' import time @@ -39,10 +39,10 @@ from utils.utils import InfiniteDataLoader as DataLoader from pytorch_msssim import MS_SSIM -low_res = 64 -is_skirt = False -is_cape = True -is_male = True +low_res = 256 +is_skirt = True +is_cape = False +is_male = False skirt_modifier = '_skirt' if is_skirt else '' cape_modifier = '_cape' if is_cape else '' gender_modifier = '' if not is_cape else ('_m' if is_male else '_f') @@ -244,6 +244,9 @@ def train(): step += 1 + if step > 100000: + break + if __name__ == '__main__': if not args['debug']: diff --git a/utils/postprocess.py b/utils/postprocess.py index 294e379..ec5f358 100644 --- a/utils/postprocess.py +++ b/utils/postprocess.py @@ -35,6 +35,8 @@ def fix_collision(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr=.004 F_fix_coll_local = F_fix_coll_orig else: raise ValueError("V must have 14475 or 3638 vertices") + if V.shape[0] == 1: + F_fix_coll_local = F_fix_coll_local[:1] mesh = Meshes(B, F_fix_coll_local) B_vn = torch.stack(mesh.verts_normals_list()) diff --git a/utils/utils.py b/utils/utils.py index 394c549..5bf1a15 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -1,9 +1,11 @@ import os import sys +import pickle import numpy as np from torch.utils.data import DataLoader from stl import mesh as mesh_stl - +from pathlib import Path +from data.utils import Frame class Accumulator(): @@ -59,4 +61,29 @@ def save_stl(file, mesh, faces): for j in range(3): obj.vectors[i][j] = mesh[f[j],:] - obj.save(file) \ No newline at end of file + obj.save(file) + + +def create_temp_frames(folder: str, frame: Frame): + with open(Path(folder) / 'frames_temp.txt', 'w') as f: + f.writelines([f'{frame}\n' for _ in range(4)]) + + +def pickle_load(file): + """ + Load a pickle file. + """ + with open(file, 'rb') as f: + loadout = pickle.load(f) + + return loadout + + +def pickle_dump(loadout, file): + """ + Dump a pickle file. Create the directory if it does not exist. + """ + os.makedirs(os.path.dirname(str(file)), exist_ok=True) + + with open(file, 'wb') as f: + pickle.dump(loadout, f) \ No newline at end of file