extrapolation train

.vscode/launch.json

@@ -10,8 +10,9 @@
             "request": "launch",
             "program": "${file}",
             "console": "integratedTerminal",
-            "justMyCode": false, // true false
-            "cwd": "${workspaceFolder}"
+            "justMyCode": true, // true false
+            "cwd": "${workspaceFolder}",
+            "args": ["--config", "configs.pyr_32", "--debug"]
         }
     ]
 }

configs/ae_body_cape_f.py

@@ -0,0 +1,73 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'body_torch_cape_f'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cape_processed/train/female'
+config['valid_dir'] = '/data-local/127/cape_processed/test/female'
+config['test_dir'] = '/data-local/127/cape_processed/test/female'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = True
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        }

configs/ae_curv_cape_f.py

@@ -0,0 +1,72 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'curv_torch_cape_f'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'curv'] 
+config['script'] = 'scripts/ae_curv_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cape_processed/train/female'
+config['valid_dir'] = '/data-local/127/cape_processed/test/female'
+config['test_dir'] = '/data-local/127/cape_processed/test/female'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = True
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = True
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_curvature',
+                        'mesh_curvature',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

configs/ae_static_cape_f.py

@@ -0,0 +1,74 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'static_torch_cape_f'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'static'] 
+config['script'] = 'scripts/ae_static_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cape_processed/train/female'
+config['valid_dir'] = '/data-local/127/cape_processed/test/female'
+config['test_dir'] = '/data-local/127/cape_processed/test/female'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = True
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0005
+config['cond_ae']['upscale_output'] = False
+config['pretrained_model'] = ''
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

configs/cape_00032_SL.py

@@ -0,0 +1,193 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'cape_00032_SL'
+config['notes'] = ''
+config['tags'] = ['iccv', 'cape_00032_SL', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = ''
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['avg_garment_path'] = 'custom_data/avg/train_cape_f_03223_SL_2_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = True
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/cape_processed/train/male'
+config['valid_dir'] = '/data-local/127/cape_processed/train/male'
+config['test_dir'] = '/data-local/127/cape_processed/train/male'
+config['train_frames'] = 'frames_00032_SL_train_2.txt'
+config['valid_frames'] = 'frames_00032_SL_test_2.txt'
+config['test_frames'] = 'frames_00032_SL_test_2.txt'
+
+config['static_model'] = 'checkpoints/iccv/torch/static_torch_cape_m'
+config['body_model'] = 'checkpoints/iccv/torch/body_torch_cape_m'
+config['curvature_model'] = 'checkpoints/iccv/torch/curv_torch_cape_m'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'uv_curvature',
+                            '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+7 if config['use_curvature'] 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_curv ############################
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['device'] = config['device']
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['latent_dim'] = config['latent_dim']
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_curv']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_curv']['encoder_block_settings'] = block_setting
+config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_curv']['activation'] = 'tanh'
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

configs/cape_03223_SL.py

@@ -0,0 +1,193 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'cape_03223_SL'
+config['notes'] = ''
+config['tags'] = ['iccv', 'cape_03223_SL', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = ''
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['avg_garment_path'] = 'custom_data/avg/train_cape_f_03223_SL_2_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = True
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/cape_processed/train/female'
+config['valid_dir'] = '/data-local/127/cape_processed/train/female'
+config['test_dir'] = '/data-local/127/cape_processed/train/female'
+config['train_frames'] = 'frames_03223_SL_train_2.txt'
+config['valid_frames'] = 'frames_03223_SL_test_2.txt'
+config['test_frames'] = 'frames_03223_SL_test_2.txt'
+
+config['static_model'] = 'checkpoints/iccv/torch/static_torch_cape_f'
+config['body_model'] = 'checkpoints/iccv/torch/body_torch_cape_f'
+config['curvature_model'] = 'checkpoints/iccv/torch/curv_torch_cape_f'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'uv_curvature',
+                            '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+7 if config['use_curvature'] 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_curv ############################
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['device'] = config['device']
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['latent_dim'] = config['latent_dim']
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_curv']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_curv']['encoder_block_settings'] = block_setting
+config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_curv']['activation'] = 'tanh'
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

configs/cape_03223_SS.py

@@ -0,0 +1,193 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'cape_03223_SS'
+config['notes'] = ''
+config['tags'] = ['iccv', 'cape_03223_SS', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = ''
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['avg_garment_path'] = 'custom_data/avg/train_cape_f_03223_SL_2_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = True
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/cape_processed/train/female'
+config['valid_dir'] = '/data-local/127/cape_processed/train/female'
+config['test_dir'] = '/data-local/127/cape_processed/train/female'
+config['train_frames'] = 'frames_03223_SS_train_2.txt'
+config['valid_frames'] = 'frames_03223_SS_test_2.txt'
+config['test_frames'] = 'frames_03223_SS_test_2.txt'
+
+config['static_model'] = 'checkpoints/iccv/torch/static_torch_cape_f'
+config['body_model'] = 'checkpoints/iccv/torch/body_torch_cape_f'
+config['curvature_model'] = 'checkpoints/iccv/torch/curv_torch_cape_f'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'uv_curvature',
+                            '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+7 if config['use_curvature'] 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_curv ############################
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['device'] = config['device']
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['latent_dim'] = config['latent_dim']
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_curv']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_curv']['encoder_block_settings'] = block_setting
+config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_curv']['activation'] = 'tanh'
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

configs/pyr_64_cape_m.py

@@ -3,7 +3,7 @@ from utils.convnext_utils import CNBlockConfig
 config = dict()
 
 config['gpu'] = '5'
-config['name'] = 'pyr_torch_64_cape_m'
+config['name'] = 'pyr_torch_64_cape_m_2'
 config['notes'] = ''
 config['tags'] = ['pyr', 'iccv', 'torch']
 config['script'] = 'scripts/cvae_pyramid.py'
@@ -44,7 +44,7 @@ if config['is_cape']:
   config['smpl_res'] = 3638
 
 config['lr'] = 0.0001
-config['starting_balancing_coeff'] = 0.001
+config['starting_balancing_coeff'] = 0.0001
 config['balancing_activation_limit'] = 0.015
 config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
 config['avg_garment_path'] = ''

configs/pyr_64_cape_m_3.py

@@ -0,0 +1,296 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'pyr_torch_64_cape_m_3'
+config['notes'] = 'without edge loss'
+config['tags'] = ['pyr', 'iccv', 'torch']
+config['script'] = 'scripts/cvae_pyramid.py'
+
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = True
+config['device'] = 'cuda'
+
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+config['use_curvature'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64]
+config['pyramid']['pretrained_pyramid'] = ''
+config['pyramid']['pretrained_pyramid_iteration'] = ''
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/data-local/127/cape_processed/train/male'
+config['valid_dir'] = '/data-local/127/cape_processed/test/male'
+config['test_dir'] = '/data-local/127/cape_processed/test/male'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['balancing_activation_limit'] = 0.015
+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_curvature',
+                            '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_curvature' + res_modifier)
+        config['signature'].add('uv_static' + res_modifier)
+        config['signature'].add('uv_mask' + res_modifier)
+
+
+config['static_model'] = 'checkpoints/iccv/torch/static_torch_cape_m'
+config['body_model'] = 'checkpoints/iccv/torch/body_torch_cape_m'
+config['curvature_model'] = 'checkpoints/iccv/torch/curv_torch_cape_m'
+config['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 7
+    level['upscale_enhancer'] = ''
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(16, 32, 3),
+        CNBlockConfig(32, 64, 3),
+        CNBlockConfig(64, 128, 3),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 3),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
+        CNBlockConfig(128, 64, 9),
+        CNBlockConfig(64, 32, 3),
+        CNBlockConfig(32, 16, 3),
+        CNBlockConfig(16, 3, 3),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['kernel_size'] = 5
+    level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_256_cape_m'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, 128, 3),
+        CNBlockConfig(128, 128, 3),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 3),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
+        CNBlockConfig(128, 128, 9),
+        CNBlockConfig(128, 65, 3),
+        CNBlockConfig(64, 3, 3),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['kernel_size'] = 5
+    level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_128_cape_m'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, 128, 9),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 9),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 9),
+        CNBlockConfig(128, 128, 9),
+        CNBlockConfig(64, 3, 9),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 3
+    level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_64_cape_m'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 9),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 64, 9),
+        CNBlockConfig(64, 3, 9),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['kernel_size'] = 3
+    level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_32_cape_m'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(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_curv ############################
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['device'] = config['device']
+config['cond_ae_curv']['is_label'] = True
+# config['cond_ae_curv']['latent_dim'] = config['latent_dim']
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_curv']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_curv']['encoder_block_settings'] = block_setting
+config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_curv']['activation'] = 'tanh'
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+# config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+# config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

data/dataset.py

@@ -45,7 +45,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 is_cape:
+            frames = [Frame(line.split('.')[0], line.split('.')[1]) for line in lines]
+        else:
+            frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines]
 
     if shuffle:
         random.shuffle(frames)

models/pyramid_model.py

@@ -154,7 +154,7 @@ class Pyramid(nn.Module):
             for param in self.body_enc.parameters(): param.requires_grad = False
             for param in self.curv_enc.parameters(): param.requires_grad = False
 
-        self.levels = [PyramidLevel({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))]
+        self.levels = nn.ModuleList([PyramidLevel({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))])
 
 
         # if len(config['pyramid']['resolutions']) == 2:

run.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.abl_all import config
+from configs.cape_03223_SS import config
 import argparse
 
 # parser = argparse.ArgumentParser(description='')
@@ -10,9 +10,9 @@ import argparse
 
 
 if __name__ == '__main__':
-    import wandb
+    # 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')
     name = config['name']
     script = config['script']
     config = [key for key in sys.modules.keys()  if key.startswith('configs.')][-1]

scripts/ae_body_train.py

@@ -198,8 +198,8 @@ if __name__ == '__main__':
         load_model(model.encoder, config['pretrained_model'] + '_enc')
         load_model(model.decoder, config['pretrained_model'] + '_dec')
 
-    ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_enc', top_n=5)
-    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_dec', top_n=5)
+    ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
+    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
 
     run_name = config['name']
     print(f'started training: {run_name}')

scripts/cvae.py

@@ -10,10 +10,12 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
-config_module = importlib.import_module(args['config'])
-config = config_module.config
 if args['debug']:
-    pass
+    from configs.cape_00032_SL import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -50,7 +52,7 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
 
 
 def init_wandb():
-    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
+    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/cvae.py')
     artifact.add_file('losses/losses.py')
@@ -457,7 +459,7 @@ def train():
             g_balancing_coeff += 1 / 10000 
 
         if step % 10000 == 0:
-            ckpt_saver(model, step, optimizer)
+            ckpt_saver(model, step, optimizer, None)
 
         step += 1
 
@@ -469,8 +471,8 @@ if __name__ == '__main__':
     smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
-    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'])
-    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
     train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
     valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
     train_accumulator = Accumulator()
@@ -478,12 +480,11 @@ if __name__ == '__main__':
 
     model = CVAE(config)
 
-
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
     g_balancing_coeff = config['starting_balancing_coeff']
 
-    ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
+    ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + config['name'], top_n=10)
 
     step = 0
     if config['pretrained_model'] != '':

scripts/cvae_pyramid.py

@@ -13,7 +13,8 @@ args = vars(parser.parse_args())
 config_module = importlib.import_module(args['config'])
 config = config_module.config
 if args['debug']:
-    from configs.pyr_64_cape_m import config
+    from configs.pyr_32 import config
+    # from configs.pyr_64_cape_m import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -440,10 +441,14 @@ def plot_everything(batch: dict, type: str):
         levels_input = group_inputs_for_levels(batch, config)
         output_levels = model(levels_input)
 
-    mesh_pred_levels = [uv2mesh(output_levels[i][0]).detach().cpu().numpy() for i in range(len(output_levels))]
-    mesh_orig_posed = batch['mesh_posed'].cpu().numpy()
+    mesh_pred_levels = [uv2mesh(output_levels[i][0]).detach().cpu() for i in range(len(output_levels))]
+    mesh_orig_posed = batch['mesh_posed'].cpu()
     body = batch['mesh_body_posed'].cpu().numpy()
     
+    # mesh_pred = torch.cat((mesh_pred[:, :, 0:1], -mesh_pred[:, :, 2:3], mesh_pred[:, :, 1:2]), dim=2)
+    # mesh_pred_levels = [torch.cat((mesh_pred_levels[i][:, :, 0:1], -mesh_pred_levels[i][:, :, 2:3], mesh_pred_levels[i][:, :, 1:2]), dim=2).numpy() for i in range(len(mesh_pred_levels))]
+    # mesh_orig_posed = torch.cat((mesh_orig_posed[:, :, 0:1], -mesh_orig_posed[:, :, 2:3], mesh_orig_posed[:, :, 1:2]), dim=2).numpy()
+
     orig_img = (torch.permute(uv_orig, [0,2,3,1]) + 1) / 2
     uv_static_img = (torch.permute(batch['uv_static'], [0,2,3,1]) + 1) / 2
     uv_body_posed_img = (torch.permute(batch['uv_body_posed'], [0,2,3,1]) + 1) / 2
@@ -477,6 +482,14 @@ def plot_everything(batch: dict, type: str):
         for i in range(mask.shape[0]):
             mesh_unposed_shaped_full[i, mask[i], :] = mesh_unposed_shaped[i]
         mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_unposed_shaped_full, device=config['device']))[0].cpu().numpy()
+    
+        # rotate gt 180 degrees along z axis
+        # gt = [np.matmul(gt[i], np.array([[np.cos(np.pi), -np.sin(np.pi), 0], [np.sin(np.pi), np.cos(np.pi), 0], [0, 0, 1]])) for i in range(mask.shape[0])]
+
+        # gt_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32)
+        # for i in range(mask.shape[0]):
+        #     gt_full[i, mask[i], :] = gt[i]
+        # mesh_orig_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(np.stack(gt_full, axis=0), device=config['device'], dtype=torch.float32))[0].cpu().numpy()
     else:
         mesh_posed = mesh_pred_levels[-1]
         gt = [mesh_orig[i][mask[i], :] for i in range(mesh_posed.shape[0])]
@@ -819,9 +832,9 @@ if __name__ == '__main__':
         # assert len(config['pyramid']['resolutions']) == 1
         step = load_model(model, config['pretrained_model'], optimizer=optimizer, scaler=scaler, iteration=config['pretrained_model_iteration'])
 
-    if not args['debug']:
-        model = torch.compile(model)
-        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+    # if not args['debug']:
+    #     model = torch.compile(model)
+    #     wandb.watch(model, log='all', log_freq=1000, log_graph=True)
 
 
     if len(config['pyramid']['resolutions']) == 1:
@@ -852,7 +865,7 @@ if __name__ == '__main__':
     # load_model(model.levels[-1], 'checkpoints/cvae_19')
 
     if not args['debug']:
-        model = torch.compile(model)
+        # model = torch.compile(model)
         wandb.watch(model, log='all', log_freq=1000, log_graph=True)
 
     train()

scripts/upscaler_train.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = '7'
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
 
 
 import time
@@ -26,6 +26,7 @@ import plotly.graph_objects as go
 import plotly.express as ex
 import torch.optim as optim
 import torch.nn.functional as F
+import torchvision
 
 from data.dataset import Dataset
 # from torch.utils.data import DataLoader
@@ -121,22 +122,26 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces):
 
 def train_step(batch):
     model.train()
-    with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
-        pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
-        pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
+    # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
+    pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
+    # pred = F.interpolate(pred, size=(512, 256), mode='bilinear', align_corners=True, antialias=True)
+    pred = resizer(pred)
     
-        l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1))
-        if low_res == 32:
-            ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
-        else:
-            ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
+    # l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1))
+    # loss = l1_loss
+    if low_res == 32:
+        ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
+    else:
+        ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
+    loss = 1e3 * (1 - ms_ssim)
 
-        loss = 1e3 * (1 - ms_ssim)
-
-    scaler.scale(loss).backward()
-    scaler.step(optimizer)
-    scaler.update()
+    # scaler.scale(loss).backward()
+    # scaler.step(optimizer)
+    # scaler.update()
+    # optimizer.zero_grad()
     optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
 
     model.eval()
     with torch.inference_mode():
@@ -144,7 +149,7 @@ def train_step(batch):
         mesh_orig = batch['mesh_unposed']
         metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'].unsqueeze(2) * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask'])
 
-    return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
+    return { 'train/loss': loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } # 'train/l1_loss': l1_loss,
         
 
 def valid_step(batch):
@@ -152,27 +157,30 @@ def valid_step(batch):
 
     with torch.inference_mode():
         pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
-        pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
+        # pred = F.interpolate(pred, size=(512, 256), mode='bilinear', align_corners=True, antialias=True)
+        pred = resizer(pred)
         
-        l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1))
+        # l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1))
+        # loss = l1_loss
         if low_res == 32:
             ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
         else:
             ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
-
         loss = 1e3 * (1 - ms_ssim)
+
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_unposed']
         metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'].unsqueeze(2) * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask'])
 
-    return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
+    return { 'valid/loss': loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, } #  'valid/l1_loss': l1_loss,
        
 
 def plot_everything(batch: dict, type: str):
     model.eval()
     with torch.inference_mode():
         pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
-        pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
+        # pred = F.interpolate(pred, size=(512, 256), mode='bilinear', align_corners=True, antialias=True)
+        pred = resizer(pred)
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_unposed']
@@ -187,7 +195,7 @@ def plot_everything(batch: dict, type: str):
     # figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
     
     orig_img = batch['uv_unposed']
-    img = torch.concat((orig_img, pred), axis=2)
+    img = torch.concat((orig_img, pred), axis=3)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
 
     return { **figs, **{ type + '/images': imgs} }
@@ -221,8 +229,8 @@ def train():
             valid_accumulator.reset()
             print(step)
 
-        if step % 10000 == 0:
-            ckpt_saver(model, step, optimizer, scaler)
+        if step % 10000 == 0 and not args['debug']:
+            ckpt_saver(model, step, optimizer, None)
 
         step += 1
 
@@ -263,10 +271,11 @@ if __name__ == '__main__':
         
     model : Upscaler = Upscaler(kernel_size=kernel_size).to('cuda')
     optimizer = optim.Adam(model.parameters(), lr=0.0001)
-    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    # scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    resizer = torchvision.transforms.Resize((512, 256), interpolation=torchvision.transforms.InterpolationMode.BILINEAR)
 
     # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
-    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[0.5, 0.2, 0.2, 0.1])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[0.8, 0.2])
     _MSSSIM_small = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[1.0])
 
     step = 0

smpl/smpl_torch.py

@@ -163,7 +163,9 @@ class SMPL(Module):
     v_shaped = torch.stack(v_shaped, axis=-1)
     
     # J = torch.matmul(self.J_regressor, v_shaped)
-    J = torch.stack([torch.matmul(self.J_regressor[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])])
+    J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
+    J = J - J0
     R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
 
     if verts is not None:
@@ -173,7 +175,7 @@ class SMPL(Module):
       I_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) + \
         torch.zeros((batch_num, R_cube.shape[1], 3, 3), dtype=torch.float32)).to(self.device)
       lrotmin = (R_cube - I_cube).reshape(batch_num, -1, 1).squeeze(dim=2)
-      v_posed = torch.stack([v_shaped[i, :, :, g.to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[g.to(torch.long)], dims=([0], [2])) for i, g in enumerate(betas[:,-1])])
+      v_posed = torch.stack([v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[torch.greater(g, 0.5).to(torch.long)], dims=([0], [2])) for i, g in enumerate(betas[:,-1])])
 
     results = []
     results.append(
@@ -204,9 +206,9 @@ class SMPL(Module):
         )
       )
     
-    # joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
+    joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
             
-    T = torch.stack([torch.tensordot(results[i], self.weights[g.to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
+    T = torch.stack([torch.tensordot(results[i], self.weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
     rest_shape_h = torch.cat(
       (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
     )
@@ -214,9 +216,9 @@ class SMPL(Module):
     v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
     result = v
 
-    # root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
-    # result = result - root_joint
-    # joints = joints - root_joint
+    root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
+    result = result - root_joint
+    joints = joints - root_joint
 
     # # plut result as a plotly 3d mesh
     # import plotly.graph_objects as go
@@ -228,4 +230,4 @@ class SMPL(Module):
     if remove and verts is None:
       result = result[:,self.to_keep,:]
 
-    return result, None # joints
+    return result, joints, v_shaped, v_posed

utils/checkpoints.py

@@ -37,7 +37,7 @@ def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=No
     checkpoint = torch.load(checkpoints[-1])
     iteration = checkpoints[-1].split('_')[-1].split('.')[0]
     model.load_state_dict(checkpoint['model'])
-    if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'])
-    if scaler is not None: scaler.load_state_dict(checkpoint['scaler'])
+    if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict())
+    if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler'])
     return int(iteration) if iteration != '' else 0