generate normals figure

configs/iccvw/abl_all.py

@@ -34,14 +34,16 @@ config['device'] = 'cuda'
 
 config['train_dir'] = '/data-local/127/cloth3d_processed'
 config['valid_dir'] = '/data-local/127/cloth3d_processed'
-config['test_dir'] = '/data-net/dataset/test_ready'
+config['test_dir'] = '/data-local/127_2/cloth3d_processed_test'
 
 if config['is_skirt']:
   config['train_frames'] = 'frames_train_skirt.txt'
   config['valid_frames'] = 'frames_valid_skirt.txt'
+  config['test_frames'] = 'frames_test_skirt.txt'
 else:
   config['train_frames'] = 'frames_train.txt'
   config['valid_frames'] = 'frames_valid.txt'
+  config['test_frames'] = 'frames_test.txt'
 
 config['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
 config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'

configs/iccvw/abl_all_3.py

@@ -33,14 +33,16 @@ config['device'] = 'cuda'
 
 config['train_dir'] = '/data-local/127/cloth3d_processed'
 config['valid_dir'] = '/data-local/127/cloth3d_processed'
-config['test_dir'] = '/data-net/dataset/test_ready'
+config['test_dir'] = '/data-local/127_2/cloth3d_processed_test'
 
 if config['is_skirt']:
   config['train_frames'] = 'frames_train_skirt.txt'
   config['valid_frames'] = 'frames_valid_skirt.txt'
+  config['test_frames'] = 'frames_test_skirt.txt'
 else:
   config['train_frames'] = 'frames_train.txt'
   config['valid_frames'] = 'frames_valid.txt'
+  config['test_frames'] = 'frames_test.txt'
 
 config['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
 config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'

configs/iccvw/abl_all_skirt.py

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

configs/iccvw/pyr_32.py

@@ -30,14 +30,16 @@ 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'
+config['test_dir'] = '/data-local/127_2/cloth3d_processed_test'
 
 if config['is_skirt']:
   config['train_frames'] = 'frames_train_skirt.txt'
   config['valid_frames'] = 'frames_valid_skirt.txt'
+  config['test_frames'] = 'frames_test_skirt.txt'
 else:
   config['train_frames'] = 'frames_train.txt'
   config['valid_frames'] = 'frames_valid.txt'
+  config['test_frames'] = 'frames_test.txt'
 
 config['dataset_in_memory'] = False
 config['resolution'] = None

configs/iccvw/pyr_64.py

@@ -30,14 +30,16 @@ 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'
+config['test_dir'] = '/data-local/127_2/cloth3d_processed_test'
 
 if config['is_skirt']:
   config['train_frames'] = 'frames_train_skirt.txt'
   config['valid_frames'] = 'frames_valid_skirt.txt'
+  config['test_frames'] = 'frames_test_skirt.txt'
 else:
   config['train_frames'] = 'frames_train.txt'
   config['valid_frames'] = 'frames_valid.txt'
+  config['test_frames'] = 'frames_test.txt'
 
 config['dataset_in_memory'] = False
 config['resolution'] = None

configs/iccvw/pyr_64_cape_f.py

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

configs/iccvw/pyr_64_cape_m.py

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

configs/iccvw/pyr_64_skirt.py

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

data/preprocessing.py

@@ -36,94 +36,6 @@ from smpl.smpl_torch import SMPL
 from custom_utils.custom_utils import _maskFaces
 
 
-
-def clean_dataset(folder):
-    ids = os.listdir(folder)
-
-    for id in tqdm(ids):
-
-        if not os.path.isdir(os.path.join(folder, id)):
-            continue
-
-        files = os.listdir(os.path.join(folder, id))
-
-        for file in files:
-            if file.startswith('glmarks') or file.find('nricp') != -1:
-                os.remove(os.path.join(folder, id, file))
-
- 
-def generate_frames_file(folder, file_name):
-    ids = os.listdir(folder)
-    frames = []
-
-    for id in tqdm(ids):
-
-        if not os.path.isdir(os.path.join(folder, id)):
-            continue
-
-        garments = [f.replace('.obj','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('.obj')]
-        info = loadInfo(os.path.join(folder, id, 'info.mat'))
-
-        for garment in garments:
-            for i in range(info['poses'].shape[1]):
-                frames.append(Frame(id, i, garment))
-
-    with open(os.path.join(folder, file_name), 'w') as f:
-        for i, frame in enumerate(frames):
-            f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n')
-
-
-def split_data(folder, ratio):
-    with open(os.path.join(folder, 'frames_all.txt'), "r") as f:
-        frames = f.readlines()
-
-    ids = sorted(list(set([frame.split()[0] for frame in frames])))
-    train_last_id = round(len(ids) * ratio)
-
-    train_frames = []
-    valid_frames = []
-
-    for frame in tqdm(frames):
-        id = frame.split()[0]
-        if id in ids[:train_last_id]:
-            train_frames.append(frame)
-        else:
-            valid_frames.append(frame)
-
-    with open(os.path.join(folder, f'frames_train_{ratio}.txt'), 'w') as f:
-        f.writelines(train_frames)
-
-    with open(os.path.join(folder, f'frames_valid_{round(1-ratio,1)}.txt'), 'w') as f:
-        f.writelines(valid_frames)
-
-
-def genereated_registered_and_mask(folder, is_skirt = False):
-    ids = os.listdir(folder)
-
-    for id in tqdm(ids):
-        garment_dir = os.path.join(folder, id)
-        if not os.path.isdir(garment_dir):
-            continue
-        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and not f.endswith('.mat') and not f.endswith('.obj') and not f.startswith('Combined')]
-        
-        assert len(garments) == 1
-
-        if os.path.exists(os.path.join(garment_dir, garments[0] + '_registered.pc16')):
-            continue
-        
-        garment = garments[0]
-
-        verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
-        map = sio.loadmat(os.path.join(garment_dir, garment + '_map.mat'))['map'].reshape((-1))
-        verts = verts[:, map, :] # ---------------> This is the garment after registration
-        raw_mask = sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].reshape((-1))
-        mask = raw_mask > 0
-        # boundaries = raw_mask == 2
-
-        writePC2(os.path.join(garment_dir, garment + '_registered.pc16'), verts, True)
-        np.save(os.path.join(garment_dir, garment + '_mask.npy'), mask)
-
-
 def create_extras_file(folder):
     # create a new pytables file
     root = '/'.join(folder.split('/')[:-1])
@@ -261,14 +173,15 @@ def create_mesh_body_posed_no_orient(folder, out_folder):
             continue
 
         info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
-    
-
         pose, shape, gender = info['poses'], info['shape'], info['gender']
+        # if pose.shape[0] == 72:
         pose = np.transpose(pose, [1, 0]).astype(np.float32)
 
         pose[:, :3] = 0
         beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
 
+        smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0])
+
         T = smpl(torch.tensor(pose, device=smpl.device), torch.tensor(np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), device=smpl.device))[0]
         T = T.cpu().numpy()
 
@@ -456,6 +369,9 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1)
 
         pose = torch.tensor(pose, dtype=torch.float32, device='cuda')
         shape = torch.tensor(shape, dtype=torch.float32, device='cuda') 
+
+        smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0])
+
         
         for garment in garments:
             if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
@@ -627,6 +543,49 @@ def create_frames_file(folder, out_folder, split=0.9):
                 f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
 
 
+def create_frames_file_test(folder, out_folder):
+    ids = os.listdir(folder)
+
+    frames_non_skirt = []
+    frames_skirt = []
+    ids_non_skirt = set()
+    ids_skirt = set()
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        frames = info['poses'].shape[1]
+
+        assert info['poses'].shape[0] == 72
+
+        for garment in garments:
+            for i in range(frames):
+                if garment == 'Skirt' or garment == 'Dress':
+                    frames_skirt.append(Frame(id, i, garment))
+                    ids_skirt.add(id)
+                else:
+                    frames_non_skirt.append(Frame(id, i, garment))
+                    ids_non_skirt.add(id)
+
+    ids_skirt = list(ids_skirt)
+    ids_non_skirt = list(ids_non_skirt)
+
+    with open(os.path.join(out_folder, 'frames_test.txt'), 'w') as f:
+        for frame in frames_non_skirt:
+            if frame.id in ids_non_skirt:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+    with open(os.path.join(out_folder, 'frames_test_skirt.txt'), 'w') as f:
+        for frame in frames_skirt:
+            if frame.id in ids_skirt:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+
+
 def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1):
     ids = os.listdir(folder)
 
@@ -674,25 +633,26 @@ def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1):
 
 
 def check_irregular_poses(folder, folder_reg, out_folder):
-    frame: Frame = Frame('01512', 50, 'Dress')
+    for i in range(100, 200, 5):
-    mesh_posed = readPC2Frame(os.path.join(out_folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
+        frame: Frame = Frame('01512', i, 'Dress')
-    # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy')
+        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_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 = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V']
-    mesh_body_posed = mesh_body_posed[int(mesh_body_posed.shape[0]/2)+frame.frame_nr]
+        mesh_body_posed = mesh_body_posed[::2][frame.frame_nr]
 
-    F = np.load('smpl/smpl_faces_skirt.npy')
+        F = np.load('smpl/smpl_faces_skirt.npy')
-    _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
+        _, 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 = 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.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(
+        fig.update_layout(scene = dict(
-        xaxis = dict(range=[-1,1],),
+            xaxis = dict(range=[-1,1],),
-                    yaxis = dict(range=[-1,1],),
+                        yaxis = dict(range=[-1,1],),
-                    zaxis = dict(range=[-1,1],),),)
+                        zaxis = dict(range=[-1,1],),),)
-    fig.update_layout(scene_aspectmode='cube')
+        fig.update_layout(scene_aspectmode='cube')
-    fig.show()
+        fig.show()
     pass
 
 
@@ -734,17 +694,88 @@ def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'):
     np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment)
 
 
+def remove_irregular_poses(folder, folder_reg, out_folder):
+    ids = sorted(os.listdir(folder))
+
+    for id in tqdm(ids):
+
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        info = loadInfo(os.path.join(folder, id, 'info.mat'))
+        
+        for garment in garments:
+       
+            verts_count = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'].shape[0]
+            if info['poses'].shape[1] != verts_count:
+                shutil.rmtree(os.path.join(folder, id))
+                if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16')):
+                    os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16'))
+                if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16')):
+                    os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'))
+                if os.path.exists(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16')):
+                    os.remove(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'))
+                if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')):
+                    os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
+                # os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.npy'))
+
+                if len(garments) == 2:
+                    if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16')):
+                        os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16'))
+                    if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16')):
+                        os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16'))
+                    if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.pc16')):
+                        os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.npy'))
+                    break
+                    # os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garments[1]}.npy'))
+
+                
+def double_check_irregular_poses(folder, folder_reg, out_folder):
+    # ids = sorted(os.listdir(folder))
+
+    # for id in tqdm(ids):
+
+    #     garment_dir = os.path.join(folder, id)
+    #     if not os.path.isdir(garment_dir):
+    #         continue
+
+    #     garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+    #     info = loadInfo(os.path.join(folder, id, 'info.mat'))
+        
+    #     for garment in garments:
+    #         body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0]
+    #         verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0]
+
+    #         assert body_count == verts_count
+
+    frames = sorted(os.listdir(out_folder + '/mesh_posed_no_orient'))
+    frames = [(id.split('_')[0], id.split('_')[1].split('.')[0]) for id in frames]
+
+    for frame in tqdm(frames):
+        id = frame[0]
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garment = frame[1]
+        body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0]
+        verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0]
+
+        assert body_count == verts_count
+
 
 if __name__ == '__main__':
 
 
-    # folder = '/data-net/dataset/cloth3d/test'
+    folder = '/data-net/dataset/cloth3d/test'
-    # folder_reg = '/data-local/127/cloth3d_registration/test'
+    folder_reg = '/data-local/127/cloth3d_registration/test'
-    # out_folder = '/data-local/127/cloth3d_processed_test'
+    out_folder = '/data-local/127_2/cloth3d_processed_test'
 
-    folder = '/data-net/dataset/cloth3d/train'
+    # folder = '/data-net/dataset/cloth3d/train'
-    folder_reg = '/data-local/127/cloth3d_registration/train'
+    # folder_reg = '/data-local/127/cloth3d_registration/train'
-    out_folder = '/data-local/127/cloth3d_processed'
+    # out_folder = '/data-local/127/cloth3d_processed'
 
     smpl = SMPL(path='smpl', device='cuda')
 
@@ -754,6 +785,8 @@ if __name__ == '__main__':
     # create_extras_file(folder)
 
     # check_irregular_poses(folder, folder_reg, out_folder)
+    # remove_irregular_poses(folder, folder_reg, out_folder)
+    # double_check_irregular_poses(folder, folder_reg, out_folder)
 
     # create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
     # verify_registration_ans_mask(folder, out_folder)
@@ -762,9 +795,10 @@ if __name__ == '__main__':
     # create_mesh_body_posed_no_orient(folder, out_folder)
     # create_mesh_normals(folder, out_folder)
     # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
-    calculate_average(folder, out_folder)
+    # calculate_average(folder, out_folder)
 
     # create_frames_file(folder, out_folder)
+    create_frames_file_test(folder, out_folder)
     # create_frames_file_partial(folder, out_folder, part=0.5)
     
     # folder = '/data-local/127_2/cape_processed/test/female'

scripts/cvae_pyramid.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_32 import config
+    from configs.iccvw.pyr_64_cape_m import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -549,7 +549,7 @@ def plot_everything(batch: dict, type: str):
         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) }
 
     batch_id = batch['id']
-    batch_frame_nr = batch['frame_nr'].cpu().numpy()
+    batch_frame_nr = batch['frame_nr']
     batch_garment = batch['garment']
     batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))]
     batch_info = text_to_html(batch_info)
@@ -855,7 +855,7 @@ if __name__ == '__main__':
 
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
-    ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + config['name'], top_n=5)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=5)
     scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
 
     g_balancing_coeff = config['starting_balancing_coeff']

scripts/generate_results.py

@@ -0,0 +1,129 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+if args['debug']:
+    from configs.iccvw.pyr_64 import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import argparse
+import wandb
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from PIL import Image
+
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from tqdm import tqdm
+from data.dataset import Dataset
+from utils.checkpoints import CheckpointSaver, load_model
+from torch.utils.data import DataLoader
+from utils.utils import text_to_html
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss, kl_loss
+from pytorch_msssim import MS_SSIM
+from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
+from smpl.smpl_torch import SMPL
+from plotly.subplots import make_subplots
+from models.pyramid_model import Pyramid
+from models.cvae_model import CVAE
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+
+def step_pyramid(batch):
+    model.eval()
+    with torch.inference_mode():
+        levels_input, conds = group_inputs_for_levels(batch, config)
+
+        out_levels = model([levels_input, conds])
+        pred, _, _ = out_levels[-1]
+
+        mesh_pred = uv2mesh(pred)
+
+        if config['use_skinning']:
+            shape = batch['pose_shape'][:, 72:]
+            gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+            shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+            mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+        else:
+            mesh_pred_posed = mesh_pred
+
+    return mesh_pred_posed
+
+
+def step(batch):
+    return None
+
+
+def generate():
+    for batch in tqdm(test_dataloader):
+        if is_pyramid:
+            mesh_pred = step_pyramid(batch)
+        else: 
+            mesh_pred = step(batch)
+
+        ids = batch['id']
+        frame_nrs = batch['frame_nr']
+        garments = batch['garment']
+
+        for i in range(len(ids)):
+            np.savez_compressed(f'{out_folder}/{ids[i]}_{frame_nrs[i]}_{garments[i]}', mesh_pred[i].cpu().numpy())
+
+
+
+
+if __name__ == '__main__':
+
+    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+    dataset_resolutions = config['pyramid']['resolutions'].copy()
+    if 512 not in dataset_resolutions:
+        dataset_resolutions.append(512)
+
+    dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file=config['test_frames'], use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
+
+    is_pyramid = config['name'].startswith('pyr')
+
+    if is_pyramid:
+        model: Pyramid = Pyramid(config)
+    else:
+        model: CVAE = CVAE(config)
+
+    it = load_model(model, 'checkpoints/iccv/torch/' + config['name'])
+    # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
+
+    name = config['name'] + '_' + str(it)
+
+    folder = '/data-local/127_2/results'
+    out_folder = f'{folder}/{name}'
+
+    os.makedirs(out_folder, exist_ok=True)
+
+    generate()

utils/generate_figures.py

@@ -3,6 +3,9 @@ import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
 import numpy as np
+import cmapy as cm
+
+from utils import remove_long_edges
 
 from custom_utils.custom_utils import save_stl, write_obj
 from data.cloth3d_io import readPC2Frame
@@ -12,24 +15,6 @@ 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():
@@ -87,7 +72,7 @@ def plot_normals_example():
     # mesh_unposed = readPC2Frame(os.path.join(folder_test, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
     # mesh_unposed_shaped = readPC2Frame(os.path.join('temp', f'{frame.id}_{frame.garment}_unposed_shaped.pc16'), frame.frame_nr, True)[mesh_mask]
     mesh_normals = readPC2Frame(os.path.join(folder_test, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
-    uv_normals = np.load(os.path.join(folder_test, 'uv_normals', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))
+    # uv_normals = np.load(os.path.join(folder_test, 'uv_normals', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))
 
     F = np.load('smpl/smpl_faces.npy')
     _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
@@ -101,12 +86,13 @@ def plot_normals_example():
     new_mask[mesh_mask] = v_mask
     _, faces = _maskFaces(np.zeros((14475, 3)), F, new_mask)
 
-    fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}]])
+    v_colors = (mesh_normals + 1.4) / 2.8
+    v_colors = v_colors[v_mask]
+    # v_colors = (v_colors - v_colors.min()) / (v_colors.max() - v_colors.min())
 
-    fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), row=1, col=1)
+    fig = go.Figure()
-    fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], 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_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], vertexcolor=v_colors, opacity=1),)
-    fig.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]),
@@ -114,11 +100,22 @@ def plot_normals_example():
         zaxis=dict(range=[-1, 1]),
     ), scene_aspectmode='cube')
 
+    helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
+    fig.update_layout(scene = helper)
+    
+    fig.show()
+
+    
+
+    pass
+
 
 
 if __name__ == '__main__':
     folder_train = '/data-local/127/cloth3d_processed'
-    folder_test = '/data-local/127/cloth3d_processed_test'
+    folder_test = '/data-local/127_2/cloth3d_processed_test'
     
     # plot_unposing_example()
     plot_normals_example()