pyramid training working

configs/archive/abl_norm.py

@@ -0,0 +1,192 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'abl_torch_norm'
+config['notes'] = ''
+config['tags'] = ['iccv', 'abl_norm', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = ''
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_normals'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_normals'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
+config['body_model'] = 'checkpoints/iccv/body_torch_1'
+config['normals_model'] = 'checkpoints/iccv/torch/norm_torch_1'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'uv_normals',
+                            'uv_static',
+                            'mesh_mask',
+                            'uv_mask',
+                        }
+
+if config['use_skinning']:
+  config['signature'].add('uv_unposed')
+  config['signature'].add('mesh_unposed')
+else:
+  config['signature'].add('uv_posed')
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+config['bottleneck_scale'] = 2
+
+config['cvae'] = dict()
+config['cvae']['device'] = config['device']
+config['cvae']['is_label'] = False
+config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+config['cvae']['kernel_size'] = 7
+config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
+config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
+
+block_setting = [
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 3),
+    CNBlockConfig(128, 256, 3),
+    CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 3),
+    CNBlockConfig(32, 3, 3),
+]
+
+config['cvae']['encoder_block_settings'] = block_setting
+config['cvae']['decoder_block_settings'] = block_setting_reverse
+
+
+
+####################### cond_ae_norm ############################
+config['cond_ae_norm'] = dict()
+config['cond_ae_norm']['device'] = config['device']
+config['cond_ae_norm']['is_label'] = True
+config['cond_ae_norm']['latent_dim'] = config['latent_dim']
+config['cond_ae_norm']['kernel_size'] = 7
+config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_norm']['encoder_block_settings'] = block_setting
+config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_norm']['activation'] = 'tanh'
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

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
 

configs/archive/ae_norm_1.py

@@ -0,0 +1,73 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'norm_torch_2'
+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/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['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',
+                        }

configs/archive/cvae_norm_1.py

@@ -0,0 +1,172 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'cvae_norm_1'
+config['notes'] = ''
+config['tags'] = ['iccv', 'cvae_norm', 'torch']
+config['script'] = 'scripts/cvae_norm.py'
+
+config['pretrained_model'] = ''
+
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
+config['body_model'] = 'checkpoints/iccv/body_torch_1'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'uv_normals',
+                            'uv_static',
+                            'mesh_mask',
+                            'uv_mask',
+                            'mesh_normals',
+                        }
+
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+config['bottleneck_scale'] = 2
+
+config['cvae'] = dict()
+config['cvae']['device'] = config['device']
+config['cvae']['is_label'] = False
+config['cvae']['input_features'] = 3+3+3
+config['cvae']['kernel_size'] = 7
+config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
+config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
+
+block_setting = [
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 3),
+    CNBlockConfig(128, 256, 3),
+    CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(176, 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 3),
+    CNBlockConfig(32, 3, 3),
+]
+
+config['cvae']['encoder_block_settings'] = block_setting
+config['cvae']['decoder_block_settings'] = block_setting_reverse
+
+
+
+####################### cond_ae_curv ############################
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['device'] = config['device']
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['latent_dim'] = config['latent_dim']
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['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_curv']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_curv']['encoder_block_settings'] = block_setting
+config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_curv']['activation'] = 'tanh'
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

configs/iccvw/abl_all.py

@@ -0,0 +1,195 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '2'
+config['name'] = 'abl_iccvw_all_2'
+config['notes'] = 'fixed collision loss'
+config['tags'] = ['iccv', 'abl_all', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = 'checkpoints/iccvw/abl_iccvw_all_1'
+
+config['use_skinning'] =              True # and deshaping
+config['use_normals'] =               True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+
+config['use_norm_loss'] =             True
+
+ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['avg_garment_path'] != '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+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['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
+config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'
+config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+config['balancing_activation_limit'] = 0.01
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'uv_normals',
+                            'uv_static',
+                            'mesh_mask',
+                            'uv_mask',
+                        }
+
+if config['use_skinning']:
+  config['signature'].add('uv_unposed')
+  config['signature'].add('mesh_unposed')
+else:
+  config['signature'].add('uv_posed')
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+config['bottleneck_scale'] = 2
+
+config['cvae'] = dict()
+config['cvae']['device'] = config['device']
+config['cvae']['is_label'] = False
+config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+config['cvae']['kernel_size'] = 7
+config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
+config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
+
+block_setting = [
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 3),
+    CNBlockConfig(128, 256, 3),
+    CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 3),
+    CNBlockConfig(32, 3, 3),
+]
+
+config['cvae']['encoder_block_settings'] = block_setting
+config['cvae']['decoder_block_settings'] = block_setting_reverse
+
+
+
+####################### cond_ae_norm ############################
+config['cond_ae_norm'] = dict()
+config['cond_ae_norm']['device'] = config['device']
+config['cond_ae_norm']['latent_dim'] = config['latent_dim']
+config['cond_ae_norm']['kernel_size'] = 7
+config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_norm']['encoder_block_settings'] = block_setting
+config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_norm']['activation'] = 'tanh'
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

configs/iccvw/abl_all_3.py

@@ -0,0 +1,194 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'abl_iccvw_all_3'
+config['notes'] = 'fixed collision loss, train from zero'
+config['tags'] = ['iccv', 'abl_all', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = ''
+
+config['use_skinning'] =              True # and deshaping
+config['use_normals'] =               True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_norm_loss'] =             True
+
+ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['avg_garment_path'] != '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+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['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
+config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'
+config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+config['balancing_activation_limit'] = 0.01
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'uv_normals',
+                            'uv_static',
+                            'mesh_mask',
+                            'uv_mask',
+                        }
+
+if config['use_skinning']:
+  config['signature'].add('uv_unposed')
+  config['signature'].add('mesh_unposed')
+else:
+  config['signature'].add('uv_posed')
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+config['bottleneck_scale'] = 2
+
+config['cvae'] = dict()
+config['cvae']['device'] = config['device']
+config['cvae']['is_label'] = False
+config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+config['cvae']['kernel_size'] = 7
+config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
+config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
+
+block_setting = [
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 3),
+    CNBlockConfig(128, 256, 3),
+    CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 3),
+    CNBlockConfig(32, 3, 3),
+]
+
+config['cvae']['encoder_block_settings'] = block_setting
+config['cvae']['decoder_block_settings'] = block_setting_reverse
+
+
+
+####################### cond_ae_norm ############################
+config['cond_ae_norm'] = dict()
+config['cond_ae_norm']['device'] = config['device']
+config['cond_ae_norm']['latent_dim'] = config['latent_dim']
+config['cond_ae_norm']['kernel_size'] = 7
+config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_norm']['encoder_block_settings'] = block_setting
+config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_norm']['activation'] = 'tanh'
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

configs/iccvw/ae_norm_1.py

@@ -0,0 +1,77 @@
+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),
+]
+
+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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

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',
+                        }

configs/iccvw/pyr_32.py

@@ -0,0 +1,300 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '1'
+config['name'] = 'pyr_iccvw_32'
+config['notes'] = ''
+config['tags'] = ['pyr', 'iccv', 'torch']
+config['script'] = 'scripts/cvae_pyramid.py'
+
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+
+config['use_skinning'] =              True # and deshaping
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_norm_loss'] =             True
+config['use_normals'] =               True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32]
+config['pyramid']['pretrained_pyramid'] = ''
+config['pyramid']['pretrained_pyramid_iteration'] = ''
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+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['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['balancing_activation_limit'] = 0.012
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+config['bottleneck_scale'] = 2
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'mesh_unposed',
+                            'uv_unposed',
+                            'uv_normals',
+                            'uv_static',
+                            'mesh_mask',
+                            'uv_mask',
+                        }
+
+config['pyramid']['is_label'] = False
+# config['pyramid']['add_noise'] = False
+# config['pyramid']['noise_amount'] = 0.0
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+    if res != 512:
+        res_modifier = '_' + str(res)
+        config['signature'].add('uv_body_posed' + res_modifier)
+        config['signature'].add('uv_unposed' + res_modifier)
+        config['signature'].add('uv_normals' + res_modifier)
+        config['signature'].add('uv_static' + res_modifier)
+        config['signature'].add('uv_mask' + res_modifier)
+
+
+config['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
+config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'
+config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1'
+config['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 7
+    level['upscale_enhancer'] = ''
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(16, 32, 3),
+        CNBlockConfig(32, 64, 3),
+        CNBlockConfig(64, 128, 3),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 3),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
+        CNBlockConfig(128, 64, 9),
+        CNBlockConfig(64, 32, 3),
+        CNBlockConfig(32, 16, 3),
+        CNBlockConfig(16, 3, 3),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['kernel_size'] = 5
+    level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, 128, 3),
+        CNBlockConfig(128, 128, 3),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 3),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
+        CNBlockConfig(128, 128, 9),
+        CNBlockConfig(128, 65, 3),
+        CNBlockConfig(64, 3, 3),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['kernel_size'] = 5
+    level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, 128, 9),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 9),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 9),
+        CNBlockConfig(128, 128, 9),
+        CNBlockConfig(64, 3, 9),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 3
+    level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 9),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 64, 9),
+        CNBlockConfig(64, 3, 9),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['kernel_size'] = 3
+    level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(level['bottleneck_size'], None, 16),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 3, 16),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+####################### cond_ae_norm ############################
+config['cond_ae_norm'] = dict()
+config['cond_ae_norm']['device'] = config['device']
+config['cond_ae_norm']['is_label'] = False
+# config['cond_ae_norm']['latent_dim'] = config['latent_dim']
+config['cond_ae_norm']['kernel_size'] = 7
+config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_norm']['encoder_block_settings'] = block_setting
+config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_norm']['activation'] = 'tanh'
+
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+# config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+# config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

configs/iccvw/pyr_64.py

@@ -0,0 +1,300 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '3'
+config['name'] = 'pyr_iccvw_64'
+config['notes'] = ''
+config['tags'] = ['pyr', 'iccv', 'torch']
+config['script'] = 'scripts/cvae_pyramid.py'
+
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+
+config['use_skinning'] =              True # and deshaping
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_norm_loss'] =             True
+config['use_normals'] =               True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64]
+config['pyramid']['pretrained_pyramid'] = ''
+config['pyramid']['pretrained_pyramid_iteration'] = ''
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+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['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['balancing_activation_limit'] = 0.012
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+config['bottleneck_scale'] = 2
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'mesh_unposed',
+                            'uv_unposed',
+                            'uv_normals',
+                            'uv_static',
+                            'mesh_mask',
+                            'uv_mask',
+                        }
+
+config['pyramid']['is_label'] = False
+# config['pyramid']['add_noise'] = False
+# config['pyramid']['noise_amount'] = 0.0
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+    if res != 512:
+        res_modifier = '_' + str(res)
+        config['signature'].add('uv_body_posed' + res_modifier)
+        config['signature'].add('uv_unposed' + res_modifier)
+        config['signature'].add('uv_normals' + res_modifier)
+        config['signature'].add('uv_static' + res_modifier)
+        config['signature'].add('uv_mask' + res_modifier)
+
+
+config['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
+config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'
+config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1'
+config['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 7
+    level['upscale_enhancer'] = ''
+    level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(16, 32, 3),
+        CNBlockConfig(32, 64, 3),
+        CNBlockConfig(64, 128, 3),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 3),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
+        CNBlockConfig(128, 64, 9),
+        CNBlockConfig(64, 32, 3),
+        CNBlockConfig(32, 16, 3),
+        CNBlockConfig(16, 3, 3),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['kernel_size'] = 5
+    level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, 128, 3),
+        CNBlockConfig(128, 128, 3),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 3),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
+        CNBlockConfig(128, 128, 9),
+        CNBlockConfig(128, 65, 3),
+        CNBlockConfig(64, 3, 3),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['kernel_size'] = 5
+    level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, 128, 9),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 9),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 9),
+        CNBlockConfig(128, 128, 9),
+        CNBlockConfig(64, 3, 9),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 3
+    level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 9),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 64, 9),
+        CNBlockConfig(64, 3, 9),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['kernel_size'] = 3
+    level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(level['bottleneck_size'], None, 16),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 3, 16),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+####################### cond_ae_norm ############################
+config['cond_ae_norm'] = dict()
+config['cond_ae_norm']['device'] = config['device']
+config['cond_ae_norm']['is_label'] = False
+# config['cond_ae_norm']['latent_dim'] = config['latent_dim']
+config['cond_ae_norm']['kernel_size'] = 7
+config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_norm']['encoder_block_settings'] = block_setting
+config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_norm']['activation'] = 'tanh'
+
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+# config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+# config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

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()

data/calculate_normals.py

@@ -0,0 +1,112 @@
+import os
+import sys
+import numpy as np
+import trimesh as tm
+import scipy.io as sio
+import plotly.graph_objects as go
+import plotly.express as px
+from tqdm import tqdm
+import torch
+from scipy.spatial.transform import Rotation as R
+import argparse
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+from data.cloth3d_io import readOBJ, readPC2
+from data.utils import Frame
+from smpl.smpl_torch import SMPL
+
+
+def calculate_normals(vertices, mesh_file, poses):
+    _, faces, _, _ = readOBJ(mesh_file)
+    rodr = smpl.rodrigues(torch.tensor(poses.T.reshape(-1, 1, 3), dtype=torch.float32)).numpy().reshape(vertices.shape[0], -1, 3, 3)
+
+    normals = []
+    for i in range(vertices.shape[0]):
+        r = R.from_matrix(rodr[i, 0, :, :])
+        # assert r.as_matrix() == rodr[frame, 0, :, :]
+        out = r.apply(vertices[i], inverse=True)
+
+        # create a plotly mesh plot of out and vertices [i]
+        # fig = go.Figure(data=[
+        #     go.Mesh3d(
+        #         x=out[:, 0], y=out[:, 2], z=out[:, 1],
+        #         # i, j and k give the vertices of triangles
+        #         i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
+        #         name='y',
+        #         showscale=False
+        #     ),
+        #     go.Mesh3d(
+        #         x=vertices[i][:, 0], y=vertices[i][:, 1], z=vertices[i][:, 2],
+        #         # i, j and k give the vertices of triangles
+        #         i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
+        #         name='x',
+        #         showscale=False
+        #     )
+        # ])
+        # fig.show()
+    
+        mesh = tm.Trimesh(vertices=out, faces=faces, process=False, validate=False)
+        normals.append(mesh.vertex_normals)
+    return np.stack(normals)
+
+
+def get_frames(folder):
+    uv_static_folder = os.path.join(folder, 'pose_shape') 
+    frames = sorted(os.listdir(uv_static_folder))
+
+    frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+    frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
+    return frames
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--folder", help = "end of range", default='/data-net/dataset/00761')
+    parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/data-net/dataset/cloth3d/train')
+    parser.add_argument("-s", "--start", help = "start of range", default=0)
+    parser.add_argument("-e", "--end", help = "end of range", default=10)
+    parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
+    args = parser.parse_args()
+
+    obj_folder = args.folder_object # '/data-net/dataset/cloth3d/train'
+    folder = args.folder #'/data-local/127/train'
+
+    os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True)
+    frames = get_frames(folder)
+    smpl = SMPL('cpu', 'smpl')
+
+    ids = {frame.id for frame in frames}
+    ids = sorted(list(ids))[int(args.start):int(args.end)]
+
+    for id in tqdm(ids):
+        # find obj files in the folder
+        garment_obj_folder = os.path.join(obj_folder, id[:5])
+        garments = [file.split('.')[0] for file in os.listdir(garment_obj_folder) if file.endswith('.obj')]
+
+        if len(id) == 5:
+            garment = garments[0]
+        elif len(id) == 6:
+            garment = garments[1]
+        else:
+            raise ValueError('Invalid frame id')
+
+        mesh_file = os.path.join(garment_obj_folder, f'{garment}.obj')
+        vertices =readPC2(os.path.join(garment_obj_folder, f'{garment}.pc16'), True)['V']
+        poses = sio.loadmat(os.path.join(garment_obj_folder, f'info.mat'))['poses']
+        normals = calculate_normals(vertices, mesh_file, poses)
+
+        map = sio.loadmat(os.path.join(garment_obj_folder, f'{garment}_map.mat'))['map'][0]
+        mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy'))
+
+        mesh_normals = normals[:,map,:]
+        mesh_normals = mesh_normals * mask[None,:, None]
+
+        for i in range(vertices.shape[0]):
+            # np.savez(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals)
+            np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals[i])
+
+        # pass
+
+        # print(frame.id, frame.frame_nr)
+

data/cloth3d_io.py

@@ -342,7 +342,6 @@ import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
 import plotly.graph_objects as go
-from custom_utils.custom_utils import _maskFaces
 
 def plot_3d(mesh, faces):
     fig = go.Figure()
@@ -370,7 +369,19 @@ def plot_3d(mesh, faces):
     return fig
 
 
+def quads2tris(F):
+    out = []
+    for f in F:
+        if len(f) == 3: out += [f]
+        elif len(f) == 4: out += [[f[0],f[1],f[2]],
+                                [f[0],f[2],f[3]]]
+        else: print("This should not happen...")
+    return np.array(out, np.int32)
+
+
+
 if __name__ == '__main__':
+    from custom_utils.custom_utils import _maskFaces
 
     folder = '/home/workspace/data/dataset/cloth3d/train'
     F = np.load("custom_utils/smpl_faces.npy")#.astype(np.int32)

data/convert_to_uv.py

@@ -12,7 +12,7 @@ import pickle
 import numpy as np
 import scipy as sp
 import trimesh as tm
-import cmapy as cm
+# import cmapy as cm
 import tables as tb
 import scipy.io as sio
 import scipy.interpolate as spi
@@ -24,7 +24,7 @@ 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.utils import save_to_h5file
+from data.cloth3d_io import readPC2
 
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "6"
@@ -34,26 +34,26 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
     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, :] #   add garment mask
-    uv_values /= uv_counter[:, np.newaxis]
+    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] #   add garment 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[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+    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[garment_mask, :]
-    rbf_dest_current = rbf_dest
+    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))
@@ -62,8 +62,8 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
         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[:,:,0]
-        rbf_dest_current = rbf_dest * 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
 
 
@@ -86,7 +86,7 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
     # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]]
 
     result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final
-    result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+    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]
 
@@ -102,26 +102,26 @@ def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None,
     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, :]
-    uv_values /= uv_counter[np.newaxis, :, np.newaxis]
+    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[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+    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] #   add garment 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[garment_mask, :]
-    rbf_dest_current = rbf_dest
+    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[np.invert(garment_mask), :]
+    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
     
@@ -133,7 +133,7 @@ def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None,
         return_garment_uv_mask = True
 
     if garment_uv_mask is not None:
-        rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
+        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
@@ -172,18 +172,18 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
     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
+    garment_mask = garment_mask[uv_map_mask[type]] #   add garment mask
 
-    rbf_dest_current = rbf_dest
+    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[np.invert(garment_mask), :]
+    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 + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
+    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
@@ -196,11 +196,11 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
         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, :]
-        uv_values /= uv_counter[:, np.newaxis]
+        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[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)
@@ -386,12 +386,12 @@ def test(verts, faces, values):
 
 
 def get_shapedirs():
-    to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
+    to_keep = sio.loadmat(os.path.join('smpl/keep_remove.mat'))['keep'].reshape(-1)
     shapedirs = []
-    with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
+    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('custom_utils/super_model_m.pkl'), 'rb') as f:
+    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
@@ -399,162 +399,139 @@ def get_shapedirs():
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/cloth3d/train')
-    parser.add_argument("-s", "--start", help = "start of range", default=70500)
-    parser.add_argument("-e", "--end", help = "end of range", default=90000)
+    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, 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("-sk", "--skirt", help = "end of range", default=0)
     args = parser.parse_args()
 
-    F = np.load("custom_utils/smpl_faces" + ("_skirt" if args.skirt == 1 else "") + ".npy")
-    uv_map = np.load('custom_utils/stretched_uv_map_' + ("skirt" if args.skirt == 1 else "3") + '.npy')
-    m2uv = np.load('custom_utils/m2uv' + ("_skirt" if args.skirt == 1 else "") + '.npy', allow_pickle=True).item()
-    m2uv = m2uv_array(m2uv)
-    uv_mask = np.load('custom_utils/mask_256x512' + ("_skirt" if args.skirt == 1 else "") + '.npy')
-    curvature_label_type = '6_25 '
+    F = {'non-skirt': np.load("smpl/smpl_faces.npy"), 'skirt': np.load("smpl/smpl_faces_skirt.npy")}
+    uv_map = {'non-skirt': np.load('custom_utils/stretched_uv_map_3.npy'), 'skirt': np.load('custom_utils/stretched_uv_map_skirt.npy')}
+    m2uv = {'non-skirt': np.load('custom_utils/m2uv.npy', allow_pickle=True).item(), 'skirt': np.load('custom_utils/m2uv_skirt.npy', allow_pickle=True).item()}
+    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_mask = (cv2.resize(uv_mask * 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 = cv2.dilate(uv_mask.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 = dict()
+    uv_map_discreet_unique = dict()
+    uv_counter = dict()
+    rbf_dest = dict()
+    uv_map_mask = dict()
+    for k,v in uv_mask.items():
+        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 = np.floor(uv_map * uv_res).astype(np.int32)
-    uv_map_mask = unique_uv_map_mask(uv_map_discreet)
+        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 = np.zeros((uv_map.shape[0]))
-    uv_counter[m2uv[:, 0]] = 1
-    uv_counter[m2uv[:, 1]] += 1
-    uv_counter = uv_counter[uv_map_mask]
+        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 = uv_map_discreet[uv_map_mask, :]
+        uv_map_discreet_unique[k] = uv_map_discreet[k][uv_map_mask[k], :]
 
-    rbf_dest = np.ones((uv_res * 2, uv_res))
-    rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0
-    rbf_dest[np.invert(uv_mask[:,:,0])] = 0
+        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
 
 
     #########################################################################################
 
-    use_no_orient = False
-    no_shape = False
-    is_skirt = True if args.skirt == 1 else False
     folder = args.folder
-    neighbourhood = 25
-    frames_file = 'frames_all.txt'
-    data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_curvature', 'body_posed'] # 'mesh_unposed',
+    data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
 
-    print("skirt" if args.skirt == 1 else "not skirt")
-    print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -skirt {args.skirt}')
-
-    skirt_modifier = '_skirt' if is_skirt else ''
-    orient_mod = '_no_orient' if use_no_orient else ''
-    shape_mod = '_no_shape' if no_shape else ''
-
-    frames = []
-    with open(os.path.join(folder, frames_file), 'r') as f:
-        lines = f.readlines()
-        lines = [line.strip() for line in lines]
-    frames = lines[int(args.start):int(args.end)]
-    ids = [frame.split()[0] for frame in frames]
-    ids = sorted(list(set(ids)))
+    ids = sorted(os.listdir(folder + '/mesh_body_posed_no_orient'))[args.start:args.end]
+    ids = [id.split('.')[0].split('_')[0] for id in ids]
 
     shapedirs = get_shapedirs()
 
     root = '/'.join(folder.split('/')[:-1])
     type = folder.split('/')[-1]
-    # h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')
-    # h5file = tb.open_file(os.path.join(root, f'{type}_extras_uv.h5'), mode='a')
 
-    extras_folder = os.path.join(root, f'{type}_extras')
-    os.makedirs(extras_folder, exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_posed' + orient_mod), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_static'), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_curvature'), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_unposed'), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_body_posed' + orient_mod), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_mask'), exist_ok=True)
+    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)
 
 
     for id in tqdm.tqdm(ids):
+        garments = glob.glob(os.path.join(folder, 'mesh_mask', f'{id}_*'))
+        garments = [garment.split('/')[-1].split('.')[0].split('_')[-1] for garment in garments]
 
+        # if os.path.exists(os.path.join(folder, 'uv_static',  f'{id}_{garment}.npz')):
+        #     continue
 
-        info = loadInfo(os.path.join(folder, id, 'info.mat'))
-        garments = list(info['outfit'].keys())
-        assert info['poses'].shape[1] <= 300
-
-        # put all this in a function and call for all garments
-
-        # also masking is not the same as just filling the empty values with neighbouring
-        # values, at least the plots are different
-
-        for garment in garments:
-        # garment = garments[0]
-
-            # if all the uv maps have been generated, skip this id
-            # if all([os.path.exists(os.path.join(extras_folder, 'uv_body_posed' + orient_mod,  f'{id}_{i}_{garment}.npz')) for i in range(info['poses'].shape[1])]):
+        for k, garment in enumerate(garments):
+            # if not ('Dress' in garments or 'Skirt' in garments):
+            #     if k != 1:
+            #         files = glob.glob(os.path.join(folder, 'uv_body_posed_no_orient',  f'{id}_*_{garment}.npz'))
+            #         for file in files:
+            #             shutil.move(file, '_'.join(file.split('_')[:-1]) + '.npz')
             #     continue
 
-            mask = read_garment_mask(folder, id, garment)
-
-            if 'mesh_posed' in data_to_convert:
-                # mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
-                mesh_posed = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], poses=info['poses'], deorient=use_no_orient, unshape=no_shape)
+            mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}_{garment}.npy'))
+            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:
-                raise NotImplementedError()
-                mesh_unposed = read_unposed_garment(folder, id, garment, mask=mask, shape=info['shape'], unshape=no_shape)
+                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 = read_garment_template(folder, id, garment, mask)
-            if 'mesh_curvature' in data_to_convert:
-                mesh_curvature = read_curvatures(folder + '_extras', id, garment, count=info['poses'].shape[1])
-            if 'body_posed' in data_to_convert:
-                body_posed = read_posed_body(folder, id, info['poses'], use_no_orient)
-
+                mesh_static = np.zeros((14475, 3), dtype=np.float32)
+                mesh_static[mask, :] = np.load(os.path.join(folder, 'mesh_static', f'{id}_{garment}.npy'))
+            if 'mesh_normals' in data_to_convert:
+                mesh_normals = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
+                mesh_normals[:,mask, :] = readPC2(os.path.join(folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V']
+            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()
+            type = 'skirt' if garment == 'Dress' or garment == 'Skirt' else 'non-skirt'
 
-            uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
-            uv_static = convert_to_uv_batch(uv_res, m2uv, uv_map, 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.astype(np.float32)
+            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, uv_map, mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+                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:
-                raise NotImplementedError()
-                uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
-            if 'mesh_curvature' in data_to_convert:
-                uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
-            if 'body_posed' in data_to_convert:
-                uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv, uv_map, body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
-
-
+                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)
 
-            # mesh_static
-            # save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
-            # save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
-            np.savez_compressed(os.path.join(extras_folder, 'uv_static',  f'{id}_{garment}.npz'), uv_static=uv_static)
-            np.savez_compressed(os.path.join(extras_folder, 'uv_mask',  f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
+            # 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(extras_folder, 'uv_posed' + orient_mod,  f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
-                    # save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
+                    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:
-                raise NotImplementedError()
-                for i in range(uv_unposed.shape[0]):
-                    save_to_h5file(h5file, id, i, garment, uv_unposed[i, :, :, :], 'uv_unposed' + orient_mod + shape_mod, 'float32')
-            if 'mesh_curvature' in data_to_convert:
-                for i in range(len(uv_curvature)):
-                    np.savez_compressed(os.path.join(extras_folder, 'uv_curvature',  f'{id}_{i}_{garment}.npz'), uv_curvature=uv_curvature[i].astype(np.uint8))
-                    # save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
-            if 'body_posed' 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 in range(len(uv_normals)):
+                    np.savez_compressed(os.path.join(folder, 'uv_normals',  f'{id}_{i}_{garment}.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(extras_folder, 'uv_body_posed' + orient_mod,  f'{id}_{i}_{garment}.npz'), uv_body_posed=uv_body_posed[i])
-                    # save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
+                    np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient',  f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
             
-            pass
-
-
-    h5file_curv.close()
+            # pass

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

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,69 +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 = 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
 
 
-        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):
@@ -214,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 
@@ -274,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)
 

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()

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()

data/investigate.py

@@ -4,11 +4,26 @@ import random
 import numpy as np
 import plotly.graph_objects as go
 from pymeshfix import MeshFix
+import trimesh
 
 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, read_curvature_frame, writeOBJ
-from custom_utils.custom_utils import _maskFaces
+# from custom_utils.custom_utils import _maskFaces
 
+def _maskFaces(V, F, mask, return_unfixed=False):
+    V_mask = mask
+    F_mask = np.take(V_mask, F.reshape(-1), axis=0)
+    F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
+
+    V_ind = np.zeros(V_mask.shape)
+    V_ind[V_mask] = np.arange(V_mask.sum())
+    F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
+    V_new = V[V_mask, :]
+
+    if return_unfixed:
+        return V_new, F_new, F[F_mask, :]
+    else:
+        return V_new, F_new
 
 def get_frames(folder, frames_file):
     frames = []
@@ -133,7 +148,7 @@ def check(folder, frame):
     v = mesh_posed[mask,:]
     is_skirt = (garment == 'Dress' or garment == 'Skirt')
     skirt_mod = '_skirt' if is_skirt else ''
-    F = np.load(f'custom_utils/smpl_faces{skirt_mod}.npy').astype(np.int32)
+    F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
     _, faces, unfixed = _maskFaces(np.zeros((14475,3)), F, mask, True)
 
     meshfix = MeshFix(v, faces)
@@ -180,7 +195,22 @@ def check(folder, frame):
 
     faces_3 = mask_faces(F, mask_2, mask_2_raw, repeated, mesh_posed)
     # writeOBJ('temp/investigate_3.obj', v_2, faces_3)
-    plot(v_2, faces_3)
+    # plot(v_2, faces_3)
+
+    # save stl
+    # mesh = trimesh.Trimesh(vertices=v_2, faces=faces_3, process=False, validate=False)
+    # mesh.export('temp/investigate_00761.stl')
+
+    # # read stl file
+    # mesh = trimesh.load('temp/fixed_00761.stl', process=False, validate=False)
+    # mesh = trimesh.load('temp/fixed_00761.ply', process=False, validate=False)
+
+    # plotly plot of mesh with v_2 vertices and faces from trimesh mesh
+    # faces_4 = faces_3 #np.asarray(mesh.faces)
+    # go.Figure(data = [go.Mesh3d(x=v_2[:, 0], y=v_2[:, 1], z=v_2[:, 2], i=faces_4[:, 0], j=faces_4[:, 1], k=faces_4[:, 2],)
+    #                      ]).show()
+    
+    
 
 
     # plot repeated
@@ -189,9 +219,9 @@ def check(folder, frame):
         r.extend(rr)
     x = mesh_posed
     y = mesh_posed[r,:]
-    go.Figure(data = [go.Scatter3d(x=x[:, 0], y=x[:, 1], z=x[:, 2], marker=dict(size=3, line=dict(width=0),), mode='markers'),
-                      go.Scatter3d(x=y[:, 0], y=y[:, 1], z=y[:, 2], marker=dict(size=5, line=dict(width=0),), mode='markers'),
-                     ]).show()
+    # go.Figure(data = [go.Scatter3d(x=x[:, 0], y=x[:, 1], z=x[:, 2], marker=dict(size=3, line=dict(width=0),), mode='markers'),
+    #                   go.Scatter3d(x=y[:, 0], y=y[:, 1], z=y[:, 2], marker=dict(size=5, line=dict(width=0),), mode='markers'),
+    #                  ]).show()
 
 
     # unq, idx, count = np.unique(v, axis=0, return_counts=True, return_index=True)
@@ -207,30 +237,49 @@ def check(folder, frame):
     # faces_4 = mask_faces(faces.astype(np.int64), mask_3, mask_2_raw, repeated, v)
     # plot(v[mask_3,:], faces_4)
 
-
+    return mask_2_raw, faces_3
     pass
-    
+
+
+def remask(folder, frame, mask, out_folder):
+    id, frame_nr = frame.split('_')
+    frame_nr = int(frame_nr)
+
+    if frame_nr == 0:
+        mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static']
+        mesh_static_new = np.zeros((14475,3), dtype=np.float32)
+        mesh_static_new[mask,:] = mesh_static[mask,:]
+        np.savez_compressed(os.path.join(out_folder, 'mesh_static', f'{id}.npz'), mesh_static=mesh_static_new)
+
+    # mesh_posed = read_posed_garment_frame(folder, id, garment, frame_nr, mask=None, shape=None, pose=info['poses'][:,frame_nr], deorient=True, unshape=None)
+    mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{frame_nr}.npz'))['mesh_posed_no_orient']
+    mesh_posed_new = np.zeros((14475,3), dtype=np.float32)
+    mesh_posed_new[mask,:] = mesh_posed[mask,:]
+    np.savez_compressed(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{frame_nr}.npz'), mesh_posed_no_orient=mesh_posed)
+
 
 
 if __name__ == "__main__":
-    folder = '/home/workspace/data/dataset/cloth3d/train'
-    frames = get_frames(folder, 'frames_all.txt')
+    folder = '/data-local/127/train'
+    frames = [f'00761_{i}' for i in range(300)]
     root = '/'.join(folder.split('/')[:-1])
     type = folder.split('/')[-1]
 
     use_no_orient = False
     no_shape = False
-    is_skirt = True
+    is_skirt = False
 
     skirt_modifier = '_skirt' if is_skirt else ''
     orient_mod = '_no_orient' if use_no_orient else ''
     shape_mod = '_no_shape' if no_shape else ''
 
+    mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
+    faces = np.load('/data-net/dataset/00761/faces/00761.npy')
 
-    while True:
-        frame = random.choice(frames)
-        frame = '05936 0 Dress'
+    for frame in frames:
+        # frame = random.choice(frames)
+        # frame = '00761 0 Tshirt'
 
-        check(folder, frame)
+        remask(folder, frame, mask, '/data-net/dataset/00761')
 
         pass

data/preprocessing.py

@@ -1,22 +1,39 @@
 import os
 import sys
 
+import argparse
+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", 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"
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+
 import math
 import shutil
 import warnings
 import tables as tb
 import numpy as np
 import scipy.io as sio
-
+import trimesh as tm
 from tqdm import tqdm
 
 from tables import NaturalNameWarning
 warnings.filterwarnings('ignore', category=NaturalNameWarning)
+import plotly.graph_objects as go
+
+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
+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
 
 
 
@@ -132,12 +149,623 @@ def create_extras_file(folder):
     h5file.close()
 
 
+def deorient(verts, pose):
+    rodr = smpl.rodrigues(torch.tensor(pose.T, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).cpu().numpy()
+    rotated = np.zeros_like(verts)
+    for frame in range(verts.shape[0]):
+        r = R.from_matrix(rodr[frame, 0, :, :])
+        rotated[frame] = r.apply(verts[frame], inverse=True)
+
+    return rotated
+
+
+def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_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)
+
+    for id in tqdm(ids):
+        # id = '01467'
+
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        info = loadInfo(os.path.join(folder, id, 'info.mat'))
+        
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue
+        
+            verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
+            # if info['poses'].shape[1] == verts.shape[0]:
+            #     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
+            # import plotly.graph_objects as go
+            # frame = 5
+            # 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 = 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)
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        info = loadInfo(os.path.join(folder, id, 'info.mat'))
+        
+        shape = np.concatenate([info['shape'], np.asarray([info['gender']])]).astype(np.float32)
+        np.save(os.path.join(out_folder, 'pose', f'{id}.npy'), info['poses'].T.astype(np.float32))
+        np.save(os.path.join(out_folder, 'shape', f'{id}.npy'), shape)
+
+
+def create_mesh_static(folder, out_folder):
+    ids = sorted(os.listdir(folder))
+    os.makedirs(os.path.join(out_folder, 'mesh_static'), exist_ok=True)
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue
+
+            verts_static, _, _, _ = readOBJ(os.path.join(garment_dir, f'{garment}.obj'))
+            reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
+            verts_static = (verts_static[reg_data['F'], :] * reg_data['W'][:, :, None]).sum(axis=1).astype(np.float32)
+
+            # rotate 90 degrees around x axis
+            verts_static = np.matmul(verts_static, np.array([[1,0,0],[0,0,-1],[0,1,0]]))
+
+            # fig = go.Figure(data=[go.Scatter3d(x=verts_static[:,0], y=verts_static[:,2], z=verts_static[:,1], mode='markers', marker=dict(size=1))])
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+            # fig.show()
+
+            np.save(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'), verts_static.astype(np.float32))
+
+
+
+def create_mesh_body_posed_no_orient(folder, out_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):
+        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'))
+    
+
+        pose, shape, gender = info['poses'], info['shape'], info['gender']
+        pose = np.transpose(pose, [1, 0]).astype(np.float32)
+
+        pose[:, :3] = 0
+        beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
+
+        T = smpl(torch.tensor(pose, device=smpl.device), torch.tensor(np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), device=smpl.device))[0]
+        T = T.cpu().numpy()
+
+        writePC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), T, True)
+
+        # verts = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
+        # # plot as scatter plot with plotly
+        # import plotly.graph_objects as go
+        # frame = 5
+        # 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()
+
+        pass
+
+
+def calculate_normals(vertices, faces):
+    normals = []
+    for i in range(vertices.shape[0]):
+        mesh = tm.Trimesh(vertices=vertices[i], faces=faces, process=False, validate=False)
+        normals.append(mesh.vertex_normals)
+    return np.stack(normals)
+
+
+def create_mesh_normals(folder, out_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)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+        
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue
+
+            verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
+            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))
+
+            reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
+            normals = (normals[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
+
+            writePC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), normals, True)
+
+            # 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[frame,:,0],
+            #     y=verts[frame,:,2],
+            #     z=verts[frame,:,1],
+            #     mode='markers',
+            #     marker=dict(
+            #         size=1,
+            #         color=normals[frame,:,:],
+            #         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[frame,:,0],
+            #     y=moved_verts[frame,:,2],
+            #     z=moved_verts[frame,:,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 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 = sorted(os.listdir(folder))
+
+    for id in ids:
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        
+        for garment in garments:
+            vertices = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
+            mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
+
+            if mask.sum() == vertices.shape[1]:
+                print(id, garment)
+
+            pass
+
+
+def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1):
+    ids = sorted(os.listdir(folder))[start:end]
+    os.makedirs(os.path.join(out_folder, 'mesh_unposed_no_shape'), exist_ok=True)
+    F = np.load('smpl/smpl_faces.npy')
+    # F_skirt = np.load('smpl/smpl_faces_skirt.npy')
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        # info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+        pose = np.load(os.path.join(out_folder, 'pose', f'{id}.npy'))
+        shape = np.tile(np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:], [pose.shape[0],1])
+
+        pose = torch.tensor(pose, dtype=torch.float32, device='cuda')
+        shape = torch.tensor(shape, dtype=torch.float32, device='cuda') 
+        
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue    
+
+            mesh_static = np.load(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'))
+            mesh_posed = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
+            mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
+            # mask_idx = np.where(mask)[0]
+
+            pose[:,:3] = 0
+            # pose = pose[:10]
+            # shape = shape[:10]
+            # mesh_posed = mesh_posed[:10]
+
+            unposed = np.zeros((14475, 3), dtype=np.float32)
+            unposed[mask] = mesh_static
+
+            unposed = torch.tensor(np.tile(unposed[None,:,:], [mesh_posed.shape[0],1,1]), dtype=torch.float32, device='cuda')
+            unposed.requires_grad = True
+            mesh_posed = torch.tensor(mesh_posed, dtype=torch.float32, device='cuda')
+
+            optimizer = torch.optim.Adam([unposed], lr=0.01)
+
+            # verts_1 = unposed[0].cpu().detach().numpy()[mask]
+            # _, faces = _maskFaces(np.zeros((14475, 3)), F, mask)
+            # fig = go.Figure(data=[go.Mesh3d(x=verts_1[:,0], y=verts_1[:,2], z=verts_1[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
+            # fig.update_layout(scene = dict(
+            #     xaxis = dict(range=[-1,1],),
+            #                 yaxis = dict(range=[-1,1],),
+            #                 zaxis = dict(range=[-1,1],),),)
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.show()
+
+            # verts_2 = mesh_posed[0].cpu().detach().numpy()
+            # fig = go.Figure(data=[go.Mesh3d(x=verts_2[:,0], y=verts_2[:,2], z=verts_2[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
+            # fig.update_layout(scene = dict(
+            #     xaxis = dict(range=[-1,1],),
+            #                 yaxis = dict(range=[-1,1],),
+            #                 zaxis = dict(range=[-1,1],),),)
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.show()
+
+            # zero_body = smpl(torch.zeros_like(pose), torch.zeros_like(shape))[0]
+
+            # verts_3 = smpl(pose, shape, zero_body)[0][0].cpu().detach().numpy()
+            # # verts_3 = zero_body[0].cpu().detach().numpy()
+            # faces_3 = F
+            # fig = go.Figure(data=[go.Mesh3d(x=verts_3[:,0], y=verts_3[:,2], z=verts_3[:,1], i=faces_3[:,0], j=faces_3[:,1], k=faces_3[:,2], color='lightpink', opacity=1)])
+            # fig.update_layout(scene = dict(
+            #     xaxis = dict(range=[-1,1],),
+            #                 yaxis = dict(range=[-1,1],),
+            #                 zaxis = dict(range=[-1,1],),),)
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.show()
+
+            for i in range(5000):
+                    optimizer.zero_grad()
+
+                    pred = smpl(pose, shape, unposed)[0]
+                    loss : torch.Tensor = ((pred[:,mask,:] - mesh_posed)**2).sum()
+
+                    # Backward pass
+                    loss.backward()
+                    optimizer.step()
+
+                    with torch.inference_mode():
+                        euclidean_error = torch.sqrt(((pred[:,mask,:] - mesh_posed)**2).sum(-1)).mean()
+
+                    # Print the loss
+                    if i % 50 == 0:
+                        print(f"it: {i} loss: {loss:.8f} euclidean_error: {euclidean_error:.8f}")
+ 
+                    if euclidean_error < 0.0000001:
+                        break
+
+                    if euclidean_error < 0.0000005 and i > 2000:
+                        break
+
+                    if euclidean_error < 0.000001 and i > 4000:
+                        break
+
+            unposed = unposed.detach().cpu().numpy()
+
+            gender = int(shape[0].cpu().numpy()[-1] > 0)
+            shape_offsets = torch.tensordot(shape[:,:10], smpl.shapedirs[gender], dims=[[1], [2]])
+            shape_offsets = shape_offsets.cpu().numpy()[:,smpl.to_keep,:]
+
+            unshaped = (unposed - shape_offsets)[:,mask,:]
+
+            writePC2(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True)
+
+            # verts = unposed[:,mask,:]
+            # verts = unshaped
+            # _, faces = _maskFaces(np.zeros((14475, 3)), F, mask)
+            # frame = 9
+            # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
+            # fig.update_layout(scene = dict(
+            #     xaxis = dict(range=[-1,1],),
+            #                 yaxis = dict(range=[-1,1],),
+            #                 zaxis = dict(range=[-1,1],),),)
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.show()
+
+            # verts = pred.detach().cpu().numpy()[:,mask,:]
+            # frame = 9
+            # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
+            # fig.update_layout(scene = dict(
+            #     xaxis = dict(range=[-1,1],),
+            #                 yaxis = dict(range=[-1,1],),
+            #                 zaxis = dict(range=[-1,1],),),)
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.show()
+            pass
+
+
+def create_frames_file(folder, out_folder, split=0.9):
+    ids = os.listdir(folder)
+
+    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 = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V']
+    mesh_body_posed = mesh_body_posed[int(mesh_body_posed.shape[0]/2)+frame.frame_nr]
+
+    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
+
+
+def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'):
+    from custom_utils.custom_utils import UV2Mesh
+    import plotly.express as ex
+    import plotly.graph_objects as go
+
+    frames = []
+    with open(os.path.join(out_folder, frames_file), 'r') as f:
+        for line in f:
+            line = line.strip()
+            id, frame_nr, garment = line.split(' ')
+            frames.append(Frame(id, int(frame_nr), garment))
+
+    uv_sum = np.zeros((512, 256, 3), dtype=np.float32)
+    mask_sum = np.zeros((512, 256), dtype=np.float32)
+
+    for frame in tqdm(frames):
+        uv = np.load(os.path.join(out_folder, 'uv_unposed_no_shape', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_unposed']
+        mask = np.load(os.path.join(out_folder, 'uv_mask', f'{frame.id}_{frame.garment}.npz'))['uv_mask']
+        uv_sum += uv
+        mask_sum += mask
+
+    mask_sum[mask_sum == 0] = 1
+    avg_garment = uv_sum / mask_sum[:, :, None]
+
+    avg_mesh = UV2Mesh()(torch.from_numpy(avg_garment.transpose(2,0,1))[None,:,:,:]).numpy()[0]
+    F = np.load('smpl/smpl_faces.npy')
+
+    fig = go.Figure(data=[go.Mesh3d(x=avg_mesh[:,0], y=avg_mesh[:,2], z=avg_mesh[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightpink', opacity=1)])
+    fig.update_layout(scene = dict(
+        xaxis = dict(range=[-1,1],),
+                    yaxis = dict(range=[-1,1],),
+                    zaxis = dict(range=[-1,1],),),)
+    fig.update_layout(scene_aspectmode='cube')
+    fig.show()
+
+    np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment)
+
+
 
 if __name__ == '__main__':
 
-    folder = '/home/workspace/data/dataset/cloth3d/train'
+
+    # 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'
+
+    smpl = SMPL(path='smpl', device='cuda')
 
     # clean_dataset(folder)
     # generate_frames_file(folder, 'frames_all.txt')
     # split_data(folder, 0.9)
-    create_extras_file(folder)
+    # 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)
+    calculate_average(folder, out_folder)
+
+    # 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)

data/registration.py

@@ -0,0 +1,461 @@
+import os
+import sys
+import random
+import numpy as np
+import plotly.graph_objects as go
+import torch
+from torch.nn.functional import mse_loss
+from pytorch3d.structures import Meshes
+from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
+from pymeshfix import MeshFix
+
+# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+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,
+    read_curvature_frame,
+    writeOBJ,
+    readOBJ,
+    quads2tris
+)
+from smpl.smpl_torch import SMPL
+# from custom_utils.custom_utils import _maskFaces
+
+device = 'cuda'
+smpl = SMPL(device, 'smpl')
+
+
+import torch.nn.functional as F
+from pytorch3d.structures import Meshes
+
+def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
+    """
+    Given two PyTorch3D meshes and a set of vertex indices for mesh1,
+    find the closest vertices from mesh2 to these vertices.
+
+    Args:
+        mesh1: A PyTorch3D Meshes object representing the first mesh.
+        mesh2: A PyTorch3D Meshes object representing the second mesh.
+        vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
+
+    Returns:
+        A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
+    """
+    # Extract the vertices of mesh1 and mesh2.
+    vertices1 = mesh1.verts_packed()
+    vertices2 = mesh2.verts_packed()
+
+    # Extract the vertices of mesh1 that correspond to the given indices.
+    vertices1_selected = vertices1[vertex_indices]
+
+    # Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
+    distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
+
+    # Find the indices of the closest vertices in mesh2.
+    closest_indices = torch.argmin(distances, dim=1)
+
+    return closest_indices
+
+
+def _maskFaces(V, F, mask, return_unfixed=False):
+    V_mask = mask
+    F_mask = np.take(V_mask, F.reshape(-1), axis=0)
+    F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
+
+    V_ind = np.zeros(V_mask.shape)
+    V_ind[V_mask] = np.arange(V_mask.sum())
+    F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
+    V_new = V[V_mask, :]
+
+    if return_unfixed:
+        return V_new, F_new, F[F_mask, :]
+    else:
+        return V_new, F_new
+
+def get_frames(folder, frames_file):
+    frames = []
+    with open(os.path.join(folder, frames_file), 'r') as f:
+        for line in f:
+            frames.append(line.strip())
+    return frames
+
+
+def plot(mesh, faces):
+    fig = go.Figure()
+
+    fig.add_trace(go.Mesh3d(
+                x=mesh[:,0],
+                y=mesh[:,1],
+                z=mesh[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),)
+    
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                    )
+    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                yaxis = dict(nticks=4, range=[-1,1],),
+                zaxis = dict(nticks=4, range=[-1,1],),)
+    fig.update_layout(scene = helper,
+                        )
+
+    fig.show()
+
+
+def build_lookup_table(repeated):
+    # iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
+
+    lookup = {}
+    for i in range(len(repeated)):
+        for j in range(len(repeated[i])):
+            if len(repeated[i]) == 1:
+                continue
+            lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
+    return lookup
+
+
+from data.cloth3d_io import readPC2Frame
+import scipy.io as sio
+import scipy
+import scipy.spatial as spatial
+
+def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
+    pc16_path = os.path.join(folder, id, garment + '.pc16')
+    verts = readPC2Frame(pc16_path, frame, True)
+
+    map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
+
+    if deorient:
+        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
+        verts = R.apply(verts)
+
+    return verts, map
+
+def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
+    pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
+    verts = readPC2Frame(pc16_path, frame, True)
+
+    if deorient:
+        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
+        verts = R.apply(verts)
+
+    return verts
+
+def edge_loss_custom(x, y, edges):
+    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
+    x_e = torch.sum(x_e ** 2, -1)
+    x_e = torch.sqrt(x_e + 1e-9)
+    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
+    y_e = torch.sum(y_e ** 2, -1)
+    y_e = torch.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    loss = torch.mean(d_e ** 2, dim=-1)[0]
+    return loss
+
+
+def crotch(mesh, body, crotch_idx, nr=75):
+    crotch_middle_on_body = crotch_idx
+    # create kdtree for body and mesh
+    body_tree = spatial.cKDTree(body)
+    mesh_tree = spatial.cKDTree(mesh)
+    # find closest point on mesh from crotch middle on body along the z axis
+
+    crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
+
+    crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
+    crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
+
+    body_verts = body[crotch_body]
+    mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
+    mesh_verts = spatial.cKDTree(mesh_verts)
+    closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
+    crotch_mesh = crotch_mesh[closest_mesh_verts]
+
+    return crotch_body, crotch_mesh
+
+
+def find_closest_point_along_z(point, verts, eps=0.01):
+    min_dist = np.inf
+    closest_point = None
+    for i, v in enumerate(verts):
+        dist = np.linalg.norm(point - v)
+        if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
+            min_dist = dist
+            closest_point = i
+    return closest_point
+
+
+def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
+    """
+    Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
+
+    Args:
+        mesh: A PyTorch3D Meshes object representing the mesh.
+
+    Returns:
+        A tensor of shape (E,) representing the lengths of all edges in the mesh,
+        where E is the total number of edges in the mesh.
+    """
+    # Extract the vertices and faces of the mesh.
+    vertices = mesh.verts_packed()
+    faces = mesh.faces_packed()
+
+    # Compute the pairwise distances between all pairs of vertices that form an edge.
+    edge_indices = faces.view(-1, 3)
+    edge_vertices = vertices[edge_indices]
+    edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
+
+    return edge_lengths
+
+
+def register(folder, frame):
+    id, frame_nr, garment = frame.split()
+    frame_nr = int(frame_nr)
+    info = loadInfo(os.path.join(folder, id, 'info.mat'))
+    raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
+    # convert it to indices
+    mask = raw_mask > 0
+    mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
+    body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
+    template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
+    faces = quads2tris(faces)
+    
+    # boundaries = raw_mask == 2
+    # # invert map and use it to invert boudaries
+    # map_inv = np.zeros(map.max()+1, dtype=np.int32)
+    # map_inv[map] = np.arange(len(map))
+    # boundaries = boundaries[map_inv]
+    # # get the boundary indices
+    
+    # mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
+
+
+    
+    # pose = torch.zeros(1, 72).to('cuda')
+    shape = torch.zeros(1, 11).to(device)
+    shape[:, :10] = torch.tensor(info['shape']).to(device)
+    shape[:, -1] = torch.tensor(info['gender']).to(device)
+    # for star pose
+    # pose[:, 5] = 3.1415/4
+    # pose[:,8] = -3.1415/4
+    STAR_POSE = np.zeros((24, 3), np.float32)
+    STAR_POSE[1,2] = .15
+    STAR_POSE[2,2] = -.15
+    pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
+    star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
+    # rotate star body with 90 degrees around x axis
+    R = spatial.transform.Rotation.from_euler('x', np.pi/2)
+    star_body = R.apply(star_body).astype(np.float32)
+
+    is_skirt = (garment == 'Dress' or garment == 'Skirt')
+    skirt_mod = '_skirt' if is_skirt else ''
+    F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
+
+    ##################### mask body #####################
+    new_crotch = 4729
+
+    # star_body = star_body[mask]
+    # new_crotch = np.zeros((14475), bool)
+    # new_crotch[4729] = True
+    # new_crotch = new_crotch[mask]
+    # new_crotch = np.where(new_crotch)[0][0]
+    # _, F = _maskFaces(np.zeros((14475, 3)), F, mask)
+
+    # plot(template, faces)
+    # v = mesh_posed[mask,:]
+    
+    ##################### boundaries #####################
+    # meshfix = MeshFix(template, faces)
+    # holes = meshfix.extract_holes()
+    # new_border = []
+    # for border in np.array(holes.points):
+    #     # find the index of the vertex in the original mesh
+    #     ind = np.argwhere(np.all(template == border, axis=1))
+    #     new_border.extend(ind.ravel())
+    # mesh1_boundaries = torch.tensor(new_border).to(device)
+    # np.save('temp/body_boundaries.npy', new_border)
+    mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
+
+    # meshfix = MeshFix(template, faces)
+    # holes = meshfix.extract_holes()
+    # new_border = []
+    # for border in np.array(holes.points):
+    #     # find the index of the vertex in the original mesh
+    #     ind = np.argwhere(np.all(template == border, axis=1))
+    #     new_border.extend(ind.ravel())
+    # mesh1_boundaries = torch.tensor(new_border).to(device)
+    # np.save('temp/body_boundaries.npy', new_border)
+    # mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
+
+    ##################### crotch #####################
+    crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
+
+
+
+
+    # export star_body as obj
+    # writeOBJ('temp/star_body.obj', star_body, F)
+
+
+    template = torch.tensor(template).to(device)
+    template_orig = template.clone()
+    faces = torch.tensor(faces).to(device)
+    star_body = torch.tensor(star_body).to(device)
+    F = torch.tensor(F).to(device)
+
+    # verts = template.cpu().numpy()
+    # star_body = star_body.cpu().numpy()
+    # faces = faces.detach().cpu().numpy()
+    # F = F.detach().cpu().numpy()
+    # fig = go.Figure(data=[
+    #     go.Mesh3d(
+    #         x=verts[:,0], y=verts[:,1], z=verts[:,2],
+    #         i=faces[:,0], j=faces[:,1], k=faces[:,2],
+    #         color='lightpink', opacity=0.50),
+    #     go.Mesh3d(
+    #         x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
+    #         i=F[:,0], j=F[:,1], k=F[:,2],
+    #         color='lightblue', opacity=0.50),
+    #     go.Scatter3d(
+    #         x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
+    #         mode='markers', marker=dict(size=2, color='red')),
+    #     go.Scatter3d(
+    #         x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
+    #         mode='markers', marker=dict(size=2, color='blue')),
+    # ])
+    # # make it cuve and limit to -1 and -1
+    # fig.update_layout(scene = dict(
+    #     xaxis = dict(range=[-1,1],),
+    #     yaxis = dict(range=[-1,1],),
+    #     zaxis = dict(range=[-1,1],),),)
+    # fig.show()
+
+    # Define the meshes to optimize
+    mesh1 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
+    mesh2 = Meshes(verts=star_body[None,:,:], faces=F[None,:,:])
+
+    mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
+
+    avg_edge_length = mesh_edge_lengths(mesh1).mean()
+
+    # Define the optimizer and learning rate
+    mesh1.verts_packed().requires_grad = True
+    optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
+    # optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
+
+    # Define the loss functions
+    def laplacian_smoothing_loss(mesh):
+        return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
+
+    def edge_loss(mesh):
+        return mesh_edge_loss(mesh, target_length=avg_edge_length)
+
+    def normal_consistency_loss(mesh):
+        return mesh_normal_consistency(mesh)
+
+    # Optimization loop
+    for i in range(10000):
+        optimizer.zero_grad()
+
+        # find boundaries with the find holes function?
+
+        # Forward pass
+        mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
+        ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
+        bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
+        # crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
+        crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
+        loss: torch.Tensor = ch_loss * 10 +  (bound_ch_loss + crotch_ch_loss) * 5
+        # loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
+        loss += laplacian_smoothing_loss(mesh1_out) * 0.05
+        loss += edge_loss(mesh1_out) * 100
+        # loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
+        loss += normal_consistency_loss(mesh1_out) * 0.1
+
+        # Backward pass
+        loss.backward()
+        optimizer.step()
+
+        # Print the loss
+        if i % 10 == 0:
+            print(f"Loss: {loss.item()} chamfer: {ch_loss.item()} ")
+
+    # create a plotly plot with the body and the mesh1_out
+    verts = mesh1_out.verts_packed().detach().cpu().numpy()
+    star_body = star_body.detach().cpu().numpy()
+    faces = faces.detach().cpu().numpy()
+    F = F.detach().cpu().numpy()
+    mesh1_boundaries = mesh1_boundaries.detach().cpu().numpy()
+    mesh2_boundaries = mesh2_boundaries.detach().cpu().numpy()
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            x=verts[:,0], y=verts[:,1], z=verts[:,2],
+            i=faces[:,0], j=faces[:,1], k=faces[:,2],
+            color='lightpink', opacity=0.80),
+        go.Mesh3d(
+            x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
+            i=F[:,0], j=F[:,1], k=F[:,2],
+            color='lightblue', opacity=0.50),
+        go.Scatter3d(
+            x=verts[mesh1_boundaries,0], y=verts[mesh1_boundaries,1], z=verts[mesh1_boundaries,2],
+            mode='markers', marker=dict(size=2, color='red')),
+        go.Scatter3d(
+            x=star_body[mesh2_boundaries,0], y=star_body[mesh2_boundaries,1], z=star_body[mesh2_boundaries,2],
+            mode='markers', marker=dict(size=2, color='blue')),   
+        go.Scatter3d(
+            x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
+            mode='markers', marker=dict(size=2, color='red')),
+        go.Scatter3d(
+            x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
+            mode='markers', marker=dict(size=2, color='blue')),
+    ])
+    # make it cuve and limit to -1 and -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/code/cloth3d/train'
+    # folder = r'D:\Work\code\cloth3d\train'
+    # frames = [f'00761_{i}' for i in range(300)]
+    # folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
+    frames = get_frames(folder, 'frames_all.txt')
+    root = '/'.join(folder.split('/')[:-1])
+    type = folder.split('/')[-1]
+
+    use_no_orient = False
+    no_shape = False
+    is_skirt = False
+
+    skirt_modifier = '_skirt' if is_skirt else ''
+    orient_mod = '_no_orient' if use_no_orient else ''
+    shape_mod = '_no_shape' if no_shape else ''
+
+    # mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
+    # faces = np.load('/data-net/dataset/00761/faces/00761.npy')
+
+    for frame in frames:
+        frame = random.choice(frames)
+        # frame = '00761 0 Tshirt'
+        frame = '01893 256 Jumpsuit'
+        
+
+        register(folder, frame)
+        break
+        pass

data/registration_post.py

@@ -0,0 +1,482 @@
+import os
+import sys
+import random
+import numpy as np
+import plotly.graph_objects as go
+import torch
+from torch.nn.functional import mse_loss
+from pytorch3d.structures import Meshes
+from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
+from pymeshfix import MeshFix
+
+# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+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,
+    read_curvature_frame,
+    writeOBJ,
+    readOBJ,
+    quads2tris
+)
+from smpl.smpl_torch import SMPL
+# from custom_utils.custom_utils import _maskFaces
+
+device = 'cuda'
+smpl = SMPL(device, 'smpl')
+
+
+import torch.nn.functional as F
+from pytorch3d.structures import Meshes
+
+def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
+    """
+    Given two PyTorch3D meshes and a set of vertex indices for mesh1,
+    find the closest vertices from mesh2 to these vertices.
+
+    Args:
+        mesh1: A PyTorch3D Meshes object representing the first mesh.
+        mesh2: A PyTorch3D Meshes object representing the second mesh.
+        vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
+
+    Returns:
+        A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
+    """
+    # Extract the vertices of mesh1 and mesh2.
+    vertices1 = mesh1.verts_packed()
+    vertices2 = mesh2.verts_packed()
+
+    # Extract the vertices of mesh1 that correspond to the given indices.
+    vertices1_selected = vertices1[vertex_indices]
+
+    # Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
+    distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
+
+    # Find the indices of the closest vertices in mesh2.
+    closest_indices = torch.argmin(distances, dim=1)
+
+    return closest_indices
+
+
+def _maskFaces(V, F, mask, return_unfixed=False):
+    V_mask = mask
+    F_mask = np.take(V_mask, F.reshape(-1), axis=0)
+    F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
+
+    V_ind = np.zeros(V_mask.shape)
+    V_ind[V_mask] = np.arange(V_mask.sum())
+    F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
+    V_new = V[V_mask, :]
+
+    if return_unfixed:
+        return V_new, F_new, F[F_mask, :]
+    else:
+        return V_new, F_new
+
+def get_frames(folder, frames_file):
+    frames = []
+    with open(os.path.join(folder, frames_file), 'r') as f:
+        for line in f:
+            frames.append(line.strip())
+    return frames
+
+
+def plot(mesh, faces):
+    fig = go.Figure()
+
+    fig.add_trace(go.Mesh3d(
+                x=mesh[:,0],
+                y=mesh[:,1],
+                z=mesh[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),)
+    
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                    )
+    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                yaxis = dict(nticks=4, range=[-1,1],),
+                zaxis = dict(nticks=4, range=[-1,1],),)
+    fig.update_layout(scene = helper,
+                        )
+
+    fig.show()
+
+
+def build_lookup_table(repeated):
+    # iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
+
+    lookup = {}
+    for i in range(len(repeated)):
+        for j in range(len(repeated[i])):
+            if len(repeated[i]) == 1:
+                continue
+            lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
+    return lookup
+
+
+from data.cloth3d_io import readPC2Frame
+import scipy.io as sio
+import scipy
+import scipy.spatial as spatial
+
+def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
+    pc16_path = os.path.join(folder, id, garment + '.pc16')
+    verts = readPC2Frame(pc16_path, frame, True)
+
+    map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
+
+    if deorient:
+        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
+        verts = R.apply(verts)
+
+    return verts, map
+
+def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
+    pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
+    verts = readPC2Frame(pc16_path, frame, True)
+
+    if deorient:
+        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
+        verts = R.apply(verts)
+
+    return verts
+
+def edge_loss_custom(x, y, edges):
+    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
+    x_e = torch.sum(x_e ** 2, -1)
+    x_e = torch.sqrt(x_e + 1e-9)
+    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
+    y_e = torch.sum(y_e ** 2, -1)
+    y_e = torch.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    loss = torch.mean(d_e ** 2, dim=-1)[0]
+    return loss
+
+
+def crotch(mesh, body, crotch_idx, nr=75):
+    crotch_middle_on_body = crotch_idx
+    # create kdtree for body and mesh
+    body_tree = spatial.cKDTree(body)
+    mesh_tree = spatial.cKDTree(mesh)
+    # find closest point on mesh from crotch middle on body along the z axis
+
+    crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
+
+    crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
+    crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
+
+    body_verts = body[crotch_body]
+    mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
+    mesh_verts = spatial.cKDTree(mesh_verts)
+    closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
+    crotch_mesh = crotch_mesh[closest_mesh_verts]
+
+    return crotch_body, crotch_mesh
+
+
+def find_closest_point_along_z(point, verts, eps=0.01):
+    min_dist = np.inf
+    closest_point = None
+    for i, v in enumerate(verts):
+        dist = np.linalg.norm(point - v)
+        if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
+            min_dist = dist
+            closest_point = i
+    return closest_point
+
+
+def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
+    """
+    Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
+
+    Args:
+        mesh: A PyTorch3D Meshes object representing the mesh.
+
+    Returns:
+        A tensor of shape (E,) representing the lengths of all edges in the mesh,
+        where E is the total number of edges in the mesh.
+    """
+    # Extract the vertices and faces of the mesh.
+    vertices = mesh.verts_packed()
+    faces = mesh.faces_packed()
+
+    # Compute the pairwise distances between all pairs of vertices that form an edge.
+    edge_indices = faces.view(-1, 3)
+    edge_vertices = vertices[edge_indices]
+    edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
+
+    return edge_lengths
+
+# try to optimize the already registered garment ??????????
+
+def register(folder, frame):
+    id, frame_nr, garment = frame.split()
+    frame_nr = int(frame_nr)
+    info = loadInfo(os.path.join(folder, id, 'info.mat'))
+    raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
+    # convert it to indices
+    mask = raw_mask > 0
+    mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
+    body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
+    template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
+    faces = quads2tris(faces)
+    
+    boundaries = raw_mask == 2
+    registered_boundaries = torch.tensor(np.where(boundaries[mask])[0]).to(device)
+    # invert map and use it to invert boudaries
+    # map_inv = np.zeros(map.max()+1, dtype=np.int32)
+    # map_inv[map] = np.arange(len(map))
+    # boundaries = boundaries[map_inv]
+    # # get the boundary indices
+    # mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
+
+
+    
+    # # pose = torch.zeros(1, 72).to('cuda')
+    # shape = torch.zeros(1, 11).to(device)
+    # shape[:, :10] = torch.tensor(info['shape']).to(device)
+    # shape[:, -1] = torch.tensor(info['gender']).to(device)
+    # # for star pose
+    # # pose[:, 5] = 3.1415/4
+    # # pose[:,8] = -3.1415/4
+    # STAR_POSE = np.zeros((24, 3), np.float32)
+    # STAR_POSE[1,2] = .15
+    # STAR_POSE[2,2] = -.15
+    # pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
+    # star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
+    # # rotate star body with 90 degrees around x axis
+    # R = spatial.transform.Rotation.from_euler('x', np.pi/2)
+    # star_body = R.apply(star_body).astype(np.float32)
+
+    is_skirt = (garment == 'Dress' or garment == 'Skirt')
+    skirt_mod = '_skirt' if is_skirt else ''
+    F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
+
+    ##################### mask body #####################
+    new_crotch = 4729
+
+    # star_body = star_body[mask]
+    new_crotch = np.zeros((14475), bool)
+    new_crotch[4729] = True
+    new_crotch = new_crotch[mask]
+    new_crotch = np.where(new_crotch)[0][0]
+    _, F = _maskFaces(np.zeros((14475, 3)), F, mask)
+    registered = template[map][mask]
+
+    # plot(template, faces)
+    # v = mesh_posed[mask,:]
+    
+    ##################### boundaries #####################
+    # meshfix = MeshFix(template, faces)
+    # holes = meshfix.extract_holes()
+    # new_border = []
+    # for border in np.array(holes.points):
+    #     # find the index of the vertex in the original mesh
+    #     ind = np.argwhere(np.all(template == border, axis=1))
+    #     new_border.extend(ind.ravel())
+    # mesh1_boundaries = torch.tensor(new_border).to(device)
+    # np.save('temp/body_boundaries.npy', new_border)
+    template_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
+
+    # meshfix = MeshFix(template, faces)
+    # holes = meshfix.extract_holes()
+    # new_border = []
+    # for border in np.array(holes.points):
+    #     # find the index of the vertex in the original mesh
+    #     ind = np.argwhere(np.all(template == border, axis=1))
+    #     new_border.extend(ind.ravel())
+    # mesh1_boundaries = torch.tensor(new_border).to(device)
+    # np.save('temp/body_boundaries.npy', new_border)
+    # mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
+
+    ##################### crotch #####################
+    # crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
+
+
+
+
+    # export star_body as obj
+    # writeOBJ('temp/star_body.obj', star_body, F)
+
+
+    template = torch.tensor(template).to(device)
+    template_orig = template.clone()
+    faces = torch.tensor(faces).to(device)
+    registered = torch.tensor(registered).to(device)
+    F = torch.tensor(F).to(device)
+
+    # verts = template.cpu().numpy()
+    # star_body = star_body.cpu().numpy()
+    # faces = faces.detach().cpu().numpy()
+    # F = F.detach().cpu().numpy()
+    # fig = go.Figure(data=[
+    #     go.Mesh3d(
+    #         x=verts[:,0], y=verts[:,1], z=verts[:,2],
+    #         i=faces[:,0], j=faces[:,1], k=faces[:,2],
+    #         color='lightpink', opacity=0.50),
+    #     go.Mesh3d(
+    #         x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
+    #         i=F[:,0], j=F[:,1], k=F[:,2],
+    #         color='lightblue', opacity=0.50),
+    #     go.Scatter3d(
+    #         x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
+    #         mode='markers', marker=dict(size=2, color='red')),
+    #     go.Scatter3d(
+    #         x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
+    #         mode='markers', marker=dict(size=2, color='blue')),
+    # ])
+    # # make it cuve and limit to -1 and -1
+    # fig.update_layout(scene = dict(
+    #     xaxis = dict(range=[-1,1],),
+    #     yaxis = dict(range=[-1,1],),
+    #     zaxis = dict(range=[-1,1],),),)
+    # fig.show()
+
+    # Define the meshes to optimize
+    mesh1 = Meshes(verts=registered[None,:,:], faces=F[None,:,:])
+    mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
+    mesh1_boundaries = registered_boundaries
+    mesh2_boundaries = template_boundaries
+
+    # mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
+
+    avg_edge_length = mesh_edge_lengths(mesh1).mean()
+
+    # Define the optimizer and learning rate
+    mesh1.verts_packed().requires_grad = True
+    optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
+    # optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
+
+    # Define the loss functions
+    def laplacian_smoothing_loss(mesh):
+        return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
+
+    def edge_loss(mesh):
+        return mesh_edge_loss(mesh, target_length=avg_edge_length)
+
+    def normal_consistency_loss(mesh):
+        return mesh_normal_consistency(mesh)
+
+    # Optimization loop
+    for i in range(10000):
+        optimizer.zero_grad()
+
+        # find boundaries with the find holes function?
+
+        # Forward pass
+        mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
+        ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
+        bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
+        # crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
+        # crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
+        loss: torch.Tensor = ch_loss * 3 +  (bound_ch_loss) * 20
+        # loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
+        lap_loss = laplacian_smoothing_loss(mesh1_out)
+        e_loss = edge_loss(mesh1_out)
+        # loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
+        norm_loss = normal_consistency_loss(mesh1_out)
+        loss += lap_loss * 0.05 + e_loss * 100 + norm_loss * 0.1
+
+        # Backward pass
+        loss.backward()
+        optimizer.step()
+
+        # Print the loss
+        if i % 10 == 0:
+            print(f"it: {i} : loss: {loss:.5f} chamfer: {ch_loss:.6f} bound_chamfer: {bound_ch_loss:.6f} laplacian: {lap_loss:.6f} edge: {e_loss:.6f} normal: {norm_loss:.5f}")
+    
+    # create a plotly plot with the body and the mesh1_out
+    registered = mesh1_out.verts_packed().detach().cpu().numpy()
+    template = mesh2.verts_packed().detach().cpu().numpy()
+    faces = faces.detach().cpu().numpy()
+    F = F.detach().cpu().numpy()
+    # faces, F = F, faces
+    # verts, star_body = star_body, verts
+    mesh1_boundaries, mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy(), mesh2_boundaries.detach().cpu().numpy()
+    # mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy()
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            x=template[:,0], y=template[:,1], z=template[:,2],
+            i=faces[:,0], j=faces[:,1], k=faces[:,2],
+            color='lightblue', opacity=0.80),
+        go.Mesh3d(
+            x=registered[:,0], y=registered[:,1], z=registered[:,2],
+            i=F[:,0], j=F[:,1], k=F[:,2],
+            color='lightpink', opacity=0.50),
+        go.Scatter3d(
+            x=registered[mesh1_boundaries,0], y=registered[mesh1_boundaries,1], z=registered[mesh1_boundaries,2],
+            mode='markers', marker=dict(size=2, color='red')),
+        go.Scatter3d(
+            x=template[mesh2_boundaries,0], y=template[mesh2_boundaries,1], z=template[mesh2_boundaries,2],
+            mode='markers', marker=dict(size=2, color='blue')),   
+        # go.Scatter3d(
+        #     x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
+        #     mode='markers', marker=dict(size=2, color='red')),
+        # go.Scatter3d(
+        #     x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
+        #     mode='markers', marker=dict(size=2, color='blue')),
+    ])
+    # make it cuve and limit to -1 and -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()
+
+    # create a plotly wireframe plot of registered
+    fig = go.Figure()
+    for edge in mesh1_out.edges_packed().detach().cpu().numpy():
+        fig.add_trace(go.Scatter3d(
+            x=registered[edge,0], y=registered[edge,1], z=registered[edge,2],
+            mode='lines', line=dict(color='red', width=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/code/cloth3d/train'
+    # folder = r'D:\Work\code\cloth3d\train'
+    # frames = [f'00761_{i}' for i in range(300)]
+    # folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
+    frames = get_frames(folder, 'frames_all.txt')
+    root = '/'.join(folder.split('/')[:-1])
+    type = folder.split('/')[-1]
+
+    use_no_orient = False
+    no_shape = False
+    is_skirt = False
+
+    skirt_modifier = '_skirt' if is_skirt else ''
+    orient_mod = '_no_orient' if use_no_orient else ''
+    shape_mod = '_no_shape' if no_shape else ''
+
+    # mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
+    # faces = np.load('/data-net/dataset/00761/faces/00761.npy')
+
+    for frame in frames:
+        frame = random.choice(frames)
+        # frame = '00761 0 Tshirt'
+        frame = '01893 256 Jumpsuit'
+        
+
+        register(folder, frame)
+        break
+        pass

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))

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

docker/docker-compose.yml

@@ -34,6 +34,29 @@ services:
       - '/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
+    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:

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

losses/losses.py

@@ -4,9 +4,9 @@ import numpy as np
 # import tensorflow_graphics as tfg
 import scipy as sp
 import scipy.io as sio
+from scipy.spatial import KDTree
 
 from pytorch3d.structures import Meshes
-from scipy.spatial import cKDTree
 from torch_cluster import knn
 # from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
 # from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
@@ -18,12 +18,12 @@ from torch_cluster import knn
 smpl_to_keep_orig = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
 smpl_to_keep_orig = smpl_to_keep_orig[smpl_to_keep_orig < 6890]
 
-to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
-F_coll = torch.tensor(np.load("smpl/smpl_faces.npy")[to_keep,:], device='cuda')
-F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:], device='cuda')
+# to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
+F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda')
+F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda')
 F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device='cuda')
-F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
-F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
+F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [4,1,1]), device='cuda')
+F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [4,1,1]), device='cuda')
 F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')
 F_norm_orig = F_norm_orig[:,smpl_to_keep_orig,:]
 
@@ -33,15 +33,26 @@ def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
     batch_size = V.shape[0]
     vertices_count = V.shape[1]
     batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(V.shape[1])
-    batch_y = torch.arange(batch_size, device=B.get_device()).repeat_interleave(B.shape[1])
     V = V.reshape(-1, V.shape[-1])
     B = B.reshape(-1, B.shape[-1])
-    idx = knn(V, B.to(V.dtype), k=1, batch_x=batch_x, batch_y=batch_y)[1]
+    idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_x)[1]
     idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count)
     return idx
 
+# def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
+#     # create an empty tensor to store the indices of the nearest neighbors
+#     batch_size = V.shape[0]
+#     vertices_count = V.shape[1]
+#     idx = torch.zeros((batch_size, vertices_count), dtype=torch.long, device=V.device)
+
+#     # loop over the batch dimension and find the nearest neighbor for each point in B
+#     for i in range(batch_size):
+#         tree = KDTree(B[i].cpu().numpy())
+#         _, nearest = tree.query(V[i].cpu().numpy(), k=1)
+#         idx[i] = torch.from_numpy(nearest)
+
+#     return idx
 
-# @tf.function
 def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004):
     if B.shape[1] == 14475:
         F_norm_local = F_norm

models/cvae_model.py

@@ -18,25 +18,24 @@ class CVAE(nn.Module):
         super(CVAE, self).__init__()
         self.size = config['resolution']
 
-        if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
-            self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
-            if config['static_model'] != '':
-                load_model(self.static_enc, config['static_model'] + '_enc')
+        self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
+        if config['static_model'] != '':
+            load_model(self.static_enc, config['static_model'] + '_enc')
 
-            self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
-            if config['body_model'] != '':
-                load_model(self.body_enc, config['body_model'] + '_enc')
+        self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
+        if config['body_model'] != '':
+            load_model(self.body_enc, config['body_model'] + '_enc')
 
-        if config['use_curvature']:
-            self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'], device=config['device'], is_label=True, nr_of_labels=7)
-            if config['curvature_model'] != '':
-                load_model(self.curv_enc, config['curvature_model'] + '_enc')
+        if config['use_normals']:
+            self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device'])
+            if config['normals_model'] != '':
+                load_model(self.norm_enc, config['normals_model'] + '_enc')
 
         if config['freeze_encoders']:
             for param in self.static_enc.parameters(): param.requires_grad = False
             for param in self.body_enc.parameters(): param.requires_grad = False
-            if config['use_curvature']:
-                for param in self.curv_enc.parameters(): param.requires_grad = False
+            if config['use_normals']:
+                for param in self.norm_enc.parameters(): param.requires_grad = False
 
         self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
         self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])
@@ -48,7 +47,7 @@ class CVAE(nn.Module):
         else:
             self.avg_garment = None
 
-        self.use_curvature = config['use_curvature']
+        self.use_normals = config['use_normals']
         self.use_offsets_over_body = not config['use_skinning']
 
         if self.use_offsets_over_body:
@@ -98,9 +97,9 @@ class CVAE(nn.Module):
     def encode_conds(self, uv_template, uv_body_posed, uv_curv):
         z_body = self.body_enc(uv_body_posed)
         z_static = self.static_enc(uv_template)
-        if self.use_curvature:
-            z_curv = self.curv_enc(uv_curv)
-            z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
+        if self.use_normals:
+            z_norm = self.norm_enc(uv_curv)
+            z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
         else:
             z_cond = torch.concat([z_static, z_body], dim=1)
 
@@ -108,13 +107,13 @@ class CVAE(nn.Module):
 
 
     def forward(self, inputs):
-        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
-        x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
+        uv_unposed, uv_template, uv_body_posed, uv_norm, uv_mask = inputs
+        x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_norm], dim=1)
         
         mean, logvar = self.encoder(x)
         z: torch.Tensor = self.reparameterize(mean, logvar)
         z = z.view((mean.shape[0], *self.latent_shape))
-        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_norm)
 
         z = torch.concat([z, z_cond], dim=1)
         x = self.decoder(z)

models/cvae_norm_model.py

@@ -0,0 +1,97 @@
+import os
+import time
+import random
+import torch
+from torch import nn
+import numpy as np
+import scipy as sp
+# import plotly.graph_objects as go
+# import plotly.express as ex
+
+from utils.checkpoints import load_model
+from models.base import AEEncoder, VariationalEncoder, Decoder
+
+
+class CVAENorm(nn.Module):
+    
+    def __init__(self, config):
+        super(CVAENorm, self).__init__()
+        self.size = config['resolution']
+
+        self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
+        if config['static_model'] != '':
+            load_model(self.static_enc, config['static_model'] + '_enc')
+
+        self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
+        if config['body_model'] != '':
+            load_model(self.body_enc, config['body_model'] + '_enc')
+
+        if config['freeze_encoders']:
+            for param in self.static_enc.parameters(): param.requires_grad = False
+            for param in self.body_enc.parameters(): param.requires_grad = False
+
+        self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
+        self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])
+
+        self.latent_shape = (config['cvae']['latent_dim'], 8, 4)
+
+
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
+        if noise is None:
+            eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
+        else:
+            eps = noise
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
+        if add_noise:
+            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
+        
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z, training=False)
+        
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        
+        if self.use_offsets_over_body:
+            x = x + uv_body_posed
+
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean: torch.Tensor, logvar: torch.Tensor, training: bool = True, fix_noise: bool = False):
+        if training == True:
+            if not fix_noise:
+                eps = torch.randn(size=mean.shape).to(mean.device)
+            else:
+                eps = torch.randn(size=(1, mean.shape[1]))
+                eps = torch.tile(eps, [mean.shape[0],1])
+            out = eps * torch.exp(logvar * 0.5) + mean
+        else:
+            out = torch.exp(logvar * 0.5) + mean
+        return out
+
+
+    def encode_conds(self, uv_template, uv_body_posed):
+        z_body = self.body_enc(uv_body_posed)
+        z_static = self.static_enc(uv_template)
+        z_cond = torch.concat([z_static, z_body], dim=1)
+        return z_cond
+
+
+    def forward(self, inputs):
+        uv_unposed, uv_template, uv_body_posed, uv_mask = inputs
+        x = torch.concat([uv_unposed, uv_template, uv_body_posed], dim=1)
+        
+        mean, logvar = self.encoder(x)
+        z: torch.Tensor = self.reparameterize(mean, logvar)
+        z = z.view((mean.shape[0], *self.latent_shape))
+        z_cond = self.encode_conds(uv_template, uv_body_posed)
+
+        z = torch.concat([z, z_cond], dim=1)
+        x = self.decoder(z)
+        x = x * uv_mask[:,None,:,:]
+
+        return x, mean, logvar

models/pyramid_model.py

@@ -12,7 +12,8 @@ import torch.nn.functional as F
 from utils.checkpoints import load_model
 from models.base import AEEncoder
 from models.pyramid_level import PyramidLevel
-
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
 
 # class Pyramid(nn.Module):
     
@@ -145,17 +146,20 @@ class Pyramid(nn.Module):
         if config['body_model'] != '':
             load_model(self.body_enc, config['body_model'] + '_enc')
 
-        self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], kernel_size=config['cond_ae_curv']['kernel_size'], activation=config['cond_ae_curv']['activation'], is_label=True, nr_of_labels=7)
-        if config['curvature_model'] != '':
-            load_model(self.curv_enc, config['curvature_model'] + '_enc')
+        self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device'])
+        if config['normals_model'] != '':
+            load_model(self.norm_enc, config['normals_model'] + '_enc')
 
         if config['freeze_encoders']:
             for param in self.static_enc.parameters(): param.requires_grad = False
             for param in self.body_enc.parameters(): param.requires_grad = False
-            for param in self.curv_enc.parameters(): param.requires_grad = False
+            for param in self.norm_enc.parameters(): param.requires_grad = False
 
         self.levels = nn.ModuleList([PyramidLevel({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))])
 
+        self.resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BILINEAR).to(config['device']) for res in config['pyramid']['resolutions'] }
+        self.resizers[512] = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device'])
+
 
         # if len(config['pyramid']['resolutions']) == 2:
         #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
@@ -187,20 +191,20 @@ class Pyramid(nn.Module):
     #     x = self.upscale_enhancer([low, mask], training=training)
     #     return x
 
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+    def encode_conds(self, uv_template, uv_body_posed, uv_norm):
         z_static = self.static_enc(uv_template)
         z_body = self.body_enc(uv_body_posed)
-        z_curv = self.curv_enc(uv_curv)
+        z_norm = self.norm_enc(uv_norm)
 
-        z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
+        z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
 
         return z_cond
 
     def forward(self, inputs):
         levels_input, conds = inputs
-        uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
+        uv_static, uv_body, uv_norm = conds[0], conds[1], conds[2]
 
-        z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
+        z_cond = self.encode_conds(uv_static, uv_body, uv_norm)
 
         # if self.add_noise == True and training == True:
         #     z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
@@ -211,13 +215,13 @@ class Pyramid(nn.Module):
         #     pred = pred + levels_input[0][2] # uv_body_posed_32
         if 1 != len(self.levels):
             pred = self.levels[0].enhance(pred, levels_input[0][-2])  # levels_input[i][-2]=uv_mask_res
-            pred = F.interpolate(pred, size=(self.levels[1].size, self.levels[1].size//2), mode='bilinear')
+            pred = self.resizers[self.levels[1].size](pred)
             out_levels.append([pred, mean, logvar])
             pred = pred * levels_input[1][-2][:,None,:,:]
         else:
             # pred = self.enhance(pred, levels_input[0][-2], training=training)
             pred = self.levels[0].enhance(pred, levels_input[0][-2])  # levels_input[i][-2]=uv_mask_res
-            pred = F.interpolate(pred, size=(512, 256), mode='bilinear') 
+            pred = self.resizers[512](pred)
             pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask
             out_levels.append([pred, mean, logvar])
 
@@ -233,13 +237,13 @@ class Pyramid(nn.Module):
             if i != len(self.levels) - 1:
                 pred = self.levels[i].enhance(pred, levels_input[i][-2])
                 out_levels.append([pred, mean, logvar])
-                pred = F.interpolate(pred, size=(self.levels[i+1].size, self.levels[i+1].size//2), mode='bilinear')
+                pred = self.resizers[self.levels[i+1].size](pred)
                 pred = pred * levels_input[i+1][-2][:,None,:,:]
             else:
                 # if self.upscale_enhancer is not None:
                     # pred = self.enhance(pred, levels_input[-1][-2], training=training)
                 pred = self.levels[-1].enhance(pred, levels_input[-1][-2])
-                pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
+                pred = self.resizers[512](pred)
                 pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask
                 out_levels.append([pred, mean, logvar])
 

run.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.pyr_64_cape_m import config
+from configs.cvae_norm_1 import config
 import argparse
 
 # parser = argparse.ArgumentParser(description='')

run_parallel.py

@@ -4,24 +4,27 @@ import os
 
 if __name__ == '__main__':
 
-    folder = '/home/workspace/data/dataset/cloth3d/train'
+    # folder = '/home/workspace/data/dataset/cloth3d/train'
     # folder = '/home/workspace/local-data/cape_processed/train/male'
     # folder = '/home/workspace/local-data/valid_skirt'
     # folder = '/home/workspace/local-data/valid'
     # ids = os.listdir(folder)
     # ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:]
     skirt = 0
+    obj_folder = '/data-net/dataset/cloth3d/train'
+    folder = '/data-local/127/train'
 
-    count = int(500 * 1e3)
-    step = 10000
+    count = 0
+    step = 100
 
     # !!!!!!! change skirt parameter in script
 
-    while count < 510 * 1e3:
+    while count <1200:
         # os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv_local.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
-        os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
+        # os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
+        os.system(f'nohup python data/calculate_normals.py -f {folder} -f_obj {obj_folder}  -s {count} -e {(count+step)} -n normals_train_calc_{count}_{(count+step)} > nohup/normals/normals_train_calc_{count}_{(count+step)}.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_valid_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_valid_skirt_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')

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()
 

scripts/ae_norm_train.py

@@ -0,0 +1,221 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import argparse
+import importlib
+
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+if args['debug']:
+    from configs.iccvw.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']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import wandb
+import random
+import numpy as np
+import torch
+from torch import nn
+# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str
+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 data.dataset import Dataset
+# from torch.utils.data import DataLoader
+from custom_utils.custom_utils import plot_3d
+from models.cond_encoder import AutoEncoder
+from utils.checkpoints import CheckpointSaver, load_model
+from custom_utils.custom_utils import Accumulator
+from utils.utils import InfiniteDataLoader as DataLoader
+
+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_norm_train.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+
+
+def train_step(batch):
+    model.train()
+    pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4
+
+    ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4))
+    loss = 1e3 * (1 - ms_ssim)
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    model.eval()
+    with torch.inference_mode():
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_normals']
+        l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_normals'] - pred)), axis=1))
+        metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1)))
+    
+    return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
+        
+
+def valid_step(batch):
+    model.eval()
+
+    with torch.inference_mode():
+        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.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)))
+    
+    return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
+       
+
+def plot_everything(batch: dict, type: str):
+    model.eval()
+    pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.4
+    
+    # mesh_pred = uv2mesh(pred)
+    # mesh_orig = batch['mesh_normals']
+    
+    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.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
+    # gt = [mesh_orig[i].cpu().numpy() for i in range(pred.shape[0])]
+    
+    # body = uv2mesh(batch['uv_body_posed']).cpu().numpy()
+    # garment_posed = garment_unposed
+    # collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    # collision_faces = [np.concatenate((F, faces + (3638 if config['is_cape'] else 14475)), axis=0) for i in range(len(garment_posed))]
+    # collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    # figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    figs = {}
+
+    return { **figs, **{ type + '/images': imgs } }
+
+
+def train():
+    valid_iterator = iter(valid_dataloader)
+
+    step = 0
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            if not args['debug']:
+                wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if step % 10000 == 0:
+            ckpt_saver_enc(model.encoder, step, optimizer, scaler=None)
+            ckpt_saver_dec(model.decoder, step, optimizer, scaler=None)
+
+        step += 1
+
+
+class Scale(torch.nn.Module):
+    def __init__(self, scale_factor):
+        super().__init__()
+        self.scale_factor = scale_factor
+
+    def forward(self, x):
+        return x * self.scale_factor
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['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()
+    valid_accumulator = Accumulator()
+
+    model : AutoEncoder = AutoEncoder(config['cond_ae'])
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+
+    # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
+    _MSSSIM = MS_SSIM(data_range=2.8, 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')
+
+    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=20000, log_graph=True)
+
+    train()
+
+
+

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()
 

scripts/cvae.py

@@ -10,10 +10,12 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
-config_module = importlib.import_module(args['config'])
-config = config_module.config
 if args['debug']:
-    pass
+    from configs.iccvw.abl_all 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']:
@@ -41,6 +43,7 @@ from data.dataset import Dataset
 from models.cond_encoder import AutoEncoder
 from utils.checkpoints import CheckpointSaver, load_model
 from utils.utils import InfiniteDataLoader as DataLoader
+from utils.utils import text_to_html
 from utils.postprocess import fix_collision
 from losses.losses import collision_loss, normal_loss
 from pytorch_msssim import MS_SSIM
@@ -50,7 +53,8 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
 
 
 def init_wandb():
-    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
+    wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
+    run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/cvae.py')
     artifact.add_file('losses/losses.py')
@@ -83,7 +87,7 @@ def train_step(batch):
         mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
 
     model.train()
-    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
+    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
     mesh_pred = uv2mesh(pred)
 
     if config['use_skinning']:
@@ -139,7 +143,7 @@ def valid_step(batch):
 
     model.eval()
     with torch.inference_mode():
-        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
+        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
         mesh_pred = uv2mesh(pred)
 
         if config['use_skinning']:
@@ -188,15 +192,16 @@ def plot_everything(batch: dict, type: str):
 
     model.eval()
     with torch.inference_mode():
-        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
+        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
     
     mesh_pred = uv2mesh(pred).detach().cpu().numpy()
     
     orig_img = torch.permute(uv_orig, [0,2,3,1])
     uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
     uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
+    uv_normals_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
     pred_img = torch.permute(pred, [0,2,3,1])
-    img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
+    img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals_img, orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
     
     F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -247,13 +252,20 @@ def plot_everything(batch: dict, type: str):
     collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
     collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
 
-    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
-    figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'])) for i in range(4) }
+    figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'])) for i in range(4) }
 
     mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0))
     logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0))
 
-    return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+    batch_id = batch['id']
+    batch_frame_nr = batch['frame_nr'].cpu().numpy()
+    batch_garment = batch['garment']
+    batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))]
+    batch_info = text_to_html(batch_info)
+    batch_info = wandb.Html(batch_info)
+
+    return { **figs, **figs_coll, **{ f'{type}/batch_info': batch_info }, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
 
 def sample(folder): 
@@ -262,10 +274,10 @@ def sample(folder):
     count = 0
     for batch in dataset_train:
         batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
-        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+        uv_normals = tf.one_hot(batch['uv_normals'], 7)
         nr_of_samples=2
 
-        pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
+        pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
         pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
         pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
 
@@ -281,7 +293,7 @@ def sample(folder):
         mask = batch['mesh_mask'].numpy().astype(bool)[0]
         _, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
 
-        pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
+        pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
         mesh_pred_zero = uv2mesh(pred_zero)
         mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
 
@@ -453,11 +465,11 @@ def train():
             valid_accumulator.reset()
             print(step)
 
-        if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
+        if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
             g_balancing_coeff += 1 / 10000 
 
-        if step % 10000 == 0:
-            ckpt_saver(model, step, optimizer)
+        if step % 25000 == 0 and not args['debug']:
+            ckpt_saver(model, step, optimizer, None)
 
         step += 1
 
@@ -483,7 +495,7 @@ if __name__ == '__main__':
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
     g_balancing_coeff = config['starting_balancing_coeff']
 
-    ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=10)
 
     step = 0
     if config['pretrained_model'] != '':
@@ -491,6 +503,6 @@ if __name__ == '__main__':
 
     if not args['debug']:
         # model = torch.compile(model)
-        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+        wandb.watch(model, log='all', log_freq=10000, log_graph=True)
 
     train()

scripts/cvae_norm.py

@@ -0,0 +1,400 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+if args['debug']:
+    from configs.cvae_norm_1 import config
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import argparse
+import wandb
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from PIL import Image
+
+if config['is_skirt']:
+    from losses.losses import edge_loss_skirt as edge_loss
+else:
+    from losses.losses import edge_loss
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from data.dataset import Dataset
+from models.cond_encoder import AutoEncoder
+from utils.checkpoints import CheckpointSaver, load_model
+from utils.utils import InfiniteDataLoader as DataLoader
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss
+from pytorch_msssim import MS_SSIM
+from models.cvae_norm_model import CVAENorm
+from smpl.smpl_torch import SMPL
+from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision
+
+
+def init_wandb():
+    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/cvae_norm.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('models/cvae_norm_model.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+
+def kl_loss(mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+    kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar)
+    kl_loss = torch.mean(kl_loss)
+    kl_loss *= -0.5
+    return kl_loss
+
+
+def train_step(batch):
+    model.train()
+    pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
+    pred = pred * 1.1
+    mesh_pred = uv2mesh(pred)
+
+
+    l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred) ** 2, axis=-1), axis=-1))
+    ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
+    kl = kl_loss(mean, logvar)
+
+    reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
+    loss = reconstr_loss + g_balancing_coeff * kl
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    model.eval()
+    with torch.inference_mode():
+        mean = torch.mean(mean)
+        logvar = torch.mean(logvar)
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+    
+    return { 'train/loss': loss, 'train/euclid_error': metric,
+             'train/reconstr_loss': reconstr_loss,
+             'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
+             'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss,
+             'train/ms_ssim': ms_ssim, 
+            }
+ 
+
+def valid_step(batch):
+    model.eval()
+    with torch.inference_mode():
+        pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
+        pred = pred * 1.1
+        mesh_pred = uv2mesh(pred)
+
+        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
+        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_pred - batch['mesh_normals']) ** 2, axis=-1), axis=-1))
+        ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_static'] + 1.1))
+        kl = kl_loss(mean, logvar)
+
+        reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
+        loss = reconstr_loss + g_balancing_coeff * kl
+
+        mean = torch.mean(mean)
+        logvar = torch.mean(logvar)
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+    
+
+    return { 'valid/loss': loss, 'valid/euclid_error': metric, 
+             'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
+             'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
+             'valid/ms_ssim': ms_ssim, 
+              }
+
+
+def plot_everything(batch: dict, type: str):
+    model.eval()
+    with torch.inference_mode():
+        pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
+    
+    pred = pred * 1.1
+    # mesh_pred = uv2mesh(pred).detach().cpu().numpy()
+    
+    uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
+    uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
+    orig_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
+    pred_img = torch.permute(pred, [0,2,3,1])
+    img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
+    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")
+    # mask = batch['mesh_mask'].cpu().numpy().astype(bool)
+
+    mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0))
+    logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0))
+
+    return { **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+
+
+def sample(folder): 
+    # run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
+    load_model(model,folder)
+    count = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+        uv_normals = tf.one_hot(batch['uv_normals'], 7)
+        nr_of_samples=2
+
+        pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
+        pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
+        pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
+
+        mesh_pred = uv2mesh(pred)
+        mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
+        mesh_orig = batch['mesh_posed'][0].numpy()
+        mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
+
+        body = batch['mesh_body_posed'][0].numpy()
+        body_color = np.ones_like(body, dtype=np.int32)*100
+
+        F = np.load("custom_utils/smpl_faces.npy")
+        mask = batch['mesh_mask'].numpy().astype(bool)[0]
+        _, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
+
+        pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
+        mesh_pred_zero = uv2mesh(pred_zero)
+        mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
+
+
+        meshes = []
+        meshes.append(go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=mesh_orig_reposed[mask,0],
+                    y=mesh_orig_reposed[mask,2],
+                    z=mesh_orig_reposed[mask,1],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ))
+
+        for i in range(nr_of_samples):
+            x = mesh_posed[i][mask, :].copy()
+            # x[:, 2] += 1 * (i + 1)
+            x[:, 2] += 1
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,2],
+                        z=x[:,1],
+                        # i, j and k give the vertices of triangles
+                        i = faces[:,0],
+                        j = faces[:,1],
+                        k = faces[:,2]
+                        ))
+
+        
+        for i in range(nr_of_samples):
+            x = mesh_posed_zero[i][mask, :].copy()
+            # x[:, 2] += 1 * (i + 1)
+            x[:, 2] += 2
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,2],
+                        z=x[:,1],
+                        # i, j and k give the vertices of triangles
+                        i = faces[:,0],
+                        j = faces[:,1],
+                        k = faces[:,2]
+                        ))
+
+        # for i in range(5):
+        #     x = body.copy()
+        #     x[:, 2] += 1 * i
+        #     meshes.append(go.Mesh3d(
+        #                 # 8 vertices of a cube
+        #                 x=x[:,0],
+        #                 y=x[:,2],
+        #                 z=x[:,1],
+        #                 # i, j and k give the vertices of triangles
+        #                 i = F[:,0],
+        #                 j = F[:,1],
+        #                 k = F[:,2],
+        #                 vertexcolor=body_color
+        #                 ))
+
+
+        fig = go.Figure(data=meshes)
+
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
+        # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
+        fig.update_layout(
+            scene = dict(
+                xaxis = dict(nticks=4, range=[-1,1],),
+                            # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
+                            yaxis = dict(nticks=4, range=[-1,3],),
+                            zaxis = dict(nticks=4, range=[-1,1],),)
+                            )
+
+        # wandb.log({'plot_' + str(count): fig})
+        fig.show()
+
+        pred = pred.numpy()
+        frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
+        frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
+               save_all=True, duration=500, loop=0)
+
+
+        if count > 10:
+            break
+
+        print(count)
+        count += 1
+
+        print()
+
+
+
+def test():
+    str_split = config['name'].split('_')
+    test_idx = str_split.index('test')
+    checkpoint_name = '_'.join(str_split[:test_idx])
+    load_model(model, 'checkpoints/' + checkpoint_name)
+
+    dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
+    static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
+   
+    step = 0
+    losses = []
+
+    for batch in tqdm.tqdm(dataset_test):
+        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+
+        test_loss = test_step(batch)
+        losses.append(test_loss)
+
+        to_log = { **{ 'current_step': step }, **test_loss }
+        if step % 1000 == 0:
+            figs = plot_everything(batch, 'test')
+            to_log = { **to_log, **figs }
+
+        wandb.log(to_log, step = step)
+        
+        # print(step)
+        step += 1
+   
+    # averages = dict()
+    path_to_plotly_html = "temp/plotly_figure.html"
+
+    histograms = dict()
+    for key in test_loss.keys():
+        accumulated = [losses[i][key].numpy() for i in range(len(losses))]
+        fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
+        fig.write_html(path_to_plotly_html, auto_play = False)
+
+        histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
+        # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
+
+    
+
+    wandb.log({ **histograms }, step = step)
+    wandb.finish()
+    pass
+
+
+
+def train():
+    global g_balancing_coeff
+    global step
+    valid_iterator = iter(valid_dataloader)
+
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            if not args['debug']:
+                wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
+            g_balancing_coeff += 1 / 10000 
+
+        if step % 10000 == 0:
+            ckpt_saver(model, step, optimizer, None)
+
+        step += 1
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+
+    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
+    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()
+    valid_accumulator = Accumulator()
+
+    model = CVAENorm(config)
+
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
+    g_balancing_coeff = config['starting_balancing_coeff']
+
+    ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
+
+    step = 0
+    if config['pretrained_model'] != '':
+        step = load_model(model, config['pretrained_model'])
+
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=10000, log_graph=True)
+
+    train()

scripts/cvae_pyramid.py

@@ -10,11 +10,12 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
-config_module = importlib.import_module(args['config'])
-config = config_module.config
 if args['debug']:
-    from configs.pyr_32 import config
-    # from configs.pyr_64_cape_m import config
+    from configs.iccvw.pyr_32 import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -42,6 +43,7 @@ import torch.nn.functional as F
 from data.dataset import Dataset
 from utils.checkpoints import CheckpointSaver, load_model
 from utils.utils import InfiniteDataLoader as DataLoader
+from utils.utils import text_to_html
 from utils.postprocess import fix_collision
 from losses.losses import collision_loss, normal_loss, kl_loss
 from pytorch_msssim import MS_SSIM
@@ -52,7 +54,8 @@ from models.pyramid_model import Pyramid
 
 
 def init_wandb():
-    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
+    wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
+    run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/cvae_pyramid.py')
     artifact.add_file('losses/losses.py')
@@ -99,8 +102,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
 
     fig.add_trace(go.Mesh3d(
             x=mesh_posed[:,0],
-            y=mesh_posed[:,1],
-            z=mesh_posed[:,2],
+            y=mesh_posed[:,2],
+            z=mesh_posed[:,1],
             i = faces[:,0],
             j = faces[:,1],
             k = faces[:,2],
@@ -270,13 +273,37 @@ def train_step(batch):
         reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
     
     kl_total += kl
-    loss = reconstr_loss_total + g_balancing_coeff * kl_total
+    loss: torch.Tensor = reconstr_loss_total + g_balancing_coeff * kl_total
 
     # scaler.scale(loss).backward()
     # scaler.step(optimizer)
     # scaler.update()
     # optimizer.zero_grad()
 
+    # b = 0
+    # from custom_utils.custom_utils import _maskFaces
+    # mesh_mask = batch['mesh_mask'][b].cpu().numpy()
+    # faces = np.load('smpl/smpl_faces_skirt.npy') if batch['garment'][b] in ['Skirt', 'Dress'] else np.load('smpl/smpl_faces.npy')
+    # faces_body = np.load('smpl/smpl_faces.npy')
+    # _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask)
+
+    # # fig = go.Figure(data=[go.Mesh3d(x=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,0], y=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,2], z=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
+    # #                         go.Mesh3d(x=batch['mesh_body_posed'].cpu().numpy()[0,:,0], y=batch['mesh_body_posed'].cpu().numpy()[0,:,2], z=batch['mesh_body_posed'].cpu().numpy()[0,:,1], i=faces_body[:,0], j=faces_body[:,1], k=faces_body[:,2], color='lightblue', opacity=1),
+    # # ])
+    
+    # fig = go.Figure(data=[go.Mesh3d(x=mesh_orig.cpu().numpy()[0,mesh_mask,0], y=mesh_orig.cpu().numpy()[0,mesh_mask,2], z=mesh_orig.cpu().numpy()[0,mesh_mask,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
+    # ])
+    # fig.update_layout(scene_aspectmode='cube')
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+    # fig.show()
+
+    # ex.imshow((uv_orig[b].cpu().numpy().transpose(1,2,0) + 1) / 2).show()
+    # upscaled = model.levels[0].enhance(levels_input[0][0], levels_input[0][-2])  # levels_input[i][-2]=uv_mask_res
+    # upscaled = model.resizers[512](upscaled)
+    # ex.imshow((levels_input[0][0][b].cpu().numpy().transpose(1,2,0) + 1) / 2).show()
+    # ex.imshow((upscaled[b].cpu().numpy().transpose(1,2,0) + 1) / 2).show()
+
+
     optimizer.zero_grad()
     loss.backward()
     optimizer.step()
@@ -452,12 +479,12 @@ def plot_everything(batch: dict, type: str):
     orig_img = (torch.permute(uv_orig, [0,2,3,1]) + 1) / 2
     uv_static_img = (torch.permute(batch['uv_static'], [0,2,3,1]) + 1) / 2
     uv_body_posed_img = (torch.permute(batch['uv_body_posed'], [0,2,3,1]) + 1) / 2
-    uv_curvature = (torch.tile(torch.argmax(torch.permute(batch['uv_curvature'], [0,2,3,1]), dim=-1)[:,:,:,None], [1,1,1,3]) / 6)
-    pred_img = torch.concatenate([torch.permute(F.interpolate(output_levels[i][0], size=(512, 256), mode='bilinear'), [0,2,3,1]) for i in range(len(output_levels))], dim=3)
+    uv_normals = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.4) / 2.8
+    pred_img = torch.concatenate([torch.permute(model.resizers[512](output_levels[i][0]), [0,2,3,1]) for i in range(len(output_levels))], dim=3)
     # pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
     pred_img = (pred_img - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None]) / (pred_img.flatten().reshape(pred_img.shape[0], -1).max(1)[0][:,None,None,None] - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None])
 
-    img = torch.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
+    img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals, orig_img, pred_img), axis=2)
     imgs = [wandb.Image(img[i].cpu().numpy()) for i in range(4)]
 
     _F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -521,7 +548,14 @@ def plot_everything(batch: dict, type: str):
         collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))]
         figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
 
-    return { **figs, **figs_coll, **{ type + '/images': imgs, } }
+    batch_id = batch['id']
+    batch_frame_nr = batch['frame_nr'].cpu().numpy()
+    batch_garment = batch['garment']
+    batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))]
+    batch_info = text_to_html(batch_info)
+    batch_info = wandb.Html(batch_info)
+
+    return { **figs, **figs_coll, **{ type + '/images': imgs, }, **{ f'{type}/batch_info': batch_info }, }
 
 
 def sample(folder):
@@ -793,7 +827,7 @@ def train():
         if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
             g_balancing_coeff += 1 / 10000 
 
-        if step % 10000 == 0:
+        if step % 25000 == 0:
             ckpt_saver(model, step, optimizer, None)
 
         step += 1
@@ -840,12 +874,12 @@ if __name__ == '__main__':
     if len(config['pyramid']['resolutions']) == 1:
         for param in model.static_enc.parameters(): param.requires_grad = False
         for param in model.body_enc.parameters(): param.requires_grad = False
-        for param in model.curv_enc.parameters(): param.requires_grad = False
+        for param in model.norm_enc.parameters(): param.requires_grad = False
     else:
         load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else ''))
         for param in model.static_enc.parameters(): param.requires_grad = False
         for param in model.body_enc.parameters(): param.requires_grad = False
-        for param in model.curv_enc.parameters(): param.requires_grad = False
+        for param in model.norm_enc.parameters(): param.requires_grad = False
 
     # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
     # model.levels[0].upscale_enhancer.trainable = False
@@ -866,6 +900,6 @@ if __name__ == '__main__':
 
     if not args['debug']:
         # model = torch.compile(model)
-        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+        wandb.watch(model, log='all', log_freq=10000, log_graph=True)
 
     train()

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()
 

smpl/smpl_torch.py

@@ -163,9 +163,12 @@ class SMPL(Module):
     v_shaped = torch.stack(v_shaped, axis=-1)
     
     # J = torch.matmul(self.J_regressor, v_shaped)
-    J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
-    J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
-    J = J - J0
+    if verts is None:
+      J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])])
+    else:
+      J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
+      J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
+      J = J - J0
     R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
 
     if verts is not None:
@@ -208,7 +211,8 @@ class SMPL(Module):
     
     joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
             
-    T = torch.stack([torch.tensordot(results[i], self.weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
+    weights = self.weights if verts is not None else self.weights_all 
+    T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) 
     rest_shape_h = torch.cat(
       (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
     )

utils/checkpoints.py

@@ -39,5 +39,6 @@ def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=No
     model.load_state_dict(checkpoint['model'])
     if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict())
     if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler'])
+    print(f'Loaded model from {checkpoints[-1]} at iteration {iteration}')
     return int(iteration) if iteration != '' else 0
 

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

utils/postprocess.py

@@ -22,7 +22,7 @@ from losses.losses import _nearest_neighbour
 """ COLLISION FUNCTIONS START """
 # F = np.load("smpl/smpl_faces.npy")
 to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
-F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
+F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [4,1,1]), device='cuda')
 smpl_to_keep_orig_post = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
 smpl_to_keep_orig_post = smpl_to_keep_orig_post[smpl_to_keep_orig_post < 6890]
 F_fix_coll_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')[:,smpl_to_keep_orig_post,:]

utils/utils.py

@@ -42,4 +42,10 @@ class InfiniteDataLoader(DataLoader):
             # Dataset exhausted, use a new fresh iterator.
             self.dataset_iterator = super().__iter__()
             batch = next(self.dataset_iterator)
-        return batch
+        return batch
+
+def text_to_html(lines):  
+    html = ''
+    for line in lines:
+        html += "<pre>" + line + "</pre>"
+    return html