diff --git a/configs/abl_all.py b/configs/archive/abl_all.py similarity index 100% rename from configs/abl_all.py rename to configs/archive/abl_all.py diff --git a/configs/abl_all_2.py b/configs/archive/abl_all_2.py similarity index 100% rename from configs/abl_all_2.py rename to configs/archive/abl_all_2.py diff --git a/configs/abl_norm.py b/configs/archive/abl_norm.py similarity index 100% rename from configs/abl_norm.py rename to configs/archive/abl_norm.py diff --git a/configs/ae_body_1.py b/configs/archive/ae_body_1.py similarity index 100% rename from configs/ae_body_1.py rename to configs/archive/ae_body_1.py diff --git a/configs/ae_body_cape_f.py b/configs/archive/ae_body_cape_f.py similarity index 100% rename from configs/ae_body_cape_f.py rename to configs/archive/ae_body_cape_f.py diff --git a/configs/ae_body_cape_m.py b/configs/archive/ae_body_cape_m.py similarity index 89% rename from configs/ae_body_cape_m.py rename to configs/archive/ae_body_cape_m.py index b7b0a94..a896f40 100644 --- a/configs/ae_body_cape_m.py +++ b/configs/archive/ae_body_cape_m.py @@ -2,7 +2,7 @@ from utils.convnext_utils import CNBlockConfig config = dict() -config['gpu'] = '4' +config['gpu'] = '0' config['name'] = 'body_torch_cape_m' config['notes'] = ' ' config['tags'] = ['torch', 'body'] @@ -11,9 +11,9 @@ config['script'] = 'scripts/ae_body_train.py' config['pretrained_model'] = '' config['batch_size'] = 4 -config['train_dir'] = '/data-local/127/cape_processed/train/male' -config['valid_dir'] = '/data-local/127/cape_processed/test/male' -config['test_dir'] = '/data-local/127/cape_processed/test/male' +config['train_dir'] = '/data-local/128/cape_processed/train/male' +config['valid_dir'] = '/data-local/128/cape_processed/test/male' +config['test_dir'] = '/data-local/128/cape_processed/test/male' config['train_frames'] = None config['valid_frames'] = None diff --git a/configs/ae_curv_1.py b/configs/archive/ae_curv_1.py similarity index 100% rename from configs/ae_curv_1.py rename to configs/archive/ae_curv_1.py diff --git a/configs/ae_curv_cape_f.py b/configs/archive/ae_curv_cape_f.py similarity index 100% rename from configs/ae_curv_cape_f.py rename to configs/archive/ae_curv_cape_f.py diff --git a/configs/ae_curv_cape_m.py b/configs/archive/ae_curv_cape_m.py similarity index 100% rename from configs/ae_curv_cape_m.py rename to configs/archive/ae_curv_cape_m.py diff --git a/configs/ae_norm_1.py b/configs/archive/ae_norm_1.py similarity index 100% rename from configs/ae_norm_1.py rename to configs/archive/ae_norm_1.py diff --git a/configs/ae_static.py b/configs/archive/ae_static.py similarity index 100% rename from configs/ae_static.py rename to configs/archive/ae_static.py diff --git a/configs/ae_static_cape_f.py b/configs/archive/ae_static_cape_f.py similarity index 100% rename from configs/ae_static_cape_f.py rename to configs/archive/ae_static_cape_f.py diff --git a/configs/ae_static_cape_m.py b/configs/archive/ae_static_cape_m.py similarity index 100% rename from configs/ae_static_cape_m.py rename to configs/archive/ae_static_cape_m.py diff --git a/configs/ae_static_orig.py b/configs/archive/ae_static_orig.py similarity index 100% rename from configs/ae_static_orig.py rename to configs/archive/ae_static_orig.py diff --git a/configs/ae_static_orig_2.py b/configs/archive/ae_static_orig_2.py similarity index 100% rename from configs/ae_static_orig_2.py rename to configs/archive/ae_static_orig_2.py diff --git a/configs/cvae_norm_1.py b/configs/archive/cvae_norm_1.py similarity index 100% rename from configs/cvae_norm_1.py rename to configs/archive/cvae_norm_1.py diff --git a/configs/pyr_32.py b/configs/archive/pyr_32.py similarity index 100% rename from configs/pyr_32.py rename to configs/archive/pyr_32.py diff --git a/configs/pyr_64_cape_m.py b/configs/archive/pyr_64_cape_m.py similarity index 100% rename from configs/pyr_64_cape_m.py rename to configs/archive/pyr_64_cape_m.py diff --git a/configs/pyr_64_cape_m_3.py b/configs/archive/pyr_64_cape_m_3.py similarity index 100% rename from configs/pyr_64_cape_m_3.py rename to configs/archive/pyr_64_cape_m_3.py diff --git a/configs/iccvw/ae_body_1.py b/configs/iccvw/ae_body_1.py new file mode 100644 index 0000000..a494bff --- /dev/null +++ b/configs/iccvw/ae_body_1.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '2' +config['name'] = 'body_iccvw_1' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_body_1_10p.py b/configs/iccvw/ae_body_1_10p.py new file mode 100644 index 0000000..84d88e5 --- /dev/null +++ b/configs/iccvw/ae_body_1_10p.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'body_iccvw_1_10p' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt_10p.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train_10p.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_body_1_25p.py b/configs/iccvw/ae_body_1_25p.py new file mode 100644 index 0000000..9c285ca --- /dev/null +++ b/configs/iccvw/ae_body_1_25p.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'body_iccvw_1_25p' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt_25p.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train_25p.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_body_1_50p.py b/configs/iccvw/ae_body_1_50p.py new file mode 100644 index 0000000..c10ee2d --- /dev/null +++ b/configs/iccvw/ae_body_1_50p.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'body_iccvw_1_50p' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt_50p.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train_50p.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_body_1_skirt.py b/configs/iccvw/ae_body_1_skirt.py new file mode 100644 index 0000000..f383361 --- /dev/null +++ b/configs/iccvw/ae_body_1_skirt.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '1' +config['name'] = 'body_iccvw_1_skirt' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = True +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_body_cape_f.py b/configs/iccvw/ae_body_cape_f.py new file mode 100644 index 0000000..af816a5 --- /dev/null +++ b/configs/iccvw/ae_body_cape_f.py @@ -0,0 +1,73 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'body_iccvw_cape_f' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/128/cape_processed/train/female' +config['valid_dir'] = '/data-local/128/cape_processed/test/female' +config['test_dir'] = '/data-local/128/cape_processed/test/female' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_body_cape_m.py b/configs/iccvw/ae_body_cape_m.py new file mode 100644 index 0000000..a896f40 --- /dev/null +++ b/configs/iccvw/ae_body_cape_m.py @@ -0,0 +1,73 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'body_torch_cape_m' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/128/cape_processed/train/male' +config['valid_dir'] = '/data-local/128/cape_processed/test/male' +config['test_dir'] = '/data-local/128/cape_processed/test/male' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_norm_1.py b/configs/iccvw/ae_norm_1.py new file mode 100644 index 0000000..f4cdb16 --- /dev/null +++ b/configs/iccvw/ae_norm_1.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '1' +config['name'] = 'norm_iccvw_1' +config['notes'] = ' ' +config['tags'] = ['torch', 'norm'] +config['script'] = 'scripts/ae_norm_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_mask', + 'uv_normals', + 'mesh_mask', + 'mesh_normals', + } \ No newline at end of file diff --git a/configs/iccvw/ae_norm_1_skirt.py b/configs/iccvw/ae_norm_1_skirt.py new file mode 100644 index 0000000..cd44125 --- /dev/null +++ b/configs/iccvw/ae_norm_1_skirt.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'norm_iccvw_1_skirt' +config['notes'] = ' ' +config['tags'] = ['torch', 'norm'] +config['script'] = 'scripts/ae_norm_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = True +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_mask', + 'uv_normals', + 'mesh_mask', + 'mesh_normals', + } \ No newline at end of file diff --git a/configs/iccvw/ae_norm_cape_f.py b/configs/iccvw/ae_norm_cape_f.py new file mode 100644 index 0000000..8cca13c --- /dev/null +++ b/configs/iccvw/ae_norm_cape_f.py @@ -0,0 +1,73 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '6' +config['name'] = 'norm_iccvw_cape_f' +config['notes'] = ' ' +config['tags'] = ['torch', 'norm'] +config['script'] = 'scripts/ae_norm_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127_2/cape_processed/train/female' +config['valid_dir'] = '/data-local/127_2/cape_processed/test/female' +config['test_dir'] = '/data-local/127_2/cape_processed/test/female' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_mask', + 'uv_normals', + 'mesh_mask', + 'mesh_normals', + } \ No newline at end of file diff --git a/configs/iccvw/ae_norm_cape_m.py b/configs/iccvw/ae_norm_cape_m.py new file mode 100644 index 0000000..0deb82e --- /dev/null +++ b/configs/iccvw/ae_norm_cape_m.py @@ -0,0 +1,74 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '6' +config['name'] = 'norm_iccvw_cape_m' +config['notes'] = ' ' +config['tags'] = ['torch', 'norm'] +config['script'] = 'scripts/ae_norm_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127_2/cape_processed/train/male' +config['valid_dir'] = '/data-local/127_2/cape_processed/test/male' +config['test_dir'] = '/data-local/127_2/cape_processed/test/male' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_mask', + 'uv_normals', + 'mesh_mask', + 'mesh_normals', + } \ No newline at end of file diff --git a/configs/iccvw/ae_static_1.py b/configs/iccvw/ae_static_1.py new file mode 100644 index 0000000..a80de57 --- /dev/null +++ b/configs/iccvw/ae_static_1.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '3' +config['name'] = 'static_iccvw_4' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] +config['script'] = 'scripts/ae_static_train.py' + +config['pretrained_model'] = 'checkpoints/iccvw/static_iccvw_3' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0005 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/configs/iccvw/ae_static_1_skirt.py b/configs/iccvw/ae_static_1_skirt.py new file mode 100644 index 0000000..127c5b8 --- /dev/null +++ b/configs/iccvw/ae_static_1_skirt.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '3' +config['name'] = 'static_iccvw_1_skirt' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] +config['script'] = 'scripts/ae_static_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = True +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0005 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/configs/iccvw/ae_static_5.py b/configs/iccvw/ae_static_5.py new file mode 100644 index 0000000..24d1fd9 --- /dev/null +++ b/configs/iccvw/ae_static_5.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'static_iccvw_5' +config['notes'] = 'lower learning rate' +config['tags'] = ['torch', 'static'] +config['script'] = 'scripts/ae_static_train.py' + +config['pretrained_model'] = 'checkpoints/iccvw/static_iccvw_3' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.00005 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/configs/iccvw/ae_static_cape_f.py b/configs/iccvw/ae_static_cape_f.py new file mode 100644 index 0000000..c7a928d --- /dev/null +++ b/configs/iccvw/ae_static_cape_f.py @@ -0,0 +1,73 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'static_iccvw_cape_f' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] +config['script'] = 'scripts/ae_static_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127_2/cape_processed/train/female' +config['valid_dir'] = '/data-local/127_2/cape_processed/test/female' +config['test_dir'] = '/data-local/127_2/cape_processed/test/female' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0005 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/configs/iccvw/ae_static_cape_m.py b/configs/iccvw/ae_static_cape_m.py new file mode 100644 index 0000000..25e2662 --- /dev/null +++ b/configs/iccvw/ae_static_cape_m.py @@ -0,0 +1,73 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'static_iccvw_cape_m' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] +config['script'] = 'scripts/ae_static_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127_2/cape_processed/train/male' +config['valid_dir'] = '/data-local/127_2/cape_processed/test/male' +config['test_dir'] = '/data-local/127_2/cape_processed/test/male' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['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']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0005 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index d9b256d..58bb990 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -9,9 +9,6 @@ from plotly.subplots import make_subplots from torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode -# from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool -# from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template - def m2uv_array(m2uv, is_smpl_orig=False): result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32) @@ -74,35 +71,6 @@ class UV2Mesh(nn.Module): return x - -class UV2MeshLabel(): - def __init__(self, res=256, is_skirt=False, is_smpl_orig=False): - super(UV2MeshLabel, self).__init__() - self.res = res - uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy') - self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32) - m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item() - m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig) - self.uv_pixels = self.uv_map_discreet[m2uv] - - # input = channels_first - # @tf.function - def uv_to_mesh(self, uv): - batch_size = uv.shape[0] - - x1_ind = tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]) - x2_ind = tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]) - - x1 = tf.gather_nd(uv, x1_ind, batch_dims=1) - x2 = tf.gather_nd(uv, x2_ind, batch_dims=1) - - mesh = (x1 + x2) / 2 - return mesh - - def __call__(self, uv): - return self.uv_to_mesh(uv) - - def _maskFaces(V, F, mask, return_unfixed=False): V_mask = mask F_mask = np.take(V_mask, F.reshape(-1), axis=0) @@ -119,20 +87,6 @@ def _maskFaces(V, F, mask, return_unfixed=False): return V_new, F_new -# def _maskFaces_tf(V, F, mask): -# V_mask = mask -# F_mask = tf.gather(V_mask, F.reshape(-1), axis=0) -# F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3 -# F_mask = tf.reduce_sum(tf.reshape(F_mask, (-1, 3)), axis=1) == 3 - -# V_ind = tf.zeros(V_mask.shape) -# V_ind[V_mask] = tf.range(V_mask.sum()) -# F_new = tf.gather(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3)) -# V_new = V[V_mask, :] - -# return V_new, F_new - - def maskFacesBatch(V, F, mask): Vs = [] Fs = [] @@ -272,25 +226,6 @@ def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_c return fig -def batch_to_dict(batch, signature): - result = dict() - for i, key in enumerate(signature.keys()): - result[key] = batch[i] - return result - -conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body'] -def nest_conditions(batch): - batch['conditions'] = {d: batch.pop(d) for d in conditions} - return batch - - -def sparse_to_tensor(L): - idx = np.mat([L.row, L.col]).transpose().astype(np.float32) - data = L.data.astype(dtype=np.float32) - shape = L.shape - return tf.SparseTensor(idx, data, shape) - - def get_summary_str(model): # from torchview import draw_graph @@ -304,6 +239,7 @@ def get_summary_str(model): return [f'all_params: {total_params}, trainable: {trainable_params}'] + def get_summary_html(model): lines = get_summary_str(model) html = '' @@ -312,88 +248,6 @@ def get_summary_html(model): return html -def repeat_laplacian(L, times): - idx = np.repeat(np.mat([L.row, L.col]), repeats=times, axis=1) - idx_offset = np.linspace(0, L.shape[0] * (times-1), times).astype(np.int32) - idx_offset = np.repeat(idx_offset, L.row.shape[0], axis=0) - idx += idx_offset[np.newaxis,:] - data = np.repeat(L.data, repeats=times, axis=0) - shape = (L.shape[0] * times, L.shape[1] * times) - return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape) - - -def merge_verts(Vs, batchDim): - V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0) - return V - - -class GraphConvolution(nn.Module): - def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None): - super(GraphConvolution, self).__init__() - self.shape = shape - self.w = self.add_weight(name=name + '/weight', shape=(2, *shape), initializer=tf.initializers.glorot_normal, trainable=True) - self.b = self.add_weight(name=name + '/bias', shape=shape[-1], initializer=tf.initializers.zeros, trainable=True) - self.activation = activation or (lambda x: x) - self.normalization = normalization - self.batchSize = batchSize - self.indices = tf.cast(tf.linspace(0, self.batchSize * L.shape[1], self.batchSize+1), dtype=tf.int32) - self.L = repeat_laplacian(L, self.batchSize) - - def gather(self): - return [self.w, self.b] - - def _single(self, X): - X0 = tf.einsum('ab,bc->ac', X, self.w[0]) # Node - X1 = tf.einsum('ab,bc->ac', X, self.w[1]) # Neigh - # Neigh conv - X1 = edge_convolution_template( - X1, sparse_to_tensor(self.L), sizes=None, - edge_function=lambda x, y: y, - reduction='weighted', - edge_function_kwargs={} - ) - X = X0 + X1 + self.b - return X - - def call(self, X, training): - X = merge_verts(X, self.batchSize) - X = self._single(X) - X = tf.stack([X[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))]) - X = self.activation(X) - if self.normalization: - X = self.normalization(X, training=training) - return X - - -def replace_arrays(string): - out = string[:1] - i = 0 - while i < len(string) and string[i+1:].find('[') != -1: - start = string[i+1:].find('[') - end = string[i+1+start:].find(']') - out += string[i+1:i+1+start] - array = string[i+1+start:i+1+start+end+1].replace('\n', '') - out += array - i = i + 1 + start + end - - out += string[i+1:] - - return out - - -def group_inputs_for_levels(batch, config): - levels_input = [] - for res in config['pyramid']['resolutions']: - res_modifier = '_' + str(res) if res != 512 else '' - if config['use_skinning']: - uv_orig = batch['uv_unposed' + res_modifier] - else: - uv_orig = batch['uv_posed' + res_modifier] - levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']]) - conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']] - return levels_input, conds - - class Accumulator(): def __init__(self): self.dict = dict() diff --git a/data/convert_to_uv_cape.py b/data/convert_to_uv_cape.py new file mode 100644 index 0000000..8fe86ee --- /dev/null +++ b/data/convert_to_uv_cape.py @@ -0,0 +1,530 @@ +import os +import sys +import glob +import tqdm +import math +import time +import random +import shutil +import collections +import argparse +import pickle +import numpy as np +import scipy as sp +import trimesh as tm +# import cmapy as cm +import tables as tb +import scipy.io as sio +import scipy.interpolate as spi +import matplotlib.pyplot as plt +import scipy.stats as stat +import cv2 +import plotly.express as ex +import plotly.graph_objects as go + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') +from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures +from data.cloth3d_io import readPC2 + + +os.environ["CUDA_VISIBLE_DEVICES"] = "6" + +def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rbf', garment_uv_mask=None): + verts_uv = np.zeros((res * 2, res, 3)) + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = verts + uv_values[m2uv[:, 1]] += verts + uv_values = uv_values[uv_map_mask[type], :] # add garment mask + uv_values /= uv_counter[type][:, np.newaxis] + return_garment_uv_mask = False + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask[type]] # add garment mask + else: + garment_mask = np.sum(uv_values, axis=-1) != 0 + + # test + verts_uv[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values + # verts_uv_save = verts_uv.copy() + + # garment_mask = np.sum(uv_values, axis=-1) != 0 + unique_uv_values = uv_values[garment_mask, :] + uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :] + rbf_dest_current = rbf_dest[type] + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (4 * 2 + 1, 4 * 2 + 1), (4, 4)), iterations=4) + garment_uv_mask_border = garment_uv_mask.astype(bool) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask + garment_uv_mask = garment_uv_mask * uv_mask[type][:,:,0] + rbf_dest_current = rbf_dest[type] * garment_uv_mask + return_garment_uv_mask = True + + + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T + if method == 'rbf': + interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values, neighbors=neighbours, kernel='linear') + elif method == 'nearest': + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values) + else: + raise NotImplementedError + + final = interpolator(rbf_dest_current) + + result = np.zeros((res * 2, res, 3)) + + # not sure, shouldn't be needed + # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]] + + result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final + result[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values + if garment_uv_mask is not None: + result = result * garment_uv_mask[:,:,np.newaxis] + + if return_garment_uv_mask: + return result, garment_uv_mask, garment_uv_mask_border + else: + return result + + +def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None): + batch = verts.shape[0] + verts_uv = np.zeros((batch, res * 2, res, 3)) + uv_values = np.zeros((batch, uv_map.shape[0], 3)) + uv_values[:, m2uv[:, 0]] = verts + uv_values[:, m2uv[:, 1]] += verts + uv_values = uv_values[:, uv_map_mask['non-skirt'], :] + uv_values /= uv_counter['non-skirt'][np.newaxis, :, np.newaxis] + return_garment_uv_mask = False + # test + verts_uv[:, uv_map_discreet_unique['non-skirt'][:, 0], uv_map_discreet_unique['non-skirt'][:, 1], :] = uv_values + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask['non-skirt']] # add garment mask + else: + garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0 + + unique_uv_values = uv_values[:, garment_mask, :] + uv_map_discreet_unique_unique = uv_map_discreet_unique['non-skirt'][garment_mask, :] + rbf_dest_current = rbf_dest['non-skirt'] + + # add back mesh vertices that are not part of the garment + uv_map_discreet_unique_unique_inv = uv_map_discreet_unique['non-skirt'][np.invert(garment_mask), :] + rbf_dest_fill = np.zeros((res * 2, res), dtype=bool) + rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + return_garment_uv_mask = True + + if garment_uv_mask is not None: + rbf_dest_current = (rbf_dest['non-skirt'] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border + + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + final = [] + for i in range(batch): + print('processing frame: ' + str(i)) + if method == 'rbf': + interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[i,:,:], neighbors=neighbours, kernel='linear') + elif method == 'nearest': + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:,:]) + else: + raise NotImplementedError + + result = np.zeros((res * 2, res, 3), np.float32) + result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current) + result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[i,:,:] + + if garment_uv_mask is not None: + result *= garment_uv_mask[:,:,np.newaxis] + + final.append(result) + + if return_garment_uv_mask: + return final, garment_uv_mask + else: + return final + + +def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None, id=None): + batch = len(verts) + + assert method == 'rbf' + + assert mask is not None + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask[type]] # add garment mask + + rbf_dest_current = rbf_dest[type] + + # add back mesh vertices that are not part of the garment + uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[type][np.invert(garment_mask), :] + rbf_dest_fill = np.zeros((res * 2, res), dtype=bool) + rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True + + assert not (garment_uv_mask is None and mask is not None) + assert garment_uv_mask is not None + rbf_dest_current = (rbf_dest[type] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border + + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + out = [] + + for i in range(batch): + print('processing frame: ' + str(i)) + + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = verts[i] + uv_values[m2uv[:, 1]] += verts[i] + uv_values = uv_values[uv_map_mask[type], :] + uv_values /= uv_counter[type][:, np.newaxis] + + unique_uv_values = uv_values[garment_mask, :] + uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :] + + interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[:,:], neighbors=neighbours, kernel='linear') + result = np.zeros((res * 2, res, 3), np.float32) + result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current) + result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[:,:] + + if garment_uv_mask is not None: + result *= garment_uv_mask[:,:,np.newaxis] + + out.append(result) + + return out + + +def convert_to_uv_batch_label(res, m2uv, uv_map, values, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None): + batch = values.shape[0] + values = (values + 1) * mask + values_uv = np.zeros((batch, res * 2, res)) + missing = np.zeros((res * 2, res)) + uv_values = np.zeros((batch, uv_map.shape[0])) + uv_values[:, m2uv[:, 0]] = values + uv_values[:, m2uv[:, 1]] = values + uv_values = uv_values[:, uv_map_mask] + return_garment_uv_mask = False + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + else: + raise Exception + + + # test + values_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values + missing[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = (uv_values[0,:] == 0) + + unique_uv_values = uv_values[:, garment_mask] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + # rbf_dest_current = rbf_dest.astype(bool) + garment_uv_mask_border * uv_mask[:,:,0] + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + return_garment_uv_mask = True + + rbf_dest_current = rbf_dest * garment_uv_mask + missing * garment_uv_mask + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T + final = [] + for i in range(batch): + print('processing frame: ' + str(i)) + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:]) + final.append(interpolator(rbf_dest_current)) + final = np.asarray(final) + + result = np.zeros((batch, res * 2, res)) + result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values + result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1]] = final + if garment_uv_mask is not None: + result = result * garment_uv_mask[np.newaxis,:,:] + + if return_garment_uv_mask: + return result, garment_uv_mask + else: + return result + + +def m2uv_array(m2uv): + if m2uv is None: + return None + result = np.zeros((3638, 2), dtype=np.int32) + for k, v in m2uv.items(): + if len(v) == 2: + result[k] = list(v) + elif len(v) == 1: + x = v.pop() + result[k] = [x, x] + else: + assert False + + return result + + +def uv_to_mesh(uv, m2uv): + res = uv.shape[0] + mesh = np.zeros((14475, 3)) + uv_pixels = uv_map_discreet[m2uv] + # counter = np.ones((14475)) + # counter[m2uv[:, 0] == m2uv[:, 1]] = 2 + mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis] + return mesh + pass + + +def verify_reconstruction(orig, reconstructed): + error = np.sqrt(np.sum((orig - reconstructed) ** 2, axis=-1)) + colors = np.array([cm.color('jet', 1 - i) for i in error]) + error_sum = np.sum(error) + + + verts_uv = np.zeros((uv_res * 2, uv_res, 3)) + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = colors + uv_values[m2uv[:, 1]] += colors + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_values /= uv_counter[:, np.newaxis] + verts_uv[uv_map_discreet[:, 0], uv_map_discreet[:, 1], :] = uv_values + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(111) + ax.scatter(uv_map[:, 0], uv_map[:, 1], s=5, c=uv_values/255) + plt.show() + + + + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(121, projection='3d') + x = orig + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5) + + ax = fig.add_subplot(122, projection='3d') + x = reconstructed#[error > 0.1, :] + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5, c=colors/255) #[error > 0.1, :] + + plt.show() + + pass + +# https://stackoverflow.com/questions/38187286/find-unique-pairs-in-list-of-pairs +def unique_pairs(a): + a.sort(axis=1) + b = np.ascontiguousarray(a).view( + np.dtype((np.void, a.dtype.itemsize * a.shape[1])) + ) + _, inv, ct = np.unique(b, return_inverse=True, return_counts=True) + return ct[inv] == 1 + + +def unique_uv_map_mask(uv_map_discreet): + mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool) + temp = np.zeros((uv_res * 2, uv_res)) + for p in uv_map_discreet: + temp[p[0], p[1]] += 1 + # temp = temp == 1 + for i, p in enumerate(uv_map_discreet): + if temp[p[0], p[1]] == 1: + mask[i] = True + elif temp[p[0], p[1]] == 2: + mask[i] = True + temp[p[0], p[1]] = 0 + return mask + + +def test(verts, faces, values): + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = values + uv_values[m2uv[:, 1]] += values + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_values /= uv_counter[:, np.newaxis] + x = uv_values + verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1) + im = ((x - x.min()) / (x.max() - x.min())) * 255 + tm.Trimesh(verts, faces, vertex_colors=im.astype(np.int32)).show() + + +def get_shapedirs(): + to_keep = sio.loadmat(os.path.join('smpl/keep_remove.mat'))['keep'].reshape(-1) + shapedirs = [] + with open(os.path.join('smpl/super_model_f.pkl'), 'rb') as f: + female = pickle.load(f, encoding="latin1") + shapedirs.append(female['shapedirs'][to_keep, :]) + with open(os.path.join('smpl/super_model_m.pkl'), 'rb') as f: + male = pickle.load(f, encoding="latin1") + shapedirs.append(male['shapedirs'][to_keep, :]) + return 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("-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("-d", "--debug", help = "end of range", default=False, action='store_true') + args = parser.parse_args() + + F = {'non-skirt': np.load("smpl/smpl_faces_orig.npy"), 'skirt': None} + uv_map = {'non-skirt': np.load('custom_utils/stretched_uv_map_3_orig.npy'), 'skirt': None} + m2uv = {'non-skirt': np.load('custom_utils/m2uv_orig.npy', allow_pickle=True).item(), 'skirt': None} + m2uv = {k: m2uv_array(v) for k, v in m2uv.items()} + uv_mask = {'non-skirt': np.load('custom_utils/mask_256x512.npy'), 'skirt': np.load('custom_utils/mask_256x512_skirt.npy')} + + + uv_res = 256 + uv_map_discreet = dict() + uv_map_discreet_unique = dict() + uv_counter = dict() + rbf_dest = dict() + uv_map_mask = dict() + for k,v in uv_mask.items(): + if k == 'skirt': + continue + uv_mask[k] = (cv2.resize(v * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis] + # uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis] + # uv_mask_border = uv_mask * np.invert(uv_mask_dilate) + uv_mask[k] = cv2.dilate(v.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis] + + uv_map_discreet[k] = np.floor(uv_map[k] * uv_res).astype(np.int32) + uv_map_mask[k] = unique_uv_map_mask(uv_map_discreet[k]) + + uv_counter[k] = np.zeros((uv_map[k].shape[0])) + uv_counter[k][m2uv[k][:, 0]] = 1 + uv_counter[k][m2uv[k][:, 1]] += 1 + uv_counter[k] = uv_counter[k][uv_map_mask[k]] + + uv_map_discreet_unique[k] = uv_map_discreet[k][uv_map_mask[k], :] + + rbf_dest[k] = np.ones((uv_res * 2, uv_res)) + rbf_dest[k][uv_map_discreet_unique[k][:, 0], uv_map_discreet_unique[k][:, 1]] = 0 + rbf_dest[k][np.invert(uv_mask[k][:,:,0])] = 0 + + + ######################################################################################### + + 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', + + frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))#[args.start:args.end] + frames = [[frame.split('.')[0], frame.split('.')[1]] for frame in frames] + ids = list(set([frame[0] for frame in frames])) + frames_by_id = {id: [frame[1] for frame in frames if frame[0] == id] for id in ids} + + shapedirs = get_shapedirs() + + root = '/'.join(folder.split('/')[:-1]) + type = 'non-skirt' + + 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) + + mask = np.ones((3638), dtype=bool) + + for id in tqdm.tqdm(ids): + + + # 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_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'] + + start = time.perf_counter() + + uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv[type], uv_map[type], mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None) + uv_static = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_static[None, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border) + 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) + + print(time.perf_counter() - start) + + # ex.imshow(uv_mask_local).show() + # ex.imshow((uv_static+1)/2).show() + # ex.imshow((uv_posed[0]+1)/2).show() + # ex.imshow((uv_unposed[0]+1)/2).show() + # ex.imshow((uv_normals[0]+1.5)/3).show() + # ex.imshow((uv_body_posed[0]+1)/2).show() + + # pass + # np.savez_compressed(os.path.join(folder, 'uv_static', f'{id}_{garment}.npz'), uv_static=uv_static) + # 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]) + + # pass diff --git a/data/dataset.py b/data/dataset.py index 99aa5e8..7ebc6d5 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -1,7 +1,7 @@ import os import sys -# os.environ["CUDA_VISIBLE_DEVICES"] = '6' +# os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3' import math import time @@ -14,13 +14,14 @@ import scipy.io as sio import plotly.express as ex import torchvision -from torch.utils.data import Dataset from torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from data.utils import Frame +from data.cloth3d_io import readPC2, readPC2Frame +from data.dataset_cape import DatasetCape def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None): @@ -60,7 +61,6 @@ class StaticCache(): self.folder = folder assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) self.resolutions = resolutions - self.filled = '' self.is_cape = is_cape self.device = device @@ -74,50 +74,65 @@ class StaticCache(): self.uv_masks[res] = down.cpu().numpy() - def read(self, id): + def read(self, id, garment): data = dict() + mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', f'{id}_{garment}.npy')); mesh_mask = mesh_mask.astype(bool) + data['mesh_mask'] = mesh_mask + for res in self.resolutions: res_string = '' if res != 512: res_string = '_' + str(res) if 'uv_static' + res_string in self.signature: - uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static + uv_static = np.load(os.path.join(self.folder, f'uv_static{res_string}', f'{id}_{garment}.npz')); uv_static = uv_static['uv_static']; data['uv_static' + res_string] = uv_static if 'uv_mask' + res_string in self.signature: if self.is_cape: data['uv_mask' + res_string] = self.uv_masks[res] else: - uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask + uv_mask = np.load(os.path.join(self.folder, f'uv_mask{res_string}', f'{id}_{garment}.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask if 'mesh_static' in self.signature: - mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32) - if 'mesh_mask' in self.signature: - if self.is_cape: - data['mesh_mask'] = np.ones((3638), np.float32) - else: - mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask + mesh_static = np.load(os.path.join(self.folder, 'mesh_static', f'{id}_{garment}.npy')) + data['mesh_static'] = np.zeros((14475, 3), dtype=np.float32) + data['mesh_static'][mesh_mask] = mesh_static.astype(np.float32) + # if 'uv_shaped_body' in self.signature: # uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body return data - def __getitem__(self, id) -> dict: - if id not in self.store: - item = self.read(id) + def __getitem__(self, frame: Frame) -> dict: + if (frame.id, frame.garment) not in self.store: + item = self.read(frame.id, frame.garment) stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()} for key, value in stored_data.items(): if key.startswith('uv_') and value.shape[-1] == 3: stored_data[key] = torch.permute(value, [2,0,1]) - self.store[id] = stored_data + self.store[(frame.id, frame.garment)] = stored_data - data = { key: value.to(self.device) for key, value in self.store[id].items()} + data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items()} return data -class Dataset(Dataset): - def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False): + +def read_frames_file(frames_file, shuffle=False): + frames = [] + with open(frames_file, 'r') as f: + lines = f.readlines() + lines = [line.strip() for line in lines] + frames = [Frame(line.split()[0], int(line.split()[1]), line.split()[2]) for line in lines] + + if shuffle: + random.shuffle(frames) + + return frames + + +class DatasetCloth3D(torch.utils.data.Dataset): + def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False): self.folder = folder self.device = device self.is_cape = is_cape @@ -125,13 +140,20 @@ class Dataset(Dataset): self.resolutions = resolutions self.transforms = tv.transforms.ToTensor() - self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True) + self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=True) - self.body_faces = np.load('smpl/smpl_faces.npy') self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device) self.cache = dict() if use_cache else None + ids = set([frame.id for frame in self.frames]) + self.pose_cache = dict() + for id in ids: + self.pose_cache[id] = np.load(os.path.join(self.folder, 'pose', f'{id}.npy')) + self.shape_cache = dict() + for id in ids: + self.shape_cache[id] = np.load(os.path.join(self.folder, 'shape', f'{id}.npy')) + def __len__(self): return len(self.frames) @@ -142,76 +164,59 @@ class Dataset(Dataset): if self.cache is not None: if frame in self.cache.keys(): data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()} - data = { **data, **self.static_cache[frame.id] } + data = { **data, **self.static_cache[frame] } return data - separator = '_' - if self.is_cape: - separator = '.' - - uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2' - if self.is_cape: - uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient' - item = dict() + mesh_mask = self.static_cache[frame]['mesh_mask'].cpu().numpy().astype(bool) + + item['id'] = frame.id + item['frame_nr'] = frame.frame_nr + item['garment'] = frame.garment - if 'id' in self.signature: - item['id'] = frame.id - if 'frame_nr' in self.signature: - item['frame_nr'] = frame.frame_nr if 'pose_shape' in self.signature: - pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32) - pose_shape[:3] = 0 - item['pose_shape'] = pose_shape + pose = self.pose_cache[frame.id][frame.frame_nr] + pose[:3] = 0 + item['pose_shape'] = np.concatenate((pose, self.shape_cache[frame.id]), axis=0) if 'mesh_body_posed' in self.signature: - mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] - item['mesh_body_posed'] = mesh_posed_body_no_orient + mesh_body_posed_no_orient = readPC2Frame(os.path.join(self.folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True) + item['mesh_body_posed'] = mesh_body_posed_no_orient if 'mesh_posed' in self.signature: - mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32) - item['mesh_posed'] = mesh_posed + mesh_posed = readPC2Frame(os.path.join(self.folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + item['mesh_posed'] = np.zeros((14475, 3), dtype=np.float32) + item['mesh_posed'][mesh_mask] = mesh_posed if 'mesh_unposed' in self.signature: - mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] - item['mesh_unposed'] = mesh_unposed_no_shape - if 'mesh_curvature' in self.signature: - mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] - item['mesh_curvature'] = mesh_curvature.astype(np.int32) + mesh_unposed_no_shape = readPC2Frame(os.path.join(self.folder, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + item['mesh_unposed'] = np.zeros((14475, 3), dtype=np.float32) + item['mesh_unposed'][mesh_mask] = mesh_unposed_no_shape if 'mesh_normals' in self.signature: - mesh_normals = np.load(os.path.join(self.folder, 'mesh_normals', frame.id + separator + frame.frame_nr + '.npz')); mesh_normals = mesh_normals['mesh_normals'] - item['mesh_normals'] = mesh_normals + mesh_normals = readPC2Frame(os.path.join(self.folder, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + item['mesh_normals'] = np.zeros((14475, 3), dtype=np.float32) + item['mesh_normals'][mesh_mask] = mesh_normals for res in self.resolutions: res_string = '' if res != 512: res_string = '_' + str(res) - if 'uv_curvature' + res_string in self.signature: - uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz')); - if res_string != '': - uv_curvature = uv_curvature['uv_curvature'] - else: - uv_curvature = uv_curvature['curvature_uv'] - uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device) - item['uv_curvature' + res_string] = uv_curvature + if 'uv_normals' + res_string in self.signature: + uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_normals = uv_normals['uv_normals'] + item['uv_normals' + res_string] = uv_normals if 'uv_unposed' + res_string in self.signature: - uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape'] + uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed'] item['uv_unposed' + res_string] = uv_unposed_no_shape if 'uv_body_posed' + res_string in self.signature: - uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz')) - uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient'] - item['uv_body_posed' + res_string] = uv_posed_body_no_orient + uv_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed'] + item['uv_body_posed' + res_string] = uv_body_posed_no_orient if 'uv_posed' + res_string in self.signature: - uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient'] + uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_posed = uv_posed['uv_posed'] item['uv_posed' + res_string] = uv_posed - if 'uv_normals' + res_string in self.signature: - uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals'] - item['uv_normals' + res_string] = uv_normals - - data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()} + data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor or type(value) == int) else value) for key, value in item.items()} for key, value in data.items(): if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7): @@ -221,38 +226,101 @@ class Dataset(Dataset): self.cache[frame] = data data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()} - data = { **data, **self.static_cache[frame.id] } + data = { **data, **self.static_cache[frame] } return data +def Dataset(folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False): + if not is_cape: + return DatasetCloth3D(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=False) + else: + return DatasetCape(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=True) + + + if __name__ == '__main__': from torch.utils.data import DataLoader - + import plotly.graph_objects as go + import plotly.express as ex + from plotly.subplots import make_subplots + from custom_utils.custom_utils import _maskFaces + + F = np.load('smpl/smpl_faces.npy') + F_skirt = np.load('smpl/smpl_faces_skirt.npy') + device = "cuda" if torch.cuda.is_available() else "cpu" - signature = { 'id', - 'frame_nr', - 'pose_shape', + signature = { 'pose_shape', 'uv_posed', 'uv_body_posed', 'mesh_body_posed', 'mesh_posed', 'mesh_unposed', 'uv_unposed', - 'uv_curvature', - 'mesh_curvature', + 'uv_normals', + 'mesh_normals', 'uv_static', 'mesh_static', 'mesh_mask', 'uv_mask', - 'uv_shaped_body', } - dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') + dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt') train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) for batch in train_dataloader: batch + b = 0 + mesh_mask = batch['mesh_mask'][b].cpu().numpy() + faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F + _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask) + + mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask] + mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask] + mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask] + mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy() + + fig = make_subplots(rows=2, cols=2, + specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}], + [{'type': 'mesh3d'}, {'type': 'mesh3d'}]]) + + fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1) + fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2) + fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=1) + fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], opacity=1), row=2, col=2) + + fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False), + yaxis=dict(range=[-1, 1], autorange=False), + zaxis=dict(range=[-1, 1], autorange=False), + aspectmode='cube'), + scene2=dict(xaxis=dict(range=[-1, 1], autorange=False), + yaxis=dict(range=[-1, 1], autorange=False), + zaxis=dict(range=[-1, 1], autorange=False), + aspectmode='cube'), + scene3=dict(xaxis=dict(range=[-1, 1], autorange=False), + yaxis=dict(range=[-1, 1], autorange=False), + zaxis=dict(range=[-1, 1], autorange=False), + aspectmode='cube'), + scene4=dict(xaxis=dict(range=[-1, 1], autorange=False), + yaxis=dict(range=[-1, 1], autorange=False), + zaxis=dict(range=[-1, 1], autorange=False), + aspectmode='cube')) + + fig.show() + + + uv_mask = batch['uv_mask'][b].cpu().numpy() + uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0) + uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2 + uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2 + uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2 + uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2 + uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3 + + img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2) + img = np.transpose(img, (1, 2, 0)) + ex.imshow(img).show() + pass pass @@ -281,7 +349,7 @@ if __name__ == '__main__': # for i, batch in enumerate(dataset): # print(i) # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) -# ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show() +# ex.imshow(batch['uv_normals'][0].numpy().argmax(-1) / 6.0).show() # sum_uv += batch['uv_unposed'].numpy().sum(axis=0) # sum_mask += batch['uv_mask'].numpy().sum(axis=0) diff --git a/data/dataset_cape.py b/data/dataset_cape.py new file mode 100644 index 0000000..4c4c370 --- /dev/null +++ b/data/dataset_cape.py @@ -0,0 +1,306 @@ +import os +import sys + +# os.environ["CUDA_VISIBLE_DEVICES"] = '6' + +import math +import time +import torch +import random +import pytorch3d +import torchvision as tv +import numpy as np +import scipy.io as sio +import plotly.express as ex +import torchvision + +from torchvision.transforms import Resize +from torchvision.transforms.functional import InterpolationMode + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +from data.utils import Frame + + +def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None): + if frames_file is None: + uv_static_folder = os.path.join(folder, 'pose_shape') + frames = sorted(os.listdir(uv_static_folder)) + + if is_cape: + frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0] + frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames] + else: + frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] + frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames] + + # if testing: + # frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) + # frames = [frame.split('.')[0] + '_4' for frame in frames] + + + if frames_file is not None: + frames = [] + with open(os.path.join(folder, frames_file), 'r') as f: + lines = f.readlines() + lines = [line.strip() for line in lines] + frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines] + + if shuffle: + random.shuffle(frames) + + return frames + + +class StaticCache(): + def __init__(self, folder, signature, resolutions=[512], is_cape=False, device='cuda') -> None: + self.store = dict() + self.signature = signature + self.folder = folder + assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) + self.resolutions = resolutions + self.filled = '' + self.is_cape = is_cape + self.device = device + + if is_cape: + self.uv_masks = dict() + uv_mask = torch.tensor(np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32), device=device) + self.uv_masks[512] = uv_mask.to(torch.float32).cpu().numpy() + for res in [32, 64, 128, 256]: + down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT)(uv_mask[None,:,:]) + down = torch.squeeze(down).to(torch.float32) + self.uv_masks[res] = down.cpu().numpy() + + + def read(self, id): + data = dict() + + for res in self.resolutions: + res_string = '' + if res != 512: + res_string = '_' + str(res) + + if 'uv_static' + res_string in self.signature: + uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static + if 'uv_mask' + res_string in self.signature: + if self.is_cape: + data['uv_mask' + res_string] = self.uv_masks[res] + else: + uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask + + if 'mesh_static' in self.signature: + mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32) + if 'mesh_mask' in self.signature: + if self.is_cape: + data['mesh_mask'] = np.ones((3638), np.float32) + else: + mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask + # if 'uv_shaped_body' in self.signature: + # uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body + + return data + + + def __getitem__(self, id) -> dict: + if id not in self.store: + item = self.read(id) + stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()} + for key, value in stored_data.items(): + if key.startswith('uv_') and value.shape[-1] == 3: + stored_data[key] = torch.permute(value, [2,0,1]) + self.store[id] = stored_data + + data = { key: value.to(self.device) for key, value in self.store[id].items()} + return data + + +class DatasetCape(torch.utils.data.Dataset): + def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False): + self.folder = folder + self.device = device + self.is_cape = is_cape + self.signature = signature + self.resolutions = resolutions + + self.transforms = tv.transforms.ToTensor() + self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True) + + self.body_faces = np.load('smpl/smpl_faces.npy') + self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device) + + self.cache = dict() if use_cache else None + + + def __len__(self): + return len(self.frames) + + + def __getitem__(self, idx): + frame : Frame = self.frames[idx] + if self.cache is not None: + if frame in self.cache.keys(): + data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()} + data = { **data, **self.static_cache[frame.id] } + return data + + separator = '_' + if self.is_cape: + separator = '.' + + uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2' + if self.is_cape: + uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient' + + item = dict() + + if 'id' in self.signature: + item['id'] = frame.id + if 'frame_nr' in self.signature: + item['frame_nr'] = frame.frame_nr + if 'pose_shape' in self.signature: + pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32) + pose_shape[:3] = 0 + item['pose_shape'] = pose_shape + if 'mesh_body_posed' in self.signature: + mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] + item['mesh_body_posed'] = mesh_posed_body_no_orient + if 'mesh_posed' in self.signature: + mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32) + item['mesh_posed'] = mesh_posed + if 'mesh_unposed' in self.signature: + mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] + item['mesh_unposed'] = mesh_unposed_no_shape + if 'mesh_curvature' in self.signature: + mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] + item['mesh_curvature'] = mesh_curvature.astype(np.int32) + if 'mesh_normals' in self.signature: + mesh_normals = np.load(os.path.join(self.folder, 'mesh_normals', frame.id + separator + frame.frame_nr + '.npz')); mesh_normals = mesh_normals['mesh_normals'] + item['mesh_normals'] = mesh_normals + + for res in self.resolutions: + res_string = '' + if res != 512: + res_string = '_' + str(res) + + if 'uv_curvature' + res_string in self.signature: + uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz')); + if res_string != '': + uv_curvature = uv_curvature['uv_curvature'] + else: + uv_curvature = uv_curvature['curvature_uv'] + uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device) + item['uv_curvature' + res_string] = uv_curvature + + if 'uv_unposed' + res_string in self.signature: + uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape'] + item['uv_unposed' + res_string] = uv_unposed_no_shape + + if 'uv_body_posed' + res_string in self.signature: + uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz')) + uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient'] + item['uv_body_posed' + res_string] = uv_posed_body_no_orient + + if 'uv_posed' + res_string in self.signature: + uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient'] + item['uv_posed' + res_string] = uv_posed + + if 'uv_normals' + res_string in self.signature: + uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals'] + item['uv_normals' + res_string] = uv_normals + + + data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()} + + for key, value in data.items(): + if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7): + data[key] = torch.permute(value, [2,0,1]) + + if self.cache is not None: + self.cache[frame] = data + + data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()} + data = { **data, **self.static_cache[frame.id] } + return data + + +if __name__ == '__main__': + from torch.utils.data import DataLoader + + device = "cuda" if torch.cuda.is_available() else "cpu" + + signature = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_posed', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'mesh_unposed', + 'uv_unposed', + 'uv_curvature', + 'mesh_curvature', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + 'uv_shaped_body', + } + + dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') + train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) + + for batch in train_dataloader: + batch + pass + + pass + + + + + + + + + + +# if __name__ == '__main__': +# os.environ["CUDA_VISIBLE_DEVICES"] = "6" +# from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces +# import plotly.graph_objects as go +# import plotly.express as ex +# # import cmapy as cm +# uv2mesh = UV2Mesh(is_skirt=True) + +# F = np.load("custom_utils/smpl_faces.npy") + +# sum_uv = np.zeros((512, 256, 3), dtype=np.float32) +# sum_mask = np.zeros((512, 256), dtype=np.float32) +# for i, batch in enumerate(dataset): +# print(i) +# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) +# ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show() +# sum_uv += batch['uv_unposed'].numpy().sum(axis=0) +# sum_mask += batch['uv_mask'].numpy().sum(axis=0) + +# sum_mask[sum_mask == 0] = 1 +# avg_garment = sum_uv / sum_mask[:, :, np.newaxis] + +# # x = np.transpose(avg_garment, [1,2,0]) +# x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0] +# go.Figure(data=[ +# go.Mesh3d( +# # 8 vertices of a cube +# x=x[:,0], +# y=x[:,1], +# z=x[:,2], +# # i, j and k give the vertices of triangles +# i = F[:,0], +# j = F[:,1], +# k = F[:,2] +# )]).show() + +# np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment) + +# print() \ No newline at end of file diff --git a/data/downscale.py b/data/downscale.py new file mode 100644 index 0000000..1924cae --- /dev/null +++ b/data/downscale.py @@ -0,0 +1,223 @@ +import argparse +import os +import sys +import tqdm +import cv2 +import torch +import scipy.interpolate as spi +import numpy as np +import plotly.graph_objects as go +import plotly.express as ex +from torchvision.transforms import Resize +from torchvision.transforms.functional import InterpolationMode + +os.environ["CUDA_VISIBLE_DEVICES"] = "4" + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') +from data.utils import Frame + + +def fill(img, mask): + coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0)) + img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0)) + img_data = img[img_idx[:,0], img_idx[:,1], :] + + interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) + # interpolator = spi.RBFInterpolator(y=img_idx, d=img_data, neighbors=5, kernel='linear') + + img_2 = interpolator(coords) + res = np.zeros((512,256,3), dtype=np.float32) + res[coords[:,0], coords[:,1], :] = img_2 + res[img_idx[:,0], img_idx[:,1], :] = img_data + + return res + + +def fill_label(img, mask): + coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0)) + img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0)) + img_data = img[img_idx[:,0], img_idx[:,1]] + + interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) + img_2 = interpolator(coords) + res = np.zeros((512,256), dtype=np.float32) + res[coords[:,0], coords[:,1]] = img_2 + res[img_idx[:,0], img_idx[:,1]] = img_data + + return res + + +def generate_downscaled_uv_map(folder, frame: Frame, size): + img = np.load(os.path.join(folder, 'uv_mask', f'{frame.id}_{frame.garment}.npz')); img = img['uv_mask'].astype(np.float32) + down = resizers_nearest[size](torch.from_numpy(img[None,:,:]).to(device)) + down = np.squeeze(down.cpu().numpy().astype(bool)) + if not os.path.exists(os.path.join(folder, 'uv_mask' + '_' + str(size))): + os.makedirs(os.path.join(folder, 'uv_mask' + '_' + str(size)), exist_ok=True) + np.savez_compressed(os.path.join(folder, 'uv_mask' + '_' + str(size), f'{frame.id}_{frame.garment}.npz'), uv_mask = down) + + +class UVMaskStore(): + def __init__(self, folder, size): + self.store = dict() + self.folder = folder + self.size = size + + def read(self, frame: Frame): + path = os.path.join(self.folder, 'uv_mask' + ('_' + str(self.size) if self.size != 512 else ''), f'{frame.id}_{frame.garment}.npz') + if not os.path.exists(path): + generate_downscaled_uv_map(self.folder, frame, self.size) + + mask = np.load(path) + mask = mask['uv_mask'] + return mask + + def __getitem__(self, frame: Frame) -> dict: + if (frame.id, frame.garment) not in self.store: + self.store[(frame.id, frame.garment)] = self.read(frame) + + return self.store[(frame.id, frame.garment)] + + +def downscale_uv_maps_posed_body(folder, type, sizes): + for size in sizes: + os.makedirs(os.path.join(folder, type + str(size)), exist_ok=True) + + frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)] + frames = [Frame(frame.split('_')[0], frame.split('_')[1].split('.')[0], '') for frame in frames] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}.npz'))['uv_body_posed'] + img = fill(img, uv_masks_body[512]) + for size in sizes: + mask = uv_masks_body[size] + down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() + down = down * mask[None,:,:] + np.savez_compressed(os.path.join(folder, type + str(size), f'{frame.id}_{frame.frame_nr}.npz'), uv_body_posed = down) + + +def downscale_uv_maps_posed_no_orient(folder, type, sizes): + for size in sizes: + 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)] + frames = [Frame(frame.split('_')[0], frame.split('_')[1], frame.split('_')[2].split('.')[0]) for frame in frames] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_posed'] + img = fill(img, uv_mask_store[512][frame]) + for size in sizes: + mask = uv_mask_store[size][frame] + 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), f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'), uv_posed = down) + + +def downscale_uv_maps_normals(folder, type, sizes): + for size in sizes: + 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)] + frames = [Frame(frame.split('_')[0], frame.split('_')[1], frame.split('_')[2].split('.')[0]) for frame in frames] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_normals'] + img = fill(img, uv_mask_store[512][frame]) + for size in sizes: + mask = uv_mask_store[size][frame] + 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), f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'), uv_normals = down) + + +def downscale_uv_maps_unposed(folder, type, sizes): + for size in sizes: + os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) + + frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)] + frames = [Frame(frame.split('_')[0], frame.split('_')[1], frame.split('_')[2].split('.')[0]) for frame in frames] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_unposed'] + img = fill(img, uv_mask_store[512][frame]) + for size in sizes: + mask = uv_mask_store[size][frame] + 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), f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'), uv_unposed = down) + + +def downscale_uv_maps_static(folder, type, sizes): + for size in sizes: + os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) + + frames = sorted(os.listdir(os.path.join(folder, type))) + import random + random.shuffle(frames) + frames = [(frame.split('_')[0], frame.split('_')[1].split('.')[0]) for frame in frames] + + for frame in tqdm.tqdm(frames): + img = np.load(os.path.join(folder, type, f'{frame[0]}_{frame[1]}.npz'))['uv_static'] + img = fill(img, uv_mask_store[512][Frame(frame[0], 0, frame[1])]) + for size in sizes: + mask = uv_mask_store[size][Frame(frame[0], 0, frame[1])] + 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), 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("-s", "--start", help = "start of range", default=0) + parser.add_argument("-e", "--end", help = "end of range", default=25000) + 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] + 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 } + + uv_masks_body = dict() + uv_mask_body = np.load('custom_utils/mask_256x512.npy') + uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2) + uv_masks_body[512] = uv_mask_body.astype(bool) + for res in [32, 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) + downscale_uv_maps_posed_no_orient(folder=folder, type='uv_posed_no_orient', sizes=sizes) + downscale_uv_maps_normals(folder=folder, type='uv_normals', sizes=sizes) + downscale_uv_maps_posed_body(folder=folder, type='uv_body_posed_no_orient', sizes=sizes) + + + print() \ No newline at end of file diff --git a/data/preprocessing.py b/data/preprocessing.py index 098285c..9a4a606 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -6,8 +6,8 @@ parser = argparse.ArgumentParser() # parser.add_argument("-f", "--folder", help = "end of range") parser.add_argument("-s", "--start", help = "start of range", type=int, default=0) parser.add_argument("-e", "--end", help = "end of range", type=int, default=1) -parser.add_argument("-g", "--gpu", help = "gpu id") -# parser.add_argument("-n", "--name", help = "end of range") +parser.add_argument("-g", "--gpu", help = "gpu id", default='0') +parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true') args = parser.parse_args() os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" @@ -30,7 +30,7 @@ import torch from scipy.spatial.transform import Rotation as R sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') -from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris +from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris, readPC2Frame from data.utils import Frame from smpl.smpl_torch import SMPL from custom_utils.custom_utils import _maskFaces @@ -160,7 +160,7 @@ def deorient(verts, pose): def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(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) @@ -179,12 +179,18 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder): continue verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] + # if info['poses'].shape[1] == verts.shape[0]: + # continue reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat')) verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2) mask_idx = sio.loadmat(os.path.join(folder_reg, id, garment + '_mask.mat'))['mask_idx'].reshape((-1)) - 1 mask = np.zeros((14475), dtype=bool) mask[mask_idx] = True + if info['poses'].shape[1] != verts.shape[0]: + assert False, 'Somethign wrong here' + info['poses'] = info['poses'][:, :verts.shape[0]] + verts = deorient(verts, info['poses']) # # plot as scatter plot with plotly @@ -193,13 +199,12 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder): # fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))]) # fig.show() - writePC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), verts, True) np.save(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'), mask) def create_pose_shape(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'pose'), exist_ok=True) os.makedirs(os.path.join(out_folder, 'shape'), exist_ok=True) @@ -216,7 +221,7 @@ def create_pose_shape(folder, out_folder): def create_mesh_static(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_static'), exist_ok=True) for id in tqdm(ids): @@ -247,7 +252,7 @@ def create_mesh_static(folder, out_folder): def create_mesh_body_posed_no_orient(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_body_posed_no_orient'), exist_ok=True) for id in tqdm(ids): @@ -288,10 +293,11 @@ def calculate_normals(vertices, faces): def create_mesh_normals(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_normals'), exist_ok=True) # F = np.load('smpl/smpl_faces.npy') # F_skirt = np.load('smpl/smpl_faces_skirt.npy') + # ids = ['00438'] for id in tqdm(ids): garment_dir = os.path.join(folder, id) @@ -306,7 +312,7 @@ def create_mesh_normals(folder, out_folder): continue verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] - verts = deorient(verts, info['poses']) + verts = deorient(verts, info['poses'][:, :verts.shape[0]]) _, faces, _, _ = readOBJ(os.path.join(garment_dir, garment + '.obj')) faces = quads2tris(faces) normals = calculate_normals(verts, faces.astype(np.int32)) @@ -356,9 +362,64 @@ def create_mesh_normals(folder, out_folder): pass +def create_mesh_normals_cape(folder, out_folder): + frames = os.listdir(os.path.join(folder, 'mesh_posed')) + os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True) + F = np.load('smpl/smpl_faces_orig.npy') + # F_skirt = np.load('smpl/smpl_faces_skirt.npy') + + for frame in tqdm(frames): + id, frame_nr = frame.split('.')[:2] + + + verts = np.load(os.path.join(folder, 'mesh_posed', f'{id}.{frame_nr}.npz'))['mesh_posed'] + normals = calculate_normals(verts[None,:,:], F)[0] + + np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}.{frame_nr}.npz'), mesh_normals=normals) + + # normals = readPC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V'] + # colors = np.clip(np.abs(normals), 0, 1) + # verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2) + # plot as scatter plot with plotly + # import plotly.graph_objects as go + # frame = 5 + # fig = go.Figure(data=[go.Scatter3d( + # x=verts[:,0], + # y=verts[:,2], + # z=verts[:,1], + # mode='markers', + # marker=dict( + # size=1, + # color=normals[:,:], + # colorscale='Viridis', + # opacity=0.8 + # ) + # )]) + # moved_verts = verts + 0.05 * normals + + # # Create the scatter plot for the moved vertices + # fig.add_trace(go.Scatter3d( + # x=moved_verts[:,0], + # y=moved_verts[:,2], + # z=moved_verts[:,1], + # mode='markers', + # marker=dict( + # size=1, + # color='red', + # opacity=0.8 + # ) + # )) + # fig.show() + + # # create a plotly mesh plot for verts and faces + # 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.show() + + pass + def verify_registration_ans_mask(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) for id in ids: garment_dir = os.path.join(folder, id) @@ -510,9 +571,136 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1) # fig.show() pass + +def create_frames_file(folder, out_folder, split=0.9): + ids = os.listdir(folder) + + frames_non_skirt = [] + frames_skirt = [] + ids_non_skirt = set() + ids_skirt = set() + + for id in tqdm(ids): + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + info = sio.loadmat(os.path.join(folder, id, 'info.mat')) + garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + frames = info['poses'].shape[1] + + assert info['poses'].shape[0] == 72 + + for garment in garments: + for i in range(frames): + if garment == 'Skirt' or garment == 'Dress': + frames_skirt.append(Frame(id, i, garment)) + ids_skirt.add(id) + else: + frames_non_skirt.append(Frame(id, i, garment)) + ids_non_skirt.add(id) + + nr_train_non_skirt = int(len(ids_non_skirt) * split) + nr_train_skirt = int(len(ids_skirt) * split) + + ids_skirt = list(ids_skirt) + ids_non_skirt = list(ids_non_skirt) + + with open(os.path.join(out_folder, 'frames_train.txt'), 'w') as f: + for frame in frames_non_skirt: + if frame.id in ids_non_skirt[:nr_train_non_skirt]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, 'frames_train_skirt.txt'), 'w') as f: + for frame in frames_skirt: + if frame.id in ids_skirt[:nr_train_skirt]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, 'frames_valid.txt'), 'w') as f: + for frame in frames_non_skirt: + if frame.id in ids_non_skirt[nr_train_non_skirt:]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, 'frames_valid_skirt.txt'), 'w') as f: + for frame in frames_skirt: + if frame.id in ids_skirt[nr_train_skirt:]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + +def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1): + ids = os.listdir(folder) + + frames_non_skirt = [] + frames_skirt = [] + ids_non_skirt = set() + ids_skirt = set() + + for id in tqdm(ids): + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + info = sio.loadmat(os.path.join(folder, id, 'info.mat')) + garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + frames = info['poses'].shape[1] + + assert info['poses'].shape[0] == 72 + + for garment in garments: + for i in range(frames): + if garment == 'Skirt' or garment == 'Dress': + frames_skirt.append(Frame(id, i, garment)) + ids_skirt.add(id) + else: + frames_non_skirt.append(Frame(id, i, garment)) + ids_non_skirt.add(id) + + nr_train_non_skirt = int(len(ids_non_skirt) * split * part) + nr_train_skirt = int(len(ids_skirt) * split * part) + + ids_skirt = list(ids_skirt) + ids_non_skirt = list(ids_non_skirt) + + with open(os.path.join(out_folder, f'frames_train_{int(part*100)}p.txt'), 'w') as f: + for frame in frames_non_skirt: + if frame.id in ids_non_skirt[:nr_train_non_skirt]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, f'frames_train_skirt_{int(part*100)}p.txt'), 'w') as f: + for frame in frames_skirt: + if frame.id in ids_skirt[:nr_train_skirt]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + + +def check_irregular_poses(folder, folder_reg, out_folder): + frame: Frame = Frame('01512', 50, 'Dress') + mesh_posed = readPC2Frame(os.path.join(out_folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy') + + mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy') + mesh_body_posed = readPC2Frame(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True) + + F = np.load('smpl/smpl_faces_skirt.npy') + _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) + + fig = go.Figure(data=[go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)]) + fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1)) + fig.update_layout(scene = dict( + xaxis = dict(range=[-1,1],), + yaxis = dict(range=[-1,1],), + zaxis = dict(range=[-1,1],),),) + fig.update_layout(scene_aspectmode='cube') + fig.show() + pass + if __name__ == '__main__': + # folder = '/data-net/dataset/cloth3d/test' + # folder_reg = '/data-local/127/cloth3d_registration/test' + # out_folder = '/data-local/127/cloth3d_processed_test' + folder = '/data-net/dataset/cloth3d/train' folder_reg = '/data-local/127/cloth3d_registration/train' out_folder = '/data-local/127/cloth3d_processed' @@ -524,10 +712,18 @@ if __name__ == '__main__': # split_data(folder, 0.9) # create_extras_file(folder) + # check_irregular_poses(folder, folder_reg, out_folder) + # create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder) # verify_registration_ans_mask(folder, out_folder) # 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) \ No newline at end of file + # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end) + + # create_frames_file(folder, out_folder) + create_frames_file_partial(folder, out_folder, part=0.5) + + # folder = '/data-local/127_2/cape_processed/test/female' + # create_mesh_normals_cape(folder, None) \ No newline at end of file diff --git a/data/utils.py b/data/utils.py index 6f25887..1e1596b 100644 --- a/data/utils.py +++ b/data/utils.py @@ -1,10 +1,10 @@ import tables as tb class Frame: - def __init__(self, id, frame_nr, garment=''): - self.id = id - self.frame_nr = frame_nr - self.garment = garment + 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 __hash__(self): return hash((self.id, self.frame_nr, self.garment)) diff --git a/docker/Dockerfile2 b/docker/Dockerfile2 index 8e35b81..ef8570c 100644 --- a/docker/Dockerfile2 +++ b/docker/Dockerfile2 @@ -1,7 +1,7 @@ # Use an existing image as the base -FROM nvcr.io/nvidia/pytorch:23.04-py3 +# FROM nvcr.io/nvidia/pytorch:23.04-py3 # FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel -# FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel +FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel # # Add the current user to the image ARG USER_NAME=lhunor @@ -10,7 +10,7 @@ ARG GROUP_NAME=guest ARG GROUP_ID=11000 # Update the system and install required packages -RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git apt-utils htop +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 @@ -25,6 +25,12 @@ 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 @@ -33,8 +39,8 @@ WORKDIR /workspace ENV PATH="$PATH:/home/$USER_NAME/.local/bin" RUN pip install --upgrade pip COPY requirements.txt /tmp/ -RUN pip install --requirement /tmp/requirements.txt RUN pip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+cu117.html +RUN pip install --requirement /tmp/requirements.txt # defining the volumes that will be mounted diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index ac5fd5e..dfb7030 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -26,6 +26,28 @@ services: container_name: lhunor_torch_nv23.04 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] + + + lhunor_iccvw: + image: lhunor_iccvw:pt2.0.1 + container_name: lhunor_iccvw_pt2.0.1 + ipc: host + tty: true ulimits: memlock: -1 stack: 67108864 diff --git a/docker/requirements.txt b/docker/requirements.txt index 37f1593..a1e9bfb 100644 --- a/docker/requirements.txt +++ b/docker/requirements.txt @@ -1,13 +1,10 @@ -torch==2.0.0 numpy -matplotlib opencv-python wandb scipy plotly pandas cmapy -pytorch3d tables numpy-stl torchmetrics diff --git a/scripts/ae_body_train.py b/scripts/ae_body_train.py index 6226991..37938cc 100644 --- a/scripts/ae_body_train.py +++ b/scripts/ae_body_train.py @@ -11,10 +11,12 @@ parser.add_argument('--config', type=str) parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) -config_module = importlib.import_module(args['config']) -config = config_module.config if args['debug']: - from configs.ae_body_cape_m import config + from configs.iccvw.ae_body_cape_m import config +else: + config_module = importlib.import_module(args['config']) + config = config_module.config + os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: @@ -32,10 +34,7 @@ import plotly.graph_objects as go import plotly.express as ex import torch.optim as optim -if config['is_skirt']: - from losses.losses import edge_loss_skirt as edge_loss -else: - from losses.losses import edge_loss +from losses.losses import edge_loss from data.dataset import Dataset # from torch.utils.data import DataLoader @@ -49,6 +48,7 @@ from pytorch_msssim import MS_SSIM def init_wandb(): + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/ae_body_train.py') @@ -64,7 +64,7 @@ def init_wandb(): artifact.add_file(file_name) run.log_artifact(artifact) -uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) +uv2mesh = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config['is_cape']) @@ -121,7 +121,7 @@ def plot_everything(batch: dict, type: str): img = torch.concat((orig_img, pred_img), axis=2) imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)] - F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") + F = np.load("smpl/smpl_faces.npy") garment_unposed, faces = mesh_pred.cpu().numpy(), F gt = [mesh_orig[i].cpu().numpy() for i in range(pred.shape[0])] @@ -175,9 +175,7 @@ if __name__ == '__main__': if not args['debug']: init_wandb() - skirt_modifier = '_skirt' if config['is_skirt'] else '' - - uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy') + 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_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device']) @@ -198,15 +196,15 @@ if __name__ == '__main__': load_model(model.encoder, config['pretrained_model'] + '_enc') load_model(model.decoder, config['pretrained_model'] + '_dec') - ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5) - ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5) + ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5) + ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5) run_name = config['name'] print(f'started training: {run_name}') if not args['debug']: # model = torch.compile(model) - wandb.watch(model, log='all', log_freq=1000, log_graph=True) + wandb.watch(model, log='all', log_freq=10000, log_graph=True) train() diff --git a/scripts/ae_norm_train.py b/scripts/ae_norm_train.py index 621b493..597e662 100644 --- a/scripts/ae_norm_train.py +++ b/scripts/ae_norm_train.py @@ -11,10 +11,12 @@ parser.add_argument('--config', type=str) parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) -config_module = importlib.import_module(args['config']) -config = config_module.config if args['debug']: - from configs.ae_norm_1 import config + from configs.iccvw.ae_norm_cape_m import config +else: + config_module = importlib.import_module(args['config']) + config = config_module.config + os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: @@ -49,7 +51,8 @@ from pytorch_msssim import MS_SSIM def init_wandb(): - run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127_2/temp', tags=config['tags']) + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) + run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/ae_norm_train.py') artifact.add_file('losses/losses.py') @@ -70,9 +73,9 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_ def train_step(batch): model.train() - pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1 + pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4 - ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1)) + ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4)) loss = 1e3 * (1 - ms_ssim) optimizer.zero_grad() @@ -93,13 +96,13 @@ def valid_step(batch): model.eval() with torch.inference_mode(): - pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1 + pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4 mesh_pred = uv2mesh(pred) mesh_orig = batch['mesh_normals'] l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_normals'] - pred)), axis=1)) - ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1)) + ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4)) loss = 1e3 * (1 - ms_ssim) metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) @@ -109,7 +112,7 @@ def valid_step(batch): def plot_everything(batch: dict, type: str): model.eval() - pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.1 + pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.4 # mesh_pred = uv2mesh(pred) # mesh_orig = batch['mesh_normals'] @@ -117,7 +120,7 @@ def plot_everything(batch: dict, type: str): orig_img = torch.permute(batch['uv_normals'], [0,2,3,1]) pred_img = torch.permute(pred, [0,2,3,1]) img = torch.concat((orig_img, pred_img), axis=2) - imgs = [wandb.Image((img[i].cpu().numpy() + 1.1) / 2.2) for i in range(4)] + imgs = [wandb.Image((img[i].cpu().numpy() + 1.4) / 2.8) for i in range(4)] # F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") # garment_unposed, faces = mesh_pred.cpu().numpy(), F @@ -196,7 +199,7 @@ if __name__ == '__main__': optimizer = optim.Adam(model.parameters(), lr=config['lr']) # _MSSSIM = MSSSIM(data_range=2).to(config['device']) - _MSSSIM = MS_SSIM(data_range=2.2, size_average=True, channel=3) + _MSSSIM = MS_SSIM(data_range=2.8, size_average=True, channel=3) if config['pretrained_model'] != '': load_model(model.encoder, config['pretrained_model'] + '_enc') diff --git a/scripts/ae_static_train.py b/scripts/ae_static_train.py index 1d03715..3e6de96 100644 --- a/scripts/ae_static_train.py +++ b/scripts/ae_static_train.py @@ -10,11 +10,13 @@ parser.add_argument('--config', type=str) parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) -config_module = importlib.import_module(args['config']) -config = config_module.config if args['debug']: - from configs.ae_static_cape_m import config - # from configs.ae_static_orig_2 import config + # from configs.ae_static_cape_m import config + from configs.iccvw.ae_static_1 import config +else: + config_module = importlib.import_module(args['config']) + config = config_module.config + os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: @@ -49,6 +51,7 @@ from pytorch_msssim import MS_SSIM def init_wandb(): + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/ae_static_train.py') @@ -91,7 +94,7 @@ def train_step(batch): mesh_orig = batch['mesh_static'] l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1)) e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } @@ -205,18 +208,20 @@ if __name__ == '__main__': _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3) if config['pretrained_model'] != '': - load_model(model.encoder, config['pretrained_model'] + '_enc') - load_model(model.decoder, config['pretrained_model'] + '_dec') + it = load_model(model.encoder, config['pretrained_model'] + '_enc', optimizer=optimizer) + load_model(model.decoder, config['pretrained_model'] + '_dec', optimizer=optimizer) + pre_name = config['pretrained_model'] + print(f'continued training from: {pre_name} at iteration: {it}') - ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5) - ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5) + ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5) + ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5) run_name = config['name'] print(f'started training: {run_name}') if not args['debug']: # model = torch.compile(model) - wandb.watch(model, log='all', log_freq=1000, log_graph=True) + wandb.watch(model, log='all', log_freq=10000, log_graph=True) train() diff --git a/scripts/upscaler_train.py b/scripts/upscaler_train.py index 7199a1d..f581ecd 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"] = '6' +os.environ["CUDA_VISIBLE_DEVICES"] = '2' import time @@ -49,26 +49,36 @@ gender_modifier = '' if not is_cape else ('_m' if is_male else '_f') smpl_res = 3638 if is_cape else 14475 use_cache = False name=f'iccv_{low_res}' + skirt_modifier + cape_modifier + gender_modifier +train_dir = '/data-local/127/cloth3d_processed' +valid_dir = '/data-local/127/cloth3d_processed' -if not is_skirt: - train_dir = '/data-local/127/train' - valid_dir = '/data-local/127/valid' +if is_skirt: + train_frames = 'frames_train_skirt.txt' + valid_frames = 'frames_valid_skirt.txt' else: - train_dir = '/data-local/128/train_skirt' - valid_dir = '/data-local/127/valid_skirt' + train_frames = 'frames_train.txt' + valid_frames = 'frames_valid.txt' + +if is_cape: + train_frames = None + valid_frames = None if is_cape: if is_male: - train_dir = '/data-local/127/cape_processed/train/male' - valid_dir = '/data-local/127/cape_processed/test/male' + train_dir = '/data-local/127_2/cape_processed/train/male' + valid_dir = '/data-local/127_2/cape_processed/test/male' else: - train_dir = '/data-local/127/cape_processed/train/female' - valid_dir = '/data-local/127/cape_processed/test/female' + train_dir = '/data-local/127_2/cape_processed/train/female' + valid_dir = '/data-local/127_2/cape_processed/test/female' +print(f'nohup python scripts/upscaler_train.py --name {name} > nohup/up/up_{name}.out 2>&1 &') +if args['debug']: + exit() def init_wandb(): - run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccv']) + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) + run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccvw']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/upscaler_train.py') artifact.add_file('losses/losses.py') @@ -255,8 +265,8 @@ if __name__ == '__main__': } # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') - dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape) - dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape) + dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=train_frames, resolutions=[512, low_res], is_cape=is_cape) + dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=valid_frames, resolutions=[512, low_res], is_cape=is_cape) train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True) valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True) train_accumulator = Accumulator() @@ -279,13 +289,13 @@ if __name__ == '__main__': _MSSSIM_small = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[1.0]) step = 0 - ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + name, top_n=5) + ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + name, top_n=5) print(f'started training: {name}') if not args['debug']: # model = torch.compile(model) - wandb.watch(model, log='all', log_freq=1000, log_graph=True) + wandb.watch(model, log='all', log_freq=10000, log_graph=True) train() diff --git a/utils/generate_figures.py b/utils/generate_figures.py new file mode 100644 index 0000000..b08abbf --- /dev/null +++ b/utils/generate_figures.py @@ -0,0 +1,126 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +import numpy as np + +from custom_utils.custom_utils import save_stl, write_obj +from data.cloth3d_io import readPC2Frame +from data.utils import Frame +from custom_utils.custom_utils import _maskFaces +import plotly.graph_objects as go +from plotly.subplots import make_subplots + + +def remove_long_edges(vertices, faces): + # Compute edge lengths + edges = np.vstack((faces[:,[0,1]], faces[:,[1,2]], faces[:,[2,0]])) + edge_lengths = np.linalg.norm(vertices[edges[:,0]] - vertices[edges[:,1]], axis=1) + + # Find edges longer than 0.1 + long_edges = edges[edge_lengths > 0.1] + + # Find vertices that are part of long edges + vertices_to_remove = np.unique(long_edges) + + # Remove vertices and update faces + mask = np.ones(vertices.shape[0], dtype=bool) + mask[vertices_to_remove] = False + new_vertices = vertices[mask] + + + return new_vertices, mask + + +def plot_unposing_example(): + frame = Frame('00438', 146, 'Jumpsuit') + mesh_static = np.load(f'{folder_test}/mesh_static/{frame.id}_{frame.garment}.npy') + + mesh_mask = np.load(f'{folder_test}/mesh_mask/{frame.id}_{frame.garment}.npy') + mesh_posed = readPC2Frame(os.path.join(folder_test, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + mesh_unposed = readPC2Frame(os.path.join(folder_test, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + mesh_unposed_shaped = readPC2Frame(os.path.join('temp', f'{frame.id}_{frame.garment}_unposed_shaped.pc16'), frame.frame_nr, True)[mesh_mask] + + + F = np.load('smpl/smpl_faces.npy') + _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) + + mesh_static, v_mask = remove_long_edges(mesh_static, faces.astype(np.int32)) + + mesh_posed = mesh_posed[v_mask] + mesh_unposed = mesh_unposed[v_mask] + mesh_unposed_shaped = mesh_unposed_shaped[v_mask] + new_mask = np.zeros_like(mesh_mask) + new_mask[mesh_mask] = v_mask + _, faces = _maskFaces(np.zeros((14475, 3)), F, new_mask) + + fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}]]) + + fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), row=1, col=1) + fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=1), row=1, col=2) + fig.add_trace(go.Mesh3d(x=mesh_unposed_shaped[:,0], y=mesh_unposed_shaped[:,2], z=mesh_unposed_shaped[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightgreen', opacity=1), row=1, col=3) + fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightyellow', opacity=1), row=1, col=4) + + fig.update_layout(scene=dict( + xaxis=dict(range=[-1, 1]), + yaxis=dict(range=[-1, 1]), + zaxis=dict(range=[-1, 1]), + ), scene_aspectmode='cube') + + # fig.show() + faces = faces.astype(np.uint32) + + save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_static.stl', mesh_static, faces) + save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_posed.stl', mesh_posed, faces) + save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_unposed.stl', mesh_unposed, faces) + save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_unposed_shaped.stl', mesh_unposed_shaped, faces) + + pass + + +def plot_normals_example(): + frame = Frame('00438', 146, 'Jumpsuit') + mesh_static = np.load(f'{folder_test}/mesh_static/{frame.id}_{frame.garment}.npy') + + mesh_mask = np.load(f'{folder_test}/mesh_mask/{frame.id}_{frame.garment}.npy') + mesh_posed = readPC2Frame(os.path.join(folder_test, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + # mesh_unposed = readPC2Frame(os.path.join(folder_test, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + # mesh_unposed_shaped = readPC2Frame(os.path.join('temp', f'{frame.id}_{frame.garment}_unposed_shaped.pc16'), frame.frame_nr, True)[mesh_mask] + mesh_normals = readPC2Frame(os.path.join(folder_test, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + uv_normals = np.load(os.path.join(folder_test, 'uv_normals', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')) + + F = np.load('smpl/smpl_faces.npy') + _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) + + mesh_static, v_mask = remove_long_edges(mesh_static, faces.astype(np.int32)) + + mesh_posed = mesh_posed[v_mask] + # mesh_unposed = mesh_unposed[v_mask] + # mesh_unposed_shaped = mesh_unposed_shaped[v_mask] + new_mask = np.zeros_like(mesh_mask) + new_mask[mesh_mask] = v_mask + _, faces = _maskFaces(np.zeros((14475, 3)), F, new_mask) + + fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}]]) + + fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), row=1, col=1) + fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=1), row=1, col=2) + fig.add_trace(go.Mesh3d(x=mesh_unposed_shaped[:,0], y=mesh_unposed_shaped[:,2], z=mesh_unposed_shaped[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightgreen', opacity=1), row=1, col=3) + fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightyellow', opacity=1), row=1, col=4) + + fig.update_layout(scene=dict( + xaxis=dict(range=[-1, 1]), + yaxis=dict(range=[-1, 1]), + zaxis=dict(range=[-1, 1]), + ), scene_aspectmode='cube') + + + +if __name__ == '__main__': + folder_train = '/data-local/127/cloth3d_processed' + folder_test = '/data-local/127/cloth3d_processed_test' + + # plot_unposing_example() + plot_normals_example() + + pass \ No newline at end of file