save all

.vscode/launch.json

@@ -12,7 +12,7 @@
             "console": "integratedTerminal",
             "justMyCode": true, // true false
             "cwd": "${workspaceFolder}",
-            "args": ["--config", "configs.pyr_32", "--debug"]
+            "args": ["--debug"]
         }
     ]
 }

configs/iccvw/abl_all_skirt_var.py

@@ -0,0 +1,197 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '1'
+config['name'] = 'abl_iccvw_all_skirt_var_2'
+config['notes'] = 'fixed collision loss, train from zero'
+config['tags'] = ['iccv', 'abl_all', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = 'checkpoints/iccvw/abl_iccvw_all_skirt_var'
+
+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['ignore_saved_optimizer'] = True
+###############
+config['starting_balancing_coeff'] = 2
+config['balancing_activation_limit'] = 0.015
+
+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_64_128_256_512_abl.py

@@ -0,0 +1,38 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['batch_size'] = 4
+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'
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'mesh_unposed',
+                            'uv_unposed',
+                            'uv_posed',
+                            'uv_normals',
+                            'uv_static',
+                            'mesh_mask',
+                            'uv_mask',
+                        }
+
+
+
+for res in [64, 128, 256]:
+    res_modifier = '_' + str(res)
+    config['signature'].add('uv_body_posed' + res_modifier)
+    config['signature'].add('uv_unposed' + res_modifier)
+    config['signature'].add('uv_posed' + res_modifier)
+    config['signature'].add('uv_normals' + res_modifier)
+    config['signature'].add('uv_static' + res_modifier)
+    config['signature'].add('uv_mask' + res_modifier)
+

configs/iccvw/pyr_64_128_256_512_e2e_var.py

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

configs/iccvw/pyr_64_128_256_512_reg_e2e.py

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

configs/iccvw/pyr_64_128_256_512_reg_skirt_e2e.py

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

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

configs/iccvw/pyr_64_128_256_512_skin_skirt_e2e.py

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

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

data/convert_to_uv.py

@@ -399,10 +399,11 @@ def get_shapedirs():
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127/cloth3d_processed')
+    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cloth3d_processed_test')
     parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
     parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
+    parser.add_argument('--debug', action='store_true')
     args = parser.parse_args()
 
     F = {'non-skirt': np.load("smpl/smpl_faces.npy"), 'skirt': np.load("smpl/smpl_faces_skirt.npy")}
@@ -442,6 +443,7 @@ if __name__ == '__main__':
     #########################################################################################
 
     folder = args.folder
+    # data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed'] # 'mesh_unposed',
     data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
 
     ids = sorted(os.listdir(folder + '/mesh_body_posed_no_orient'))[args.start:args.end]
@@ -460,6 +462,8 @@ if __name__ == '__main__':
     os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
     os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
 
+    # ids = ['00438']
+
 
     for id in tqdm.tqdm(ids):
         garments = glob.glob(os.path.join(folder, 'mesh_mask', f'{id}_*'))

data/convert_to_uv_cape.py

@@ -401,7 +401,7 @@ def get_shapedirs():
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/test/female')
+    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/female')
     parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
     parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
@@ -448,11 +448,11 @@ if __name__ == '__main__':
 
     folder = args.folder
     # data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
-    data_to_convert = ['mesh_static', 'mesh_normals'] # 'mesh_unposed',
+    data_to_convert = ['mesh_posed', 'mesh_static'] # 'mesh_unposed',
 
-    frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))#[args.start:args.end]
+    frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))
     frames = [[frame.split('.')[0], frame.split('.')[1]] for frame in frames]
-    ids = list(set([frame[0] for frame in frames]))
+    ids = sorted(list(set([frame[0] for frame in frames])))[args.start:args.end]
     frames_by_id = {id: [frame[1] for frame in frames if frame[0] == id] for id in ids}
 
     shapedirs = get_shapedirs()
@@ -462,11 +462,11 @@ if __name__ == '__main__':
 
     os.makedirs(folder, exist_ok=True)
     os.makedirs(os.path.join(folder, 'uv_posed_no_orient'), exist_ok=True)
-    os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True)
-    os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True)
-    os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True)
-    os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
-    os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
+    # os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True)
+    # os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True)
+    # os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True)
+    # os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
+    # os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
 
     mask = np.ones((3638), dtype=bool)
 
@@ -476,15 +476,14 @@ if __name__ == '__main__':
         # x = readPC2(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
         # mesh_posed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
         # mesh_posed[:,mask, :] = x
-        if 'mesh_unposed' in data_to_convert:
-            mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
-            mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V']
+        # if 'mesh_unposed' in data_to_convert:
+        #     mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
+        #     mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V']
         if 'mesh_static' in data_to_convert:
             mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static']
-        if 'mesh_normals' in data_to_convert:
-            mesh_normals = np.stack([np.load(os.path.join(folder, 'mesh_normals', f'{id}.{int(i):06d}.npz'))['mesh_normals'] for i in frames_by_id[id]], axis=0)
-        if 'body_posed' in data_to_convert and k != 1:
-            body_posed = readPC2(os.path.join(folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
+        if 'mesh_posed' in data_to_convert:
+            mesh_posed = np.stack([np.load(os.path.join(folder, 'mesh_posed_no_orient', f'{id}.{int(i):06d}.npz'))['mesh_posed_no_orient'] for i in frames_by_id[id]], axis=0)
+
 
         start = time.perf_counter()
 
@@ -493,13 +492,13 @@ if __name__ == '__main__':
         uv_static = uv_static[0].astype(np.float32)
 
         if 'mesh_posed' in data_to_convert:
-            uv_posed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
-        if 'mesh_unposed' in data_to_convert:
-            uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
-        if 'mesh_normals' in data_to_convert:
-            uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
-        if 'body_posed' in data_to_convert and k != 1:
-            uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
+            uv_posed_no_orient = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+        # if 'mesh_unposed' in data_to_convert:
+        #     uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+        # if 'mesh_normals' in data_to_convert:
+        #     uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+        # if 'body_posed' in data_to_convert and k != 1:
+        #     uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
 
         print(time.perf_counter() - start)
 
@@ -515,16 +514,16 @@ if __name__ == '__main__':
         # np.savez_compressed(os.path.join(folder, 'uv_mask',  f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
 
         if 'mesh_posed' in data_to_convert:
-            for i in range(len(uv_posed)):
-                np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient',  f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
-        if 'mesh_unposed' in data_to_convert:
-            for i in range(len(uv_unposed)):
-                np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape',  f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i])
-        if 'mesh_normals' in data_to_convert:
-            for i, frame_nr in enumerate(frames_by_id[id]):
-                np.savez_compressed(os.path.join(folder, 'uv_normals',  f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i])
-        if 'body_posed' in data_to_convert and k != 1:
-            for i in range(len(uv_body_posed)):
-                np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient',  f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
+            for k, i in enumerate(frames_by_id[id]):
+                np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient',  f'{id}.{int(i):06d}.npz'), uv_posed_no_orient=uv_posed_no_orient[k])
+        # if 'mesh_unposed' in data_to_convert:
+        #     for i in range(len(uv_unposed)):
+        #         np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape',  f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i])
+        # if 'mesh_normals' in data_to_convert:
+        #     for i, frame_nr in enumerate(frames_by_id[id]):
+        #         np.savez_compressed(os.path.join(folder, 'uv_normals',  f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i])
+        # if 'body_posed' in data_to_convert and k != 1:
+        #     for i in range(len(uv_body_posed)):
+        #         np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient',  f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
         
         # pass

data/dataset.py

@@ -131,7 +131,7 @@ def read_frames_file(frames_file, shuffle=False):
     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]
+    frames = [Frame(frame_str=line) for line in lines]
 
     if shuffle:
         random.shuffle(frames)
@@ -276,7 +276,7 @@ if __name__ == '__main__':
                   'uv_mask',
                   }
     
-    dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt')
+    dataset = Dataset('/data-local/127_2/cloth3d_processed_test', signature, resolutions=[512], device=device, frames_file='frames_temp.txt')
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
     smpl = SMPL(orig_smpl=False, path='smpl', device='cuda')
 
@@ -293,6 +293,16 @@ if __name__ == '__main__':
     step = 0
     for batch in train_dataloader:
         batch
+
+        from custom_utils.custom_utils import save_stl
+        fix_data = np.load(os.path.join(dataset.folder, 'mesh_fixed', f'{batch["id"][0]}_{batch["garment"][0]}.npz'))
+        faces = fix_data['faces']
+        mask = fix_data['mask']
+
+        save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_posed.stl', batch['mesh_posed'][0].cpu().numpy()[mask], faces)
+        save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_unposed.stl', batch['mesh_unposed'][0].cpu().numpy()[mask], faces)
+        
+
         print(batch['id'][0])
         with profiler:
             m = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], batch['mesh_unposed'])[0]

data/dataset_cape.py

@@ -44,7 +44,10 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
         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 not is_cape:
+            frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines]
+        else:
+            frames = [Frame(line.split('.')[0], line.split('.')[1], 'cape') for line in lines]
 
     if shuffle:
         random.shuffle(frames)
@@ -160,6 +163,7 @@ class DatasetCape(torch.utils.data.Dataset):
             item['garment'] = 'cape'
         if 'pose_shape' in self.signature:
             pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
+            item['global_orient'] = pose_shape[:3].copy()
             pose_shape[:3] = 0
             item['pose_shape'] = pose_shape
         if 'mesh_body_posed' in self.signature:

data/downscale.py

@@ -11,7 +11,7 @@ import plotly.express as ex
 from torchvision.transforms import Resize
 from torchvision.transforms.functional import InterpolationMode
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "4"
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from data.utils import Frame
@@ -165,32 +165,19 @@ def downscale_uv_maps_static(folder, type, sizes):
             np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame[0]}_{frame[1]}.npz'), uv_static = down)
 
 
-def downscale_uv_maps_mask(folder, type, size):
-    os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
-
-    frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
-
-    for frame in tqdm.tqdm(frames):
-        id, _ = frame.split('.')
-        img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
-        down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
-        down = np.squeeze(down.numpy().astype(bool))
-        np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down)
-
-
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127/cloth3d_processed')
+    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cloth3d_processed_test')
     parser.add_argument("-s", "--start", help = "start of range", default=0)
-    parser.add_argument("-e", "--end", help = "end of range", default=25000)
+    parser.add_argument("-e", "--end", help = "end of range", default=300000)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
     parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
     args = parser.parse_args()
 
     # sizes = [32, 64, 128, 256]
     device = 'cuda'
-    sizes = [32, 64, 128, 256]
+    sizes = [64, 128, 256]
     folder = args.folder
     resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
     resizers_nearest = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT).to(device) for res in sizes }
@@ -199,19 +186,14 @@ if __name__ == '__main__':
     uv_mask_body = np.load('custom_utils/mask_256x512.npy')
     uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
     uv_masks_body[512] = uv_mask_body.astype(bool)
-    for res in [32, 64, 128, 256]:
+    for res in [64, 128, 256]:
         down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:].astype(np.int32)).to(device))
         down = np.squeeze(down.cpu().numpy().astype(bool))
         uv_masks_body[res] = down
         
-
     uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes }
     uv_mask_store[512] = UVMaskStore(folder, size=512)
 
-    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32)
-    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64)
-    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128)
-    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
     # downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
 
     downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)

data/downscale_cape.py

@@ -23,7 +23,7 @@ def get_img_idxs():
     img_idxs[512] = img_idx
     img_coords[512] = coords
     uv_masks[512] = uv_mask.astype(np.int32)
-    for res in [32, 64, 128, 256]:
+    for res in [64, 128, 256]:
         down = resizers_nearest[res](torch.from_numpy(uv_mask[None,:,:]).to(device)).cpu().numpy()
         down = np.squeeze(down.astype(bool))
         coords = np.transpose(np.stack(np.nonzero((1 - down))), (1,0))
@@ -37,7 +37,7 @@ def get_img_idxs():
     uv_mask_body = np.load('custom_utils/mask_256x512.npy')
     uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
     uv_masks_body[512] = uv_mask_body
-    for res in [32, 64, 128, 256]:
+    for res in [64, 128, 256]:
         down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:]).to(device)).cpu().numpy()
         down = np.squeeze(down.astype(bool))
         uv_masks_body[res] = down
@@ -113,6 +113,8 @@ def downscale_uv_normals(folder, type, sizes):
     frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
 
     for frame in tqdm.tqdm(frames):
+        if not frame.startswith('03223_shortshort_chicken_wings'):
+            continue
         img = np.load(os.path.join(folder, type, frame)); img = img['uv_normals']
         img = fill(img, uv_masks_body[512])
         for size in sizes:
@@ -140,6 +142,23 @@ def downscale_uv_maps_unposed(folder, type, sizes):
             np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down)
 
 
+def downscale_uv_maps_posed(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_no_orient']
+        mask_full = uv_mask_store[512][None]
+        img = fill(img, mask_full)
+        for size in sizes:
+            mask = uv_mask_store[size][None]
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_no_orient = down)
+
+
 def downscale_uv_maps_static(folder, type, sizes):
     for size in sizes:
         os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
@@ -173,7 +192,7 @@ def downscale_uv_maps_mask(folder, type, size):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/male')
+    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/female')
     parser.add_argument("-s", "--start", help = "start of range", default=0)
     parser.add_argument("-e", "--end", help = "end of range", default=100000)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
@@ -182,7 +201,7 @@ if __name__ == '__main__':
 
     device = 'cuda'
 
-    sizes = [32, 64, 128, 256]
+    sizes = [64, 128, 256]
     folder = args.folder
 
     resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
@@ -199,9 +218,10 @@ if __name__ == '__main__':
     # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
     # downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
 
-    downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
-    downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes)
-    downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
+    # downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
+    downscale_uv_maps_posed(folder=folder, type='uv_posed_no_orient', sizes=sizes)
+    # downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes)
+    # downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
 
 
     print()

data/preprocessing.py

@@ -12,6 +12,7 @@ args = parser.parse_args()
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 import math
 import shutil
@@ -76,6 +77,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
     os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
     os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
 
+    ids = ['00021']
+
     for id in tqdm(ids):
         # id = '01467'
 
@@ -87,8 +90,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
         info = loadInfo(os.path.join(folder, id, 'info.mat'))
         
         for garment in garments:
-            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
-                continue
+            # if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+            #     continue
         
             verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
             # if info['poses'].shape[1] == verts.shape[0]:
@@ -488,6 +491,176 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1)
             pass
 
 
+F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda')
+F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda')
+F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [1,1,1]), device='cuda')
+F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [1,1,1]), device='cuda')
+
+from scipy.spatial import KDTree
+from pytorch3d.structures import Meshes
+from torch_cluster import knn
+
+def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
+    # flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
+    batch_size = V.shape[0]
+    vertices_count = V.shape[1]
+    batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(vertices_count)
+    batch_y = torch.arange(batch_size, device=V.get_device()).repeat_interleave(14475)
+    V = V.reshape(-1, V.shape[-1])
+    B = B.reshape(-1, B.shape[-1])
+    idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_y, batch_y=batch_x)[1]
+    idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count)
+    return idx
+
+
+def collision_loss(V: torch.Tensor, B: torch.Tensor, thr: float = .004):
+    F_norm_local = F_norm
+    mesh = Meshes(B, F_norm_local)
+    B_vn = torch.stack(mesh.verts_normals_list())
+   
+    # compute correspondences
+    idx = _nearest_neighbour(V, B)
+    # compute loss
+    D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
+    b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
+    dot = torch.einsum('abc,abc->ab', D, b_vn)
+    thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))
+    loss = torch.sum(thresholded ** 2, axis=-1)
+
+    return loss
+
+
+def area_loss(pred, gt, faces):
+    """
+    Calculates the area loss between predicted and ground truth meshes.
+
+    Args:
+        pred (torch.Tensor): Predicted mesh vertices of shape (batch_size, num_vertices, 3).
+        gt (torch.Tensor): Ground truth mesh vertices of shape (batch_size, num_vertices, 3).
+        faces (torch.Tensor): Mesh faces of shape (num_faces, 3).
+
+    Returns:
+        torch.Tensor: Area loss between predicted and ground truth meshes.
+    """
+    # Calculate the area of each face for predicted and ground truth meshes
+    pred_faces = pred[:, faces, :]
+    gt_faces = gt[:, faces, :]
+    pred_edges1 = pred_faces[:, :, 1, :] - pred_faces[:, :, 0, :]
+    pred_edges2 = pred_faces[:, :, 2, :] - pred_faces[:, :, 0, :]
+    gt_edges1 = gt_faces[:, :, 1, :] - gt_faces[:, :, 0, :]
+    gt_edges2 = gt_faces[:, :, 2, :] - gt_faces[:, :, 0, :]
+    pred_cross = torch.cross(pred_edges1, pred_edges2, dim=2)
+    gt_cross = torch.cross(gt_edges1, gt_edges2, dim=2)
+    pred_area = 0.5 * torch.norm(pred_cross, dim=2)
+    gt_area = 0.5 * torch.norm(gt_cross, dim=2)
+
+    # Calculate the area loss
+    area_loss = torch.mean(torch.abs(pred_area - gt_area))
+
+    return area_loss
+
+
+
+def create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=0, end=1):
+    ids = sorted(os.listdir(folder))[start:end]
+    os.makedirs(os.path.join(out_folder, 'mesh_static_unposed_no_shape'), exist_ok=True)
+    F = np.load('smpl/smpl_faces.npy')
+    # F_skirt = np.load('smpl/smpl_faces_skirt.npy')
+
+    ids = ['00005']
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        shape = np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:]
+
+        shape = torch.tensor(shape, dtype=torch.float32, device='cuda') 
+
+        smpl = SMPL(path='smpl', device='cuda', batch_size=1)
+        pose = np.zeros((24, 3), np.float32)
+        pose[1,2] = .15
+        pose[2,2] = -.15
+        pose = torch.tensor(pose.flatten()[None,:], dtype=torch.float32, device='cuda')
+        
+
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue    
+
+            mesh_static, faces, _, _ = readOBJ(os.path.join(folder, id, f'{garment}.obj'))
+            faces = quads2tris(faces)
+            mesh_static = np.matmul(mesh_static, np.array([[1,0,0],[0,0,-1],[0,1,0]]))
+
+            unposed = mesh_static
+
+            unposed = torch.tensor(unposed[None,:,:], dtype=torch.float32, device='cuda')
+            unposed.requires_grad = True
+            mesh_static = torch.tensor(mesh_static[None,:,:], dtype=torch.float32, device='cuda')
+
+            optimizer = torch.optim.Adam([unposed], lr=0.0001)
+
+            for i in range(10000):
+                    optimizer.zero_grad()
+
+                    body = smpl(pose, shape)[0]
+                    c_loss = collision_loss(unposed, body)
+                    a_loss = area_loss(unposed, mesh_static, faces)
+                    loss : torch.Tensor = c_loss + a_loss
+
+                    # Backward pass
+                    loss.backward()
+                    optimizer.step()
+
+                    pose = pose / 1.0005
+
+                    # Print the loss
+                    if i % 50 == 0:
+                        pose_sum = np.abs(pose.detach().cpu().numpy()).sum().item()
+                        loss_str = loss.detach().cpu().numpy()[0]
+                        c_loss_str = c_loss.detach().cpu().numpy()[0]
+                        a_loss_str = a_loss.detach().cpu().numpy()
+                        print(f"it: {i} loss: {loss_str:.10f} c_loss: {c_loss_str:.10f}, a_loss: {a_loss_str:.8f}, pose: {pose_sum:.8f}")
+
+
+            unposed = unposed.detach().cpu().numpy()
+            body = body.detach().cpu().numpy()
+
+            gender = int(shape[0].cpu().numpy()[-1] > 0)
+            shape_offsets = torch.tensordot(shape[:,:10], smpl.shapedirs[gender], dims=[[1], [2]])
+            shape_offsets = shape_offsets.cpu().numpy()[:,smpl.to_keep,:]
+
+            unshaped = (unposed - shape_offsets)
+
+            writePC2(os.path.join(out_folder, 'mesh_static_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True)
+
+            verts = unposed
+            # verts = unshaped
+            frame = 0
+            fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
+            # add body mesh
+            fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,2], z=body[frame,:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1))
+            fig.update_layout(scene = dict(
+                xaxis = dict(range=[-1,1],),
+                            yaxis = dict(range=[-1,1],),
+                            zaxis = dict(range=[-1,1],),),)
+            fig.update_layout(scene_aspectmode='cube')
+            fig.show()
+
+            # verts = pred.detach().cpu().numpy()[:,mask,:]
+            # frame = 9
+            # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
+            # fig.update_layout(scene = dict(
+            #     xaxis = dict(range=[-1,1],),
+            #                 yaxis = dict(range=[-1,1],),
+            #                 zaxis = dict(range=[-1,1],),),)
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.show()
+            pass
+
+
 def create_frames_file(folder, out_folder, split=0.9):
     ids = os.listdir(folder)
 
@@ -879,13 +1052,13 @@ def generate_fixed_masks(folder, out_folder):
 if __name__ == '__main__':
 
 
-    # folder = '/data-net/dataset/cloth3d/test'
-    # folder_reg = '/data-local/127/cloth3d_registration/test'
-    # out_folder = '/data-local/127_2/cloth3d_processed_test'
+    folder = '/data-net/dataset/cloth3d/test'
+    folder_reg = '/data-local/127/cloth3d_registration/test'
+    out_folder = '/data-local/127_2/cloth3d_processed_test'
 
-    folder = '/data-net/dataset/cloth3d/train'
-    folder_reg = '/data-local/127/cloth3d_registration/train'
-    out_folder = '/data-local/127/cloth3d_processed'
+    # folder = '/data-net/dataset/cloth3d/train'
+    # folder_reg = '/data-local/127/cloth3d_registration/train'
+    # out_folder = '/data-local/127/cloth3d_processed'
 
     smpl = SMPL(path='smpl', device='cuda')
 
@@ -898,13 +1071,14 @@ if __name__ == '__main__':
     # remove_irregular_poses(folder, folder_reg, out_folder)
     # double_check_irregular_poses(folder, folder_reg, out_folder)
 
-    # create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
+    create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
     # verify_registration_ans_mask(folder, out_folder)
     # create_pose_shape(folder, out_folder)
     # create_mesh_static(folder, out_folder)
     # create_mesh_body_posed_no_orient(folder, out_folder)
     # create_mesh_normals(folder, out_folder)
     # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
+    # create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=args.start, end=args.end)
     # calculate_average(folder, out_folder)
 
     # create_frames_file(folder, out_folder)
@@ -915,4 +1089,4 @@ if __name__ == '__main__':
     # folder = '/data-local/127_2/cape_processed/test/female'
     # create_mesh_normals_cape(folder, None)
 
-    generate_fixed_masks(folder, out_folder)
+    # generate_fixed_masks(folder, out_folder)

data/utils.py

@@ -1,10 +1,26 @@
 import tables as tb
 
 class Frame:
-    def __init__(self, id: str, frame_nr: int, garment: str=''):
-        self.id: str = id
-        self.frame_nr: int = frame_nr
-        self.garment: str = garment
+    def __init__(self, id: str ='', frame_nr: int= 0, garment: str='', frame = None, frame_str: str = ''):
+        if frame is not None:
+            self.id: str = frame.id
+            self.frame_nr: int = frame.frame_nr
+            self.garment: str = frame.garment
+            assert id == '' and frame_nr == 0 and garment == '' and frame_str == ''
+        elif frame_str != '':
+            if ' ' in frame_str:
+                self.id: str = frame_str.split()[0]
+                self.frame_nr: int = int(frame_str.split()[1])
+                self.garment: str = frame_str.split()[2]
+            else:
+                self.id: str = frame_str.split('_')[0]
+                self.frame_nr: int = int(frame_str.split('_')[1])
+                self.garment: str = frame_str.split('_')[2]
+            assert id == '' and frame_nr == 0 and garment == '' and frame is None
+        else:
+            self.id: str = id
+            self.frame_nr: int = frame_nr
+            self.garment: str = garment
 
     def __hash__(self):
         return hash((self.id, self.frame_nr, self.garment))

docker/Dockerfile2

@@ -9,8 +9,11 @@ ARG USER_ID=11017
 ARG GROUP_NAME=guest
 ARG GROUP_ID=11000
 
+# fix for tzdata asking for user input
+ARG DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC
+
 # Update the system and install required packages
-RUN apt-get update && apt-get install -y sudo git apt-utils htop
+RUN apt-get update && apt-get install -y sudo git apt-utils htop ffmpeg libsm6 libxext6
 
 # Create a new group with the same GID as the host
 RUN groupadd --gid $GROUP_ID $GROUP_NAME

docker/Dockerfile_hood

@@ -0,0 +1,48 @@
+# Use an existing image as the base
+# FROM nvcr.io/nvidia/pytorch:23.04-py3
+# FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
+# 
+# Add the current user to the image
+ARG USER_NAME=lhunor
+ARG USER_ID=11017
+ARG GROUP_NAME=guest
+ARG GROUP_ID=11000
+
+# Update the system and install required packages
+RUN apt-get update && apt-get install -y sudo git apt-utils htop
+
+# Create a new group with the same GID as the host
+RUN groupadd --gid $GROUP_ID $GROUP_NAME
+
+# Create a new user with the same UID and GID as the host
+RUN useradd --uid $USER_ID --gid $GROUP_ID --shell /bin/bash --create-home $USER_NAME
+
+# don't require password with sudo, for convenience
+# not the safest thing to do, but hopefully okay inside the container
+RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
+
+# Add the new user to the sudo group
+RUN usermod -aG sudo $USER_NAME
+
+# install conda stuff still as root
+RUN conda update -n base conda
+RUN conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y
+RUN conda install -c bottler nvidiacub -y
+RUN conda install pytorch3d -c pytorch3d -y
+
+# Set the default user to the new user
+USER $USER_NAME
+
+WORKDIR /workspace
+
+ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
+RUN pip install --upgrade pip
+
+# defining the volumes that will be mounted
+VOLUME [ "/data-net" ]
+VOLUME [ "/data-local" ]
+VOLUME [ "/home-net" ]
+
+# Start the container with a bash shell
+CMD ["/bin/bash"]

docker/docker-compose.yml

@@ -104,4 +104,25 @@ services:
       resources:
         reservations:
           devices:
-            - capabilities: [gpu]
+            - capabilities: [gpu]
+
+  lhunor_hood:
+    image: lhunor_hood:pt2.0.1
+    container_name: lhunor_hood_pt2.0.1
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '/home-net/lhunor:/home-net'
+      - '/data-net/guest/lhunor/data:/data-net'
+      - '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/127-data2/lhunor/data:/data-local/127_2'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]

losses/losses.py

@@ -32,7 +32,7 @@ def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
     # flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
     batch_size = V.shape[0]
     vertices_count = V.shape[1]
-    batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(V.shape[1])
+    batch_x = torch.arange(batch_size, device=V.device).repeat_interleave(V.shape[1])
     V = V.reshape(-1, V.shape[-1])
     B = B.reshape(-1, B.shape[-1])
     idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_x)[1]
@@ -58,7 +58,9 @@ def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: fl
         F_norm_local = F_norm
     elif B.shape[1] == 3638:
         F_norm_local = F_norm_orig
-    mesh = Meshes(B, F_norm_local)
+    if V.device != F_norm_local.device:
+        F_norm_local = F_norm_local.to(V.device)
+    mesh = Meshes(B, F_norm_local[:B.shape[0],:,:])
     B_vn = torch.stack(mesh.verts_normals_list())
    
     # compute correspondences
@@ -74,6 +76,54 @@ def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: fl
     return loss, ratio
 
 
+def _nearest_neighbour_error(V: torch.Tensor, B: torch.Tensor):
+    # flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
+    batch_size = V.shape[0]
+    vertices_count = V.shape[1]
+    batch_x = torch.arange(batch_size, device=V.device).repeat_interleave(B.shape[1])
+    batch_y = torch.arange(batch_size, device=V.device).repeat_interleave(V.shape[1])
+    V = V.reshape(-1, V.shape[-1])
+    B = B.reshape(-1, B.shape[-1])
+    idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_y)[1]
+    idx = (idx - batch_y * vertices_count).reshape(batch_size, vertices_count)
+    return idx
+
+# def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
+#     # create an empty tensor to store the indices of the nearest neighbors
+#     batch_size = V.shape[0]
+#     vertices_count = V.shape[1]
+#     idx = torch.zeros((batch_size, vertices_count), dtype=torch.long, device=V.device)
+
+#     # loop over the batch dimension and find the nearest neighbor for each point in B
+#     for i in range(batch_size):
+#         tree = KDTree(B[i].cpu().numpy())
+#         _, nearest = tree.query(V[i].cpu().numpy(), k=1)
+#         idx[i] = torch.from_numpy(nearest)
+
+#     return idx
+
+def collision_error(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = 0.0005):
+    if B.shape[1] == 14475:
+        F_norm_local = F_norm
+    elif B.shape[1] == 3638:
+        F_norm_local = F_norm_orig
+    if V.device != F_norm_local.device:
+        F_norm_local = F_norm_local.to(V.device)
+    mesh = Meshes(B, F_norm_local[:B.shape[0],:,:])
+    B_vn = torch.stack(mesh.verts_normals_list())
+   
+    # compute correspondences
+    idx = _nearest_neighbour_error(V, B)
+    # compute loss
+    D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
+    b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
+    dot = torch.einsum('abc,abc->ab', D, b_vn)
+    thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))
+    loss = torch.sum(thresholded ** 2, axis=-1)
+    ratio = torch.sum((thresholded < 0).to(torch.float32)) / V.shape[1]
+
+    return loss, ratio
+
 
 """ COLLISION FUNCTIONS END """
 
@@ -218,9 +268,9 @@ def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
     else:
         raise NotImplementedError
 
-    mesh = Meshes(pred, f_norm)
+    mesh = Meshes(pred, f_norm[:pred.shape[0], :, :])
     pred_norm = torch.stack(mesh.verts_normals_list())
-    mesh = Meshes(gt, f_norm)
+    mesh = Meshes(gt, f_norm[:pred.shape[0], :, :])
     gt_norm = torch.stack(mesh.verts_normals_list())
 
     return torch.mean(torch.mean((torch.sum(torch.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))

models/cvae_model.py

@@ -55,28 +55,24 @@ class CVAE(nn.Module):
         if self.avg_garment is not None:
             assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
 
-    # @tf.function
-    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
+    
+    def sample(self, uv_template, uv_body_posed, uv_normals, uv_mask, noise=None):
         if noise is None:
-            eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
+            eps = torch.randn(uv_template.shape[0], self.latent_shape[0], self.latent_shape[1], self.latent_shape[2]).to(uv_template.device)
         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)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
+        z = torch.cat([eps, z_cond], dim=1)
+        x = self.decoder(z)
 
-        x = self.decoder(z, training=False)
-        
         if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-        
+            x = x + self.avg_garment[None,:,:,:]
+    
         if self.use_offsets_over_body:
             x = x + uv_body_posed
 
-        x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,None,:,:]
 
         return x
 
@@ -94,11 +90,11 @@ class CVAE(nn.Module):
         return out
 
 
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+    def encode_conds(self, uv_template, uv_body_posed, uv_normals):
         z_body = self.body_enc(uv_body_posed)
         z_static = self.static_enc(uv_template)
         if self.use_normals:
-            z_norm = self.norm_enc(uv_curv)
+            z_norm = self.norm_enc(uv_normals)
             z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
         else:
             z_cond = torch.concat([z_static, z_body], dim=1)

models/pyramid_level.py

@@ -39,22 +39,16 @@ class PyramidLevel(nn.Module):
 
     def sample(self, z_cond, uv_mask, noise=None):
         if noise is None:
-            eps = torch.normal(shape=(uv_mask.shape[0], self.latent_dim))
+            eps = torch.normal(mean=0.0, std=1.0, size=z_cond.shape).to(uv_mask.device)
         else:
             eps = noise
 
-        z = torch.concat([eps, z_cond], dim=-1)
-
+        z = torch.cat([eps, z_cond], dim=1)
         x = self.decoder(z)
-        
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[None,:,:,:]
-        
-        x = x * uv_mask[:,None:,:]
+        x = x * uv_mask[:,None,:,:]
 
         return x
 
-
     def reparameterize(self, mean, logvar, fix_noise=False):
         if self.training == True:
             if not fix_noise:

models/pyramid_model.py

@@ -10,6 +10,29 @@ from models.pyramid_level import PyramidLevel
 from torchvision.transforms import Resize
 from torchvision.transforms.functional import InterpolationMode
 
+# class SampleParallel(nn.Module):
+#     def __init__(self, levels):
+#         super().__init__()
+#         self.levels = nn.ModuleList(levels)
+
+#     def forward(self, z_cond, uv_masks, noise=None):
+#         level_outs = nn.parallel.data_parallel(self._forward_level, [(self.levels[i], z_cond, uv_masks[i], noise) for i in range(len(uv_masks))], dim=0, device_ids=[0, 0, 0, 0], output_device=0)
+#         # level_outs = nn.parallel.data_parallel(self.levels, [(z_cond, uv_masks[i], noise) for i in range(len(uv_masks))], device_ids=[0, 0, 0, 0], output_device=0)
+#         out = torch.zeros((uv_masks[0].shape[0], 3, 64, 32)).to(uv_masks[0].device)
+#         for i in range(len(self.levels)):
+#             level_out = level_outs[i]
+#             if i != len(self.levels) - 1:
+#                 out = out + level_out
+#                 out = self.levels[i].enhance(out, uv_masks[i])
+#                 out = self.resizers[self.levels[i+1].size](out)
+#                 out = out * uv_masks[i+1][:,None,:,:]
+#             else:
+#                 out = out + level_out
+#         out = out * uv_masks[0][:,None,:,:]
+#         return out
+    
+#     def _forward_level(self, level, z_cond, uv_mask, noise=None):
+#         return level(z_cond, uv_mask, noise=noise)
 
 class Pyramid(nn.Module):
     
@@ -42,6 +65,58 @@ class Pyramid(nn.Module):
 
         self.use_normals = config['use_normals']
 
+
+
+    
+    def sample(self, uv_template, uv_body_posed, uv_normals, uv_mask, noise=None):
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
+        # uv_masks = [uv_mask[64], uv_mask[128], uv_mask[256], uv_mask[512]]
+
+        # sampleP = SampleParallel(self.levels)
+        # out = sampleP(z_cond, uv_masks, noise=noise)
+        # return out
+    
+        assert len(self.levels) == 4
+        out = torch.zeros((uv_template.shape[0], 3, 64, 32)).to(uv_template.device)
+
+        for i in range(len(self.levels)):
+            level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise[i] if noise is not None else None)
+            if i != len(self.levels) -  1:
+                out = out + level_out 
+                out = self.levels[i].enhance(out, uv_mask[self.levels[i].size])
+                out = self.resizers[self.levels[i+1].size](out)
+                out = out * uv_mask[self.levels[i+1].size][:,None,:,:]
+            else:
+                out = out + level_out
+
+        out = out * uv_mask[512][:,None,:,:]
+
+        return out
+
+
+    def sample_level(self, uv_template, uv_body_posed, uv_normals, uv_mask, prev, level, noise=None):
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
+    
+        assert len(self.levels) == 4
+
+        i = level
+        out = prev
+
+        level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise[i] if noise is not None else None)
+        if i != len(self.levels) -  1:
+            out = out + level_out 
+            out = self.levels[i].enhance(out, uv_mask[self.levels[i].size])
+            out = self.resizers[self.levels[i+1].size](out)
+            out = out * uv_mask[self.levels[i+1].size][:,None,:,:]
+        else:
+            out = out + level_out
+            out = out * uv_mask[512][:,None,:,:]
+
+        return out
+
+
     def encode_conds(self, uv_template, uv_body_posed, uv_norm):
         z_static = self.static_enc(uv_template)
         z_body = self.body_enc(uv_body_posed)

run.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.cvae_norm_1 import config
+from configs.iccvw.pyr_64_128_256_512_skirt_e2e_var import config
 import argparse
 
 # parser = argparse.ArgumentParser(description='')
@@ -12,14 +12,15 @@ import argparse
 if __name__ == '__main__':
     import wandb
     # wandb.login(key='******************************************')
-    wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838')
+    # wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838')
+    # wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
     name = config['name']
     script = config['script']
     config = [key for key in sys.modules.keys()  if key.startswith('configs.')][-1]
     current_file = '/'.join(__file__.split('/')[-2:])
     
-    print(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &')
-    if args['debug']:
-        os.system(f'python {script} --name lhunor_{name} --config {config}')
-    else:
-        os.system(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &')
+    print(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/iccvw/{name}.out 2>&1 &')
+    # if args['debug']:
+    #     os.system(f'python {script} --name lhunor_{name} --config {config}')
+    # else:
+    #     os.system(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &')

run_parallel.py

@@ -12,19 +12,24 @@ if __name__ == '__main__':
     # ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:]
     skirt = 0
     obj_folder = '/data-net/dataset/cloth3d/train'
-    folder = '/data-local/127/train'
+    folder = '/data-local/127_2/cape_processed/train/male'
+    # folder = '/data-local/127_2/cloth3d_processed_test'
 
     count = 0
-    step = 100
+    step = 35
 
     # !!!!!!! change skirt parameter in script
 
-    while count <1200:
+    while count < 550:
         # os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv_local.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
-        os.system(f'nohup python data/calculate_normals.py -f {folder} -f_obj {obj_folder}  -s {count} -e {(count+step)} -n normals_train_calc_{count}_{(count+step)} > nohup/normals/normals_train_calc_{count}_{(count+step)}.out' + ' 2>&1 &')
+        # os.system(f'nohup python data/convert_to_uv_cape.py -f {folder} -s {count} -e {(count+step)} -n to_uv_cape_m_test_{count}_{(count+step)} > nohup/to_uv/cape_m_test_{count}_{(count+step)}.out' + ' 2>&1 &')
+        # os.system(f'nohup python data/convert_to_uv.py -f {folder} -s {count} -e {(count+step)} -n to_uv_{count}_{(count+step)} > nohup/to_uv/test_{count}_{(count+step)}.out' + ' 2>&1 &')
+        # os.system(f'nohup python data/downscale.py -f {folder} -s {count} -e {(count+step)} -n down_test_{count}_{(count+step)} > nohup/down/test_{count}_{(count+step)}.out' + ' 2>&1 &')
+        # os.system(f'nohup python data/downscale_cape.py -f {folder} -s {count} -e {(count+step)} -n down_train_cape_m_{count}_{(count+step)} > nohup/down/train_cape_m_{count}_{(count+step)}.out' + ' 2>&1 &')
+        os.system(f'nohup python scripts/calculate_chamfer.py -s {count} -e {(count+step)} -n hood_130k_{count}_{(count+step)} > nohup/hood/hood_130k_{count}_{(count+step)}.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_valid_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_valid_skirt_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')

scripts/calculate.py

@@ -0,0 +1,34 @@
+import numpy as np
+
+error_1 = dict()
+error_2 = dict()
+#end2end with reg
+# error_1[0] = {'Jumpsuit': 0.00888, 'Trousers': 0.01193, 'Tshirt': 0.01065, 'Top': 0.00725}
+# error_1[1] = {'Jumpsuit': 0.00491, 'Trousers': 0.00704, 'Tshirt': 0.00531, 'Top': 0.00393}
+# error_1[2] = {'Jumpsuit': 0.00379, 'Trousers': 0.00547, 'Tshirt': 0.00404, 'Top': 0.00309}
+# error_1[3] = {'Jumpsuit': 0.00349, 'Trousers': 0.00502, 'Tshirt': 0.00382, 'Top': 0.00281}
+# error_2[0] = {'Dress': 0.01009, 'Skirt': 0.01379}
+# error_2[1] = {'Dress': 0.0055, 'Skirt': 0.0077}
+# error_2[2] = {'Dress': 0.00418, 'Skirt': 0.00582}
+# error_2[3] = {'Dress': 0.00359, 'Skirt': 0.00496}
+
+#end2end without reg
+error_1[0] = {'Jumpsuit': 0.00927, 'Trousers': 0.01229, 'Tshirt': 0.01113, 'Top': 0.0076}
+error_1[1] = {'Jumpsuit': 0.00503, 'Trousers': 0.00721, 'Tshirt': 0.0054, 'Top': 0.004}
+error_1[2] = {'Jumpsuit': 0.00402, 'Trousers': 0.0059, 'Tshirt': 0.00421, 'Top': 0.00319}
+error_1[3] = {'Jumpsuit': 0.00347, 'Trousers': 0.00528, 'Tshirt': 0.00365, 'Top': 0.00271}
+error_2[0] = {'Dress': 0.01005, 'Skirt': 0.01395}
+error_2[1] = {'Dress': 0.0056, 'Skirt': 0.00787}
+error_2[2] = {'Dress': 0.00441, 'Skirt': 0.00634}
+error_2[3] = {'Dress': 0.00379, 'Skirt': 0.0054}
+
+for i in range(4):
+    errors_mean = dict()
+    error = {**error_1[i], **error_2[i]}
+    for category in error.keys():
+        errors_mean[category] = round(np.mean(error[category]), 5)
+    print(f'Level {i}: {errors_mean}')
+
+    mean = round(np.mean(list(errors_mean.values())), 5)
+
+    print(f'Final level {i}: {mean}')

scripts/calculate_area_diff.py

@@ -0,0 +1,47 @@
+import trimesh
+import numpy as np
+
+cases = [
+    # ('06461_237_Trousers', 'temp/sampling_pyramid/06461_237_Trousers', [0,8,12,13]),
+    # ('01171_213_Dress', 'temp/sampling_pyramid/01171_213_Dress', [0,6,10,15]),
+    ('01852_70_Tshirt_lvl_0', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,10,1,5,13]),
+    ('01852_70_Tshirt_lvl_1', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,14,3,5,0 ]),
+    ('01852_70_Tshirt_lvl_2', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,12,0,8,6]),
+    ('01852_70_Tshirt_lvl_3', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,2,8,1,0]),
+]
+
+for case in cases:
+    # Load the original mesh
+    folder = case[1]
+    original_mesh = trimesh.load(f'{folder}/{case[0]}_{case[2][0]}.stl')
+
+    # Load the variation meshes
+    variation_meshes = []
+    for i in case[2][1:]:
+        variation_meshes.append(trimesh.load(f'{folder}/{case[0]}_{i}.stl'))
+
+    # Calculate area of original mesh
+    # original_area = original_mesh.area
+
+    # Calculate area difference for each variation
+    area_diffs = []
+    for variation_mesh in variation_meshes:
+        # calculate per vertex euclid distance
+        euclid = np.linalg.norm(original_mesh.vertices - variation_mesh.vertices, axis=1) * 1000
+        euclid_diff = np.mean(euclid)
+
+
+        # area_diff = abs(variation_mesh.area - original_area) / original_area
+        area_diffs.append(euclid_diff)
+
+    # Average area differences
+    # avg_area_diff = np.mean(area_diffs)
+
+    variance = np.var(area_diffs)
+
+    # Print results
+    print(f'Case: {case[0]}')
+    # print('Original area:', original_area)
+    print('Euclidean differences:', area_diffs)
+    print('Euclidean variance:', variance)
+    # print('Average Euclidean difference:', avg_area_diff)

scripts/calculate_chamfer.py

@@ -0,0 +1,263 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+
+
+import time
+import argparse
+
+import tqdm
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+
+
+import torch
+from tqdm import tqdm
+from data.utils import Frame
+from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
+from utils.utils import pickle_load
+from scipy.spatial import cKDTree
+from smpl.smpl_torch import SMPL
+from scipy.spatial.transform import Rotation as R
+import trimesh as tm
+from losses.losses import collision_loss, normal_loss, kl_loss, collision_error
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
+parser.add_argument("-e", "--end", help = "end of range", default=550, type=int)
+parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
+parser.add_argument('--debug', action='store_true')
+args = parser.parse_args()
+
+def chamfer(points1, points2, points2_tree: cKDTree):
+    tree1 = cKDTree(points1)
+    dist1, _ = tree1.query(points2, workers=10)
+    dist2, _ = points2_tree.query(points1, workers=10)
+    chamfer_dist = np.mean(dist1) + np.mean(dist2)
+    return chamfer_dist
+
+
+smpl = SMPL(path='smpl', device='cpu')
+def deorient(verts, pose, inverse=True):
+    assert pose.shape[1] == 72
+    rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy()
+    rotated = np.zeros_like(verts)
+    for frame in range(verts.shape[0]):
+        r = R.from_matrix(rodr[frame, 0, :, :])
+        rotated[frame] = r.apply(verts[frame], inverse=inverse)
+
+    return rotated
+
+
+def get_collision_depth(m1, m2):
+    cm = tm.collision.CollisionManager()
+    cm.add_object('m1', m1)
+    cm.add_object('m2', m2)
+    data = cm.in_collision_internal(return_data=True)[1]
+    # c = np.asarray([d.index('m2') for d in data])
+    # count = np.unique(c).shape[0]
+    depth_sum = np.sum([d.depth for d in data])
+    return depth_sum
+
+
+def get_area_diff(m1, m2):
+    return np.abs(m1.area - m2.area)
+
+device = 'cuda'
+def get_collision_loss(pred, body, mask):
+    c_loss, c_ratio = collision_error(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device))
+    return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
+
+def get_collision_loss_2(pred, body, mask):
+    _pred = np.zeros((14475, 3), dtype=np.float32)
+    _pred[mask] = pred
+    c_loss, c_ratio = collision_loss(torch.from_numpy(_pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005)
+    return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
+
+def calculate():
+    errors_hood = []
+    collisions_hood = []
+    areas_hood = []
+    collision_ratios_hood = []
+    if calculate_ours:
+        errors_ours = []
+        collisions_ours = []
+        areas_ours = []
+        collision_ratios_ours = []
+
+    garments = sorted(os.listdir(hood_folder))[args.start:args.end]
+    garments = [g[:-4] for g in garments]
+    body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
+
+    for garment in tqdm(garments):
+        picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
+        id, garment_type = garment.split('_')
+
+        pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl')
+        pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces']
+        gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
+        template, template_faces, _, _ = readOBJ(os.path.join(cloth3d_folder, id, garment_type + '.obj'))
+        mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
+        pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:]
+        gt = deorient(gt, pose, True)
+        body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:]
+        fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
+        faces = fix_data['faces']
+
+        m_template = tm.Trimesh(vertices=template, faces=template_faces, process=False, validate=False)
+
+        for i in range(pred_hood.shape[0]):
+            gt_tree = cKDTree(gt[i])
+
+            try:
+                if calculate_ours:
+                    # pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
+                    pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0']
+            except:
+                print('missed one')
+                continue
+
+            # m_body = tm.Trimesh(vertices=body[i], faces=body_faces, process=False, validate=False)
+            
+            errors_hood.append((f'{id}_{i}_{garment_type}', chamfer(pred_hood[i], gt[i], gt_tree)))
+            m_hood = tm.Trimesh(vertices=pred_hood[i], faces=pred_hood_faces, process=False, validate=False)
+            # collision_hood = get_collision_depth(m_body, m_hood)
+            # collisions_hood.append((f'{id}_{i}_{garment_type}', collision_hood))
+            area_hood = get_area_diff(m_template, m_hood)
+            areas_hood.append((f'{id}_{i}_{garment_type}', area_hood))
+            collision_ratio_hood = get_collision_loss(pred_hood[i], body[i], mask)[1]
+            collision_ratios_hood.append((f'{id}_{i}_{garment_type}', collision_ratio_hood))
+            
+            if calculate_ours:
+                errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask], gt[i], gt_tree)))
+                m_ours = tm.Trimesh(vertices=pred_ours[mask], faces=faces, process=False, validate=False)
+                # collision_ours = get_collision_depth(m_body, m_ours)
+                # collisions_ours.append((f'{id}_{i}_{garment_type}', collision_ours))
+                area_ours = get_area_diff(m_template, m_ours)
+                areas_ours.append((f'{id}_{i}_{garment_type}', area_ours))
+                collision_ratio_ours = get_collision_loss(pred_ours[mask], body[i], mask)[1]
+                # collision_ratio_ours_2 = get_collision_loss_2(pred_ours[mask], body[i], mask)[1]
+                collision_ratios_ours.append((f'{id}_{i}_{garment_type}', collision_ratio_ours))
+
+
+            # fig = go.Figure(data=[go.Mesh3d(x=pred_ours[mask][:,0], y=pred_ours[mask][:,2], z=pred_ours[mask][:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
+            # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
+            # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+            # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+            # fig.show()
+
+            # fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
+            # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
+            # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+            # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+            # fig.show()
+
+            # fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
+            # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
+            # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+            # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+            # fig.show()
+            # pass
+
+            # fig = go.Figure()
+            # fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
+            # fig.add_trace(go.Scatter3d(x=pred_hood[i][:, 0], y=pred_hood[i][:, 2], z=pred_hood[i][:, 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()
+            # pass
+
+            # fig = go.Figure()
+            # fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
+            # fig.add_trace(go.Scatter3d(x=pred_ours[i,mask,:][:, 0], y=pred_ours[i,mask,:][:, 2], z=pred_ours[i,mask,:][:, 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()
+
+
+    # save errors to a csv file
+    os.makedirs(f'{results_folder}/hood/{name}_hood', exist_ok=True)
+    os.makedirs(f'{results_folder}/hood/{name}_hood/chamfer', exist_ok=True)
+    os.makedirs(f'{results_folder}/hood/{name}_hood/collision', exist_ok=True)
+    os.makedirs(f'{results_folder}/hood/{name}_hood/area', exist_ok=True)
+    os.makedirs(f'{results_folder}/hood/{name}_hood/collision_ratio', exist_ok=True)
+
+    for frame, error in errors_hood:
+        np.save(f'{results_folder}/hood/{name}_hood/chamfer/{frame}.npy', error)
+    for frame, error in collisions_hood:
+        np.save(f'{results_folder}/hood/{name}_hood/collision/{frame}.npy', error)
+    for frame, error in areas_hood:
+        np.save(f'{results_folder}/hood/{name}_hood/area/{frame}.npy', error)
+    for frame, error in collision_ratios_hood:
+        np.save(f'{results_folder}/hood/{name}_hood/collision_ratio/{frame}.npy', error)
+
+    os.makedirs(f'{results_folder}/hood/{name}_ours', exist_ok=True)
+    os.makedirs(f'{results_folder}/hood/{name}_ours/chamfer', exist_ok=True)
+    os.makedirs(f'{results_folder}/hood/{name}_ours/collision', exist_ok=True)
+    os.makedirs(f'{results_folder}/hood/{name}_ours/area', exist_ok=True)
+    os.makedirs(f'{results_folder}/hood/{name}_ours/collision_ratio', exist_ok=True)
+
+    for frame, error in errors_ours:
+        np.save(f'{results_folder}/hood/{name}_ours/chamfer/{frame}.npy', error)
+    for frame, error in collisions_ours:
+        np.save(f'{results_folder}/hood/{name}_ours/collision/{frame}.npy', error)
+    for frame, error in areas_ours:
+        np.save(f'{results_folder}/hood/{name}_ours/area/{frame}.npy', error)
+    for frame, error in collision_ratios_ours:
+        np.save(f'{results_folder}/hood/{name}_ours/collision_ratio/{frame}.npy', error)
+
+    # with open(f'{results_folder}/hood/{name}/chamfer_hood.csv', 'w') as f:
+    #     for frame, error in errors_hood:
+    #         frame_str = ','.join(str(frame).split('_'))
+    #         f.write(f'{frame_str},{error}\n')
+
+    # with open(f'{results_folder}/hood/{name}/collision_hood.csv', 'w') as f:
+    #     for frame, error in collisions_hood:
+    #         frame_str = ','.join(str(frame).split('_'))
+    #         f.write(f'{frame_str},{error}\n')
+
+    # with open(f'{results_folder}/hood/{name}/area_hood.csv', 'w') as f:
+    #     for frame, error in areas_hood:
+    #         frame_str = ','.join(str(frame).split('_'))
+    #         f.write(f'{frame_str},{error}\n')
+
+    # if calculate_ours:
+    #     with open(f'{results_folder}/hood/{name}_chamfer_ours.csv', 'w') as f:
+    #         for frame, error in errors_ours:
+    #             frame_str = ','.join(str(frame).split('_'))
+    #             f.write(f'{frame_str},{error}\n')
+
+    #     with open(f'{results_folder}/hood/{name}/collision_ours.csv', 'w') as f:
+    #         for frame, error in collisions_ours:
+    #             frame_str = ','.join(str(frame).split('_'))
+    #             f.write(f'{frame_str},{error}\n')
+
+    #     with open(f'{results_folder}/hood/{name}/area_ours.csv', 'w') as f:
+    #         for frame, error in areas_ours:
+    #             frame_str = ','.join(str(frame).split('_'))
+    #             f.write(f'{frame_str},{error}\n')
+
+
+
+if __name__ == '__main__':
+
+    name = 'hood_3_var'
+
+    hood_folder = '/data-local/127_2/results/hood/all_3'
+    ours_folder = dict()
+    ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_var_3_450000'
+    ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_var_2_600000'
+    cloth3d_folder = '/data-net/dataset/cloth3d/test'
+
+    results_folder = '/data-local/127_2/results'
+
+    calculate_ours = True
+
+
+
+    calculate()

scripts/calculate_error.py

@@ -9,11 +9,11 @@ parser.add_argument('--name', type=str)
 parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 parser.add_argument('--iteration', type=str, default='')
-parser.add_argument('-g', '--gpu', type=str)
+parser.add_argument('-g', '--gpu', type=str, default='1')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_cape_f_2 import config_1
+    from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_1
 else:
     config_module = importlib.import_module(args['config'])
     config_1 = config_module.config
@@ -65,19 +65,32 @@ def calculate():
         # fig.update_layout(scene_aspectmode='cube')
         # fig.show()
 
-        pred = np.load(f'{out_folder}/{frame}.npz')['arr_0']
+        pred = np.load(f'{out_folder}/{frame}.npz', allow_pickle=True)['arr_0']
+
+        try:
+            pred.item()['lvl3']
+            all_levels = True
+        except:
+            all_levels = False
+
         gt = readPC2Frame(f'{config_1["test_dir"]}/mesh_posed_no_orient/{frame.id}_{frame.garment}.pc16', frame.frame_nr, True)
-        # mask = np.load(f'{config["test_dir"]}/mesh_mask/{frame.id}_{frame.garment}.npy')
 
         fix_data = np.load(f'{config_1["test_dir"]}/mesh_fixed/{frame.id}_{frame.garment}.npz')
         mask_convert = fix_data['mas_convert']
         new_mask = fix_data['mask']
-
-        pred = pred[new_mask]
         gt = gt[mask_convert]
 
-        error = np.linalg.norm(pred - gt, axis=1).mean()
-        errors.append((str(frame), error))
+        if not all_levels:
+            pred = pred[new_mask]
+
+            error = np.linalg.norm(pred - gt, axis=1).mean()
+            errors.append((str(frame), error))
+        else:
+            for i in range(4):
+                _pred = pred.item()[f'lvl{i}'][new_mask]
+                error = np.linalg.norm(_pred - gt, axis=1).mean()
+                errors.append((f'{frame}_lvl{i}', error))
+
 
     # save errors to a csv file
     with open(f'{results_folder}/errors/{name}.csv', 'w') as f:

scripts/calculate_error_all.py

@@ -0,0 +1,263 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('--iteration', type=str, default='200000')
+parser.add_argument('-g', '--gpu', type=str, default='0')
+parser.add_argument('-f', '--final', action='store_true')
+args = vars(parser.parse_args())
+
+if args['debug']:
+    from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_1
+else:
+    config_module = importlib.import_module(args['config'])
+    config_1 = config_module.config
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = args['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 trimesh as tm
+
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from tqdm import tqdm
+from data.dataset import Dataset
+from data.utils import Frame
+from data.cloth3d_io import readPC2Frame, readPC2
+from 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 custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
+from plotly.subplots import make_subplots
+
+if config_1['is_skirt']:
+    from losses.losses import edge_loss_skirt as edge_loss
+else:
+    from losses.losses import edge_loss
+from losses.losses import collision_loss, normal_loss, kl_loss
+
+device = 'cuda'
+
+def get_edge_loss(pred, gt, mask):
+    e_loss = torch.mean(edge_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(gt).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device)))
+    return e_loss.cpu().numpy()
+
+def get_normal_loss(pred, gt, mask):
+    n_loss = normal_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(gt).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
+    return n_loss.cpu().numpy()
+
+def get_collision_loss(pred, body, mask):
+    c_loss, c_ratio = collision_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005)
+    return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
+
+
+def calculate():
+    errors_collision = []
+    errors_collision_ratio = []
+    errors_normal = []
+    errors_edge = []
+
+    for frame in tqdm(frames):
+
+        # import plotly.graph_objects as go
+        # mesh_pred = mesh_pred.cpu().numpy()
+        # mask = batch['mesh_mask'].cpu().numpy()
+        # i = 3
+        # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,mask[i],0], y=mesh_pred[i,mask[i],2], z=mesh_pred[i,mask[i],1], mode='markers', marker=dict(size=1))])
+        # fig.update_layout(scene_aspectmode='cube')
+        # fig.show()
+
+        pred = np.load(f'{out_folder}/{frame}.npz', allow_pickle=True)['arr_0']
+
+        try:
+            pred.item()['lvl3']
+            all_levels = True
+        except:
+            all_levels = False
+
+        gt = readPC2Frame(f'{config_1["test_dir"]}/mesh_posed_no_orient/{frame.id}_{frame.garment}.pc16', frame.frame_nr, True)
+        body = readPC2Frame(f'{config_1["test_dir"]}/mesh_body_posed_no_orient/{frame.id}.pc16', frame.frame_nr, True)
+
+        fix_data = np.load(f'{config_1["test_dir"]}/mesh_fixed/{frame.id}_{frame.garment}.npz')
+        mask_convert = fix_data['mas_convert']
+        new_mask = fix_data['mask']
+        gt = gt[mask_convert]
+        gt_full = np.zeros((14475, 3), dtype=np.float32)
+        gt_full[new_mask] = gt
+        gt = gt_full
+
+        if not all_levels:
+            c_loss, c_ratio = get_collision_loss(pred, body, new_mask)
+            errors_collision.append((str(frame), c_loss))
+            errors_collision_ratio.append((str(frame), c_ratio))
+            errors_normal.append((str(frame), get_normal_loss(pred, gt, new_mask)))
+            errors_edge.append((str(frame), get_edge_loss(pred, gt, new_mask)))
+
+        else:
+            levels = [0,1,2,3] if not args['final'] else [3]
+            for i in levels:
+                _pred = pred.item()[f'lvl{i}']
+                c_loss, c_ratio = get_collision_loss(_pred, body, new_mask)
+
+                errors_collision.append((f'{frame}_lvl{i}', c_loss))
+                errors_collision_ratio.append((f'{frame}_lvl{i}', c_ratio))
+                errors_normal.append((f'{frame}_lvl{i}', get_normal_loss(_pred, gt, new_mask)))
+                errors_edge.append((f'{frame}_lvl{i}', get_edge_loss(_pred, gt, new_mask)))
+
+                # faces = fix_data['faces']
+                # body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
+                # m1 = tm.Trimesh(vertices=body, faces=body_faces, process=False, validate=False)
+                # m2 = tm.Trimesh(vertices=_pred[new_mask], faces=faces, process=False, validate=False)
+                # # m2 = tm.Trimesh(vertices=gt[new_mask], faces=faces, process=False, validate=False)
+
+                # cm = tm.collision.CollisionManager()
+                # cm.add_object('m1', m1)
+                # cm.add_object('m2', m2)
+                # data = cm.in_collision_internal(return_data=True)[1]
+                # c = np.asarray([d.index('m2') for d in data])
+                # count = np.unique(c).shape[0]
+                # depth_sum = np.sum([d.depth for d in data])
+
+                # _pred = gt
+
+                # fig = go.Figure(data=[go.Mesh3d(x=_pred[new_mask,0], y=_pred[new_mask,2], z=_pred[new_mask,1], i=faces[:,0], j=faces[:,1], k=faces[:,2])])
+                # fig.add_trace(go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2]))
+                # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+                # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+                # fig.show()
+
+
+
+
+    # save errors to a csv file
+    with open(f'{results_folder}/errors/{name}_collision.csv', 'w') as f:
+        for frame, error in errors_collision:
+            frame_str = ','.join(str(frame).split('_'))
+            f.write(f'{frame_str},{error}\n')
+
+    with open(f'{results_folder}/errors/{name}_collision_ratio.csv', 'w') as f:
+        for frame, error in errors_collision_ratio:
+            frame_str = ','.join(str(frame).split('_'))
+            f.write(f'{frame_str},{error}\n')
+
+    with open(f'{results_folder}/errors/{name}_normal.csv', 'w') as f:
+        for frame, error in errors_normal:
+            frame_str = ','.join(str(frame).split('_'))
+            f.write(f'{frame_str},{error}\n')
+
+    with open(f'{results_folder}/errors/{name}_edge.csv', 'w') as f:
+        for frame, error in errors_edge:
+            frame_str = ','.join(str(frame).split('_'))
+            f.write(f'{frame_str},{error}\n')
+
+
+def calculate_cape():
+    from smpl.smpl_torch import SMPL
+    errors = []
+
+    clothing_verts_idx = np.load(os.path.join('custom_data', 'clothing_verts_idx.npy'))
+    cape_mask = np.zeros((6890), bool)
+    cape_mask[clothing_verts_idx] = True
+
+    smpl = SMPL(path='smpl', orig_smpl=True, device='cpu')
+
+    for frame in tqdm(frames):
+
+        # import plotly.graph_objects as go
+        # mesh_pred = mesh_pred.cpu().numpy()
+        # mask = batch['mesh_mask'].cpu().numpy()
+        # i = 3
+        # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,mask[i],0], y=mesh_pred[i,mask[i],2], z=mesh_pred[i,mask[i],1], mode='markers', marker=dict(size=1))])
+        # fig.update_layout(scene_aspectmode='cube')
+        # fig.show()
+
+        if not os.path.exists(f'{out_folder}/{frame}.npz'):
+            print('missed one')
+            continue
+
+
+
+        pred = np.load(f'{out_folder}/{frame}.npz')['arr_0']
+        gt = np.load(f'{config_1["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
+        pose_shape = np.load(f'{config_1["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32)
+        pose, shape = pose_shape[:72], pose_shape[72:]
+        gender = int(shape[10] >= 0.5)
+        shape_offsets = torch.tensordot(torch.from_numpy(shape[:10][None, :]), smpl.shapedirs[gender], dims=[[1], [2]])[0,smpl.to_keep,:].numpy()
+
+        gt_shaped = gt + shape_offsets 
+        
+        min_body = np.load(os.path.join('/data-local/127_2/cape_processed', 'minimal_body_shape', frame.id.split('_')[0], frame.id.split('_')[0] + '_minimal.npy'))
+        offsets_pred = pred[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:]
+        offsets_orig =  gt_shaped[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:]
+
+        error = np.linalg.norm(offsets_pred - offsets_orig, axis=1).mean()
+        errors.append((str(frame), error))
+
+        # import plotly.graph_objects as go
+        # fig = go.Figure(data=[go.Scatter3d(x=pred[:,0], y=pred[:,2], z=pred[:,1], mode='markers', marker=dict(size=1))])
+        # fig.add_trace(go.Scatter3d(x=gt[:,0], y=gt[:,2], z=gt[:,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()
+
+    # save errors to a csv file
+    with open(f'{results_folder}/errors/{name}.csv', 'w') as f:
+        for frame, error in errors:
+            split = str(frame).split('_')
+            if not config_1['is_cape']:
+                frame_str = ','.join(split)
+            else:
+                id = '_'.join(split[:-2])
+                frame_str = f'{id},{split[-2]},{split[-1]}'
+            f.write(f'{frame_str},{error}\n')
+
+
+if __name__ == '__main__':
+    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
+
+    name = config_1['name']
+
+    results_folder = '/data-local/127_2/results'
+    if args['iteration'] == '':
+        folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config_1['name'])]
+        assert len(folders) == 1
+        name = folders[0]
+    else:
+        name = f'{name}_{args["iteration"]}'
+    out_folder = f'{results_folder}/{name}'
+
+    os.makedirs(f'{results_folder}/errors', exist_ok=True)
+
+    if config_1['is_cape']:
+        from data.dataset_cape import read_frames
+        frames = read_frames(config_1['test_dir'], shuffle=False, is_cape=True)
+        calculate_cape()
+    else:
+        frames_file = config_1['test_frames']
+        with open(os.path.join(config_1['test_dir'], frames_file), 'r') as f:
+            frames = f.readlines()
+        frames = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
+
+        calculate()

scripts/cvae.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.abl_all_4 import config
+    from configs.iccvw.abl_all_5 import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -30,6 +30,7 @@ import numpy as np
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
+from pathlib import Path
 
 if config['is_skirt']:
     from losses.losses import edge_loss_skirt as edge_loss
@@ -44,7 +45,7 @@ from data.dataset import Dataset
 from models.cond_encoder import AutoEncoder
 from utils.checkpoints import CheckpointSaver, load_model
 from utils.utils import InfiniteDataLoader as DataLoader
-from utils.utils import text_to_html
+from utils.utils import text_to_html, save_stl, create_temp_frames
 from utils.postprocess import fix_collision
 from losses.losses import collision_loss, normal_loss
 from pytorch_msssim import MS_SSIM
@@ -277,174 +278,128 @@ def plot_everything(batch: dict, type: str):
     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): 
-    # run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
-    load_model(model,folder)
+
+
+def sample():
+    from data.utils import Frame
+    from utils.utils import save_stl
+
+    def create_temp_frames(folder: str, frame: Frame):
+        with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+            f.writelines([f'{frame}\n' for _ in range(4)])
+
+    model.eval()
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
     count = 0
-    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()
+    frames = ['01852_70_Tshirt']
+    frames = [Frame(frame_str=frame) for frame in frames]
 
 
-        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 frame in frames:
+        create_temp_frames(config['train_dir'], frame)
+        dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[512], is_cape=config['is_cape'])
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+        batch = next(iter(dataloader))
 
+        uv_static = batch['uv_static']
+        uv_body_posed = batch['uv_body_posed']
+        uv_mask = batch['uv_mask']
+        uv_normals = batch['uv_normals']
+
+        batch_size = uv_static.size(0)
+
+
+        for i in range(batch_size):
+            counter = 0
+            frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
+            fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+            mask = fix_data['mask']
+            faces = fix_data['faces']
+
+            gt = batch['mesh_posed'][i].cpu().numpy()[mask]
+            save_stl(Path('temp') / 'sampling' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
+
+            uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_mask_tiled = uv_mask[i][None,:,:].repeat(4, 1, 1)
+            uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
+            pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
         
-        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 _ in range(4):
+                with torch.inference_mode():
+                    pred = model.sample(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled)
+                
+                mesh_pred = uv2mesh(pred)
 
-        # 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
-        #                 ))
+                shape = pose_shape[:, 72:]
+                gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+                shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+                mesh_pred_posed = smpl(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
+
+                for k in range(4):
+                    posed = mesh_pred_posed.cpu().numpy()[k, mask, :]
+                    save_stl(Path('temp') / 'sampling' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', posed, faces)
+                    counter += 1
+
+            pass
 
 
-        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:
+        if count > 100:
             break
 
-        print(count)
-        count += 1
 
-        print()
+def generate():
+    from data.utils import Frame
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
 
+    frame = '01852_70_Tshirt'
+    # frame = '03849_27_Dress' # 03849_27_Dress 00054_1_Dress
+    frame = Frame(frame_str=frame)
+    create_temp_frames(config['train_dir'], frame)
 
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'])
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
 
-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)
+    batch = next(iter(train_dataloader))
 
-    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 = []
+    if config['use_skinning']:
+        mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
+    else:
+        mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
 
-    for batch in tqdm.tqdm(dataset_test):
-        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+    # with record_function("forward"):
+    model.eval()
+    with torch.inference_mode():
+        pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
+    mesh_pred = uv2mesh(pred)
 
-        test_loss = test_step(batch)
-        losses.append(test_loss)
+    if config['use_skinning']:
+        shape = batch['pose_shape'][:, 72:]
+        gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+        shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+        # shape_offsets = torch.stack([torch.matmul(shape[i, :10], smpl.shapedirs[gender[i]][:, smpl.to_keep]) for i in range(shape.shape[0])])
 
-        to_log = { **{ 'current_step': step }, **test_loss }
-        if step % 1000 == 0:
-            figs = plot_everything(batch, 'test')
-            to_log = { **to_log, **figs }
+        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
 
-        wandb.log(to_log, step = step)
-        
-        # print(step)
-        step += 1
-   
-    # averages = dict()
-    path_to_plotly_html = "temp/plotly_figure.html"
+    fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+    mask = fix_data['mask']
+    faces = fix_data['faces']
 
-    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()
+    pred = mesh_pred_posed.cpu().numpy()[0, mask, :]
+    body = batch['mesh_body_posed'].cpu().numpy()[0]
+    body_faces = smpl.f
+    save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_mesh.stl', pred, faces)
+    save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_body.stl', body, body_faces)
     pass
 
 
-
 def train():
     global g_balancing_coeff
     global step
@@ -500,6 +455,11 @@ if __name__ == '__main__':
 
     model = CVAE(config)
 
+    generate()
+    exit()
+
+    sample()
+    exit()
 
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)

scripts/cvae_pyramid.py

@@ -11,7 +11,8 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_cape_f_skin import config
+    # from configs.iccvw.pyr_64_128_256_512 import config
+    from configs.iccvw.pyr_64_128_256_512_e2e import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -24,7 +25,7 @@ if args['debug']:
 import time
 import argparse
 import wandb
-import tqdm
+from tqdm import tqdm
 import random
 import numpy as np
 import plotly.graph_objects as go
@@ -501,6 +502,380 @@ def speed_step(levels_input, conds):
     return model.sample([levels_input, conds], training=False)
 
 
+from scipy.spatial import KDTree
+import trimesh
+
+def connect_unconnected_vertices(mesh):
+    # Find unconnected vertices
+    unconnected_verts = np.where(np.array(mesh.vertex_degree) == 0)[0]
+
+    # Create KD-tree from mesh vertices
+    tree = KDTree(mesh.vertices)
+
+    new_faces = mesh.faces.copy()
+
+    # Connect unconnected vertices to their two closest neighbors
+    for vert in unconnected_verts:
+        # Find the two closest vertices
+        _, closest_verts = tree.query(mesh.vertices[vert], k=3)
+
+        # Create a new face with the unconnected vertex and its two closest neighbors
+        new_face = [vert, closest_verts[1], closest_verts[2]]
+        new_faces = np.vstack((new_faces, new_face))
+
+    mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=new_faces, process=False)
+
+    return mesh
+
+
+def sample():
+    from data.utils import Frame
+    from utils.utils import save_stl
+    from pathlib import Path
+    import trimesh as tm
+
+    def create_temp_frames(folder: str, frame: Frame):
+        with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+            f.writelines([f'{frame}\n' for _ in range(4)])
+
+    model.eval()
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
+    count = 0
+
+    frames = ['06461_237_Trousers'] # 03528_132_Tshirt 01852_70_Tshirt 06461_237_Trousers  01893_173_Jumpsuit
+    frames = ['01171_213_Dress'] # 01039_118_Dress 01171_213_Dress
+    frames = [Frame(frame_str=frame) for frame in frames]
+
+    for frame in frames:
+        create_temp_frames(config['train_dir'], frame)
+        dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+        batch = next(iter(dataloader))
+
+        uv_static = batch['uv_static']
+        uv_body_posed = batch['uv_body_posed']
+        uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
+        uv_normals = batch['uv_normals']
+
+        batch_size = uv_static.size(0)
+        body = batch['mesh_body_posed'][0].cpu().numpy()
+        body_faces = smpl.f
+
+
+        for i in range(batch_size):
+            counter = 0
+            frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
+            fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+            mask = fix_data['mask']
+            faces = fix_data['faces']
+
+            os.makedirs(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True)
+
+            gt = batch['mesh_posed'][i].cpu().numpy()[mask]
+            save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
+            save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces)
+
+            uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)}
+            uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
+            pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
+            body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
+            mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1)
+
+            noise = []
+        
+            for _ in range(4):
+                with torch.inference_mode():
+                    pred = model.sample(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled)
+                
+                mesh_pred = uv2mesh(pred)
+
+                shape = pose_shape[:, 72:]
+                gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+                shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+                mesh_pred_posed = smpl(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
+
+                mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, mask_tiled).cpu().numpy()
+
+                for k in range(4):
+                    _mesh_posed_fixed = mesh_posed_fixed[k, mask, :]
+                    m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False)
+                    m = connect_unconnected_vertices(m)
+                    
+                    tm.smoothing.filter_humphrey(m, iterations=2)
+                    verts = m.vertices
+                    save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces)
+                    save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}_body.stl', body, body_faces)
+                    counter += 1
+
+            pass
+
+
+        if count > 100:
+            break
+
+
+
+def sample_pro():
+    from data.utils import Frame
+    from utils.utils import save_stl
+    from pathlib import Path
+    import trimesh as tm
+
+    def create_temp_frames(folder: str, frame: Frame):
+        with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+            f.writelines([f'{frame}\n' for _ in range(4)])
+
+    model.eval()
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
+    count = 0
+
+    # frames = ['06461_237_Trousers'] # 03528_132_Tshirt 01852_70_Tshirt 06461_237_Trousers  01893_173_Jumpsuit
+    frames = ['01852_70_Tshirt'] # 01039_118_Dress 01171_213_Dress 01852_70_Tshirt
+    frames = [Frame(frame_str=frame) for frame in frames]
+
+    for frame in frames:
+        create_temp_frames(config['train_dir'], frame)
+        dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+        batch = next(iter(dataloader))
+
+        uv_static = batch['uv_static']
+        uv_body_posed = batch['uv_body_posed']
+        uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
+        uv_normals = batch['uv_normals']
+
+        batch_size = uv_static.size(0)
+        body = batch['mesh_body_posed'][0].cpu().numpy()
+        body_faces = smpl.f
+
+
+        i=0
+        frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
+        fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+        mask = fix_data['mask']
+        faces = fix_data['faces']
+
+        os.makedirs(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True)
+
+        gt = batch['mesh_posed'][i].cpu().numpy()[mask]
+        save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
+        save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces)
+
+        uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
+        uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
+        uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)}
+        uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
+        pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
+        body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
+        mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1)
+    
+        out = torch.zeros((uv_static_tiled.shape[0], 3, 64, 32)).to(uv_static_tiled.device)
+        for l in range(4):
+            with torch.inference_mode():
+                for b in range(4):
+                    pred_level = model.sample_level(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled, out, l)
+                    pred = pred_level                     
+
+                    mesh_pred = uv2mesh(pred)
+
+                    shape = pose_shape[:, 72:]
+                    gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+                    shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+                    mesh_pred_posed = smpl(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
+
+                    mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, mask_tiled).cpu().numpy()
+
+                    for k in range(4):
+                        _mesh_posed_fixed = mesh_posed_fixed[k, mask, :]
+                        m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False)
+                        m = connect_unconnected_vertices(m)
+                        
+                        tm.smoothing.filter_humphrey(m, iterations=2)
+                        verts = m.vertices
+                        save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_lvl_{l}_{b*4+k}.stl', verts, faces)
+                        save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_lvl_{l}_{b*4+k}_body.stl', body, body_faces)
+
+                    #replace the contentent of out with its own first element tiled
+                    if b == 3:
+                        out = pred_level[3,:,:,:].repeat(4, 1, 1, 1)
+
+        pass
+
+
+        # if count > 100:
+        #     break
+
+
+
+
+def pred_pro():
+    from data.utils import Frame
+    from utils.utils import save_stl
+    from pathlib import Path
+    import trimesh as tm
+    from typing import List
+    from models.cvae_norm_model import CVAENorm
+    from data.downscale import fill
+    from torchvision.transforms import Resize
+    from torchvision.transforms.functional import InterpolationMode
+
+    def create_temp_frames(folder: str, frames: List[Frame]):
+        with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+            f.writelines([f'{frame}\n' for frame in frames])
+
+    model.eval()
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
+    count = 0
+
+    from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_skirt
+    model_skirt: Pyramid = Pyramid(config_skirt)
+    load_model(model_skirt, f'checkpoints/iccvw/{config_skirt["name"]}')
+    models = {'non-skirt': model, 'skirt': model_skirt}
+    from configs.iccvw.cvae_norm_5 import config as config_norm
+    model_norm = CVAENorm(config_norm)
+    model_norm.eval()
+    load_model(model_norm, f'checkpoints/iccvw/{config_norm["name"]}')
+
+    uv2mesh_skirt = UV2Mesh(res=256, is_skirt=config_skirt['is_skirt'], is_smpl_orig=config['is_cape'])
+    uv2meshes = {'non-skirt': uv2mesh, 'skirt': uv2mesh_skirt}
+
+    # '01320_30_Dress'
+    frames = ['01852_70_Tshirt', '04717_276_Trousers', '07532_152_Tshirt', '02792_282_Jumpsuit'] # 01039_118_Dress 01171_213_Dress
+    frames = [Frame(frame_str=frame) for frame in frames]
+
+    create_temp_frames(config['train_dir'], frames)
+    dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
+    dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False, )
+    batch_orig = next(iter(dataloader))
+    
+    batch_size = 4
+    body_faces = smpl.f
+
+    masks = []
+    faces_orig = []
+    for frame in frames:
+        fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+        masks.append(fix_data['mask'])
+        faces_orig.append(fix_data['faces'])
+    masks = np.stack(masks, axis=0)
+    # faces = np.stack(faces, axis=0)
+
+    resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to('cuda') for res in [64, 128, 256, 512] }
+
+    mode = 'body'
+
+    for i in range(batch_size):
+
+        batch = batch_orig.copy()
+        uv_static = batch['uv_static']
+        uv_body_posed = batch['uv_body_posed']
+        uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
+        uv_normals = batch['uv_normals']
+
+        counter = 0
+        frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
+        # fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+        # mask = fix_data['mask']
+        # faces = fix_data['faces']
+        type = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'non-skirt'
+
+        os.makedirs(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True)
+
+        # gt = batch['mesh_posed'][i].cpu().numpy()[mask]
+        # save_stl(Path('temp') / 'pred_pro' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
+        # save_stl(Path('temp') / 'pred_pro' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces)
+
+        uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
+        uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
+        uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)}
+        uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
+        pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
+        body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
+        # masks_tiled = torch.from_numpy(masks[i][None,:]).to(body_posed.device).repeat(4, 1)
+
+        if mode == 'static':
+            batch['uv_body_posed'] = uv_body_posed_tiled
+            batch['mesh_body_posed'] = body_posed
+            batch['pose_shape'] = pose_shape
+            masks_tiled = masks
+            for res in [64, 128, 256]:
+                batch[f'uv_body_posed_{res}'] = batch[f'uv_body_posed_{res}'][i][None,:,:,:].repeat(4, 1, 1, 1)
+            # orig_img = uv_body_posed_tiled.clone()
+            # downscaled = {64: [], 128: [], 256: []}
+            # for i in range(4):
+            #     filled = fill(orig_img[i].cpu().numpy().transpose(1,2,0), uv_mask[512][i].cpu().numpy())
+            #     filled = torch.from_numpy(filled.transpose(2,0,1)).to(orig_img.device)
+            #     for res in [64, 128, 256]:
+            #         downscaled[res].append(resizers[res](filled) * uv_mask[res][i])
+            # for res in [64, 128, 256]:
+            #     batch[f'uv_body_posed_{res}'] = torch.stack(downscaled[res], dim=0)
+            faces = faces_orig
+        elif mode == 'body':
+            batch['uv_static'] = uv_static_tiled
+            for res in [64, 128, 256]:
+                batch[f'uv_static_{res}'] = batch[f'uv_static_{res}'][i][None,:,:,:].repeat(4, 1, 1, 1)
+            batch['mesh_mask'] = batch['mesh_mask'][i][None,:].repeat(4, 1)
+            for key in uv_mask_tiled.keys():
+                batch[f'uv_mask_{key}'] = uv_mask_tiled[key]
+            batch['uv_mask'] = uv_mask_tiled[512]
+            masks_tiled = torch.from_numpy(masks[i][None,:]).to(body_posed.device).repeat(4, 1).cpu().numpy()
+            faces = [faces_orig[i], faces_orig[i], faces_orig[i], faces_orig[i]]
+        else:
+            raise Exception('mode not supported')
+
+        body_posed = batch['mesh_body_posed']
+        pose_shape = batch['pose_shape']
+
+        with torch.inference_mode():
+            pred_norm = model_norm.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']) * 1.4 
+        batch['uv_normals'] = pred_norm
+        orig_img = pred_norm.clone()
+        downscaled = {64: [], 128: [], 256: []}
+        for i in range(4):
+            filled = fill(orig_img[i].cpu().numpy().transpose(1,2,0), uv_mask[512][i].cpu().numpy())
+            filled = torch.from_numpy(filled.transpose(2,0,1)).to(orig_img.device)
+            for res in [64, 128, 256]:
+                downscaled[res].append(resizers[res](filled) * uv_mask[res][i])
+        for res in [64, 128, 256]:
+            batch[f'uv_normals_{res}'] = torch.stack(downscaled[res], dim=0)
+
+        levels_input, conds = group_inputs_for_levels(batch, config)
+
+        with torch.inference_mode():
+            out_levels = models[type]([levels_input, conds])
+            pred, _, _ = out_levels[-1]
+
+        mesh_pred = uv2meshes[type](pred)
+
+        shape = pose_shape[:, 72:]
+        gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+        shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+        mesh_pred_posed = smpl(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
+
+        # mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, torch.from_numpy(masks).to(body_posed.device)).cpu().numpy()
+        mesh_posed_fixed = mesh_pred_posed.cpu().numpy()
+        for k in range(4):
+            _mesh_posed_fixed = mesh_posed_fixed[k, masks_tiled[k], :]
+            m = tm.Trimesh(_mesh_posed_fixed, faces[k], process=False, validate=False)
+            m = connect_unconnected_vertices(m)
+            
+            tm.smoothing.filter_humphrey(m, iterations=2)
+            verts = m.vertices
+            save_stl(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces[k])
+            save_stl(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}_body.stl', batch['mesh_body_posed'][k].cpu().numpy(), body_faces)
+            counter += 1
+
+    pass
+
+
 
 def train():
     global g_balancing_coeff
@@ -540,9 +915,44 @@ def train():
         step += 1
 
 
+def speed_test():
+    from data.utils import Frame
+    from utils.utils import save_stl
+    from pathlib import Path
+    import trimesh as tm
+
+    def create_temp_frames(folder: str, frame: Frame):
+        with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+            f.writelines([f'{frame}\n' for _ in range(4)])
+
+    model.eval()
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
+
+    frames = ['01852_70_Tshirt']
+    frames = [Frame(frame_str=frame) for frame in frames]
+
+    create_temp_frames(config['train_dir'], frames[0])
+    dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
+    dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+    batch = next(iter(dataloader))
+
+    uv_static = batch['uv_static']
+    uv_body_posed = batch['uv_body_posed']
+    uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
+    uv_normals = batch['uv_normals']
+    
+    for i in tqdm(range(1000)):
+        with torch.inference_mode():
+            if i ==5:
+                start = time.time()
+            model.sample(uv_static, uv_body_posed, uv_normals, uv_mask)
+                
+    print(f'per batch: {(time.time() - start) / 1000}')
+
+
 if __name__ == '__main__':
-    if not args['debug']:
-        init_wandb()
+    # if not args['debug']:
+    #     init_wandb()
 
     use_amp = True
     smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
@@ -560,6 +970,16 @@ if __name__ == '__main__':
 
     model: Pyramid = Pyramid(config)
 
+    # count number of parameters
+    total_params = sum(p.numel() for p in model.parameters())
+
+    # speed_test()
+    # exit()
+
+    # sample_pro()
+    pred_pro()
+    exit()
+
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
     ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
@@ -569,23 +989,23 @@ if __name__ == '__main__':
     reg_coeff = 1
 
     step = 0
-    if config['pretrained_model'] != '':
-        if 'ignore_saved_optimizer' in config.keys()  and config['ignore_saved_optimizer']:
-            step = load_model(model, config['pretrained_model'])
-        else:
-            step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration'])
+    # if config['pretrained_model'] != '':
+    #     if 'ignore_saved_optimizer' in config.keys()  and config['ignore_saved_optimizer']:
+    #         step = load_model(model, config['pretrained_model'])
+    #     else:
+    #         step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration'])
 
-    if len(config['pyramid']['resolutions']) == 1:
-        for param in model.static_enc.parameters(): param.requires_grad = False
-        for param in model.body_enc.parameters(): param.requires_grad = False
-        if config['use_normals']:
-            for param in model.norm_enc.parameters(): param.requires_grad = False
-    else:
-        load_model(model, config['pyramid']['pretrained_pyramid'], iteration=config['pyramid']['pretrained_pyramid_iteration'])
-        for param in model.static_enc.parameters(): param.requires_grad = False
-        for param in model.body_enc.parameters(): param.requires_grad = False
-        if config['use_normals']:
-            for param in model.norm_enc.parameters(): param.requires_grad = False
+    # if len(config['pyramid']['resolutions']) == 1:
+    #     for param in model.static_enc.parameters(): param.requires_grad = False
+    #     for param in model.body_enc.parameters(): param.requires_grad = False
+    #     if config['use_normals']:
+    #         for param in model.norm_enc.parameters(): param.requires_grad = False
+    # else:
+    #     load_model(model, config['pyramid']['pretrained_pyramid'], iteration=config['pyramid']['pretrained_pyramid_iteration'])
+    #     for param in model.static_enc.parameters(): param.requires_grad = False
+    #     for param in model.body_enc.parameters(): param.requires_grad = False
+    #     if config['use_normals']:
+    #         for param in model.norm_enc.parameters(): param.requires_grad = False
 
     # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
     # model.levels[0].upscale_enhancer.trainable = False
@@ -608,4 +1028,5 @@ if __name__ == '__main__':
         # model = torch.compile(model)
         wandb.watch(model, log='all', log_freq=10000, log_graph=True)
 
+    model.eval()
     train()

scripts/cvae_pyramid_e2e.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_256_512_cape_m_e2e import config
+    from configs.iccvw.pyr_64_128_256_512_reg_e2e import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -246,34 +246,49 @@ def train_step(batch):
         if res != 512:
             pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
 
-        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) # levels_input[i][0] = uv_orig
+        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2]))
         kl = kl_loss(mean, logvar)
 
-        mesh_pred = uv2mesh(pred)
+        if config['use_extra_losses']:
+            mesh_pred = uv2mesh(pred)
 
-        if config['use_skinning']:
-            shape = batch['pose_shape'][:, 72:]
-            gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
-            shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+            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_fix']['mesh_mask'][:,:,None]
-            mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
+                mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+            else:
+                mesh_pred_posed = mesh_pred
+
+            e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
+            c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
+            c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+            n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+            l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+            
+            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
+            if not config['is_cape']:
+                reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] 
+                if config['use_norm_loss'] and step > 2000:
+                    reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
+
+            output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss, 
+                            'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
+                            'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, 'train/normal_loss': n_loss,
+                            'train/l2_mesh_loss': l2_mesh_loss, 'train/collision_ratio': c_ratio } }
         else:
-            mesh_pred_posed = mesh_pred
+            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
 
-        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
-        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
-        c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
-        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
-        
-        reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
-        if not config['is_cape']:
-            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] 
-            if config['use_norm_loss'] and step > 2000:
-                reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * n_loss * config['reg_loss_coeff']
-        
         kl_total += kl
+
+        model.eval()
+        with torch.inference_mode():
+            mean = torch.mean(mean)
+            logvar = torch.mean(logvar)
+
+    
     loss: torch.Tensor = reconstr_loss_total + g_balancing_coeff * kl_total
 
     optimizer.zero_grad()
@@ -282,17 +297,45 @@ def train_step(batch):
 
     model.eval()
     with torch.inference_mode():
-        mean = torch.mean(mean)
-        logvar = torch.mean(logvar)
+        for i, res in enumerate(config['pyramid']['resolutions']):
+            # i = -1
+            res = config['pyramid']['resolutions'][i]
+            res_modifier = '_' + str(res) if res != 512 else ''
+            pred, mean, logvar = out_levels[i]
+            if res != 512:
+                pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
 
-        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
-        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
+            mean = torch.mean(mean)
+            logvar = torch.mean(logvar)
 
-    output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
-                                'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
-                                'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, 'train/normal_loss': n_loss,
-                                'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss, 'train/ms_ssim' + res_modifier: ms_ssim,
-                                'train/collision_ratio': c_ratio } } # 
+            if not config['use_extra_losses']:
+                mesh_pred = uv2mesh(pred)
+
+                if config['use_skinning']:
+                    shape = batch['pose_shape'][:, 72:]
+                    gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+                    shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
+                    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+                else:
+                    mesh_pred_posed = mesh_pred
+
+                e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
+                c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
+                c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+                l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+                n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+                
+                output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,
+                                            'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
+                                            'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, 'train/normal_loss': n_loss,
+                                            'train/l2_mesh_loss': l2_mesh_loss, 'train/collision_ratio': c_ratio } }
+
+            ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
+            metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
+
+            output = { **output, **{ 'train/euclid_error' + res_modifier: metric, 'train/ms_ssim' + res_modifier: ms_ssim } }
 
 
     return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output }
@@ -356,18 +399,21 @@ def valid_step(batch):
                 reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
 
             kl_total += kl
+
+            mean = torch.mean(mean)
+            logvar = torch.mean(logvar)
+
+            ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
+            metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
+
+            output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss, 
+                                        'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
+                                        'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss, 'valid/normal_loss': n_loss,
+                                        'valid/euclid_error' + res_modifier: metric, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/ms_ssim' + res_modifier: ms_ssim, } } # 
+
+
         loss = reconstr_loss_total + g_balancing_coeff * kl_total
 
-        mean = torch.mean(mean)
-        logvar = torch.mean(logvar)
-
-        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
-        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
-
-        output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 
-                                    'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
-                                    'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss, 'valid/normal_loss': n_loss,
-                                    'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/ms_ssim' + res_modifier: ms_ssim, } } # 
 
 
 
@@ -501,7 +547,7 @@ def train():
         valid_accumulator.update(valid_loss)
 
         if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
-            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step, 'balancing_coeff': g_balancing_coeff } })
             if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
                 train_figs = plot_everything(batch, 'train')
                 valid_figs = plot_everything(data_valid, 'valid')
@@ -515,10 +561,9 @@ def train():
             valid_accumulator.reset()
             print(step)
 
-        # if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
-        #     g_balancing_coeff += 1 / 50000 
+        g_balancing_coeff += 1 / 100000 
 
-        if step % 5000 == 0 and not args['debug']:
+        if step % 10000 == 0 and not args['debug']:
             ckpt_saver(model, step, optimizer, None)
 
         step += 1
@@ -547,7 +592,7 @@ if __name__ == '__main__':
     # optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
-    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=20)
     scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
     resizer = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device'])
 

scripts/generate_results.py

@@ -10,10 +10,11 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--iteration', type=str, default='')
 parser.add_argument('--debug', action='store_true')
 parser.add_argument('-g', '--gpu', type=str)
+parser.add_argument('-l', '--level', type=int, default=-1)
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_cape_f_2 import config_1
+    from configs.iccvw.pyr_64_128_256_512_e2e_var import config_1
 else:
     config_module = importlib.import_module(args['config'])
     config_1 = config_module.config
@@ -62,7 +63,7 @@ def step_pyramid(batch):
         levels_input, conds = group_inputs_for_levels(batch, config_1)
 
         out_levels = model([levels_input, conds])
-        pred, _, _ = out_levels[-1]
+        pred, _, _ = out_levels[args['level']]
 
         mesh_pred = uv2mesh(pred)
 
@@ -167,7 +168,10 @@ if __name__ == '__main__':
     model.eval()
     # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
 
-    name = config_1['name'] + '_' + str(it)
+    if args['level'] == -1:
+        name = config_1['name'] + '_' + str(it)
+    else:
+        name = config_1['name'] + '_lvl' + str(args['level']) + '_' + str(it)
 
     folder = '/data-local/127_2/results'
     out_folder = f'{folder}/{name}'

scripts/generate_results_all.py

@@ -10,6 +10,7 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--iteration', type=str, default='')
 parser.add_argument('--debug', action='store_true')
 parser.add_argument('-g', '--gpu', type=str, default='1')
+parser.add_argument('-fr', '--frames_file', type=str, default='')
 args = vars(parser.parse_args())
 
 if args['debug']:
@@ -117,7 +118,8 @@ if __name__ == '__main__':
     # remove all values from config['signature] that start with mesh_
     config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
 
-    dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
+    frames_file = args['frames_file'] if args['frames_file'] != '' else config_1['test_frames']
+    dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=frames_file, use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
     # dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
     test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
 

scripts/get_error.py

@@ -10,10 +10,14 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 parser.add_argument('--iteration', type=str, default='')
 parser.add_argument('-g', '--gpu', type=str)
+parser.add_argument('-n', '--normal', action='store_true')
+parser.add_argument('-e', '--edge', action='store_true')
+parser.add_argument('-cr', '--collision_ratio', action='store_true')
+parser.add_argument('-a', '--all', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_cape_f_2 import config_1
+    from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
 else:
     config_module = importlib.import_module(args['config'])
     config_1 = config_module.config
@@ -31,28 +35,47 @@ from data.utils import Frame
 
 
 
-def calculate():
-    errors = dict()
-    errors_mean = dict()
+def calculate(type=''):
+    all_levels = args['all']
+    if all_levels:
+        errors = {0: dict(), 1: dict(), 2: dict(), 3: dict()}
+        errors_mean = {0: dict(), 1: dict(), 2: dict(), 3: dict()}
+    else:
+        errors = dict()
+        errors_mean = dict()
 
-    with open(f'{results_folder}/errors/{name}.csv', 'r') as f:
+    with open(f'{results_folder}/errors/{name}{type}.csv', 'r') as f:
         lines = f.readlines()
 
     for line in lines:
         line = line.split(',')
         category = line[2]
-        if category not in errors:
-            errors[category] = []
-        errors[category].append(float(line[3]))
+        if all_levels:
+            level = int(line[3][-1])
+            if category not in errors[level].keys():
+                errors[level][category] = []
+            errors[level][category].append(float(line[4]))
+        else:
+            if category not in errors.keys():
+                errors[category] = []
+            errors[category].append(float(line[3]))
 
-    for category in errors:
-        errors_mean[category] = np.mean(errors[category])
+    if all_levels:
+        for i in range(4):
+            for category in errors[i]:
+                errors_mean[i][category] = round(np.mean(errors[i][category]), 5)
+            print(f'Level {i} error{type}: {errors_mean[i]}')
 
-    print(errors_mean) 
+            mean = np.mean(list(errors_mean[i].values()))
+            print(f'Final error{type} level {i}: {mean}')
+        pass
+    else:
+        for category in errors:
+            errors_mean[category] = round(np.mean(errors[category]), 5)
+        print(errors_mean) 
 
-    mean = np.mean(list(errors_mean.values()))
-
-    print(f'Final error: {mean}')
+        mean = round(np.mean(list(errors_mean.values())), 5)
+        print(f'Final error{type}: {mean}')
 
 
 
@@ -78,5 +101,10 @@ if __name__ == '__main__':
         name = f'{name}_{args["iteration"]}'
     out_folder = f'{results_folder}/{name}'
 
-
+    if args['normal']:
+        calculate('_normal')
+    if args['edge']:
+        calculate('_edge')
+    if args['collision_ratio']:
+        calculate('_collision_ratio')
     calculate()

scripts/get_error_all.py

@@ -0,0 +1,102 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('--iteration', type=str, default='190000')
+parser.add_argument('-g', '--gpu', type=str)
+parser.add_argument('-a', '--all', action='store_true')
+args = vars(parser.parse_args())
+
+if args['debug']:
+    from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
+else:
+    config_module = importlib.import_module(args['config'])
+    config_1 = config_module.config
+
+
+import time
+import argparse
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from data.utils import Frame
+
+
+
+
+def calculate():
+    all_levels = args['all']
+    if all_levels:
+        errors = {0: dict(), 1: dict(), 2: dict(), 3: dict()}
+        errors_mean = {0: dict(), 1: dict(), 2: dict(), 3: dict()}
+    else:
+        errors = dict()
+        errors_mean = dict()
+
+
+    with open(f'{results_folder}/errors/{name}.csv', 'r') as f:
+        lines = f.readlines()
+
+    for line in lines:
+        line = line.split(',')
+        category = line[2]
+        if all_levels:
+            level = int(line[3][-1])
+            if category not in errors:
+                errors[level][category] = []
+            errors[level][category].append(float(line[4]))
+        else:
+            if category not in errors:
+                errors[category] = []
+            errors[category].append(float(line[3]))
+
+    if all_levels:
+        for i in range(4):
+            for category in errors[i]:
+                errors_mean[i][category] = round(np.mean(errors[i][category]), 5)
+            print(f'Level {i}: {errors_mean[i]}')
+
+            mean = np.mean(list(errors_mean[i].values()))
+            print(f'Final error level {i}: {mean}')
+    else:
+        for category in errors:
+            errors_mean[category] = round(np.mean(errors[category]), 5)
+        print(errors_mean) 
+
+        mean = np.mean(list(errors_mean.values()))
+        print(f'Final error: {mean}')
+
+
+
+if __name__ == '__main__':
+    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
+
+    name = config_1['name']
+    if not config_1['is_cape']:
+        frames_file = config_1['test_frames']
+        with open(os.path.join(config_1['test_dir'], frames_file), 'r') as f:
+            frames = f.readlines()
+        frames = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
+    else:
+        from data.dataset_cape import read_frames
+        frames = read_frames(config_1['test_dir'], shuffle=False, is_cape=True)
+
+    results_folder = '/data-local/127_2/results'
+    if args['iteration'] == '':
+        folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config_1['name'])]
+        assert len(folders) == 1
+        name = folders[0]
+    else:
+        name = f'{name}_{args["iteration"]}'
+    out_folder = f'{results_folder}/{name}'
+
+
+    calculate()

scripts/get_error_hood.py

@@ -0,0 +1,66 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import time
+import argparse
+from tqdm import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from data.utils import Frame
+
+
+
+
+def calculate(type='chamfer'):
+    errors = dict()
+    errors_mean = dict()
+
+    files_hood = sorted(os.listdir(f'{folder}/{run}_hood/{type}'))
+    errors_hood = [np.load(f'{folder}/{run}_hood/{type}/{file}') for file in tqdm(files_hood)]
+    files_ours = sorted(os.listdir(f'{folder}/hood_3_var_ours/{type}'))
+    errors_ours = [np.load(f'{folder}/hood_3_var_ours/{type}/{file}') for file in tqdm(files_ours)]
+
+    for error, frame in zip(errors_hood, files_hood):
+        frame_id, frame_nr, frame_garment = frame.split('_')
+        category = frame_garment
+
+        if category not in errors:
+            errors[category] = []
+        errors[category].append(error)
+
+    for category in errors:
+        errors_mean[category] = round(np.mean(errors[category]), 5)
+    print(errors_mean) 
+
+    mean = round(np.mean(list(errors_mean.values())), 5)
+    print(f'Final HOOD error_{type}: {mean}')
+
+    # errors = dict()
+    # errors_mean = dict()
+
+    # for error, frame in zip(errors_ours, files_ours):
+    #     frame_id, frame_nr, frame_garment = frame.split('_')
+    #     category = frame_garment
+
+    #     if category not in errors:
+    #         errors[category] = []
+    #     errors[category].append(error)
+
+    # for category in errors:
+    #     errors_mean[category] = round(np.mean(errors[category]), 5)
+    # print(errors_mean) 
+
+    # mean = round(np.mean(list(errors_mean.values())), 5)
+    # print(f'Final OURS error_{type}: {mean}')
+
+
+if __name__ == '__main__':
+
+    folder = '/data-local/127_2/results/hood'
+    run = 'hood_3'
+
+    for type in ['chamfer', 'area', 'collision_ratio']: # ['chamfer', 'collision', 'area', 'collision_ratio']:
+        calculate(type)

scripts/get_final_cape.py

@@ -13,10 +13,10 @@ parser.add_argument('-g', '--gpu', type=str, default='1')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
+    from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_f
 else:
     config_module = importlib.import_module(args['config'])
-    config_1 = config_module.config
+    config_f = config_module.config
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = args['gpu']
@@ -54,13 +54,13 @@ from models.cvae_model import CVAE
 from data.utils import Frame
 
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
+uv2mesh = UV2Mesh(res=256, is_skirt=config_f['is_skirt'], is_smpl_orig=config_f['is_cape'])
 
 
 def step_pyramid(batch):
     model.eval()
     with torch.inference_mode():
-        levels_input, conds = group_inputs_for_levels(batch, config_1)
+        levels_input, conds = group_inputs_for_levels(batch, config_f)
 
         out = []
 
@@ -70,16 +70,13 @@ def step_pyramid(batch):
 
             mesh_pred = uv2mesh(pred)
 
-            if config_1['use_skinning']:
+            if config_f['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]
-                if not config_1['is_cape']:
-                    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
-                else:
-                    mesh_pred_posed = mesh_pred_shaped
+                mesh_pred_posed = mesh_pred_shaped
             else:
                 mesh_pred_posed = mesh_pred
 
@@ -94,8 +91,8 @@ cape_mask = np.zeros((6890), bool)
 cape_mask[clothing_verts_idx] = True
 
 def get_error(pred, frame: Frame):
-    gt = np.load(f'{config_1["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
-    pose_shape = np.load(f'{config_1["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32)
+    gt = np.load(f'{config_f["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
+    pose_shape = np.load(f'{config_f["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32)
     pose, shape = pose_shape[:72], pose_shape[72:]
     gender = int(shape[10] >= 0.5)
     shape_offsets = torch.tensordot(torch.from_numpy(shape[:10][None, :]), smpl.shapedirs[gender].to('cpu'), dims=[[1], [2]])[0,smpl.to_keep,:].numpy()
@@ -155,30 +152,30 @@ def generate():
 
 if __name__ == '__main__':
 
-    smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
-    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
+    smpl = SMPL(orig_smpl=config_f['is_cape'], path='smpl', device=config_f['device'])
+    skirt_modifier = '_skirt' if config_f['is_skirt'] else ''
     dataset_resolutions = []
-    if 'pyramid' in config_1:
-        dataset_resolutions = config_1['pyramid']['resolutions'].copy()
+    if 'pyramid' in config_f:
+        dataset_resolutions = config_f['pyramid']['resolutions'].copy()
     if 512 not in dataset_resolutions:
         dataset_resolutions.append(512)
 
     # remove all values from config['signature] that start with mesh_
-    config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
+    config_f['signature'] = set([x for x in config_f['signature'] if not x.startswith('mesh_')])
 
-    dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
+    dataset_test = Dataset(config_f['test_dir'], config_f['signature'], device=config_f['device'], frames_file=config_f['test_frames'], use_cache=config_f['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_f['is_cape'])
     # dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
     test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
 
-    is_pyramid = config_1['name'].startswith('pyr')
+    is_pyramid = config_f['name'].startswith('pyr')
 
-    model: Pyramid = Pyramid(config_1)
+    model: Pyramid = Pyramid(config_f)
 
-    it = load_model(model, 'checkpoints/iccvw/' + config_1['name'], iteration=args['iteration'])
+    it = load_model(model, 'checkpoints/iccvw/' + config_f['name'], iteration=args['iteration'])
     model.eval()
     # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
 
-    name = config_1['name'] + '_' + str(it)
+    name = config_f['name'] + '_' + str(it)
 
     folder = '/data-local/127_2/results'
     out_folder = f'{folder}/{name}'

scripts/plot_ablation.py

@@ -0,0 +1,177 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '3'
+
+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 torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
+from pathlib import Path
+import importlib
+
+from data.utils import Frame
+from data.dataset import read_frames_file
+from utils.utils import save_stl, create_temp_frames
+from utils.checkpoints import CheckpointSaver, load_model
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss, kl_loss
+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 data.dataset import Dataset
+from utils.checkpoints import CheckpointSaver, load_model
+from torch.utils.data import DataLoader
+import trimesh as tm
+from zipfile import ZipFile
+
+from configs.iccvw.pyr_64_128_256_512_abl import config as config_1
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config_1['is_cape'])
+uv2mesh_skirt = UV2Mesh(res=256, is_skirt=True, is_smpl_orig=config_1['is_cape'])
+
+
+def step_pyramid(model, batch, config):
+    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 not config['is_skirt'] else uv2mesh_skirt(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]
+            if not config['is_cape']:
+                mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+            else:
+                mesh_pred_posed = mesh_pred_shaped
+        else:
+            mesh_pred_posed = mesh_pred
+
+    return mesh_pred_posed
+
+
+if __name__ == '__main__':
+    results_folder = Path('/data-local/127_2/results')
+    out_folder = Path('temp/supplementary')
+    os.makedirs(out_folder, exist_ok=True)
+
+    smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
+    dataset_resolutions = []
+    dataset_resolutions = [64, 128, 256, 512]
+
+    experiments = {
+                #    'inc': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) },
+                   'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 250000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 200000) },
+                #    'skin': { 'nonskirt': ('pyr_64_128_256_512_skin', 225000), 'skirt': ('pyr_64_128_256_512_skin_skirt', 200000) },
+                #    'norm': { 'nonskirt': ('pyr_64_128_256_512_norm', 200000), 'skirt': ('pyr_64_128_256_512_norm_skirt', 200000) },
+                #    'reg': { 'nonskirt': ('pyr_64_128_256_512_reg', 200000), 'skirt': ('pyr_64_128_256_512_reg_skirt', 200000) },
+    }
+
+    configs = dict()
+    models = dict()
+
+    for experiment in experiments.keys():
+        configs[experiment] = dict()
+        configs[experiment]['nonskirt'] = importlib.import_module('configs.iccvw.' + experiments[experiment]['nonskirt'][0]).config
+        configs[experiment]['skirt'] = importlib.import_module('configs.iccvw.' + experiments[experiment]['skirt'][0]).config
+
+        models[experiment] = dict()
+        models[experiment]['nonskirt'] = Pyramid(configs[experiment]['nonskirt'])
+        models[experiment]['skirt'] = Pyramid(configs[experiment]['skirt'])
+
+        load_model(models[experiment]['nonskirt'], 'checkpoints/iccvw/' + configs[experiment]['nonskirt']['name'], iteration=experiments[experiment]['nonskirt'][1])
+        load_model(models[experiment]['skirt'], 'checkpoints/iccvw/' + configs[experiment]['skirt']['name'], iteration=experiments[experiment]['skirt'][1])
+
+        models[experiment]['nonskirt'].eval()
+        models[experiment]['skirt'].eval()
+
+
+    # frames = ['00054_1_Dress', '00005_1_Jumpsuit']
+    # frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt']
+    # frames = ['01852_70_Tshirt'] # , '03849_27_Dress', '00054_1_Dress'
+    # frames = [Frame(frame_str=frame) for frame in frames]
+            
+    frames_file = configs['e2e']['nonskirt']['train_frames']
+    with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f:
+        frames = f.readlines()
+    frames_1 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
+    
+    frames_file = configs['e2e']['skirt']['train_frames']
+    with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f:
+        frames = f.readlines()
+    frames_2 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
+    
+    frames = frames_1 + frames_2
+    random.shuffle(frames)
+
+    body_faces = smpl.f
+
+    for frame in frames:
+        create_temp_frames(config_1['train_dir'], frame)
+        dataset = Dataset(config_1['train_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=False)
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+
+        model_picker = 'skirt' if frame.garment in ['Dress', 'Skirt'] else 'nonskirt'
+        
+        fix_data = np.load(Path(dataset.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+        mask = fix_data['mask']
+        faces = fix_data['faces']
+
+        batch = next(iter(dataloader))
+        body = batch['mesh_body_posed'][0].cpu().numpy()
+        gt = batch['mesh_posed'][0].cpu().numpy()
+
+        # save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces)
+        # save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces)
+
+        for experiment in experiments.keys():
+
+            pred = step_pyramid(models[experiment][model_picker], batch, configs[experiment][model_picker])
+
+            mesh_posed_fixed = fix_collision(pred, batch['mesh_body_posed'], batch['mesh_mask']).cpu().numpy()
+            mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
+            try:
+                m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
+                tm.smoothing.filter_humphrey(m, iterations=2)
+                verts = m.vertices
+            except:
+                verts = mesh_posed_fixed
+                print('failed to apply smoothing')
+
+            # save_stl(out_folder / f'{frame}_{experiment}_body.stl', body, body_faces)
+            # save_stl(out_folder / f'{frame}_{experiment}_mesh.stl', verts, faces)
+
+            fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
+                                  go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
+                                  ])
+            fig.update_layout(scene=dict(aspectmode='cube'))
+            fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+            fig.update_layout(title_text=f'{frame} {experiment}')
+            fig.show()
+            pass
+
+        pass
+
+
+
+

scripts/plot_ablation_pyramid.py

@@ -0,0 +1,152 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+
+import time
+import argparse
+import wandb
+from tqdm 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 torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
+from pathlib import Path
+
+from data.utils import Frame
+from data.dataset import read_frames_file
+from utils.utils import save_stl, create_temp_frames
+from utils.checkpoints import CheckpointSaver, load_model
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss, kl_loss
+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 data.dataset import Dataset
+from utils.checkpoints import CheckpointSaver, load_model
+from torch.utils.data import DataLoader
+import trimesh as tm
+from zipfile import ZipFile
+
+from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
+from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_2
+
+uv2mesh = dict()
+uv2mesh['nonskirt'] = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config_1['is_cape'])
+uv2mesh['skirt'] = UV2Mesh(res=256, is_skirt=True, is_smpl_orig=config_1['is_cape'])
+
+
+def step_pyramid(model, batch, config):
+    with torch.inference_mode():
+        levels_input, conds = group_inputs_for_levels(batch, config)
+
+        out = []
+        out_levels = model([levels_input, conds])
+        for i in range(len(out_levels)):
+            pred, _, _ = out_levels[i]
+
+            mesh_pred = uv2mesh[picker](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]
+                if not config['is_cape']:
+                    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+                else:
+                    mesh_pred_posed = mesh_pred_shaped
+            else:
+                mesh_pred_posed = mesh_pred
+
+            out.append(mesh_pred_posed)
+
+    return out
+
+
+
+
+
+if __name__ == '__main__':
+    results_folder = Path('/data-local/127_2/results')
+    out_folder = Path('temp/pyramid_ablation')
+    iteration = {'nonskirt': 250000, 'skirt': 200000 }
+
+    smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
+    dataset_resolutions = []
+    if 'pyramid' in config_1:
+        dataset_resolutions = config_1['pyramid']['resolutions'].copy()
+    if 512 not in dataset_resolutions:
+        dataset_resolutions.append(512)
+
+    model = { 'nonskirt': Pyramid(config_1), 'skirt': Pyramid(config_2) }
+    load_model(model['nonskirt'], 'checkpoints/iccvw/' + config_1['name'], iteration=iteration['nonskirt'])
+    load_model(model['skirt'], 'checkpoints/iccvw/' + config_2['name'], iteration=iteration['skirt'])
+    model['nonskirt'].eval()
+    model['skirt'].eval()
+
+    frames = ['01070_13_Jumpsuit', '00054_1_Dress', '03849_27_Dress', '06057_152_Dress',
+              '03179_274_Dress', '04247_102_Tshirt', '01852_70_Tshirt']
+    frames = ['00027_1_Dress']
+
+    frames_corrupt = ['00300_69_Jumpsuit', '07608_77_Jumpsuit', '06955_117_Dress']
+    frames = [Frame(frame_str=frame) for frame in frames]
+    body_faces = smpl.f
+
+    config = { 'nonskirt': config_1, 'skirt': config_2 }
+
+    for frame in tqdm(frames):
+        create_temp_frames(config_1['test_dir'], frame)
+        dataset = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+
+        picker = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'nonskirt'
+
+        batch = next(iter(dataloader))
+        pred = step_pyramid(model[picker], batch, config[picker])
+        
+        fix_data = np.load(Path(dataset.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+        mask = fix_data['mask']
+        faces = fix_data['faces']
+
+        body = batch['mesh_body_posed'][0].cpu().numpy()
+        gt = batch['mesh_posed'][0].cpu().numpy()
+
+        save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces)
+        save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces)
+
+        for i, _pred in enumerate(pred):
+            mesh_posed_fixed = fix_collision(_pred, batch['mesh_body_posed'], batch['mesh_mask']).cpu().numpy()
+
+            mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
+            m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
+            tm.smoothing.filter_humphrey(m, iterations=2)
+            verts = m.vertices
+
+            # save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
+            # save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)
+
+            fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
+                                  go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
+                                  ])
+            fig.update_layout(scene=dict(aspectmode='cube'))
+            fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+            fig.show()
+            pass
+
+        pass
+
+
+
+

scripts/plot_cape.py

@@ -0,0 +1,196 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '3'
+
+import time
+import argparse
+import wandb
+from tqdm 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 torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
+from pathlib import Path
+import importlib
+from scipy.spatial.transform import Rotation as R
+
+from data.utils import Frame
+from utils.utils import save_stl
+from utils.checkpoints import CheckpointSaver, load_model
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss, kl_loss
+from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
+from smpl.smpl_torch import SMPL
+from plotly.subplots import make_subplots
+from models.pyramid_model import Pyramid
+from data.dataset import Dataset
+from utils.checkpoints import CheckpointSaver, load_model
+from torch.utils.data import DataLoader
+import trimesh as tm
+from zipfile import ZipFile
+from custom_utils.custom_utils import _maskFaces
+
+from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_f
+from configs.iccvw.pyr_64_128_256_512_cape_m_e2e_2 import config as config_m
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config_f['is_cape'])
+
+
+def step_pyramid(model, batch, config):
+    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]
+            if not config['is_cape']:
+                mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+            else:
+                mesh_pred_posed = mesh_pred_shaped
+        else:
+            mesh_pred_posed = mesh_pred
+
+    return mesh_pred_posed
+
+def read_frames_file(file):
+    with open(file, 'r') as f:
+        lines = f.readlines()
+    frames = ['_'.join(line.strip().split('.')[:-1]) for line in lines]
+
+    return frames
+
+def create_temp_frames(folder: str, frame: Frame):
+    with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+        f.writelines([f'{frame.id}.{frame.frame_nr}\n' for _ in range(4)])
+
+def add_global_orientation(mesh, pose, inverse=False):
+    rodr = smpl.rodrigues(torch.tensor(pose.T, device=smpl.device).reshape(-1, 1, 3)).reshape(4, -1, 3, 3).cpu().numpy()
+    r = R.from_matrix(rodr[0, 0, :, :])
+    out = r.apply(mesh, inverse=inverse)
+
+    return out
+
+
+def read_obj(path):
+    with open(path, 'r') as f:
+        obj = tm.exchange.obj.load_obj(f, skip_materials=True, maintain_order=True)
+    verts = obj['vertices']
+    return verts
+
+
+if __name__ == '__main__':
+    results_folder = Path('/data-local/127_2/results')
+    out_folder = Path('temp/cape')
+
+    folder = dict()
+    folder['cape_f'] = '/data-net/code/StyleGAN2_tf2_conditioned/CAPE/resultsOld/CAPE-affineconv_nz18_pose24_clotype8_female/test_reconstruction_objs_dataset_female_4clotypes'
+    folder['cape_m'] = '/data-net/code/StyleGAN2_tf2_conditioned/CAPE/resultsOld/CAPE-affineconv_nz18_pose24_clotype8_male/test_reconstruction_objs_dataset_male_4clotypes'
+
+    frames_file = dict()
+    frames_file['cape_f'] = '/data-net/dataset/cape_processed/datasets/dataset_female_4clotypes_frames.txt'
+    frames_file['cape_m'] = '/data-net/dataset/cape_processed/datasets/dataset_male_4clotypes_frames.txt'
+    
+    frames = dict()
+    frames_to_i = dict()
+    frames['cape_f'] = read_frames_file(frames_file['cape_f'])
+    frames_to_i['cape_f'] = { frame: i for i, frame in enumerate(frames['cape_f']) } 
+    frames['cape_m'] = read_frames_file(frames_file['cape_m'])
+    frames_to_i['cape_m'] = { frame: i for i, frame in enumerate(frames['cape_m']) } 
+    random.shuffle(frames['cape_f'])
+    random.shuffle(frames['cape_m'])
+
+    smpl = SMPL(orig_smpl=True, path='smpl', device=config_f['device'])
+    J0 = dict()
+    J0['cape_f'] = np.load('smpl/root_joint_f.npy') #smpl(torch.zeros((4, 72), device=config_f['device']), torch.zeros((4, 11), device=config_f['device']))[1][0].cpu().numpy()
+    J0['cape_m'] = np.load('smpl/root_joint_m.npy')
+    clothing_verts_idx = np.load(os.path.join('custom_data', 'clothing_verts_idx.npy'))
+    cape_mask = np.zeros((6890), bool)
+    cape_mask[clothing_verts_idx] = True
+    cape_mask = cape_mask[smpl.to_keep]
+    _, faces = _maskFaces(np.zeros((3638, 3)), np.load('smpl/smpl_faces_orig.npy'), cape_mask)
+    faces = faces.astype(np.int32)
+    cape_template_body = read_obj('/data-net/code/StyleGAN2_tf2_conditioned/CAPE/data/template_mesh.obj')
+
+
+    dataset_resolutions = []
+    dataset_resolutions = [64, 128, 256, 512]
+
+    # cape_exp = { 'cape_f': ('pyr_64_128_256_512_cape', 100000), 'cape_m': ('pyr_64_128_256_512_skirt', 200000) }
+    ours_exp = { 'cape_f': ('pyr_iccvw_64_128_256_512_cape_f_e2e_fine_3_lvl2', 170000), 'cape_m': ('pyr_iccvw_64_128_256_512_cape_m_e2e_lvl2', 230000) }
+
+    picker = 'cape_m'
+    config = dict()
+    config['cape_f'] = config_f
+    config['cape_m'] = config_m
+
+    frames_override = ['00122_shortshort_ballerina_spin_000153', '00127_shortlong_ballerina_spin_000141',
+                       '00127_shortlong_soccer_000306', 
+                       '00134_shortlong_ski_trial2_000230',] #'00134_longlong_climb_trial2_000040', '00134_shortlong_soccer_trial2_000059'
+
+    _frames = frames_override if frames_override is not None else frames[picker]
+    for _frame in tqdm(_frames):
+        if _frame in frames['cape_f']:
+            picker = 'cape_f'
+        else:
+            picker = 'cape_m'
+        frame = Frame('_'.join(_frame.split('_')[:-1]), _frame.split('_')[-1], 'cape')
+        create_temp_frames(config[picker]['test_dir'], frame)
+        dataset = Dataset(config[picker]['test_dir'], config[picker]['signature'], device=config[picker]['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=True)
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+
+        batch = next(iter(dataloader))
+        gt = batch['mesh_posed'][0].cpu().numpy()[cape_mask]
+        pose = batch['pose_shape'][0, :72].cpu().numpy()
+        pose[:3] = batch['global_orient'][0].cpu().numpy()
+
+        save_stl(out_folder / f'{frame}_gt.stl', gt, faces)
+
+        pred = np.load(results_folder / f'{ours_exp[picker][0]}_{ours_exp[picker][1]}' / f'{frame}.npz')['arr_0']
+        pred = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.from_numpy(pred[None,:,:]).repeat(4,1,1).to('cuda'))[0].cpu().numpy()
+        pred = pred[0, cape_mask, :]
+
+        cape = read_obj(os.path.join(folder[picker], '{:0>4d}.obj'.format(frames_to_i[picker][f'{frame.id}_{frame.frame_nr}'])))[clothing_verts_idx, :]  - J0[picker][np.newaxis, :]
+        cape = add_global_orientation(cape, pose[np.newaxis, :], inverse=True)
+        min_body = np.load(os.path.join('/data-local/127_2/cape_processed', 'minimal_body_shape', frame.id.split('_')[0], frame.id.split('_')[0] + '_minimal.npy'))
+        cape = cape - cape_template_body[clothing_verts_idx, :] + min_body[clothing_verts_idx, :]
+
+        save_stl(out_folder / f'{frame}_ours.stl', pred, faces)
+        save_stl(out_folder / f'{frame}_cape.stl', cape, faces)
+
+        # pred[:,0] = pred[:,0] + 1.5
+        # cape[:,0] = cape[:,0] + 0.5
+        # gt[:,0] = gt[:,0] - 0.5
+        # fig = go.Figure(data=[
+        #     go.Mesh3d(x=gt[:,0], y=gt[:,2], z=gt[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightgreen', opacity=1),
+        #     go.Mesh3d(x=pred[:,0], y=pred[:,2], z=pred[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
+        #     go.Mesh3d(x=cape[:,0], y=cape[:,2], z=cape[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=1)
+        #                         ])
+        # fig.update_layout(scene=dict(aspectmode='cube'))
+        # fig.update_layout(scene=dict(xaxis=dict(range=[-2, 2], ), yaxis=dict(range=[-2, 2], ), zaxis=dict(range=[-2, 2], ), ), )
+        # fig.update_layout(title_text=f'{frame} ')
+        # fig.show()
+        pass
+
+        pass
+
+
+
+

scripts/plot_cloth3d.py

@@ -0,0 +1,148 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
+
+import time
+import argparse
+import random
+import tqdm
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from utils.postprocess import fix_collision
+
+
+import torch
+from tqdm import tqdm
+from data.utils import Frame
+from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
+from utils.utils import pickle_load, save_stl
+from scipy.spatial import cKDTree
+from smpl.smpl_torch import SMPL
+from scipy.spatial.transform import Rotation as R
+import trimesh as tm
+from plotly.subplots import make_subplots
+
+
+smpl = SMPL(path='smpl', device='cpu')
+def deorient(verts, pose, inverse=True):
+    assert pose.shape[1] == 72
+    rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy()
+    rotated = np.zeros_like(verts)
+    for frame in range(verts.shape[0]):
+        r = R.from_matrix(rodr[frame, 0, :, :])
+        rotated[frame] = r.apply(verts[frame], inverse=inverse)
+
+    return rotated
+
+
+def readBIN(fname):
+    if '.' not in fname: fname += '.bin'
+    elif '.' in fname and not fname.endswith('.bin'): 
+        print("File name extension should be '.bin'")
+        return
+        
+    with open(fname, 'rb') as f:
+        F = np.frombuffer(f.read(), dtype=np.uint16).astype(np.int32)
+        return F.reshape((-1,3))
+    
+
+def calculate():
+
+    garments = sorted(os.listdir('/data-local/127_2/cloth3d_processed_test/mesh_mask'))
+    garments = [g[:-4] for g in garments]
+    # random.shuffle(garments)
+    # garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)]
+    body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
+
+    for garment in tqdm(garments):
+        # garment = garment_data[0]
+        
+        picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
+        id, garment_type = garment.split('_')
+
+        if not os.path.exists(f'temp/cloth3d/{id}/outfit.pc16'):
+            continue
+        if not os.path.exists(os.path.join(cloth3d_folder, id, garment_type + '.pc16')):
+            continue
+
+        pred_cloth3d = readPC2(f'temp/cloth3d/{id}/outfit.pc16', True)['V']
+        pred_cloth3d_faces = readBIN(f'temp/cloth3d/{id}/outfit.bin')
+        gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V']
+        mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
+        pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')
+        gt = deorient(gt, pose, True)
+        # gt = deorient(gt, pose, True)
+        body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V']
+        fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
+        faces = fix_data['faces']
+
+
+        i = np.random.randint(0, len(pred_cloth3d))
+        # i = garment_data[1]
+
+        try:
+            if calculate_ours:
+                pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
+        except:
+            print('missed one')
+            continue
+
+        mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy()
+        mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
+        try:
+            m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
+            tm.smoothing.filter_humphrey(m, iterations=2)
+            pred_ours = m.vertices
+        except:
+            pred_ours = mesh_posed_fixed
+            print('failed to apply smoothing')
+
+        # fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
+        # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
+        # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+        # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+        # fig.show()
+
+        fig = go.Figure(data=[go.Mesh3d(x=pred_cloth3d[i][:,0], y=pred_cloth3d[i][:,1], z=pred_cloth3d[i][:,2], i=pred_cloth3d_faces[:,0], j=pred_cloth3d_faces[:,1], k=pred_cloth3d_faces[:,2], color='blue')])
+        fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
+        fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+        fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+        fig.show()
+
+        # fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
+        # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
+        # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+        # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+        # fig.show()
+        pass
+
+        # save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces)
+        # save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces)
+        # save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces)
+        # save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces)
+        # save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces)
+        # save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces)
+ 
+
+
+if __name__ == '__main__':
+
+    hood_folder = '/data-local/127_2/results/hood/all_2_105k'
+    ours_folder = dict()
+    ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
+    ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
+    cloth3d_folder = '/data-net/dataset/cloth3d/test'
+
+    results_folder = '/data-local/127_2/results'
+
+    calculate_ours = True
+
+
+
+    calculate()

scripts/plot_hood.py

@@ -0,0 +1,131 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
+
+import time
+import argparse
+import random
+import tqdm
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from utils.postprocess import fix_collision
+
+
+import torch
+from tqdm import tqdm
+from data.utils import Frame
+from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
+from utils.utils import pickle_load, save_stl
+from scipy.spatial import cKDTree
+from smpl.smpl_torch import SMPL
+from scipy.spatial.transform import Rotation as R
+import trimesh as tm
+from plotly.subplots import make_subplots
+
+
+smpl = SMPL(path='smpl', device='cpu')
+def deorient(verts, pose, inverse=True):
+    assert pose.shape[1] == 72
+    rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy()
+    rotated = np.zeros_like(verts)
+    for frame in range(verts.shape[0]):
+        r = R.from_matrix(rodr[frame, 0, :, :])
+        rotated[frame] = r.apply(verts[frame], inverse=inverse)
+
+    return rotated
+
+
+
+def calculate():
+
+    # garments = sorted(os.listdir(hood_folder))
+    # garments = [g[:-4] for g in garments]
+    # random.shuffle(garments)
+    garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)]
+    body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
+
+    for garment_data in tqdm(garments):
+        garment = garment_data[0]
+        picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
+        id, garment_type = garment.split('_')
+
+        pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl')
+        pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces']
+        gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
+        mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
+        pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:]
+        gt = deorient(gt, pose, True)
+        body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:]
+        fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
+        faces = fix_data['faces']
+
+
+        # i = np.random.randint(0, len(pred_hood))
+        i = garment_data[1]
+
+        try:
+            if calculate_ours:
+                pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
+        except:
+            print('missed one')
+            continue
+
+        mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy()
+        mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
+        try:
+            m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
+            tm.smoothing.filter_humphrey(m, iterations=2)
+            pred_ours = m.vertices
+        except:
+            pred_ours = mesh_posed_fixed
+            print('failed to apply smoothing')
+
+        # fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
+        # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
+        # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+        # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+        # fig.show()
+
+        fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
+        fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
+        fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+        fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+        fig.show()
+
+        # fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
+        # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
+        # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+        # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
+        # fig.show()
+        pass
+
+        # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces)
+        # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces)
+        # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces)
+        # save_stl(f'temp/hood/ne2w/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces)
+        # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces)
+        # save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces)
+ 
+
+
+if __name__ == '__main__':
+
+    hood_folder = '/data-local/127_2/results/hood/all_2_105k'
+    ours_folder = dict()
+    ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
+    ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
+    cloth3d_folder = '/data-net/dataset/cloth3d/test'
+
+    results_folder = '/data-local/127_2/results'
+
+    calculate_ours = True
+
+
+
+    calculate()

scripts/upscaler_train.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = '2'
+os.environ["CUDA_VISIBLE_DEVICES"] = '3'
 
 
 import time
@@ -39,10 +39,10 @@ from utils.utils import InfiniteDataLoader as DataLoader
 from pytorch_msssim import MS_SSIM
 
 
-low_res = 64
-is_skirt = False
-is_cape = True
-is_male = True
+low_res = 256
+is_skirt = True
+is_cape = False
+is_male = False
 skirt_modifier = '_skirt' if is_skirt else ''
 cape_modifier = '_cape' if is_cape else ''
 gender_modifier = '' if not is_cape else ('_m' if is_male else '_f')
@@ -244,6 +244,9 @@ def train():
 
         step += 1
 
+        if step > 100000:
+            break
+
 
 if __name__ == '__main__':
     if not args['debug']:

smpl/smpl_torch.py

@@ -174,7 +174,7 @@ class SMPL(Module):
     
     # J = torch.matmul(self.J_regressor, v_shaped)
     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])])
+      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])])
     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 = J[:,0,:].view((-1, 1, 3))
@@ -222,6 +222,8 @@ class SMPL(Module):
             
     weights = self.weights if verts is not None else self.weights_all 
     rest_shape_h_ones_local = self.rest_shape_h_ones if verts is not None else self.rest_shape_h_ones_full
+    if weights.shape[1] == 6890:
+      rest_shape_h_ones_local = torch.ones((batch_num, 6890, 1), dtype=torch.float32).to(self.device)
     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, rest_shape_h_ones_local), dim=2

utils/postprocess.py

@@ -35,6 +35,8 @@ def fix_collision(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr=.004
         F_fix_coll_local = F_fix_coll_orig
     else:
         raise ValueError("V must have 14475 or 3638 vertices")
+    if V.shape[0] == 1:
+        F_fix_coll_local = F_fix_coll_local[:1]
     mesh = Meshes(B, F_fix_coll_local)
     B_vn = torch.stack(mesh.verts_normals_list())
    

utils/utils.py

@@ -1,8 +1,11 @@
 import os
 import sys
+import pickle
 import numpy as np
 from torch.utils.data import DataLoader
-
+from stl import mesh as mesh_stl
+from pathlib import Path
+from data.utils import Frame
 
 
 class Accumulator():
@@ -48,4 +51,39 @@ def text_to_html(lines):
     html = ''
     for line in lines:
         html += "<pre>" + line + "</pre>"
-    return html
+    return html
+
+
+
+def save_stl(file, mesh, faces):
+    obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype))
+    for i, f in enumerate(faces):
+        for j in range(3):
+            obj.vectors[i][j] = mesh[f[j],:]    
+
+    obj.save(file)
+
+
+def create_temp_frames(folder: str, frame: Frame):
+    with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+        f.writelines([f'{frame}\n' for _ in range(4)])
+
+
+def pickle_load(file):
+    """
+    Load a pickle file.
+    """
+    with open(file, 'rb') as f:
+        loadout = pickle.load(f)
+
+    return loadout
+
+
+def pickle_dump(loadout, file):
+    """
+    Dump a pickle file. Create the directory if it does not exist.
+    """
+    os.makedirs(os.path.dirname(str(file)), exist_ok=True)
+
+    with open(file, 'wb') as f:
+        pickle.dump(loadout, f)