vto conversion, some progress, not working properl

.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/abl_all.py

@@ -0,0 +1,192 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '6'
+config['name'] = 'abl_torch_all_2'
+config['notes'] = ''
+config['tags'] = ['iccv', 'abl_all', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = 'checkpoints/iccv/abl_torch_all'
+
+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_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
+config['body_model'] = 'checkpoints/iccv/body_torch_1'
+config['curvature_model'] = 'checkpoints/iccv/curv_torch_1'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'uv_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/abl_all_2.py

@@ -0,0 +1,192 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '6'
+config['name'] = 'abl_torch_all_2'
+config['notes'] = ''
+config['tags'] = ['iccv', 'abl_all', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = 'checkpoints/iccv/abl_torch_all'
+
+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_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
+config['body_model'] = 'checkpoints/iccv/body_torch_1'
+config['curvature_model'] = 'checkpoints/iccv/curv_torch_1'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'uv_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/ae_body_1.py

@@ -0,0 +1,73 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '1'
+config['name'] = 'body_torch_1'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 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_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_body_cape_m.py

@@ -0,0 +1,73 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'body_torch_cape_m'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/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['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_1.py

@@ -0,0 +1,72 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '2'
+config['name'] = 'curv_torch_1'
+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/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = 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_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_curv_cape_m.py

@@ -0,0 +1,72 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'curv_torch_cape_m'
+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/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['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/ae_static_cape_m.py

@@ -0,0 +1,74 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'static_torch_cape_m'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'static'] 
+config['script'] = 'scripts/ae_static_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/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['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/ae_static_orig.py

@@ -7,6 +7,8 @@ config['name'] = 'static_torch_orig_1'
 config['notes'] = ' ' 
 config['tags'] = ['torch', 'static'] 
 
+config['pretrained_model'] = 'checkpoints/iccv/static_torch_orig_1'
+
 config['batch_size'] = 4
 config['train_dir'] = '/data-local/127/train'
 config['valid_dir'] = '/data-local/127/valid'
@@ -60,7 +62,6 @@ config['cond_ae']['activation'] = 'tanh'
 
 config['lr'] = 0.0005
 config['cond_ae']['upscale_output'] = False
-config['pretrained_model'] = ''
 
 config['signature'] = { 'id',
                         'frame_nr',

configs/ae_static_orig_2.py

@@ -7,6 +7,8 @@ config['name'] = 'static_torch_orig_2'
 config['notes'] = ' ' 
 config['tags'] = ['torch', 'static'] 
 
+config['pretrained_model'] = 'checkpoints/iccv/static_torch_orig_2'
+
 config['batch_size'] = 4
 config['train_dir'] = '/data-local/127/train'
 config['valid_dir'] = '/data-local/127/valid'

configs/direct_1.py

@@ -0,0 +1,77 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'direct_1'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'direct'] 
+config['script'] = 'scripts/direct_train.py'
+
+config['pretrained_model'] = ''
+config['use_curvature'] = False
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+config['input_features'] = 6
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['is_label'] = False
+config['kernel_size'] = 7
+config['bottleneck_size'] = 64 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 128, 3),
+    CNBlockConfig(128, 256, 9),
+    CNBlockConfig(256, config['bottleneck_size'], 9),
+    CNBlockConfig(config['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['bottleneck_size'], 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['encoder_block_settings'] = block_setting
+config['decoder_block_settings'] = block_setting_reverse
+config['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['reg_loss_coeff'] = 0.1
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        'mesh_posed',
+                        'uv_unposed',
+                        'mesh_unposed',
+                        'uv_static',
+                        'uv_mask',
+                        }

configs/direct_2.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'direct_2'
+config['notes'] = 'with curvature input' 
+config['tags'] = ['torch', 'direct'] 
+config['script'] = 'scripts/direct_train.py'
+
+config['pretrained_model'] = ''
+config['use_curvature'] = True
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+config['input_features'] = 10
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['is_label'] = False
+config['kernel_size'] = 7
+config['bottleneck_size'] = 64 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 128, 3),
+    CNBlockConfig(128, 256, 9),
+    CNBlockConfig(256, config['bottleneck_size'], 9),
+    CNBlockConfig(config['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['bottleneck_size'], 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['encoder_block_settings'] = block_setting
+config['decoder_block_settings'] = block_setting_reverse
+config['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['reg_loss_coeff'] = 0.1
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_curvature',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        'mesh_posed',
+                        'uv_unposed',
+                        'mesh_unposed',
+                        'uv_static',
+                        'uv_mask',
+                        'uv_body_posed',
+                        }

configs/pyr_32.py

@@ -0,0 +1,296 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'pyr_torch_32'
+config['notes'] = ''
+config['tags'] = ['pyr', 'iccv', 'torch']
+config['script'] = 'scripts/cvae_pyramid.py'
+
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+
+config['use_skinning'] =              True # and deshaping
+config['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'] = [32]
+config['pyramid']['pretrained_pyramid'] = ''
+config['pyramid']['pretrained_pyramid_iteration'] = ''
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['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.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/static_torch_orig_2'
+config['body_model'] = 'checkpoints/iccv/body_torch_1'
+config['curvature_model'] = 'checkpoints/iccv/curv_torch_1'
+config['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 7
+    level['upscale_enhancer'] = ''
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(16, 32, 3),
+        CNBlockConfig(32, 64, 3),
+        CNBlockConfig(64, 128, 3),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 3),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
+        CNBlockConfig(128, 64, 9),
+        CNBlockConfig(64, 32, 3),
+        CNBlockConfig(32, 16, 3),
+        CNBlockConfig(16, 3, 3),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['kernel_size'] = 5
+    level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_256'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, 128, 3),
+        CNBlockConfig(128, 128, 3),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 3),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
+        CNBlockConfig(128, 128, 9),
+        CNBlockConfig(128, 65, 3),
+        CNBlockConfig(64, 3, 3),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['kernel_size'] = 5
+    level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_128'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, 128, 9),
+        CNBlockConfig(128, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 9),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 128, 9),
+        CNBlockConfig(128, 128, 9),
+        CNBlockConfig(64, 3, 9),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 3
+    level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_64'
+    level['upscale_enhancer_freeze'] = True
+    level['bottleneck_size'] = 64 * config['bottleneck_scale']
+    level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3
+
+    block_setting = [
+        CNBlockConfig(64, level['bottleneck_size'], 9),
+        CNBlockConfig(level['bottleneck_size'], None, 9),
+    ]
+
+    block_setting_reverse = [
+        CNBlockConfig(level['bottleneck_size'] * 2, 64, 9),
+        CNBlockConfig(64, 3, 9),
+    ]
+
+    level['encoder_block_settings'] = block_setting
+    level['decoder_block_settings'] = block_setting_reverse
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['kernel_size'] = 3
+    level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_32'
+    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'

configs/pyr_64_cape_m.py

@@ -0,0 +1,296 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'pyr_torch_64_cape_m_2'
+config['notes'] = ''
+config['tags'] = ['pyr', 'iccv', 'torch']
+config['script'] = 'scripts/cvae_pyramid.py'
+
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = True
+config['device'] = 'cuda'
+
+
+config['use_skinning'] =              True # and deshaping
+config['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'

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'

custom_utils/custom_utils.py

@@ -40,9 +40,10 @@ class UV2Mesh(nn.Module):
 
 
     def uv_to_mesh(self, uv):
-        if uv.shape[1] == 3:
+        if uv.shape[1] == 3 or uv.shape[1] == 7:
             uv = torch.permute(uv, [0,2,3,1])
         if uv.shape[1] != 512:
+            raise NotImplementedError
             uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR)
             
         # batch_size = uv.shape[0]
@@ -65,7 +66,7 @@ class UV2Mesh(nn.Module):
         # p1 = tf.stack((x1, y1, z1), axis=-1)
         # p2 = tf.stack((x2, y2, z2), axis=-1)
 
-        mesh = torch.mean(uv[:, self.uv_pixels[:,:,0], self.uv_pixels[:,:,1]], dim=2)
+        mesh = torch.mean(uv[:, self.uv_pixels[:,:,0], self.uv_pixels[:,:,1]].float(), dim=2)
         return mesh
 
 

data/convert_vto.py

@@ -0,0 +1,259 @@
+import numpy as np
+import os
+import sys
+import pickle
+import plotly.express as ex
+import plotly.graph_objects as go
+import torch
+import open3d as o3d
+import trimesh
+# from vedo import trimesh as trimesh_plot
+from scipy.spatial import KDTree
+from scipy.sparse import lil_matrix
+from scipy.optimize import minimize
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+from smpl.smpl_torch import SMPL
+
+# import mask faces
+from custom_utils.custom_utils import _maskFaces
+
+smpl = SMPL(path='smpl')
+
+def energy_function(params, source_points, target_points, indices):
+    num_points = source_points.shape[0]
+    num_dims = source_points.shape[1]
+    
+    # Reshape the parameters
+    weights = params[:num_points * num_dims].reshape((num_points, num_dims))
+    translation = params[num_points * num_dims:num_points * num_dims + num_dims]
+    control_points = params[num_points * num_dims + num_dims:].reshape((-1, num_dims))
+    
+    # Compute the transformed source points
+    transformed_points = source_points + np.dot(weights, control_points.T)
+    
+    # Compute the TPS energy
+    tps_energy = 0.5 * np.sum(np.square(target_points[indices] - transformed_points))
+    
+    return tps_energy
+
+def non_rigid(source_points, target_points):
+    tree = KDTree(target_points)
+    distances, indices = tree.query(source_points)
+
+    # Initialize the optimization parameters
+    num_points = source_points.shape[0]
+    num_dims = source_points.shape[1]
+    initial_params = np.zeros(num_points * num_dims + num_dims + num_points)
+
+    # Define the optimization bounds
+    bounds = [(None, None)] * (num_points * num_dims)
+    bounds += [(None, None)] * num_dims
+    bounds += [(None, None)] * (num_points * num_dims)
+
+    # Perform the optimization
+    result = minimize(energy_function, initial_params, method='L-BFGS-B',
+                    args=(source_points, target_points, indices), bounds=bounds)
+
+    # Get the optimized parameters
+    optimized_params = result.x
+
+    # Reshape the parameters
+    optimized_weights = optimized_params[:num_points * num_dims].reshape((num_points, num_dims))
+    optimized_translation = optimized_params[num_points * num_dims:num_points * num_dims + num_dims]
+    optimized_control_points = optimized_params[num_points * num_dims + num_dims:].reshape((-1, num_dims))
+
+    transformed_points = source_points + np.dot(optimized_weights, optimized_control_points.T)
+    return transformed_points
+
+
+def register_smpl():
+    # generate body with zero pose and shape
+    pose = torch.zeros((1, 72))
+    shape = torch.zeros((1, 11))
+    smpl_out = smpl(pose, shape)[0][0]
+
+    with open('/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl', "rb") as f:
+        data = pickle.load(f)
+        garment = data['vertices'][100]
+        faces = data['faces']
+
+    # rotate garment x degrees along y axis
+    # convert degrees to radians
+    degrees = 107
+    theta = degrees * np.pi / 180
+    rot = np.array([[np.cos(theta), 0, np.sin(theta)],
+                    [0, 1, 0],
+                    [-np.sin(theta), 0, np.cos(theta)]])
+    garment = np.dot(rot, garment.T).T
+
+    # rotate garment x degrees along x axis
+    # convert degrees to radians
+    degrees = 5
+    theta = degrees * np.pi / 180
+    rot = np.array([[1, 0, 0],
+                    [0, np.cos(theta), -np.sin(theta)],
+                    [0, np.sin(theta), np.cos(theta)]])
+    garment = np.dot(rot, garment.T).T
+
+    # rotate garment x degrees along z axis
+    # convert degrees to radians
+    degrees = -5
+    theta = degrees * np.pi / 180
+    rot = np.array([[np.cos(theta), -np.sin(theta), 0],
+                    [np.sin(theta), np.cos(theta), 0],
+                    [0, 0, 1]])
+    garment = np.dot(rot, garment.T).T
+
+
+    garment[:,0] -= 0.06
+    garment[:,1] -= 0.03
+    garment[:,2] -= 0.015
+    
+    # visualize the body and garment
+    # fig = go.Figure(data=[go.Mesh3d(x=smpl_out[:,0], y=smpl_out[:,1], z=smpl_out[:,2], i=smpl.f[:,0], j=smpl.f[:,1], k=smpl.f[:,2])])
+    # fig.add_trace(go.Mesh3d(x=garment[:,0], y=garment[:,1], z=garment[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2]))
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+    # fig.show()
+
+    # non_rigid(garment, smpl_out.numpy())
+
+    # Load the source and target meshes
+    source_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(garment), o3d.utility.Vector3iVector(faces))
+    target_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(smpl_out), o3d.utility.Vector3iVector(smpl.f))
+    # flip source and target mesh
+    source_mesh, target_mesh = target_mesh, source_mesh
+
+    source_mesh.compute_vertex_normals()
+    source_mesh = source_mesh.filter_smooth_laplacian(number_of_iterations=10)
+    # plot the smoothed mesh with plotly
+    # fig = go.Figure(data=[go.Mesh3d(x=np.asarray(source_mesh.vertices)[:,0], y=np.asarray(source_mesh.vertices)[:,1], z=np.asarray(source_mesh.vertices)[:,2], i=np.asarray(source_mesh.triangles)[:,0], j=np.asarray(source_mesh.triangles)[:,1], k=np.asarray(source_mesh.triangles)[:,2])])
+    # fig.show()
+
+    # Perform initial alignment (translation and/or scaling)
+    # You can modify this step based on your specific needs
+    # For example, you can manually set an initial transformation
+    init_transformation = [[1.0, 0.0, 0.0, 0.0],
+                        [0.0, 1.0, 0.0, 0.0],
+                        [0.0, 0.0, 1.0, 0.0],
+                        [0.0, 0.0, 0.0, 1.0]]
+
+    # Create the ICP object and set the parameters
+    icp = o3d.pipelines.registration.TransformationEstimationPointToPoint()
+    criteria = o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=2000)
+        
+    # Perform ICP registration
+    reg_result = o3d.pipelines.registration.registration_icp(o3d.geometry.PointCloud(source_mesh.vertices),
+                                                             o3d.geometry.PointCloud(target_mesh.vertices),
+                                                             0.075,
+                                                             init_transformation, icp, criteria)
+
+    # Get the aligned source mesh after registration
+    aligned_mesh = source_mesh.transform(reg_result.transformation)
+
+    # Visualize the aligned mesh and target mesh with plotly
+    fig = go.Figure(data=[go.Mesh3d(x=np.asarray(aligned_mesh.vertices)[:,0], y=np.asarray(aligned_mesh.vertices)[:,1], z=np.asarray(aligned_mesh.vertices)[:,2], i=np.asarray(aligned_mesh.triangles)[:,0], j=np.asarray(aligned_mesh.triangles)[:,1], k=np.asarray(aligned_mesh.triangles)[:,2])])
+    fig.add_trace(go.Mesh3d(x=np.asarray(target_mesh.vertices)[:,0], y=np.asarray(target_mesh.vertices)[:,1], z=np.asarray(target_mesh.vertices)[:,2], i=np.asarray(target_mesh.triangles)[:,0], j=np.asarray(target_mesh.triangles)[:,1], k=np.asarray(target_mesh.triangles)[:,2]))
+    fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+    fig.show()
+
+    # Get the vertices of the aligned source mesh and target mesh
+    aligned_source_vertices = np.asarray(aligned_mesh.vertices)
+    target_vertices = np.asarray(target_mesh.vertices)
+
+    # Perform nearest neighbor search to find corresponding vertex indices
+    kd_tree = o3d.geometry.KDTreeFlann(target_mesh)
+    correspondences = []
+    matched = set()
+    for i in range(len(aligned_source_vertices)):
+        k = 1
+        _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
+        # while idx[-1] in matched:
+        #     k += 1
+        #     _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
+        # matched.add(idx[-1])
+        correspondences.append((i, idx[-1]))
+
+    # correspondences is a list of tuples where each tuple contains the corresponding vertex indices
+    # For example, correspondences = [(0, 10), (1, 5), (2, 15), ...]
+
+    # Visualize the correspondences with plotly
+    # fig = go.Figure(data=[go.Scatter3d(x=aligned_source_vertices[:,0], y=aligned_source_vertices[:,1], z=aligned_source_vertices[:,2], mode='markers', marker=dict(color='red', size=1))])
+    # fig.add_trace(go.Scatter3d(x=target_vertices[:,0], y=target_vertices[:,1], z=target_vertices[:,2], mode='markers', marker=dict(color='blue', size=1)))
+    # for corr in correspondences:
+    #     fig.add_trace(go.Scatter3d(x=[aligned_source_vertices[corr[0],0], target_vertices[corr[1],0]], y=[aligned_source_vertices[corr[0],1], target_vertices[corr[1],1]], z=[aligned_source_vertices[corr[0],2], target_vertices[corr[1],2]], mode='lines', line=dict(color='green', width=1)))
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+    # fig.show()
+
+    registered = np.zeros((14475, 3))
+    mask = np.zeros((14475), dtype=bool)
+    for i in range(len(correspondences)):
+        registered[correspondences[i][1]] = aligned_source_vertices[correspondences[i][0]]
+        mask[correspondences[i][1]] = True
+
+    # F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)
+
+    # registered_masked, masked_faces = _maskFaces(registered, np.asarray(target_mesh.triangles), mask)
+    registered_masked, masked_faces = registered, np.asarray(target_mesh.triangles)
+
+    # plot aligned source mesh with the target topology
+    fig = go.Figure(data=[go.Mesh3d(x=registered_masked[:,0], y=registered_masked[:,1], z=registered_masked[:,2], i=masked_faces[:,0], j=masked_faces[:,1], k=masked_faces[:,2])])
+    fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+    fig.show()
+
+    return smpl
+
+
+def convert_vto(folder, out_folder):
+    os.makedirs(out_folder, exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'mesh_posed_body_no_orient'), exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'pose_shape'), exist_ok=True)
+    F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)
+
+    files = os.listdir(folder)
+    # files = ['/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl']
+    for file in files:
+        # read pickle files and get vetices and faces from them
+        with open(os.path.join(folder, file), "rb") as f:
+            data = pickle.load(f)
+            vertices = data['vertices']
+            faces = data['faces']
+            pose = data['pose']
+            shape = data['shape']
+            trans = data['translation']
+            id = file.split('.')[0]
+            # pose[:, :3] = 0
+            id = id.replace('_', '-')
+            subject = id.split('-')[2]
+            sequence = '_'.join(id.split('-')[2:])
+            # cmu_data = np.load(f'/data-net/code/vto-garment-collisions/assets/CMU/{subject}/{sequence}_poses.npz')
+            # shape = np.tile(cmu_data['betas'], (shape.shape[0], 1))
+            # pose = cmu_data['poses']
+
+            body = smpl(torch.tensor(pose), torch.tensor(np.concatenate((shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)))[0].numpy()
+            pose_shape = np.concatenate((pose, shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)
+            
+            # for i in range(vertices.shape[0]):
+            #     np.savez_compressed(os.path.join(out_folder, 'mesh_posed_body_no_orient', f'{id}_{i}.npz'), mesh_posed_body_no_orient=body[i])
+            #     np.save(os.path.join(out_folder, 'pose_shape', f'{id}_{i}.npy'), pose_shape[i])
+
+
+            # plot mesh with plotly
+            frame = 0
+            fig = go.Figure(data=[go.Mesh3d(x=vertices[frame,:,0], y=vertices[frame,:,1], z=vertices[frame,:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2])])
+            fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,1], z=body[frame,:,2], i=F[:,0], j=F[:,1], k=F[:,2]))
+            fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+            fig.show()
+
+            
+            pass
+    pass
+
+
+if __name__ == '__main__':
+    # register_smpl()
+    convert_vto('/data-net/code/vto-dataset/dress/simulations', '/data-net/code/vto-dataset/dress_converted')
+    pass

data/dataset.py

@@ -1,7 +1,7 @@
 import os
 import sys
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+# os.environ["CUDA_VISIBLE_DEVICES"] = '6'
 
 import math
 import time
@@ -30,14 +30,15 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
 
         if is_cape:
             frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
+            frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames]
         else:
             frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+            frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
 
-        if testing:
-            frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
-            frames = [frame.split('.')[0] + '_4' for frame in frames]
+        # if testing:
+        #     frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
+        #     frames = [frame.split('.')[0] + '_4' for frame in frames]
 
-        frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
 
     if frames_file is not None:
         frames = []
@@ -65,12 +66,12 @@ class StaticCache():
         
         if is_cape:
             self.uv_masks = dict()
-            uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32)
-            self.uv_masks[512] = uv_mask.astype(np.float32)
+            uv_mask = torch.tensor(np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32), device=device)
+            self.uv_masks[512] = uv_mask.to(torch.float32).cpu().numpy()
             for res in [32, 64, 128, 256]:
-                down = Resize(uv_mask[:,:,None], size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST)
-                down = np.squeeze(down.numpy().astype(np.float32))
-                self.uv_masks[res] = down
+                down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT)(uv_mask[None,:,:])
+                down = torch.squeeze(down).to(torch.float32)
+                self.uv_masks[res] = down.cpu().numpy()
 
 
     def read(self, id):
@@ -127,7 +128,7 @@ class Dataset(Dataset):
         self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
 
         self.body_faces = np.load('smpl/smpl_faces.npy')
-        self.static_cache = StaticCache(folder, signature, is_cape=is_cape)
+        self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device)
 
         self.cache = dict() if use_cache else None
 
@@ -141,7 +142,7 @@ class Dataset(Dataset):
         if self.cache is not None:
             if frame in self.cache.keys():
                 data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()}
-                data = { **item, **self.static_cache[frame.id] }
+                data = { **data, **self.static_cache[frame.id] }
                 return data
 
         separator = '_'
@@ -159,7 +160,7 @@ class Dataset(Dataset):
         if 'frame_nr' in self.signature:
             item['frame_nr'] = frame.frame_nr
         if 'pose_shape' in self.signature:
-            pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy'))
+            pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
             pose_shape[:3] = 0
             item['pose_shape'] = pose_shape
         if 'mesh_body_posed' in self.signature:
@@ -186,7 +187,7 @@ class Dataset(Dataset):
                     uv_curvature = uv_curvature['uv_curvature'] 
                 else:
                     uv_curvature = uv_curvature['curvature_uv']
-                uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).to(self.device)
+                uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device)
                 item['uv_curvature' + res_string] = uv_curvature
 
             if 'uv_unposed' + res_string in self.signature:
@@ -206,14 +207,14 @@ class Dataset(Dataset):
         data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
 
         for key, value in data.items():
-            if key.startswith('uv_') and value.shape[-1] == 3:
+            if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7):
                 data[key] = torch.permute(value, [2,0,1])
 
         if self.cache is not None:
             self.cache[frame] = data
         
         data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
-        data = { **item, **self.static_cache[frame.id] }
+        data = { **data, **self.static_cache[frame.id] }
         return data
 
 

data/dataset_vto.py

@@ -0,0 +1,40 @@
+import os
+import sys
+
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+if args['debug']:
+    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
+import math
+import time
+import torch
+import random
+import pytorch3d
+import torchvision as tv
+import numpy as np
+import scipy.io as sio
+import plotly.express as ex
+from torch.utils.data import Dataset
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+from data.utils import Frame
+
+class Dataset(Dataset):
+    def __init__(self, folder):
+        pass
+
+    def __len__(self):
+        return len(self.frames)
+
+
+    def __getitem__(self, idx):
+        frame : Frame = self.frames[idx]
+
+
+if __name__ == '__main__':
+    from torch.utils.data import DataLoader

docker/Dockerfile2

@@ -0,0 +1,46 @@
+# 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.0-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 upgrade -y && 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
+
+# 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
+COPY requirements.txt /tmp/
+RUN pip install --requirement /tmp/requirements.txt
+RUN pip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+cu117.html
+
+
+# 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

@@ -0,0 +1,42 @@
+name: "cvpr_2023"
+services:
+  lhunor_torch:
+    image: lhunor_torch:v1.13
+    container_name: lhunor_torch_v1.13
+    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/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]
+
+
+  lhunor_torch_2:
+    image: lhunor_torch:nv23.04
+    container_name: lhunor_torch_nv23.04
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '/home-net/lhunor:/home-net'
+      - '/data-net/guest/lhunor/data:/data-net'
+      - '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]

docker/requirements.txt

@@ -1,3 +1,4 @@
+torch==2.0.0    
 numpy
 matplotlib
 opencv-python
@@ -10,3 +11,5 @@ pytorch3d
 tables
 numpy-stl
 torchmetrics
+pytorch_msssim
+importlib

losses/losses.py

@@ -7,6 +7,7 @@ import scipy.io as sio
 
 from pytorch3d.structures import Meshes
 from scipy.spatial import cKDTree
+from torch_cluster import knn
 # from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
 # from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
 
@@ -14,60 +15,54 @@ from scipy.spatial import cKDTree
 
 # from hbertiche/PBNS
 """ COLLISION FUNCTIONS START """
-F_coll = np.load("smpl/smpl_faces.npy")
-F_coll_skirt = np.load("smpl/smpl_faces_skirt.npy")
-F_coll_orig = np.load("smpl/smpl_faces_orig.npy")
-F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1])
-F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1])
-F_norm_orig = np.tile(F_coll_orig[np.newaxis,:,:], [4,1,1])
+smpl_to_keep_orig = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
+smpl_to_keep_orig = smpl_to_keep_orig[smpl_to_keep_orig < 6890]
+
+to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
+F_coll = torch.tensor(np.load("smpl/smpl_faces.npy")[to_keep,:], device='cuda')
+F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:], device='cuda')
+F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device='cuda')
+F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
+F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
+F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')
+F_norm_orig = F_norm_orig[:,smpl_to_keep_orig,:]
+
+
+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_y = torch.arange(batch_size, device=B.get_device()).repeat_interleave(B.shape[1])
+    V = V.reshape(-1, V.shape[-1])
+    B = B.reshape(-1, B.shape[-1])
+    idx = knn(V, B.to(V.dtype), k=1, batch_x=batch_x, batch_y=batch_y)[1]
+    idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count)
+    return idx
 
-def _nearest_neighbour(V, B):
-	tree = cKDTree(B.numpy())
-	return tree.query(V.numpy(), workers=-1, distance_upper_bound=0.2)[1]
 
 # @tf.function
-# def _tf_nn_parallel(V, B):
-# 	return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0])
-
-# @tf.function
-def collision_loss(V, B, mask, thr=.004):
-    mesh = Meshes(B, torch.tile(F_coll[None], [B.shape[0], 1, 1]))
+def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004):
+    if B.shape[1] == 14475:
+        F_norm_local = F_norm
+    elif B.shape[1] == 3638:
+        F_norm_local = F_norm_orig
+    mesh = Meshes(B, F_norm_local)
     B_vn = torch.stack(mesh.verts_normals_list())
    
     # compute correspondences
-    idx = torch.stack([_nearest_neighbour(V[i], B[i]) for i in range(V.shape[0])])
-    N, M = V.shape[:2]
-    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
-    _idx = torch.ones((N, M), dtype=torch.int32, device=V.get_device()) * np.array(list(range(N)), np.int32)[:,None]
-    idx = torch.stack((_idx, idx), axis=-1)
+    idx = _nearest_neighbour(V, B)
     # compute loss
-    D = V - B[idx]
-    b_vn = B_vn[idx]
+    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, 0)  * mask
+    thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))  * mask
     loss = torch.sum(thresholded ** 2, axis=-1)
     ratio = torch.sum((thresholded < 0).to(torch.float32)) / torch.sum(mask)
 
     return loss, ratio
 
 
-# def collision_ratio(V, B, mask, thr=.004):
-    # B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
-   
-    # # compute correspondences
-    # idx = _tf_nn_parallel(V, B)
-    # N, M = V.shape[:2]
-    # # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
-    # _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
-    # idx = tf.stack((_idx, idx), axis=-1)
-    # # compute loss
-    # D = V - tf.gather_nd(B, idx)
-    # b_vn = tf.gather_nd(B_vn, idx)
-    # dot = tf.einsum('abc,abc->ab', D, b_vn)
-    # loss = tf.minimum(dot - thr, 0)  * mask # / tf.reduce_sum(mask, axis=-1)
-    # ratio = tf.reduce_sum(tf.cast(loss < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
-
-    # return ratio
 
 """ COLLISION FUNCTIONS END """
 
@@ -78,18 +73,26 @@ edges: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())
 laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo()
 edges_skirt: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda')
 
+laplacian_loss_orig = laplacian_loss.todense()[:6890,:6890][smpl_to_keep_orig,:][:,smpl_to_keep_orig]
+edges_cape = torch.tensor(np.stack(np.nonzero(laplacian_loss_orig)).transpose(), device='cuda')
 
-def get_edge_mask(masks: torch.Tensor):
+def get_edge_mask(edges: torch.Tensor, masks: torch.Tensor):
     if edges.get_device() != masks.get_device():
         edges.to(masks.get_device())
     return torch.min(masks[:,edges], dim=-1)[0]
 
 def edge_loss(x, y, mask):
-    e_mask = get_edge_mask(mask)
-    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
+    if x.shape[1] == 14475:
+        edges_local = edges
+    elif x.shape[1] == 3638:
+        edges_local = edges_cape
+    else:
+        raise ValueError("Invalid number of vertices")
+    e_mask = get_edge_mask(edges_local, mask)
+    x_e = x[:,edges_local[:,0],:] - x[:,edges_local[:,1],:]
     x_e = torch.sum(x_e ** 2, -1)
     x_e = torch.sqrt(x_e + 1e-9)
-    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
+    y_e = y[:,edges_local[:,0],:] - y[:,edges_local[:,1],:]
     y_e = torch.sum(y_e ** 2, -1)
     y_e = torch.sqrt(y_e + 1e-9)
     d_e = y_e - x_e
@@ -116,41 +119,6 @@ def edge_loss_skirt(x, y, mask):
     return loss
 
 
-# class PairwiseDistance(tf.keras.layers.Layer):
-#     def __init__(self, **kwargs):
-#         super(PairwiseDistance, self).__init__(**kwargs)
-
-#     def call(self, A, B):
-#         rA = tf.reduce_sum(tf.square(A), axis=-1)
-#         rB = tf.reduce_sum(tf.square(B), axis=-1)
-#         transpose_axes = [0, 2, 1] 
-#         distances = - 2*tf.matmul(A, tf.transpose(B, transpose_axes)) + rA[:, :, tf.newaxis] + rB[:, tf.newaxis, :]
-#         return distances
-
-
-# class NearestNeighbour(tf.keras.layers.Layer):
-#     def __init__(self, **kwargs):
-#         super(NearestNeighbour, self).__init__(**kwargs)
-
-#     def call(self, A, B):
-#         distances = PairwiseDistance(dtype=self.dtype)(A, B)
-#         nearest_neighbour = tf.argmin(distances, axis=-1)
-#         return tf.cast(nearest_neighbour, dtype=tf.int32)
-
-#
-# @tf.function
-def collision_penalty(va, vb, mask, eps=1e-3):
-    nb = vertex_normals(vb, tf.tile(F[None], [vb.shape[0], 1, 1]))
-    batch_size = tf.cast(tf.shape(va)[0], va.dtype)
-
-    closest_vertices = NearestNeighbour(dtype=va.dtype)(va, vb)
-    vb = tf.gather(vb, closest_vertices, batch_dims=1)
-    nb = tf.gather(nb, closest_vertices, batch_dims=1)
-
-    distance = tf.reduce_sum(nb*(va - vb), axis=-1) 
-    interpenetration = tf.maximum(eps - distance * mask, 0)
-
-    return tf.reduce_sum(interpenetration**3) / batch_size
 
 whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy')
 
@@ -222,8 +190,8 @@ def log_normal_pdf(sample, mean, logvar, raxis=1):
 
 
 def kl_loss(mean, logvar):
-    kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
-    kl_loss = tf.reduce_mean(kl_loss)
+    kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar)
+    kl_loss = torch.mean(kl_loss)
     kl_loss *= -0.5
     return kl_loss
 
@@ -239,6 +207,9 @@ def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
     else:
         raise NotImplementedError
 
-    pred_norm = vertex_normals(pred, f_norm)
-    gt_norm = vertex_normals(gt, f_norm)
-    return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))
+    mesh = Meshes(pred, f_norm)
+    pred_norm = torch.stack(mesh.verts_normals_list())
+    mesh = Meshes(gt, f_norm)
+    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/base.py

@@ -1,27 +1,29 @@
 import torch
 from torch import nn
-from models.convnext import ConvNeXt, CNBlock, CNBlockConfig
+from models.convnext_orig import ConvNeXt, CNBlock, CNBlockConfig
 from typing import Any, Callable, List, Optional, Sequence
 
 
 class VariationalEncoder(nn.Module):
     
-    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7):
+    def __init__(self, block_setting: List[CNBlockConfig], kernel_size: int = 7, input_features: int = 16, device: str = 'cuda'):
         super(VariationalEncoder, self).__init__()
 
-        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size)
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
 
-        self.mean_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
-        self.var_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.mean_branch = nn.Sequential(*layers).to(device)
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.var_branch = nn.Sequential(*layers).to(device)
 
         self.flatten_mean = nn.Flatten()
         self.flatten_var = nn.Flatten()
 
 
-    def call(self, inputs, training):
+    def forward(self, inputs):
         x = inputs
 
-        x = self.encoder(x, training=training)
+        x = self.encoder(x)
         mean = self.mean_branch(x)
         log_var = self.var_branch(x)
 
@@ -33,10 +35,10 @@ class VariationalEncoder(nn.Module):
 
 class AEEncoder(nn.Module):
     
-    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', device='cuda'):
+    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda', input_features=3):
         super(AEEncoder, self).__init__()
 
-        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device)
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device, input_features=input_features)
         if activation == 'tanh':
             self.act = nn.Tanh()
         elif activation == 'sigmoid':
@@ -81,10 +83,10 @@ class Decoder(nn.Module):
         x = self.decoder(x)
 
         if self.is_label:
-            x = x * mask[:,:,:,None]
+            x = x * mask[:,None,:,:]
 
-            background = torch.ones((mask.shape[0],512,256,1), device=mask.get_device()) * (1 - mask[:,:,:,None])
-            x = torch.concat((background, x), axis=-1)
+            background = torch.ones((mask.shape[0],1,512,256), device=mask.get_device()) * (1 - mask[:,None,:,:])
+            x = torch.concat((background, x), dim=1)
             x = self.softmax(x)
 
         if len(inputs) == 2:

models/cond_encoder.py

@@ -8,8 +8,8 @@ class AutoEncoder(nn.Module):
     def __init__(self, config):
         super(AutoEncoder, self).__init__()
 
-        self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], device=config['device'])
-        self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'])
+        self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], is_label=config['is_label'], device=config['device'], nr_of_labels=7, input_features=config['input_features'])
+        self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'], nr_of_classes=6)
 
 
     def forward(self, input):

models/convnext_orig.py

@@ -64,6 +64,9 @@ class ConvNeXtBase(nn.Module):
         norm_layer: Optional[Callable[..., nn.Module]] = None,
         device: str = 'cuda',
         kernel_size: int = 7,
+        is_label: bool = False,
+        nr_of_labels: Optional[int] = None,
+        input_features: int = 3,
     ) -> None:
         super().__init__()
 
@@ -83,7 +86,7 @@ class ConvNeXtBase(nn.Module):
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
         self.stem = Conv2dNormActivation(
-                3,
+                input_features if not is_label else nr_of_labels,
                 firstconv_output_channels,
                 kernel_size=4,
                 stride=4,
@@ -227,6 +230,7 @@ def _convnext(
     is_label: bool,
     nr_of_labels: Optional[int],
     kernel_size: int,
+    input_features: int,
     device='cuda',
     **kwargs: Any,
 ) -> ConvNeXtBase:
@@ -234,17 +238,17 @@ def _convnext(
     if reversed:
         model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
     else:
-        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, device=device, **kwargs)
+        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, input_features=input_features, device=device, **kwargs)
 
 
     return model
 
 
 class ConvNeXt(nn.Module):
-    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, is_label=False, nr_of_labels=None, kernel_size=7, device='cuda', **kwargs) -> None:
+    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, is_label=False, nr_of_labels=None, kernel_size=7, input_features=3, device='cuda', **kwargs) -> None:
         super().__init__()
 
-        self.model = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, device=device, **kwargs)
+        self.model: nn.Module = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, input_features=input_features, device=device, **kwargs)
 
     def forward(self, x: Tensor) -> Tensor:
         x = self.model(x)

models/cvae_model.py

@@ -1,63 +1,54 @@
 import os
 import time
 import random
-import tensorflow as tf
+import torch
+from torch import nn
 import numpy as np
 import scipy as sp
 # import plotly.graph_objects as go
 # import plotly.express as ex
 
-from custom_utils.custom_utils import load_model
+from utils.checkpoints import load_model
 from models.base import AEEncoder, VariationalEncoder, Decoder
-from models.base import AEEncoderNew, VariationalEncoderNew, DecoderNew
 
 
-class CVAE(tf.keras.Model):
+class CVAE(nn.Module):
     
     def __init__(self, config):
         super(CVAE, self).__init__()
-        self.latent_dim = config['cvae']['latent_dim']
         self.size = config['resolution']
 
         if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
-            self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
+            self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
             if config['static_model'] != '':
                 load_model(self.static_enc, config['static_model'] + '_enc')
 
-            self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
+            self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
             if config['body_model'] != '':
                 load_model(self.body_enc, config['body_model'] + '_enc')
 
         if config['use_curvature']:
-            self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
+            self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'], device=config['device'], is_label=True, nr_of_labels=7)
             if config['curvature_model'] != '':
                 load_model(self.curv_enc, config['curvature_model'] + '_enc')
 
         if config['freeze_encoders']:
-            if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
-                self.static_enc.trainable = False
-                self.body_enc.trainable = False
+            for param in self.static_enc.parameters(): param.requires_grad = False
+            for param in self.body_enc.parameters(): param.requires_grad = False
             if config['use_curvature']:
-                self.curv_enc.trainable = False
+                for param in self.curv_enc.parameters(): param.requires_grad = False
 
-        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
-                                                 tf.keras.layers.Dense(256, activation='relu'),
-                                                 tf.keras.layers.Dense(128, activation='relu'),
-                                                #  tf.keras.layers.Lambda(lambda x: x / 1000),
-                                                 tf.keras.layers.Dense(config['latent_dim'], activation='tanh')])
+        self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
+        self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])
 
-        self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
-        self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
+        self.latent_shape = (config['cvae']['latent_dim'], 8, 4)
 
         if config['avg_garment_path'] != '':
-            self.avg_garment = np.load(config['avg_garment_path'])
+            self.avg_garment = torch.permute(torch.tensor(np.load(config['avg_garment_path']), device=config['device']), [2,0,1])
         else:
             self.avg_garment = None
 
-        self.add_noise = config['cvae']['add_noise']
-        self.noise_amount = config['cvae']['noise_amount']
         self.use_curvature = config['use_curvature']
-        self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
         self.use_offsets_over_body = not config['use_skinning']
 
         if self.use_offsets_over_body:
@@ -68,7 +59,7 @@ class CVAE(tf.keras.Model):
     # @tf.function
     def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
         if noise is None:
-            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+            eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
         else:
             eps = noise
 
@@ -91,195 +82,49 @@ class CVAE(tf.keras.Model):
         return x
 
 
-    def reparameterize(self, mean, logvar, training, fix_noise=False):
+    def reparameterize(self, mean: torch.Tensor, logvar: torch.Tensor, training: bool = True, fix_noise: bool = False):
         if training == True:
             if not fix_noise:
-                eps = tf.random.normal(shape=mean.shape)
+                eps = torch.randn(size=mean.shape).to(mean.device)
             else:
-                eps = tf.random.normal(shape=(1, mean.shape[1]))
-                eps = tf.tile(eps, [mean.shape[0],1])
-            out = eps * tf.exp(logvar * .5) + mean
+                eps = torch.randn(size=(1, mean.shape[1]))
+                eps = torch.tile(eps, [mean.shape[0],1])
+            out = eps * torch.exp(logvar * 0.5) + mean
         else:
-            out = tf.exp(logvar * .5) + mean
+            out = torch.exp(logvar * 0.5) + mean
         return out
 
 
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
-        if self.only_curvatures:
-            z_cond = self.curv_enc(uv_curv, training=training)
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+        z_body = self.body_enc(uv_body_posed)
+        z_static = self.static_enc(uv_template)
+        if self.use_curvature:
+            z_curv = self.curv_enc(uv_curv)
+            z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
         else:
-            z_body = self.body_enc(uv_body_posed, training=training)
-            z_static = self.static_enc(uv_template, training=training)
-            if self.use_curvature:
-                z_curv = self.curv_enc(uv_curv, training=training)
-                z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-            else:
-                z_cond = tf.concat([z_static, z_body], axis=-1)
-
-        z_cond = self.cond_encoder(z_cond, training=training)
+            z_cond = torch.concat([z_static, z_body], dim=1)
 
         return z_cond
 
 
-    def call(self, inputs, training):
+    def forward(self, inputs):
         uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
-        if self.only_curvatures:
-            x = tf.concat([uv_unposed, uv_curv], axis=-1)
-        else:
-            x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+        x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
         
-        mean, logvar = self.encoder(x, training=training)
-        z = self.reparameterize(mean, logvar, training)
-        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
+        mean, logvar = self.encoder(x)
+        z: torch.Tensor = self.reparameterize(mean, logvar)
+        z = z.view((mean.shape[0], *self.latent_shape))
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
 
-        if self.add_noise and training == True:
-            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
-
-
-        z = tf.concat([z, z_cond], axis=-1)
-        x = self.decoder(z, training=training)
+        z = torch.concat([z, z_cond], dim=1)
+        x = self.decoder(z)
 
         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, mean, logvar
-
-
-
-
-class CVAENew(tf.keras.Model):
-    
-    def __init__(self, config):
-        super(CVAENew, self).__init__()
-        self.size = config['resolution']
-
-        if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
-            self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'])
-            if config['static_model'] != '':
-                load_model(self.static_enc, config['static_model'] + '_enc')
-
-            self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size'])
-            if config['body_model'] != '':
-                load_model(self.body_enc, config['body_model'] + '_enc')
-
-        if config['use_curvature']:
-            self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['encoder_depths'], config['cond_ae_curv']['encoder_dims'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'])
-            if config['curvature_model'] != '':
-                load_model(self.curv_enc, config['curvature_model'] + '_enc')
-
-        if config['freeze_encoders']:
-            if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
-                self.static_enc.trainable = False
-                self.body_enc.trainable = False
-            if config['use_curvature']:
-                self.curv_enc.trainable = False
-
-        self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
-        self.decoder = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
-
-        self.reshape = tf.keras.layers.Reshape((8, 4, 64))
-
-        if config['avg_garment_path'] != '':
-            self.avg_garment = np.load(config['avg_garment_path'])
-        else:
-            self.avg_garment = None
-
-        self.add_noise = config['cvae']['add_noise']
-        self.noise_amount = config['cvae']['noise_amount']
-        self.use_curvature = config['use_curvature']
-        self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
-        self.use_offsets_over_body = not config['use_skinning']
-
-        if self.use_offsets_over_body:
-            assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
-        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):
-        if noise is None:
-            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
-        else:
-            eps = noise
-
-        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
-        if add_noise:
-            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
-        
-        z = tf.concat([eps, z_cond], axis=-1)
-
-        x = self.decoder(z, training=False)
-        
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-        
-        if self.use_offsets_over_body:
-            x = x + uv_body_posed
-
-        x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x
-
-
-    def reparameterize(self, mean, logvar, training, fix_noise=False):
-        if training == True:
-            if not fix_noise:
-                eps = tf.random.normal(shape=mean.shape)
-            else:
-                eps = tf.random.normal(shape=(1, mean.shape[1]))
-                eps = tf.tile(eps, [mean.shape[0],1])
-            out = eps * tf.exp(logvar * .5) + mean
-        else:
-            out = tf.exp(logvar * .5) + mean
-        return out
-
-
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
-        if self.only_curvatures:
-            z_cond = self.curv_enc(uv_curv, training=training)
-        else:
-            z_body = self.body_enc(uv_body_posed, training=training)
-            z_static = self.static_enc(uv_template, training=training)
-            if self.use_curvature:
-                z_curv = self.curv_enc(uv_curv, training=training)
-                z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-            else:
-                z_cond = tf.concat([z_static, z_body], axis=-1)
-
-        return z_cond
-
-
-    def call(self, inputs, training):
-        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
-        if self.only_curvatures:
-            x = tf.concat([uv_unposed, uv_curv], axis=-1)
-        else:
-            x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
-        
-        mean, logvar = self.encoder(x, training=training)
-        z = self.reparameterize(mean, logvar, training)
-        z = self.reshape(z)
-        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
-
-        if self.add_noise and training == True:
-            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
-
-
-        z = tf.concat([z, z_cond], axis=-1)
-        x = self.decoder(z, training=training)
-
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-    
-        if self.use_offsets_over_body:
-            x = x + uv_body_posed
-
-        x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x, mean, logvar
-

models/direct.py

@@ -0,0 +1,29 @@
+import torch
+from torch import nn
+from models.convnext_orig import ConvNeXt, CNBlock, CNBlockConfig
+from typing import Any, Callable, List, Optional, Sequence
+
+
+class DirectModel(nn.Module):
+    
+    def __init__(self, block_setting: List[CNBlockConfig], kernel_size: int = 7, input_features: int = 16, device: str = 'cuda'):
+        super(DirectModel, self).__init__()
+
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
+        self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device, reversed=True)
+
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.mean_branch = nn.Sequential(*layers).to(device)
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.var_branch = nn.Sequential(*layers).to(device)
+
+        self.flatten_mean = nn.Flatten()
+        self.flatten_var = nn.Flatten()
+
+
+    def forward(self, inputs):
+        body, template = inputs
+
+        x = self.encoder(x)
+
+        return mean

models/pyramid_level.py

@@ -1,176 +1,168 @@
 import os
 import time
 import random
-import tensorflow as tf
+import torch
+from torch import nn
 import numpy as np
 import scipy as sp
 # import plotly.graph_objects as go
 # import plotly.express as ex
 
-from custom_utils.custom_utils import load_model
-from models.base import AEEncoder, VariationalEncoder, Decoder, VariationalEncoderNew, DecoderNew
-from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
+from utils.checkpoints import load_model
+from models.base import AEEncoder, VariationalEncoder, Decoder
+from models.upscaler import Upscaler
 
 
-class PyramidLevel(tf.keras.Model):
+
+# class PyramidLevel(nn.Module):
+    
+#     def __init__(self, config):
+#         super(PyramidLevel, self).__init__()
+#         self.latent_dim = config['latent_dim']
+#         self.size = config['resolution']
+
+#         self.encoder: VariationalEncoder = VariationalEncoder(config['encoder_block_settings'], config['input_features'],  kernel_size=config['kernel_size'])
+#         self.decoder: Decoder = Decoder(config['decoder_block_settings'], config['kernel_size'])
+
+#         self.upscale_enhancer = None
+#         if config['upscale_enhancer'] != '':
+#             k_size = config['upscale_enhancer_kernel_size']
+#             self.upscale_enhancer = Upscaler(config['resolution'], kernel_size=k_size)
+#             load_model(self.upscale_enhancer, config['upscale_enhancer'])
+#             for param in self.upscale_enhancer.parameters(): param.requires_grad = False
+
+#         if config['avg_garment_path'] != '':
+#             self.avg_garment = np.load(config['avg_garment_path'])
+#         else:
+#             self.avg_garment = None
+
+
+#     def enhance(self, low, mask, training=None):
+#         x = self.upscale_enhancer([low, mask], training=training)
+#         return x
+
+
+#     def sample(self, z_cond, uv_mask, noise=None):
+#         if noise is None:
+#             eps = torch.normal(shape=(uv_mask.shape[0], self.latent_dim))
+#         else:
+#             eps = noise
+
+#         z = torch.concat([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,:,:]
+
+#         return x
+
+
+#     def reparameterize(self, mean, logvar, training, fix_noise=False):
+#         if training == True:
+#             if not fix_noise:
+#                 eps = torch.normal(mean=0.0, std=1.0, size=mean.shape)
+#             else:
+#                 eps = torch.normal(mean=0.0, std=1.0, size=mean.shape)
+#                 eps = torch.tile(eps, [mean.shape[0],1])
+#             out = eps * torch.exp(logvar * .5) + mean
+#         else:
+#             out = torch.exp(logvar * .5) + mean
+#         return out
+
+
+#     def call(self, inputs, z_cond, training=None):
+#         uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
+
+#         x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
+#         mean, logvar = self.encoder(x, training=training)
+#         z = self.reparameterize(mean, logvar, training)
+
+#         z = torch.concat([z, z_cond], dim=-1)
+#         x = self.decoder(z, training=training)
+
+#         if self.avg_garment is not None:
+#             x = x + self.avg_garment[None,:,:,:]
+
+#         x = x * uv_mask[:,None,:,:]
+
+#         # if self.upscale_enhancer is not None:
+#         #     pred_low = self.enhance(x, uv_mask, training=training)
+#         #     x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        
+#         # x = x * uv_mask_high[:,:,:,tf.newaxis]
+
+#         return x, mean, logvar
+
+
+
+class PyramidLevel(nn.Module):
     
     def __init__(self, config):
         super(PyramidLevel, self).__init__()
-        self.latent_dim = config['latent_dim']
         self.size = config['resolution']
 
-        self.encoder = VariationalEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
-        self.decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
+        self.encoder: VariationalEncoder = VariationalEncoder(config['encoder_block_settings'], input_features=config['input_features'],  kernel_size=config['kernel_size'])
+        self.decoder: Decoder = Decoder(config['decoder_block_settings'], config['kernel_size'])
 
         self.upscale_enhancer = None
         if config['upscale_enhancer'] != '':
-            k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
-            self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
+            k_size = config['kernel_size']
+            self.upscale_enhancer = Upscaler(config['resolution'], kernel_size=k_size).to('cuda')
             load_model(self.upscale_enhancer, config['upscale_enhancer'])
-            self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
-
-        if config['avg_garment_path'] != '':
-            self.avg_garment = np.load(config['avg_garment_path'])
-        else:
-            self.avg_garment = None
+            for param in self.upscale_enhancer.parameters(): param.requires_grad = False
 
 
-    def enhance(self, low, mask, training=None):
-        x = self.upscale_enhancer([low, mask], training=training)
+    def enhance(self, low, mask):
+        x = self.upscale_enhancer([low, mask])
         return x
 
 
     def sample(self, z_cond, uv_mask, noise=None):
         if noise is None:
-            eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
+            eps = torch.normal(shape=(uv_mask.shape[0], self.latent_dim))
         else:
             eps = noise
 
-        z = tf.concat([eps, z_cond], axis=-1)
+        z = torch.concat([eps, z_cond], dim=-1)
 
         x = self.decoder(z)
         
         if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
+            x = x + self.avg_garment[None,:,:,:]
         
-        x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,None:,:]
 
         return x
 
 
-    def reparameterize(self, mean, logvar, training, fix_noise=False):
-        if training == True:
+    def reparameterize(self, mean, logvar, fix_noise=False):
+        if self.training == True:
             if not fix_noise:
-                eps = tf.random.normal(shape=mean.shape)
+                eps = torch.normal(mean=0.0, std=1.0, size=mean.shape).to(mean.device)
             else:
-                eps = tf.random.normal(shape=(1, mean.shape[1]))
-                eps = tf.tile(eps, [mean.shape[0],1])
-            out = eps * tf.exp(logvar * .5) + mean
+                eps = torch.normal(mean=0.0, std=1.0, size=mean.shape).to(mean.device)
+                eps = torch.tile(eps, [mean.shape[0],1])
+            out = eps * torch.exp(logvar * .5) + mean
         else:
-            out = tf.exp(logvar * .5) + mean
+            out = torch.exp(logvar * .5) + mean
         return out
 
 
-    def call(self, inputs, z_cond, training=None):
+    def forward(self, inputs, z_cond):
         uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
 
-        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
-        mean, logvar = self.encoder(x, training=training)
-        z = self.reparameterize(mean, logvar, training)
+        x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
+        mean, logvar = self.encoder(x)
+        z = self.reparameterize(mean, logvar)
+        z = torch.reshape(z, (mean.shape[0], -1, 8, 4))
 
-        z = tf.concat([z, z_cond], axis=-1)
-        x = self.decoder(z, training=training)
+        z = torch.concat([z, z_cond], dim=1)
+        x = self.decoder(z)
 
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-
-        x = x * uv_mask[:,:,:,tf.newaxis]
-
-        # if self.upscale_enhancer is not None:
-        #     pred_low = self.enhance(x, uv_mask, training=training)
-        #     x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        
-        # x = x * uv_mask_high[:,:,:,tf.newaxis]
-
-        return x, mean, logvar
-
-
-
-class PyramidLevelNew(tf.keras.Model):
-    
-    def __init__(self, config):
-        super(PyramidLevelNew, self).__init__()
-        self.size = config['resolution']
-
-        self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
-        self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
-
-        self.upscale_enhancer = None
-        if config['upscale_enhancer'] != '':
-            k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
-            self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
-            load_model(self.upscale_enhancer, config['upscale_enhancer'])
-            self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
-
-        if config['avg_garment_path'] != '':
-            self.avg_garment = np.load(config['avg_garment_path'])
-        else:
-            self.avg_garment = None
-
-        self.reshape = tf.keras.layers.Reshape((8, 4, 64))
-
-
-    def enhance(self, low, mask, training=None):
-        x = self.upscale_enhancer([low, mask], training=training)
-        return x
-
-
-    def sample(self, z_cond, uv_mask, noise=None):
-        if noise is None:
-            eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
-        else:
-            eps = noise
-
-        z = tf.concat([eps, z_cond], axis=-1)
-
-        x = self.decoder(z)
-        
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-        
-        x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x
-
-
-    def reparameterize(self, mean, logvar, training, fix_noise=False):
-        if training == True:
-            if not fix_noise:
-                eps = tf.random.normal(shape=mean.shape)
-            else:
-                eps = tf.random.normal(shape=(1, mean.shape[1]))
-                eps = tf.tile(eps, [mean.shape[0],1])
-            out = eps * tf.exp(logvar * .5) + mean
-        else:
-            out = tf.exp(logvar * .5) + mean
-        return out
-
-
-    def call(self, inputs, z_cond, training=None):
-        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
-
-        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
-        mean, logvar = self.encoder(x, training=training)
-        z = self.reparameterize(mean, logvar, training)
-        z = self.reshape(z)
-
-        z = tf.concat([z, z_cond], axis=-1)
-        x = self.decoder(z, training=training)
-
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-
-        x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,None,:,:]
 
         # if self.upscale_enhancer is not None:
         #     pred_low = self.enhance(x, uv_mask, training=training)

models/pyramid_model.py

@@ -0,0 +1,246 @@
+import os
+import time
+import random
+import torch
+from torch import nn
+import torch.nn.functional as F
+# import numpy as np
+# import scipy as sp
+# import plotly.graph_objects as go
+# import plotly.express as ex
+
+from utils.checkpoints import load_model
+from models.base import AEEncoder
+from models.pyramid_level import PyramidLevel
+
+
+# class Pyramid(nn.Module):
+    
+#     def __init__(self, config):
+#         super(Pyramid, self).__init__()
+
+#         self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], kernel_size=config['cond_ae_static']['kernel_size'], activation=config['cond_ae_static']['activation'])
+#         if config['pyramid']['static_model'] != '':
+#             load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
+
+#         self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], kernel_size=config['cond_ae_body']['kernel_size'], activation=config['cond_ae_body']['activation'])
+#         if config['pyramid']['body_model'] != '':
+#             load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
+
+#         self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], kernel_size=config['cond_ae_curv']['kernel_size'], activation=config['cond_ae_curv']['activation'])
+#         if config['pyramid']['curvature_model'] != '':
+#             load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
+
+#         if config['freeze_encoders']:
+#             for param in self.static_enc.parameters(): param.requires_grad = False
+#             for param in self.body_enc.parameters(): param.requires_grad = False
+#             for param in self.curv_enc.parameters(): param.requires_grad = False
+
+
+#         self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae_static': config['cond_ae_static'], 'cond_ae_body': config['cond_ae_body'], 'cond_ae_curv': config['cond_ae_curv'], } }) for i in range(len(config['pyramid']['levels']))]
+
+#         # if config['is_cape']:
+#         #     if len(config['pyramid']['resolutions']) == 2:
+#         #         self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+#         #     if len(config['pyramid']['resolutions']) == 3:
+#         #         self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+#         #         self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+#         #     if len(config['pyramid']['resolutions']) == 4:
+#         #         self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+#         #         self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+#         #         self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+#         # else:
+
+#         # if len(config['pyramid']['resolutions']) == 2:
+#         #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+#         # if len(config['pyramid']['resolutions']) == 3:
+#         #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+#         #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+#         # if len(config['pyramid']['resolutions']) == 4:
+#         #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
+#         #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+#         #     self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+#         # if len(config['pyramid']['resolutions']) == 5:
+#         #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
+#         #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
+#         #     self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+#         #     self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+
+#         # self.add_noise = config['pyramid']['add_noise']
+#         # self.noise_amount = config['pyramid']['noise_amount']
+#         self.use_offsets = config['use_offsets']
+ 
+#         # self.upscale_enhancer = None
+#         # if config['pyramid']['final_upscale_enhancer'] != '':
+#         #     k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
+#         #     self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
+#         #     load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
+#         #     self.upscale_enhancer.trainable = False
+
+#     def enhance(self, low, mask):
+#         x = self.upscale_enhancer([low, mask])
+#         return x
+
+#     def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+#         z_static = self.static_enc(uv_template)
+#         z_body = self.body_enc(uv_body_posed)
+#         z_curv = self.curv_enc(uv_curv)
+
+#         z_cond = torch.concat([z_static, z_body, z_curv], dim=-1)
+
+#         return z_cond
+
+#     def call(self, inputs):
+#         levels_input, conds = inputs
+#         uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
+
+#         z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
+
+#         # if self.add_noise == True and training == True:
+#         #     z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
+
+#         out_levels = []
+#         pred, mean, logvar = self.levels[0](levels_input[0], z_cond)
+#         if self.use_offsets:
+#             pred = pred + levels_input[0][2] # uv_body_posed_32
+#         if 1 != len(self.levels):
+#             pred = self.levels[0].enhance(pred, levels_input[0][-2])  # levels_input[i][-2]=uv_mask_res
+#             pred = F.interpolate(pred, size=(self.levels[1].size, self.levels[1].size//2), mode='bilinear')
+#             out_levels.append([pred, mean, logvar])
+#             pred = pred * levels_input[1][-2][:,None,:,:]
+#         else:
+#             pred = self.enhance(pred, levels_input[0][-2])
+#             pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
+#             pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask
+#             out_levels.append([pred, mean, logvar])
+
+#         for i in range(1, len(self.levels)):
+#             pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond)
+#             pred = pred + pred_level
+#             if i != len(self.levels) - 1:
+#                 pred = self.levels[i].enhance(pred, levels_input[i][-2])
+#                 out_levels.append([pred, mean, logvar])
+#                 pred = F.interpolate(pred, size=(self.levels[i+1].size, self.levels[i+1].size//2), mode='bilinear')
+#                 pred = pred * levels_input[i+1][-2][:,None,:,:]
+#             else:
+#                 pred = self.enhance(pred, levels_input[i][-1])
+#                 pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask
+#                 pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
+#                 out_levels.append([pred, mean, logvar])
+
+#         return out_levels
+
+
+
+class Pyramid(nn.Module):
+    
+    def __init__(self, config):
+        super(Pyramid, self).__init__()
+
+        self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], kernel_size=config['cond_ae_static']['kernel_size'], activation=config['cond_ae_static']['activation'])
+        if config['static_model'] != '':
+            load_model(self.static_enc, config['static_model'] + '_enc')
+
+        self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], kernel_size=config['cond_ae_body']['kernel_size'], activation=config['cond_ae_body']['activation'])
+        if config['body_model'] != '':
+            load_model(self.body_enc, config['body_model'] + '_enc')
+
+        self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], kernel_size=config['cond_ae_curv']['kernel_size'], activation=config['cond_ae_curv']['activation'], is_label=True, nr_of_labels=7)
+        if config['curvature_model'] != '':
+            load_model(self.curv_enc, config['curvature_model'] + '_enc')
+
+        if config['freeze_encoders']:
+            for param in self.static_enc.parameters(): param.requires_grad = False
+            for param in self.body_enc.parameters(): param.requires_grad = False
+            for param in self.curv_enc.parameters(): param.requires_grad = False
+
+        self.levels = nn.ModuleList([PyramidLevel({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))])
+
+
+        # if len(config['pyramid']['resolutions']) == 2:
+        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        # if len(config['pyramid']['resolutions']) == 3:
+        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+        #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        # if len(config['pyramid']['resolutions']) == 4:
+        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
+        #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+        #     self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        # if len(config['pyramid']['resolutions']) == 5:
+        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
+        #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
+        #     self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+        #     self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+
+        # self.add_noise = config['pyramid']['add_noise']
+        # self.noise_amount = config['pyramid']['noise_amount']
+        # self.use_offsets = config['use_offsets']
+ 
+        # self.upscale_enhancer = None
+        # if config['pyramid']['final_upscale_enhancer'] != '':
+        #     k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
+        #     self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
+        #     load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
+        #     self.upscale_enhancer.trainable = False
+
+    # def enhance(self, low, mask, training):
+    #     x = self.upscale_enhancer([low, mask], training=training)
+    #     return x
+
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(uv_body_posed)
+        z_curv = self.curv_enc(uv_curv)
+
+        z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
+
+        return z_cond
+
+    def forward(self, inputs):
+        levels_input, conds = inputs
+        uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
+
+        z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
+
+        # if self.add_noise == True and training == True:
+        #     z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
+
+        out_levels = []
+        pred, mean, logvar = self.levels[0](levels_input[0], z_cond)
+        # if self.use_offsets:
+        #     pred = pred + levels_input[0][2] # uv_body_posed_32
+        if 1 != len(self.levels):
+            pred = self.levels[0].enhance(pred, levels_input[0][-2])  # levels_input[i][-2]=uv_mask_res
+            pred = F.interpolate(pred, size=(self.levels[1].size, self.levels[1].size//2), mode='bilinear')
+            out_levels.append([pred, mean, logvar])
+            pred = pred * levels_input[1][-2][:,None,:,:]
+        else:
+            # pred = self.enhance(pred, levels_input[0][-2], training=training)
+            pred = self.levels[0].enhance(pred, levels_input[0][-2])  # levels_input[i][-2]=uv_mask_res
+            pred = F.interpolate(pred, size=(512, 256), mode='bilinear') 
+            pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask
+            out_levels.append([pred, mean, logvar])
+
+        for i in range(1, len(self.levels)):
+            input = levels_input[i]
+            # input[0] = input[0] - pred
+            # input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3])
+            # input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None]
+            # input[0] = input[0] * input[-2][:,None,:,:] # uv_mask_res
+            pred_level, mean, logvar = self.levels[i](input, z_cond)
+            # pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res
+            pred = pred + pred_level
+            if i != len(self.levels) - 1:
+                pred = self.levels[i].enhance(pred, levels_input[i][-2])
+                out_levels.append([pred, mean, logvar])
+                pred = F.interpolate(pred, size=(self.levels[i+1].size, self.levels[i+1].size//2), mode='bilinear')
+                pred = pred * levels_input[i+1][-2][:,None,:,:]
+            else:
+                # if self.upscale_enhancer is not None:
+                    # pred = self.enhance(pred, levels_input[-1][-2], training=training)
+                pred = self.levels[-1].enhance(pred, levels_input[-1][-2])
+                pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
+                pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask
+                out_levels.append([pred, mean, logvar])
+
+        return out_levels

models/pyramid_model_optimized.py

@@ -1,365 +0,0 @@
-import os
-import time
-import random
-import tensorflow as tf
-import numpy as np
-# import plotly.graph_objects as go
-# import plotly.express as ex
-
-from custom_utils.custom_utils import load_model
-from models.base import AEEncoder, AEEncoderNew
-from models.pyramid_level import PyramidLevel, PyramidLevelNew
-from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
-
-
-class Pyramid(tf.keras.Model):
-    
-    def __init__(self, config):
-        super(Pyramid, self).__init__()
-
-        self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
-        if config['pyramid']['static_model'] != '':
-            load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
-
-        self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
-        if config['pyramid']['body_model'] != '':
-            load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
-
-        self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
-        if config['pyramid']['curvature_model'] != '':
-            load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
-
-        if config['pyramid']['freeze_encoders']:
-            self.static_enc.trainable = False
-            self.body_enc.trainable = False
-            self.curv_enc.trainable = False
-
-        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
-                                                 tf.keras.layers.Dense(256, activation='relu'),
-                                                 tf.keras.layers.Dense(128, activation='relu'),
-                                                 tf.keras.layers.Dense(config['pyramid']['latent_dim'], activation='tanh'),])
-
-        self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
-
-        # if config['is_cape']:
-        #     if len(config['pyramid']['resolutions']) == 2:
-        #         self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        #     if len(config['pyramid']['resolutions']) == 3:
-        #         self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        #         self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        #     if len(config['pyramid']['resolutions']) == 4:
-        #         self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        #         self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        #         self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        # else:
-        if len(config['pyramid']['resolutions']) == 2:
-            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        if len(config['pyramid']['resolutions']) == 3:
-            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
-            self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        if len(config['pyramid']['resolutions']) == 4:
-            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
-            self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
-            self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        if len(config['pyramid']['resolutions']) == 5:
-            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
-            self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
-            self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
-            self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-
-        self.add_noise = config['pyramid']['add_noise']
-        self.noise_amount = config['pyramid']['noise_amount']
-        self.use_offsets = config['use_offsets']
- 
-        self.upscale_enhancer = None
-        if config['pyramid']['final_upscale_enhancer'] != '':
-            k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
-            self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
-            load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
-            self.upscale_enhancer.trainable = False
-
-    def enhance(self, low, mask, training):
-        x = self.upscale_enhancer([low, mask], training=training)
-        return x
-
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
-        z_static = self.static_enc(uv_template, training)
-        z_body = self.body_enc(uv_body_posed, training)
-        z_curv = self.curv_enc(uv_curv, training)
-
-        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-        z_cond = self.cond_encoder(z_cond, training=training)
-
-        return z_cond
-
-    def call(self, inputs, training):
-        levels_input, conds = inputs
-        uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
-
-        z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
-
-        if self.add_noise == True and training == True:
-            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
-
-        out_levels = []
-        pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
-        if self.use_offsets:
-            pred = pred + levels_input[0][2] # uv_body_posed_32
-        if 1 != len(self.levels):
-            pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
-            pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
-            out_levels.append([pred, mean, logvar])
-            pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
-        else:
-            pred = self.enhance(pred, levels_input[0][-2], training=training)
-            pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-            pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
-            out_levels.append([pred, mean, logvar])
-
-        for i in range(1, len(self.levels)):
-            pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
-            pred = pred + pred_level
-            if i != len(self.levels) - 1:
-                pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
-                out_levels.append([pred, mean, logvar])
-                pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
-                pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
-            else:
-                if self.upscale_enhancer is not None:
-                    pred = self.enhance(pred, levels_input[-1][-2], training=training)
-                    pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-                pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
-                out_levels.append([pred, mean, logvar])
-
-        return out_levels
-
-        # # 32
-        # pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
-        # pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
-        # pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        # pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
-        # return [ [pred_0, mean_0, logvar_0] ]
-
-        # # 32 64
-        # pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
-        # pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
-        # pred_0 = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
-        # pred = pred_0 * levels_input[1][-2][:,:,:,tf.newaxis]
-
-        # pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
-        # pred_1 = pred + pred_1
-        # pred_1 = self.enhance(pred_1, levels_input[-1][-2], training=training)
-        # pred_1 = tf.image.resize(pred_1, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        # pred_1 = pred_1 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
-
-        # return [ [pred_0, mean_0, logvar_0], [pred_1, mean_1, logvar_1] ]
-
-
-# class Pyramid(tf.keras.Model):
-    
-#     def __init__(self, config):
-#         super(Pyramid, self).__init__()
-
-#         self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
-#         if config['pyramid']['static_model'] != '':
-#             load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
-
-#         self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
-#         if config['pyramid']['body_model'] != '':
-#             load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
-
-#         self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
-#         if config['pyramid']['curvature_model'] != '':
-#             load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
-
-#         if config['pyramid']['freeze_encoders']:
-#             self.static_enc.trainable = False
-#             self.body_enc.trainable = False
-#             self.curv_enc.trainable = False
-
-#         self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
-#                                                  tf.keras.layers.Dense(256, activation='relu'),
-#                                                  tf.keras.layers.Dense(128, activation='relu'),
-#                                                  tf.keras.layers.Dense(config['pyramid']['latent_dim']),])
-
-#         self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
-
- 
-#         self.upscale_enhancer = None
-#         if config['pyramid']['final_upscale_enhancer'] != '':
-#             self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
-#             load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
-#             self.upscale_enhancer.trainable = False
-
-#     def enhance(self, low, mask, training):
-#         x = self.upscale_enhancer([low, mask], training=training)
-#         return x
-
-#     def encode_conds(self, uv_template, uv_body_posed, uv_curv):
-#         z_static = self.static_enc(uv_template)
-#         z_body = self.body_enc(uv_body_posed)
-#         z_curv = self.curv_enc(uv_curv)
-
-#         z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-#         z_cond = self.cond_encoder(z_cond)
-
-#         return z_cond
-
-#     def call(self, inputs, training):
-#         levels_input, conds = inputs
-#         uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
-#         z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
-
-#         out_levels = []
-#         pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
-#         pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
-#         # pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
-            
-#         ##################    
-#         # if more than 1 level
-#         # x = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
-#         # x = x * levels_input[1][-2][:,:,:,tf.newaxis]
-
-#         # if only 1 level
-#         pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-#         ##################
-
-
-#         # pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
-#         # pred_1 = pred_0 + pred_1
-#         # pred_1 = self.levels[1].enhance(pred, levels_input[1][-2], training=training)
-
-#         # for i in range(1, len(self.levels)):
-#         #     pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
-#         #     pred = pred + pred_level
-#         #     if i != len(self.levels) - 1:
-#         #         pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
-#         #         out_levels.append([pred, mean, logvar])
-#         #         pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
-#         #         pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
-#         #     else:
-#         #         if self.upscale_enhancer is not None:
-#         #             pred = self.enhance(pred, levels_input[-1][-2], training=training)
-#         #             pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-#         #         out_levels.append([pred, mean, logvar])
-#         #         pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
-
-#         out_levels = [
-#             [pred_0, mean_0, logvar_0]
-#         ]
-
-#         return out_levels
-
-
-class PyramidNew(tf.keras.Model):
-    
-    def __init__(self, config):
-        super(PyramidNew, self).__init__()
-
-        self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'])
-        if config['pyramid']['static_model'] != '':
-            load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
-
-        self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size'])
-        if config['pyramid']['body_model'] != '':
-            load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
-
-        self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['encoder_depths'], config['cond_ae_curv']['encoder_dims'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'])
-        if config['pyramid']['curvature_model'] != '':
-            load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
-
-        if config['pyramid']['freeze_encoders']:
-            self.static_enc.trainable = False
-            self.body_enc.trainable = False
-            self.curv_enc.trainable = False
-
-        self.levels = [PyramidLevelNew({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))]
-
-
-        # if len(config['pyramid']['resolutions']) == 2:
-        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        # if len(config['pyramid']['resolutions']) == 3:
-        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
-        #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        # if len(config['pyramid']['resolutions']) == 4:
-        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
-        #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
-        #     self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        # if len(config['pyramid']['resolutions']) == 5:
-        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
-        #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
-        #     self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
-        #     self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-
-        self.add_noise = config['pyramid']['add_noise']
-        self.noise_amount = config['pyramid']['noise_amount']
-        self.use_offsets = config['use_offsets']
- 
-        self.upscale_enhancer = None
-        if config['pyramid']['final_upscale_enhancer'] != '':
-            k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
-            self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
-            load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
-            self.upscale_enhancer.trainable = False
-
-    def enhance(self, low, mask, training):
-        x = self.upscale_enhancer([low, mask], training=training)
-        return x
-
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
-        z_static = self.static_enc(uv_template, training)
-        z_body = self.body_enc(uv_body_posed, training)
-        z_curv = self.curv_enc(uv_curv, training)
-
-        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-
-        return z_cond
-
-    def call(self, inputs, training):
-        levels_input, conds = inputs
-        uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
-
-        z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
-
-        if self.add_noise == True and training == True:
-            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
-
-        out_levels = []
-        pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
-        if self.use_offsets:
-            pred = pred + levels_input[0][2] # uv_body_posed_32
-        if 1 != len(self.levels):
-            pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
-            pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR)
-            out_levels.append([pred, mean, logvar])
-            pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
-        else:
-            # pred = self.enhance(pred, levels_input[0][-2], training=training)
-            pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
-            pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
-            pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
-            out_levels.append([pred, mean, logvar])
-
-        for i in range(1, len(self.levels)):
-            input = levels_input[i]
-            input[0] = input[0] - pred
-            input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3])
-            input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None]
-            input[0] = input[0] * input[-2][:,:,:,None] # uv_mask_res
-            pred_level, mean, logvar = self.levels[i](input, z_cond, training=training)
-            pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res
-            pred = pred + pred_level
-            if i != len(self.levels) - 1:
-                pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
-                out_levels.append([pred, mean, logvar])
-                pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.BILINEAR)
-                pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
-            else:
-                # if self.upscale_enhancer is not None:
-                    # pred = self.enhance(pred, levels_input[-1][-2], training=training)
-                pred = self.levels[-1].enhance(pred, levels_input[-1][-2], training=training)
-                pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
-                pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
-                out_levels.append([pred, mean, logvar])
-
-        return out_levels

models/upscaler.py

@@ -1,172 +1,35 @@
-import tensorflow as tf
-from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip
+import torch
+import torch.nn as nn
 
-
-class Bottleneck(tf.keras.Model):
-    
-    def __init__(self, dim, kernel_size=7):
-        super(Bottleneck, self).__init__()
-
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0, kernel_size=kernel_size,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-
-    def call(self, inputs, training):
-        x = inputs
-
-        x = self.bottleneck(x, training=training)
-        
-        return x
-
-
-
-class Upscaler(tf.keras.Model):
-    
-    def __init__(self, bottleneck_size):
+class Upscaler(nn.Module):
+    def __init__(self, low_res=32, kernel_size=3):
         super(Upscaler, self).__init__()
 
-        depths = [27, 27]
-        dims = [512, 512]
-        reverse_dims = [16] + dims[:-1]
-        self.encoder = ConvNeXtWithSkip(depths=depths, dims=dims, kernel_size=3)
-
-        self.bottleneck = Bottleneck(dims[-1], kernel_size=3)
-
-        self.decoder = ConvNeXtReverseWithSkip(depths=depths, dims=reverse_dims, kernel_size=3)
-
-    def call(self, inputs, training):
-        img, mask = inputs
-
-        img = tf.transpose(img, [0,2,3,1])
-
-        x, skip_features = self.encoder(img, training=training)
-
-        x = self.bottleneck(x, training=training)
-
-        x = self.decoder([x, skip_features], training=training)
-
-        x = x * (1 - mask[:,:,:,tf.newaxis])
-        x = x + img
-        x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
-
-        return x
-
-
-class UpscalerBasic(tf.keras.Model):
-    
-    def __init__(self, bottleneck_size, low_res=32, high_res=64):
-        super(UpscalerBasic, self).__init__()
-
-        self.low_res = low_res
-        self.high_res = high_res
-        self.flatten = tf.keras.layers.Flatten()
-
-        self.seq = tf.keras.Sequential([
-                    tf.keras.layers.Dense(1024, activation='relu'),
-                    tf.keras.layers.BatchNormalization(),
-                    tf.keras.layers.Dense(1024, activation='relu'),
-                    tf.keras.layers.BatchNormalization(),
-                    tf.keras.layers.Dense(1024, activation='relu'),
-                    tf.keras.layers.BatchNormalization(),
-                    tf.keras.layers.Dense(low_res * (low_res//2) * 3),
-        ])
-
-        self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3))
-
-    def call(self, inputs, training):
-        img, mask = inputs
-
-        x = self.flatten(img)
-        x = self.seq(x, training=training)
-        x = self.reshape(x)
-
-        x = x * (1 - mask[:,:,:,tf.newaxis])
-        x = x + img
-        x = tf.image.resize(x, size=(self.high_res, self.high_res//2), method=tf.image.ResizeMethod.LANCZOS5)
-
-        return x
-
-
-
-
-class UpscalerBasicNew(tf.keras.Model):
-    
-    def __init__(self, low_res=32):
-        super(UpscalerBasicNew, self).__init__()
-
-        self.low_res = low_res
-        self.flatten = tf.keras.layers.Flatten(name='upscale_enhancer_flattent')
-
-        self.seq = tf.keras.Sequential([
-                    tf.keras.layers.Dense(4096, activation='relu', name='upscale_enhancer_dense_1'),
-                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
-                    tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_2'),
-                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
-                    tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_3'),
-                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
-                    tf.keras.layers.Dense(low_res * (low_res//2) * 3, activation='tanh', name='upscale_enhancer_dense_final'),
-        ])
-
-        self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3), name='upscale_enhancer_reshape')
-
-    def call(self, inputs, training):
-        img, mask = inputs
-
-        x = self.flatten(img)
-        x = self.seq(x, training=training)
-        x = self.reshape(x)
-
-        x = x * (1 - mask[:,:,:,tf.newaxis])
-        x = x + img
-
-        return x
-
-
-class UpscalerBasicConv(tf.keras.Model):
-    
-    def __init__(self, low_res=32, kernel_size=3):
-        super(UpscalerBasicConv, self).__init__()
-
         self.low_res = low_res
 
-        self.seq = tf.keras.Sequential([
-                    tf.keras.layers.Conv2D(256, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
-                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
-                    tf.keras.layers.Conv2D(128, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
-                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
-                    tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
-                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
-                    tf.keras.layers.Conv2D(32, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
-                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'),
-                    tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'),
-        ])
+        self.seq = nn.Sequential(
+            nn.Conv2d(3, 256, kernel_size=kernel_size, padding='same'),
+            nn.ReLU(),
+            nn.BatchNorm2d(256),
+            nn.Conv2d(256, 128, kernel_size=kernel_size, padding='same'),
+            nn.ReLU(),
+            nn.BatchNorm2d(128),
+            nn.Conv2d(128, 64, kernel_size=kernel_size, padding='same'),
+            nn.ReLU(),
+            nn.BatchNorm2d(64),
+            nn.Conv2d(64, 32, kernel_size=kernel_size, padding='same'),
+            nn.ReLU(),
+            nn.BatchNorm2d(32),
+            nn.Conv2d(32, 3, kernel_size=kernel_size, padding='same'),
+            nn.Tanh()
+        )
 
-
-    def call(self, inputs, training):
+    def forward(self, inputs) -> torch.Tensor:
         img, mask = inputs
 
-        x = self.seq(img, training=training)
+        x = self.seq(img)
 
-        x = x * (1 - mask[:,:,:,tf.newaxis])
+        x = x * (1 - mask.unsqueeze(1))
         x = x + img
 
         return x
-
-# class Upscaler(tf.keras.Model):
-    
-#     def __init__(self, bottleneck_size):
-#         super(Upscaler, self).__init__()
-
-#         self.decoder = ConvNeXtReverse(depths=list(reversed([9, 9, 9])), dims=list(reversed([bottleneck_size, 128, 32, 16])))
-
-#     def call(self, inputs, training):
-#         x, mask = inputs
-
-#         x = tf.transpose(x, [0,2,3,1])
-#         x = self.decoder(x)
-
-#         x = x * mask[:,:,:,tf.newaxis]
-
-#         return x
-

run.py

@@ -1,9 +1,25 @@
 import os
-from configs.ae_static import config
+import sys
+from configs.direct_1 import config
+import argparse
+
+# parser = argparse.ArgumentParser(description='')
+# parser.add_argument('--debug', action='store_true')
+# args = vars(parser.parse_args())
+
 
 
 if __name__ == '__main__':
-
+    import wandb
+    # wandb.login(key='******************************************')
+    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]
     current_file = '/'.join(__file__.split('/')[-2:])
-    os.system(f'nohup python scripts/ae_static_train.py --name lhunor_{name} > nohup/torch/{name}.out 2>&1 &')
+    
+    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 &')

scripts/ae_body_train.py

@@ -0,0 +1,214 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import argparse
+import importlib
+
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+if args['debug']:
+    from configs.ae_body_cape_m import config
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import wandb
+import random
+import numpy as np
+import torch
+import cv2
+# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str
+import plotly.graph_objects as go
+import plotly.express as ex
+import torch.optim as optim
+
+if config['is_skirt']:
+    from losses.losses import edge_loss_skirt as edge_loss
+else:
+    from losses.losses import edge_loss
+
+from data.dataset import Dataset
+# from torch.utils.data import DataLoader
+from custom_utils.custom_utils import plot_3d
+from models.cond_encoder import AutoEncoder
+from utils.checkpoints import CheckpointSaver, load_model
+from custom_utils.custom_utils import Accumulator
+from utils.utils import InfiniteDataLoader as DataLoader
+
+from pytorch_msssim import MS_SSIM
+
+
+def init_wandb():
+    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/ae_body_train.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+
+
+def train_step(batch):
+    model.train()
+    pred = model([batch['uv_body_posed'], uv_mask_body])
+
+    ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_body_posed'] + 1))
+    loss = 1e3 * (1 - ms_ssim)
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    model.eval()
+    with torch.inference_mode():
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_body_posed']
+        l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_body_posed'] - pred)), axis=1))
+        e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, torch.ones((pred.shape[0], 3638 if config['is_cape'] else 14475), device=pred.get_device())))
+        metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1)))
+    
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
+        
+
+def valid_step(batch):
+    model.eval()
+
+    with torch.inference_mode():
+        pred = model([batch['uv_body_posed'], uv_mask_body])
+
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_body_posed']
+
+        e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, torch.ones((pred.shape[0], 3638 if config['is_cape'] else 14475), device=pred.get_device())))
+        l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_body_posed'] - pred)), axis=1))
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_body_posed'] + 1))
+        loss = 1e3 * (1 - ms_ssim)
+
+        metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1)))
+    
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
+       
+
+def plot_everything(batch: dict, type: str):
+    model.eval()
+    pred = model([batch['uv_body_posed'], uv_mask_body]).detach()
+    
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_body_posed']
+    
+    orig_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
+    pred_img = torch.permute(pred, [0,2,3,1])
+    img = torch.concat((orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
+
+    F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
+    garment_unposed, faces = mesh_pred.cpu().numpy(), F
+    gt = [mesh_orig[i].cpu().numpy() for i in range(pred.shape[0])]
+    
+    body = uv2mesh(batch['uv_body_posed']).cpu().numpy()
+    garment_posed = garment_unposed
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces + (3638 if config['is_cape'] else 14475)), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    return { **figs, **{ type + '/images': imgs } }
+
+
+def train():
+    valid_iterator = iter(valid_dataloader)
+
+    step = 0
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            if not args['debug']:
+                wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if step % 10000 == 0:
+            ckpt_saver_enc(model.encoder, step, optimizer, scaler=None)
+            ckpt_saver_dec(model.decoder, step, optimizer, scaler=None)
+
+        step += 1
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy')
+    uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
+    uv_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device'])
+
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True)
+    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model : AutoEncoder = AutoEncoder(config['cond_ae'])
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+
+    # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
+
+    if config['pretrained_model'] != '':
+        load_model(model.encoder, config['pretrained_model'] + '_enc')
+        load_model(model.decoder, config['pretrained_model'] + '_dec')
+
+    ckpt_saver_enc = CheckpointSaver('checkpoints/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}')
+
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+
+    train()
+
+
+

scripts/ae_curv_train.py

@@ -0,0 +1,223 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import argparse
+import importlib
+
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+if args['debug']:
+    from configs.ae_curv_cape_m import config
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import wandb
+import random
+import numpy as np
+import torch
+import cv2
+# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str
+import plotly.graph_objects as go
+import plotly.express as ex
+import torch.optim as optim
+import cmapy as cm
+
+if config['is_skirt']:
+    from losses.losses import edge_loss_skirt as edge_loss
+else:
+    from losses.losses import edge_loss
+
+from data.dataset import Dataset
+# from torch.utils.data import DataLoader
+from custom_utils.custom_utils import plot_3d
+from models.cond_encoder import AutoEncoder
+from utils.checkpoints import CheckpointSaver, load_model
+from custom_utils.custom_utils import Accumulator
+from utils.utils import InfiniteDataLoader as DataLoader
+from torchmetrics.classification import MulticlassJaccardIndex
+from torchmetrics import Dice
+from pytorch_msssim import MS_SSIM
+
+
+def init_wandb():
+    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/ae_curv_train.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+
+def train_step(batch):
+    model.train()
+    pred = model([batch['uv_curvature'], batch['uv_mask']])
+    gt_dense = torch.argmax(batch['uv_curvature'], dim=1)
+
+    ce_uv_loss = _ce_uv(pred, gt_dense)
+    mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,None]
+    mesh_orig = uv2mesh(gt_dense) * batch['mesh_mask']
+    ce_mesh_loss = _ce_mesh(torch.transpose(mesh_pred, 1, 2), mesh_orig.long())
+
+    loss = ce_uv_loss + ce_mesh_loss
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    model.eval()
+    with torch.inference_mode():
+        ms_ssim = torch.mean(_MSSSIM(torch.argmax(pred, dim=1)[:,None,:,:].float(), gt_dense[:,None,:,:].float()))
+        mean_iou = _mean_iou(pred, gt_dense)
+        dice = _dice(pred, gt_dense)
+
+    return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou,
+             'train/dice': dice, 'train/ce_mesh_loss': ce_mesh_loss }
+        
+
+def valid_step(batch):
+    model.eval()
+
+    with torch.inference_mode():
+        pred = model([batch['uv_curvature'], batch['uv_mask']])
+        gt_dense = torch.argmax(batch['uv_curvature'], dim=1)
+
+        ce_uv_loss = _ce_uv(pred, gt_dense)
+        mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,None]
+        mesh_orig = uv2mesh(gt_dense) * batch['mesh_mask']
+        ce_mesh_loss = _ce_mesh(torch.transpose(mesh_pred, 1, 2), mesh_orig.long())
+
+        loss = ce_uv_loss + ce_mesh_loss
+
+        ms_ssim = torch.mean(_MSSSIM(torch.argmax(pred, dim=1)[:,None,:,:].float(), gt_dense[:,None,:,:].float()))
+        mean_iou = _mean_iou(pred, gt_dense)
+        dice = _dice(pred, gt_dense)
+
+    return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou,
+             'valid/dice': dice, 'valid/ce_mesh_loss': ce_mesh_loss }
+       
+
+def plot_everything(batch: dict, type: str):
+    model.eval()
+    pred = model([batch['uv_curvature'], batch['uv_mask']])
+    # gt_dense = torch.argmax(batch['uv_curvature'], dim=1)
+
+    # mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,None]
+    # mesh_orig = uv2mesh(gt_dense) * batch['mesh_mask']
+    
+    orig_img = torch.argmax(batch['uv_curvature'], dim=1)
+    pred_img = torch.argmax(pred, dim=1)
+    img = torch.concat((orig_img, pred_img), axis=2).cpu().numpy()
+    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+
+    colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
+    color_img = np.zeros((img.shape[0], img.shape[1], img.shape[2], 3), np.uint8)
+    for i in range(7):
+        color_img[img == i, :] = colors[i]
+    color_img = [wandb.Image(color_img[i]) for i in range(4)]
+
+    imgs = { type + '/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])], type + '/color_images': color_img }
+
+    return imgs
+
+
+def train():
+    valid_iterator = iter(valid_dataloader)
+
+    step = 0
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            if not args['debug']:
+                wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if step % 10000 == 0:
+            ckpt_saver_enc(model.encoder, step, optimizer, None)
+            ckpt_saver_dec(model.decoder, step, optimizer, None)
+
+        step += 1
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy')
+    uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
+    uv_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device'])
+
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True)
+    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model : AutoEncoder = AutoEncoder(config['cond_ae'])
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+
+    # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=1)
+    _mean_iou = MulticlassJaccardIndex(num_classes=7).to(config['device'])
+    _dice = Dice(average='macro', num_classes=7).to(config['device'])
+    _ce_uv = torch.nn.CrossEntropyLoss()
+    _ce_mesh = torch.nn.CrossEntropyLoss()
+
+    if config['pretrained_model'] != '':
+        load_model(model.encoder, config['pretrained_model'] + '_enc')
+        load_model(model.decoder, config['pretrained_model'] + '_dec')
+
+    ckpt_saver_enc = CheckpointSaver('checkpoints/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}')
+
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+
+    train()
+
+
+

scripts/ae_static_train.py

@@ -2,10 +2,23 @@ import os
 import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
-from configs.ae_static import config
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+if args['debug']:
+    from configs.ae_static_cape_m import config
+    # from configs.ae_static_orig_2 import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
-
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 import time
 import argparse
@@ -57,6 +70,7 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_
 
 def train_step(batch):
     model.train()
+    # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
     pred = model([batch['uv_static'], batch['uv_mask']])
 
     ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
@@ -66,6 +80,10 @@ def train_step(batch):
     optimizer.zero_grad()
     loss.backward()
     optimizer.step()
+    # scaler.scale(loss).backward()
+    # scaler.step(optimizer)
+    # scaler.update()
+    # optimizer.zero_grad()
 
     model.eval()
     with torch.inference_mode():
@@ -98,11 +116,16 @@ def valid_step(batch):
        
 
 def plot_everything(batch: dict, type: str):
-    pred = model([batch['uv_static'], batch['uv_mask']]).detach()
+    model.eval()
+    with torch.inference_mode():
+        pred = model([batch['uv_static'], batch['uv_mask']]).detach()
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_static']
 
+    mesh_pred = torch.cat((mesh_pred[:, :, 0:1], -mesh_pred[:, :, 2:3], mesh_pred[:, :, 1:2]), dim=2)
+    mesh_orig = torch.cat((mesh_orig[:, :, 0:1], -mesh_orig[:, :, 2:3], mesh_orig[:, :, 1:2]), dim=2)
+
     orig_img = torch.permute(batch['uv_static'], [0,2,3,1])
     pred_img = torch.permute(pred, [0,2,3,1])
     img = torch.concat((orig_img, pred_img), axis=2)
@@ -146,34 +169,29 @@ def train():
             if step == 0:
                 to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-            wandb.log(to_log, step=step)
+            if not args['debug']:
+                wandb.log(to_log, step=step)
             train_accumulator.reset()
             valid_accumulator.reset()
             print(step)
 
         if step % 10000 == 0:
-            ckpt_saver_enc(model.encoder, step, optimizer)
-            ckpt_saver_dec(model.decoder, step, optimizer)
+            ckpt_saver_enc(model.encoder, step, optimizer, scaler)
+            ckpt_saver_dec(model.decoder, step, optimizer, scaler)
 
         step += 1
 
 
 if __name__ == '__main__':
-    init_wandb()
-    # device = "cuda" if torch.cuda.is_available() else "cpu"
+    if not args['debug']:
+        init_wandb()
     
-    parser = argparse.ArgumentParser(description='')
-    parser.add_argument('--name', type=str)
-    args = vars(parser.parse_args())
-    
-    resolution_modifier = ''
-    if config['resolution'] != 512:
-        resolution_modifier = '_' + str(config['resolution'])
+    use_amp = True
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
-    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()
@@ -181,6 +199,7 @@ if __name__ == '__main__':
 
     model : AutoEncoder = AutoEncoder(config['cond_ae'])
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
 
     # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
@@ -189,11 +208,16 @@ 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}')
+
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+
     train()
 
 

scripts/cvae.py

@@ -1,21 +1,23 @@
 import os
 import sys
-from configs.config_abl_all_post import config
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+if args['debug']:
+    pass
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
-import tensorflow as tf
-
-gpus = tf.config.experimental.list_physical_devices('GPU')
-if gpus:
-    try:
-        # Currently, memory growth needs to be the same across GPUs
-        for gpu in gpus:
-            tf.config.experimental.set_memory_growth(gpu, True)
-        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
-        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
-    except RuntimeError as e:
-        # Memory growth must be set before GPUs have been initialized
-        print(e)
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 import time
 import argparse
@@ -26,256 +28,215 @@ import numpy as np
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
-from utils.postprocess import fix_collision
 
 if config['is_skirt']:
-    from utils.losses import edge_loss_skirt as edge_loss
+    from losses.losses import edge_loss_skirt as edge_loss
 else:
-    from utils.losses import edge_loss
+    from losses.losses import edge_loss
 
-from utils.losses import collision_ratio, kl_loss, normal_loss, collision_loss
-from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
-from data.dataset import DatasetFactory, StaticDataset
-from custom_utils.smpl import SMPL
-from models.cvae_model import CVAENew as CVAE
-from utils.utils import Accumulator
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from data.dataset import Dataset
+from models.cond_encoder import AutoEncoder
+from utils.checkpoints import CheckpointSaver, load_model
+from utils.utils import InfiniteDataLoader as DataLoader
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss
+from pytorch_msssim import MS_SSIM
+from models.cvae_model import CVAE
+from smpl.smpl_torch import SMPL
+from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision
 
 
 def init_wandb():
-    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags'])
+    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
-    config_file_names = [key for key in sys.modules.keys()  if (key.startswith('configs.') or key.startswith('custom_utils.') or key.startswith('utils.') or key.startswith('data.') or key.startswith('models.') or key.startswith('scripts.'))]
+    artifact.add_file('scripts/cvae.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('models/cvae_model.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
     for config_file_name in config_file_names:
         file_name = '/'.join(config_file_name.split('.')) + '.py'
         if os.path.exists(file_name):
             artifact.add_file(file_name)
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
 
 
-def kl_loss(mean, logvar):
-    kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
-    kl_loss = tf.reduce_mean(kl_loss)
+def kl_loss(mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+    kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar)
+    kl_loss = torch.mean(kl_loss)
     kl_loss *= -0.5
     return kl_loss
 
 
-# 3d loss on posed !!!!!!!!!!!!
-
-# @tf.function
 def train_step(batch):
+    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']
 
-    with tf.GradientTape() as tape:
-        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']
+    model.train()
+    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
+    mesh_pred = uv2mesh(pred)
 
-        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
+    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])])
 
-        mesh_pred = uv2mesh(pred)
+        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
 
-        if config['use_skinning']:
-            shape = batch['pose_shape'][:, 72:]
-            gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-            shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
 
-            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-            mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
-        else:
-            mesh_pred_posed = mesh_pred
+    e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+    c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+    l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
+    l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
+    n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+    kl = kl_loss(mean, logvar)
 
-        # is the output already masked?
+    reconstr_loss = l2_loss
+    if config['use_extra_losses']:
+        reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
 
-        e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
-        c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
-        l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
-        l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
-        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-        # conc_loss = concept_loss(pred, batch['uv_unposed'])
-        kl = kl_loss(mean, logvar)
+    loss = reconstr_loss + g_balancing_coeff * kl
 
-        reconstr_loss = l2_loss
-        if config['use_extra_losses']:
-            if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True:
-                reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff']
-            else:
-                reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
-        if config['use_norm_loss'] and step > 2000:
-            reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
 
-        loss = reconstr_loss + g_balancing_coeff * kl
+    model.eval()
+    with torch.inference_mode():
+        mean = torch.mean(mean)
+        logvar = torch.mean(logvar)
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
     
-        mean = tf.reduce_mean(mean)
-        logvar = tf.reduce_mean(logvar)
-
-    # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    gradients = tape.gradient(loss, model.trainable_variables)
-    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
              'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
              'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
              'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
-             # 'train/ms_ssim': ms_ssim, 
+             'train/ms_ssim': ms_ssim, 
             }
  
 
-@tf.function
 def valid_step(batch):
     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']
 
-    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
+    model.eval()
+    with torch.inference_mode():
+        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
+        mesh_pred = uv2mesh(pred)
 
-    mesh_pred = uv2mesh(pred)
+        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])])
 
-    if config['use_skinning']:
-        shape = batch['pose_shape'][:, 72:]
-        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+            mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+        else:
+            mesh_pred_posed = mesh_pred
 
-        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
-    else:
-        mesh_pred_posed = mesh_pred
+        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+        c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
+        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
+        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+        kl = kl_loss(mean, logvar)
 
-    # is the output already masked?
+        reconstr_loss = l2_loss
+        if config['use_extra_losses']:
+            reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
 
-    e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
-    c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
-    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
-    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
-    n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-    # conc_loss = concept_loss(pred, batch['uv_unposed'])
-    kl = kl_loss(mean, logvar)
+        loss = reconstr_loss + g_balancing_coeff * kl
 
-    reconstr_loss = l2_loss
-    if config['use_extra_losses']:
-        reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
-    if config['use_norm_loss']:
-        reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
+        mean = torch.mean(mean)
+        logvar = torch.mean(logvar)
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
     
-    loss = reconstr_loss + g_balancing_coeff * kl
-
-    mean = tf.reduce_mean(mean)
-    logvar = tf.reduce_mean(logvar)
-
-    # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    # c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
 
     return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric, 
              'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/collision_ratio': c_ratio,
              'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
-              'valid/normal_loss': n_loss,# 'valid/ms_ssim': ms_ssim, 
+              'valid/normal_loss': n_loss, 'valid/ms_ssim': ms_ssim, 
               }
 
 
-@tf.function
-def test_step(batch):
-    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']
-    
-    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
-
-    mesh_pred = uv2mesh(pred)
-
-    if config['use_skinning']:
-        shape = batch['pose_shape'][:, 72:]
-        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
-
-        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
-        mesh_orig_shaped = mesh_orig + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-        mesh_orig_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_orig_shaped)
-    else:
-        mesh_pred_posed = mesh_pred
-        mesh_orig_posed = mesh_orig
-    
-    mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
-
-
-    # is the output already masked?
-
-    e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
-    c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
-    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
-    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
-    # conc_loss = concept_loss(pred, batch['uv_unposed'])
-    kl = kl_loss(mean, logvar)
-
-    reconstr_loss = l2_loss
-    if config['use_extra_losses']:
-        reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
-            
-    loss = reconstr_loss + g_balancing_coeff * kl
-
-    mean = tf.reduce_mean(mean)
-    logvar = tf.reduce_mean(logvar)
-
-    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, uv_orig + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    euclid_posed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    euclid_posed_fixed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed_fixed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-
-    return { 'test/loss': loss, 'test/edge_loss': e_loss, 'test/collision_loss': c_loss, 'test/euclid_error': metric, 'test/ms_ssim': ms_ssim,
-             'test/reconstr_loss': reconstr_loss, 'test/collision_ratio': c_ratio, 'test/l2_mesh_loss': l2_mesh_loss, 
-             'test/mean': mean, 'test/logvar': logvar, 'test/kl_loss': kl, 'test/var': np.e ** logvar,
-             'test/euclid_posed': euclid_posed, 'test/euclid_posed_fixed': euclid_posed_fixed }
-
-
 def plot_everything(batch: dict, type: str):
     if config['use_skinning']:
-        mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
+        mesh_orig, uv_orig = batch['mesh_unposed'].cpu().numpy(), batch['uv_unposed']
     else:
-        mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
+        mesh_orig, uv_orig = batch['mesh_posed'].cpu().numpy(), batch['uv_posed']
 
-    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
+    model.eval()
+    with torch.inference_mode():
+        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
     
-    mesh_pred = uv2mesh(pred).numpy()
+    mesh_pred = uv2mesh(pred).detach().cpu().numpy()
     
-    orig_img = uv_orig
-    uv_static_img = batch['uv_static']
-    uv_body_posed_img = batch['uv_body_posed']
-    pred_img = pred
-    img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
-    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+    orig_img = torch.permute(uv_orig, [0,2,3,1])
+    uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
+    uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
+    pred_img = torch.permute(pred, [0,2,3,1])
+    img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
     
-    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
-    mask = batch['mesh_mask'].numpy().astype(bool)
+    F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
+    mask = batch['mesh_mask'].cpu().numpy().astype(bool)
 
     if config['use_skinning']:
         shape = batch['pose_shape'][:, 72:]
-        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+        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])]).cpu().numpy()
 
         garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
-        garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
-        gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(pred.shape[0])]
+        garment_unposed = [garment_unposed[i] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
+        gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(pred.shape[0])]
     
         for i in range(mesh_orig.shape[0]):
-            mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
-        mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+            mesh_pred[i, mask[i], :] += shape_offsets[i, mask[i], :]
+        mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_pred, device=config['device']))[0].cpu().numpy()
     else:
-        mesh_posed = mesh_pred
-        gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
-        garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
+        mesh_posed = mesh_pred.cpu().numpy()
+        gt = [mesh_orig[i][mask[i], :] for i in range(pred.shape[0])]
+        garment_unposed, faces = maskFacesBatch(mesh_pred.cpu().numpy(), F, mask)
 
-    gt_posed = [batch['mesh_posed'].numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+    gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])]
 
+    # # plot garment_unposed, gt and mesh_posed next to eachother in plotly subplots
+    # from plotly.subplots import make_subplots
+    # fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'surface'}, {'type': 'surface'}, {'type': 'surface'}, {'type': 'surface'}]])
+    # fig.add_trace(go.Mesh3d(x=gt[0][:,0], y=gt[0][:,2], z=gt[0][:,1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='red', opacity=0.5), row=1, col=1)
+    # fig.add_trace(go.Mesh3d(x=garment_unposed[0][:,0], y=garment_unposed[0][:,2], z=garment_unposed[0][:,1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='blue', opacity=0.5), row=1, col=2)
+    # # add new subplot for gt_posed
+    # fig.add_trace(go.Mesh3d(x=gt_posed[0][:,0], y=gt_posed[0][:,2], z=gt_posed[0][:,1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='red', opacity=0.5), row=1, col=3)
+    # fig.add_trace(go.Mesh3d(x=mesh_posed[0][mask[0],0], y=mesh_posed[0][mask[0],2], z=mesh_posed[0][mask[0],1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='green', opacity=0.5), row=1, col=4)
+    # # make each subplot between -1 and 1
+    # fig.update_layout(scene_aspectmode='cube', scene2_aspectmode='cube', scene3_aspectmode='cube', scene4_aspectmode='cube')
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])), scene2=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])), scene3=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])), scene4=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
 
-    body = batch['mesh_body_posed'].numpy()
-    mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
+    # fig.show()
+
+    mesh_posed_fixed = fix_collision(torch.tensor(mesh_posed, device=config['device']), batch['mesh_body_posed'], torch.tensor(mask, device=config['device'])).cpu().numpy()
+    body = batch['mesh_body_posed'].cpu().numpy()
 
     garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
     collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
@@ -289,8 +250,8 @@ 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], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
     figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
 
-    mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0))
-    logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0))
+    mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0))
+    logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0))
 
     return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
@@ -414,6 +375,7 @@ def sample(folder):
         print()
 
 
+
 def test():
     str_split = config['name'].split('_')
     test_idx = str_split.index('test')
@@ -462,108 +424,73 @@ def test():
 
 
 
-if __name__ == '__main__':
-    # init_wandb()
-    parser = argparse.ArgumentParser(description='')
-    parser.add_argument('--name', type=str)
-    args = vars(parser.parse_args())
+def train():
+    global g_balancing_coeff
+    global step
+    valid_iterator = iter(valid_dataloader)
 
-    smpl = SMPL(orig_smpl=config['is_cape'])
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            if not args['debug']:
+                wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
+            g_balancing_coeff += 1 / 10000 
+
+        if step % 10000 == 0:
+            ckpt_saver(model, step, optimizer)
+
+        step += 1
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+
+    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
-    factory = DatasetFactory()
-
-    dataset_train = factory.get_dataset(folder=config['train_dir'],
-                                        batch_size=config['batch_size'], 
-                                        signature=config['signature_dynamic'],
-                                        resolutions=[config['resolution']], 
-                                        in_memory=config['dataset_in_memory'],
-                                        use_one_hot=config['use_one_hot'],
-                                        use_no_orient=config['remove_global_orientation'],
-                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None, 
-                                        is_cape=config['is_cape'])
-    static_dataset_train = StaticDataset(folder=config['train_dir'],
-                                         signature=config['signature_static'], 
-                                         resolutions=[config['resolution']], 
-                                         in_memory=True, 
-                                        is_cape=config['is_cape'])
-
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'],
-                                        batch_size=config['batch_size'], 
-                                        signature=config['signature_dynamic'], 
-                                        resolutions=[config['resolution']], 
-                                        in_memory=config['dataset_in_memory'], 
-                                        use_one_hot=config['use_one_hot'],
-                                        use_no_orient=config['remove_global_orientation'], 
-                                        frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None, 
-                                        is_cape=config['is_cape'])
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'],
-                                         signature=config['signature_static'], 
-                                         resolutions=[config['resolution']], 
-                                         in_memory=True, 
-                                        is_cape=config['is_cape'])
-
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
+    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
     train_accumulator = Accumulator()
     valid_accumulator = Accumulator()
 
     model = CVAE(config)
 
-    optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
-    balancing_activated_step = 0
+
+    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']
 
-    # sample('checkpoints/cvae_13')
-    # test()
+    ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
 
-    ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
-
-    if config['pretrained_model'] != '':
-        load_model(model, config['pretrained_model'])
-
-    # load_model(model, 'checkpoints/' + config['name'])
- 
-    valid_iterator = iter(dataset_valid)
     step = 0
-    for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+    if config['pretrained_model'] != '':
+        step = load_model(model, config['pretrained_model'])
 
-        train_loss = train_step(batch)
-        train_accumulator.update(train_loss)
-
-        data_valid = valid_iterator.get_next()
-        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
-        valid_loss = valid_step(data_valid)
-        valid_accumulator.update(valid_loss)
-
-
-        if train_accumulator.get()['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
-            if balancing_activated_step == 0:
-                balancing_activated_step = step
-            g_balancing_coeff = g_balancing_coeff + 1 / 10000     
-
-
-        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
-            to_log = { **train_accumulator.get(), **valid_accumulator.get(), **{ 'current_step': step} }
-
-            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
-                train_figs = plot_everything(batch, 'train')
-                valid_figs = plot_everything(data_valid, 'valid')
-                to_log = { **to_log, **train_figs, **valid_figs }
-
-            if step == 0:
-                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
-
-            to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
-        
-            wandb.log(to_log, step = step)
-            train_accumulator.reset()
-            valid_accumulator.reset()
-            
-        if step % 5000 == 0:
-            manager.save(checkpoint_number=step)
-
-        print(step)
-
-        step += 1
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
 
+    train()

scripts/cvae_pyramid.py

@@ -1,65 +1,68 @@
 import os
 import sys
-from configs.config_abl_pyr_32_skirt import config
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+if args['debug']:
+    from configs.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']
-import tensorflow as tf
-
-gpus = tf.config.experimental.list_physical_devices('GPU')
-if gpus:
-    try:
-        # Currently, memory growth needs to be the same across GPUs
-        for gpu in gpus:
-            tf.config.experimental.set_memory_growth(gpu, True)
-        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
-        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
-    except RuntimeError as e:
-        # Memory growth must be set before GPUs have been initialized
-        print(e)
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 import time
-import tqdm
 import argparse
 import wandb
+import tqdm
 import random
 import numpy as np
-import scipy as sp
-from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, get_summary_str, plot_3d_collision, group_inputs_for_levels
 import plotly.graph_objects as go
 import plotly.express as ex
-from plotly.subplots import make_subplots
 from PIL import Image
 
 if config['is_skirt']:
-    from losses import edge_loss_skirt as edge_loss
+    from losses.losses import edge_loss_skirt as edge_loss
 else:
-    from losses import edge_loss
+    from losses.losses import edge_loss
 
-from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio, normal_loss
-# from losses import collision_penalty as collision_penalty
-from losses import collision_loss as collision_penalty
-from custom_utils.custom_utils import batch_to_dict, load_model
-from dataset import DatasetFactory, StaticDataset
-from custom_utils.smpl import SMPL
-from models.pyramid_model_optimized import Pyramid
-from postprocessing import fix_collision
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from data.dataset import Dataset
+from utils.checkpoints import CheckpointSaver, load_model
+from utils.utils import InfiniteDataLoader as DataLoader
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss, kl_loss
+from pytorch_msssim import MS_SSIM
+from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
+from smpl.smpl_torch import SMPL
+from plotly.subplots import make_subplots
+from models.pyramid_model import Pyramid
 
 
 def init_wandb():
-    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', 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('cvae_pyramid.py')
-    artifact.add_file('models/cvae_model.py')
-    artifact.add_file('models/pyramid_model_optimized.py')
-    artifact.add_file('models/pyramid_level.py')
-    artifact.add_file('models/upscaler.py')
-    artifact.add_file('losses.py')
+    artifact.add_file('scripts/cvae_pyramid.py')
+    artifact.add_file('losses/losses.py')
     artifact.add_file('models/base.py')
     artifact.add_file('models/cond_encoder.py')
-    artifact.add_file('dataset.py')
-    artifact.add_file('convnext/models/convnext_tf.py')
-    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs')]
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('models/pyramid_model.py')
+    artifact.add_file('models/pyramid_level.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
     for config_file_name in config_file_names:
         file_name = '/'.join(config_file_name.split('.')) + '.py'
         if os.path.exists(file_name):
@@ -210,86 +213,12 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
     return fig
 
 
-# @tf.function
 def train_step(batch):
+    model.train()
+    # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
     levels_input, conds = group_inputs_for_levels(batch, config)
 
-    with tf.GradientTape() as tape:
-        out_levels = model([levels_input, conds], training=True)
-
-        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']
-
-        output = dict()
-        reconstr_loss_total = 0
-        kl_total = 0
-
-        #for i, res in enumerate(config['pyramid']['resolutions']):
-        i = -1
-        res = config['pyramid']['resolutions'][-1]
-        res_modifier = '_' + str(res) if res != 512 else ''
-        pred, mean, logvar = out_levels[i]
-
-        mesh_pred = uv2mesh(pred)
-
-        if config['use_skinning']:
-            shape = batch['pose_shape'][:, 72:]
-            gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-            shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
-
-            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-            mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
-        else:
-            mesh_pred_posed = mesh_pred
-
-        # is the output already masked?
-
-        e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-        c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
-        c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
-        l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
-        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-        # ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
-        l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
-        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'])
-        kl = kl_loss(mean, logvar)
-
-        if config['use_extra_losses']:
-            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'] 
-        else:
-            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
-
-        # if config['use_norm_loss'] and step > 2000:
-        #     reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
-        
-        kl_total += kl
-
-        mean = tf.reduce_mean(mean)
-        logvar = tf.reduce_mean(logvar)
-
-        metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-
-        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,
-
-        loss = reconstr_loss_total + balancing_coeff * kl_total
-
-    gradients = tape.gradient(loss, model.trainable_variables)
-    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-    return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output }
-
-
-@tf.function
-def valid_step(batch):
-    levels_input, conds = group_inputs_for_levels(batch, config)
-    out_levels = model([levels_input, conds], training=False)
-
+    out_levels = model([levels_input, conds])
 
     if config['use_skinning']:
         mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
@@ -310,23 +239,24 @@ def valid_step(batch):
 
     if config['use_skinning']:
         shape = batch['pose_shape'][:, 72:]
-        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+        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'][:,:,tf.newaxis]
-        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+        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
-    # is the output already masked?
 
-    e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
-    c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
-    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
-    # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-    # ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
-    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
-    n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'])
+    # is the output already masked?
+    # ms_ssim = torch.mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+    # ms_ssim = torch.mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+
+    e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+    c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+    l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2]))
+    l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['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_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
     kl = kl_loss(mean, logvar)
 
     if config['use_extra_losses']:
@@ -336,34 +266,120 @@ def valid_step(batch):
     else:
         reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
 
-    # if config['use_norm_loss'] and step > 2000:
-    #     reconstr_loss_total = reconstr_loss_total + n_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']
     
     kl_total += kl
+    loss = reconstr_loss_total + g_balancing_coeff * kl_total
 
-    mean = tf.reduce_mean(mean)
-    logvar = tf.reduce_mean(logvar)
+    # scaler.scale(loss).backward()
+    # scaler.step(optimizer)
+    # scaler.update()
+    # optimizer.zero_grad()
 
-    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    model.eval()
+    with torch.inference_mode():
+        mean = torch.mean(mean)
+        logvar = torch.mean(logvar)
+
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+
+    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, } } # 
+
+
+    return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output }
+
+
+def valid_step(batch):
+    model.eval()
+    with torch.inference_mode():
+        levels_input, conds = group_inputs_for_levels(batch, config)
+
+        out_levels = model([levels_input, conds])
+
+        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']
+
+        output = dict()
+        reconstr_loss_total = 0
+        kl_total = 0
+
+        #for i, res in enumerate(config['pyramid']['resolutions']):
+        i = -1
+        res = config['pyramid']['resolutions'][-1]
+        res_modifier = '_' + str(res) if res != 512 else ''
+        pred, mean, logvar = out_levels[i]
+
+        mesh_pred = uv2mesh(pred)
+
+        if config['use_skinning']:
+            shape = batch['pose_shape'][:, 72:]
+            gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+            shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+            mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+        else:
+            mesh_pred_posed = mesh_pred
+
+        # is the output already masked?
+        # ms_ssim = torch.mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+        # ms_ssim = torch.mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+
+        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+        c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2]))
+        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['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_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+        kl = kl_loss(mean, logvar)
+
+        if config['use_extra_losses']:
+            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'] 
+        else:
+            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
+
+        if config['use_norm_loss'] and step > 2000:
+            reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
+        
+        kl_total += kl
+        loss = reconstr_loss_total + g_balancing_coeff * kl_total
+
+        mean = torch.mean(mean)
+        logvar = torch.mean(logvar)
+
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['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, } } # 
 
-    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,
 
-    loss = reconstr_loss_total + balancing_coeff * kl_total
 
     return { **{ 'valid/loss': loss, 'valid/reconstr_loss_total': reconstr_loss_total, 'valid/kl_total': kl_total }, **output }
 
 
-@tf.function
 def test_step(batch):
     levels_input, conds = group_inputs_for_levels(batch, config)
     out_levels = model([levels_input, conds], training=False)
 
     shape = batch['pose_shape'][:, 72:]
     gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, dim=0) for i in range(shape.shape[0])])
     mesh_orig = batch['mesh_unposed']
 
     output = dict()
@@ -377,18 +393,18 @@ def test_step(batch):
     pred, mean, logvar = out_levels[i]
 
     mesh_pred = uv2mesh(pred)
-    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
     mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
 
     # is the output already masked?
 
-    e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
     c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
-    c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
-    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
-    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
-    # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-    # ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+    c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+    l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), dim=-1), dim=[1, 2]))
+    l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, dim=-1), dim=-1))
+    # ms_ssim = torch.mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+    # ms_ssim = torch.mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
     kl = kl_loss(mean, logvar)
 
     reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
@@ -397,11 +413,11 @@ def test_step(batch):
     # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
     kl_total += kl
 
-    mean = tf.reduce_mean(mean)
-    logvar = tf.reduce_mean(logvar)
+    mean = torch.mean(mean)
+    logvar = torch.mean(logvar)
 
-    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    euclid_error_per_garment = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1)), axis=1) / tf.reduce_sum(batch['mesh_mask'], axis=1)
+    metric = torch.sum(tf.math.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask'])
+    euclid_error_per_garment = torch.sum(tf.math.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, dim=-1)), dim=1) / torch.sum(batch['mesh_mask'], dim=1)
 
     output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,'test/euclid_error' + res_modifier: metric, 
                                 'test/mean' + res_modifier: mean, 'test/logvar' + res_modifier: logvar, 'test/kl_loss' + res_modifier: kl,
@@ -416,35 +432,42 @@ def test_step(batch):
 
 def plot_everything(batch: dict, type: str):
     if config['use_skinning']:
-        mesh_orig, uv_orig = batch['mesh_unposed'].numpy(), batch['uv_unposed'].numpy()
+        mesh_orig, uv_orig = batch['mesh_unposed'].cpu().numpy(), batch['uv_unposed']
     else:
-        mesh_orig, uv_orig = batch['mesh_posed'].numpy(), batch['uv_posed'].numpy()
+        mesh_orig, uv_orig = batch['mesh_posed'].cpu().numpy(), batch['uv_posed']
     
-    levels_input = group_inputs_for_levels(batch, config)
-    output_levels = model(levels_input, training=False)
+    model.eval()
+    with torch.inference_mode():
+        levels_input = group_inputs_for_levels(batch, config)
+        output_levels = model(levels_input)
 
-    mesh_pred_levels = [uv2mesh(output_levels[i][0]).numpy() for i in range(len(output_levels))]
-    mesh_orig_posed = batch['mesh_posed'].numpy()
-    body = batch['mesh_body_posed'].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()
     
-    orig_img = (uv_orig + 1) / 2
-    uv_static_img = (batch['uv_static'].numpy() + 1) / 2
-    uv_body_posed_img = (batch['uv_body_posed'].numpy() + 1) / 2
-    uv_curvature = (np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 6)
-    pred_img = np.concatenate([tf.image.resize(output_levels[i][0], (512, 256), tf.image.ResizeMethod.LANCZOS5) for i in range(len(output_levels))], axis=2)
-    pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
-    img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
-    imgs = [wandb.Image(img[i]) for i in range(4)]
+    # 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()
 
-    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
-    mask = batch['mesh_mask'].numpy().astype(bool)
-    _, faces = maskFacesBatch(np.ones_like(mesh_orig), F, mask)
+    orig_img = (torch.permute(uv_orig, [0,2,3,1]) + 1) / 2
+    uv_static_img = (torch.permute(batch['uv_static'], [0,2,3,1]) + 1) / 2
+    uv_body_posed_img = (torch.permute(batch['uv_body_posed'], [0,2,3,1]) + 1) / 2
+    uv_curvature = (torch.tile(torch.argmax(torch.permute(batch['uv_curvature'], [0,2,3,1]), dim=-1)[:,:,:,None], [1,1,1,3]) / 6)
+    pred_img = torch.concatenate([torch.permute(F.interpolate(output_levels[i][0], size=(512, 256), mode='bilinear'), [0,2,3,1]) for i in range(len(output_levels))], dim=3)
+    # pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
+    pred_img = (pred_img - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None]) / (pred_img.flatten().reshape(pred_img.shape[0], -1).max(1)[0][:,None,None,None] - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None])
+
+    img = torch.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
+    imgs = [wandb.Image(img[i].cpu().numpy()) for i in range(4)]
+
+    _F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
+    mask = batch['mesh_mask'].cpu().numpy().astype(bool)
+    _, faces = maskFacesBatch(np.ones_like(mesh_orig), _F, mask)
     
     if config['use_skinning']:
         shape = batch['pose_shape'][:, 72:]
-        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()
-
+        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])]).cpu().numpy()
 
         gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
 
@@ -458,8 +481,15 @@ def plot_everything(batch: dict, type: str):
         mesh_unposed_shaped_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32)
         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()
     
-        mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_unposed_shaped_full).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])]
@@ -467,7 +497,7 @@ def plot_everything(batch: dict, type: str):
 
 
     gt_posed = [mesh_orig_posed[i, mask[i], :] for i in range(mask.shape[0])]
-    mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
+    mesh_posed_fixed = fix_collision(torch.Tensor(mesh_posed).to('cuda'), torch.Tensor(body).to('cuda'), torch.Tensor(mask).to('cuda')).cpu().numpy()
 
     mesh_posed = [mesh_posed[j, mask[j], :] for j in range(mask.shape[0])]
     
@@ -477,7 +507,7 @@ def plot_everything(batch: dict, type: str):
         collision_color = [np.ones_like(mesh_posed[i], dtype=np.int32)*230 for i in range(len(mesh_posed))]
     else:
         collision = [np.concatenate((body[i], mesh_posed[i]), axis=0) for i in range(len(mesh_posed))]
-        collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(mesh_posed))]
+        collision_faces = [np.concatenate((_F, faces[i] + config['smpl_res']), axis=0) for i in range(len(mesh_posed))]
         collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(mesh_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(mesh_posed))]
 
     figs = { type + "/garment_" + str(i): plot_3d(gt[i], mesh_unposed_shaped_levels[i], faces[i], gt_posed[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
@@ -487,7 +517,7 @@ def plot_everything(batch: dict, type: str):
     else:
         garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
         collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(mesh_posed))]
-        gt_posed = [batch['mesh_posed'].numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+        gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])]
         collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))]
         figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
 
@@ -505,7 +535,7 @@ def sample(folder):
 
         pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[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])
+        pose, shape = tf.tile(pose[None,:], [nr_of_samples,1]), tf.tile(shape[None,:], [nr_of_samples,1])
 
         mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
         mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
@@ -706,13 +736,13 @@ def generate_results():
 
         shape = batch['pose_shape'][:, 72:]
         gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
 
         i = -1
         pred = out_levels[i][0]
 
         mesh_pred = uv2mesh(pred)
-        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
         mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
 
         for i in range(batch['id'].shape[0]):
@@ -725,157 +755,117 @@ def generate_results():
 
     pass
 
-@tf.function
 def speed_step(levels_input, conds):
 
     return model.sample([levels_input, conds], training=False)
 
 
 
+def train():
+    global g_balancing_coeff
+    global step
+    valid_iterator = iter(valid_dataloader)
+
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            if not args['debug']:
+                wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
+            g_balancing_coeff += 1 / 10000 
+
+        if step % 10000 == 0:
+            ckpt_saver(model, step, optimizer, None)
+
+        step += 1
+
+
 if __name__ == '__main__':
-    # init_wandb()
-    parser = argparse.ArgumentParser(description='')
-    parser.add_argument('--name', type=str, default='lhunor')
-    args = vars(parser.parse_args())
+    if not args['debug']:
+        init_wandb()
 
-    # config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
-    # assert len(config_file_names) == 1
-    # config_file_name = config_file_names[0]
-    # # assert config_file_name.split('.')[1].endswith(args.name)
-    # assert config_file_name.split('.')[1].endswith(config['name'])
-
-    smpl = SMPL(orig_smpl=config['is_cape'])
+    use_amp = True
+    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
     dataset_resolutions = config['pyramid']['resolutions'].copy()
     if 512 not in dataset_resolutions:
         dataset_resolutions.append(512)
 
-    factory = DatasetFactory()
-    dataset_train = factory.get_dataset(folder=config['train_dir'], 
-                                        batch_size=config['batch_size'], 
-                                        # batch_size=1, 
-                                        signature=config['signature_dynamic'], 
-                                        resolutions=dataset_resolutions, 
-                                        in_memory=config['dataset_in_memory'], 
-                                        is_cape=config['is_cape'],
-                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None)
-    static_dataset_train = StaticDataset(folder=config['train_dir'], 
-                                        signature=config['signature_static'], 
-                                        resolutions=dataset_resolutions, 
-                                        in_memory=True,
-                                        is_cape=config['is_cape'])
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], 
-                                        batch_size=config['batch_size'], 
-                                        signature=config['signature_dynamic'], 
-                                        resolutions=dataset_resolutions, 
-                                        in_memory=config['dataset_in_memory'], 
-                                        is_cape=config['is_cape'],
-                                        frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None)
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'], 
-                                        signature=config['signature_static'], 
-                                        resolutions=dataset_resolutions, 
-                                        in_memory=True, 
-                                        is_cape=config['is_cape'])
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, 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'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
+    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
 
     model: Pyramid = Pyramid(config)
 
-    optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
-    balancing_coeff = config['starting_balancing_coeff']
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
+    ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + config['name'], top_n=5)
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+
+    g_balancing_coeff = config['starting_balancing_coeff']
     reg_coeff = 1
-    valid_iterator = iter(dataset_valid)
-    gamma = 1
-    balancing_activated_step = 0
-
-    # sample('checkpoints/cvae_13')
-    # test()
-
-    # generate_results()
-
-    ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
-
-    # load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000')
 
+    step = 0
     if config['pretrained_model'] != '':
         # assert len(config['pyramid']['resolutions']) == 1
-        load_model(model, config['pretrained_model'], iteration=config['pretrained_model_iteration'])
+        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 len(config['pyramid']['resolutions']) == 1:
-        model.static_enc.trainable = False
-        model.body_enc.trainable = False
-        model.curv_enc.trainable = False
+        for param in model.static_enc.parameters(): param.requires_grad = False
+        for param in model.body_enc.parameters(): param.requires_grad = False
+        for param in model.curv_enc.parameters(): param.requires_grad = False
     else:
         load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else ''))
-        model.static_enc.trainable = False
-        model.body_enc.trainable = False
-        model.curv_enc.trainable = False
-        model.cond_encoder.trainable = False
+        for param in model.static_enc.parameters(): param.requires_grad = False
+        for param in model.body_enc.parameters(): param.requires_grad = False
+        for param in model.curv_enc.parameters(): param.requires_grad = False
 
     # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
     # model.levels[0].upscale_enhancer.trainable = False
 
+    # for i in range(len(model.levels)-1):
+    #     load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer'])
+    #     model.levels[i].upscale_enhancer.trainable = False
+    # if config['pyramid']['resolutions'][-1] != 512:
+    #     load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
+    #     model.upscale_enhancer.trainable = False
 
     for i in range(len(model.levels)-1):
-        load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer'])
-        model.levels[i].upscale_enhancer.trainable = False
-    if config['pyramid']['resolutions'][-1] != 512:
-        load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
-        model.upscale_enhancer.trainable = False
-
-    for i in range(len(model.levels)-1):
-        model.levels[i].trainable = False
+        for param in model.levels[i].parameters(): param.requires_grad = False
 
     # load_model(model.cond_encoder, config['pyramid']['pretrained_pyramid'])
     # load_model(model.levels[0], config['pyramid']['pretrained_pyramid'])
     # load_model(model.levels[-1], 'checkpoints/cvae_19')
 
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
 
-    step = 0
-    times = []
-    for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
-
-        # levels_input, conds = group_inputs_for_levels(batch, config)
-
-        # start = time.perf_counter()
-        # out = speed_step(levels_input, conds)
-        # t = time.perf_counter() - start
-        # print(t)
-        # times.append(t)
-
-        # if step > 200:
-        #     break
-
-
-        train_loss = train_step(batch)
-
-        data_valid = valid_iterator.get_next()
-        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
-        valid_loss = valid_step(data_valid)
-
-        to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
-        if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
-            train_figs = plot_everything(batch, 'train')
-            valid_figs = plot_everything(data_valid, 'valid')
-            to_log = { **to_log, **train_figs, **valid_figs }
-
-        if step % 5000 == 0:
-            manager.save(checkpoint_number=step)
-        if step == 0:
-            to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
-
-        if train_loss['train/euclid_error'] < 0.004:
-            if balancing_activated_step == 0:
-                balancing_activated_step = step
-            balancing_coeff = balancing_coeff + 1 / 20000
-
-
-        to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
-        
-        wandb.log(to_log, step = step)
-
-        print(step)
-
-        step += 1
-    # print('mean:')
-    # print(np.mean(times[1:]))
+    train()

scripts/direct_train.py

@@ -0,0 +1,281 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+if args['debug']:
+    from configs.direct_1 import config
+    # from configs.ae_static_orig_2 import config
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import argparse
+import wandb
+import random
+import numpy as np
+import torch
+# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, plot_3d_collision
+import plotly.graph_objects as go
+import plotly.express as ex
+import torch.optim as optim
+
+if config['is_skirt']:
+    from losses.losses import edge_loss_skirt as edge_loss
+else:
+    from losses.losses import edge_loss
+
+from data.dataset import Dataset
+# from torch.utils.data import DataLoader
+from losses.losses import collision_loss, normal_loss
+from utils.postprocess import fix_collision
+from custom_utils.custom_utils import plot_3d
+from utils.checkpoints import CheckpointSaver, load_model
+from custom_utils.custom_utils import Accumulator
+from utils.utils import InfiniteDataLoader as DataLoader
+from models.cond_encoder import AutoEncoder as Model
+from smpl.smpl_torch import SMPL
+
+from pytorch_msssim import MS_SSIM
+
+
+def init_wandb():
+    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/direct_train.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+
+
+def train_step(batch):
+    model.train()
+    # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
+    if config['use_curvature']:
+        input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
+    else:
+        input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
+    pred = model([input, batch['uv_mask']])
+
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_unposed']
+    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])])
+    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]
+ 
+    e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+    c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+    # l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
+    n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+    loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+    # scaler.scale(loss).backward()
+    # scaler.step(optimizer)
+    # scaler.update()
+    # optimizer.zero_grad()
+
+    model.eval()
+    with torch.inference_mode():
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+    
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/euclid_error': metric, 'train/collision_ratio': c_ratio,
+             'train/ms_ssim': ms_ssim, 'train/collision_loss': c_loss, 'train/l2_mesh_loss': l2_mesh_loss,
+             'train/normal_loss': n_loss }
+        
+
+def valid_step(batch):
+    model.eval()
+
+    with torch.inference_mode():
+        if config['use_curvature']:
+            input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
+        else:
+            input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
+        pred = model([input, batch['uv_mask']])
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_unposed']
+        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])])
+        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]
+    
+        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+        c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+        # l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
+        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
+        loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
+
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+    
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/euclid_error': metric, 'valid/collision_ratio': c_ratio,
+             'valid/ms_ssim': ms_ssim, 'valid/collision_loss': c_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
+             'valid/normal_loss': n_loss }
+       
+
+def plot_everything(batch: dict, type: str):
+    model.eval()
+    with torch.inference_mode():
+        if config['use_curvature']:
+            input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
+        else:
+            input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
+        pred = model([input, batch['uv_mask']]).detach()
+
+    mesh_pred = uv2mesh(pred).cpu().numpy()
+    mesh_orig = batch['mesh_unposed'].cpu().numpy()
+    
+    orig_img = torch.permute(batch['uv_unposed'], [0,2,3,1])
+    uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
+    uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
+    pred_img = torch.permute(pred, [0,2,3,1])
+    img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
+    
+    F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
+    mask = batch['mesh_mask'].cpu().numpy().astype(bool)
+
+    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])]).cpu().numpy()
+
+    garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
+    garment_unposed = [garment_unposed[i] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
+    gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(pred.shape[0])]
+
+    for i in range(mesh_orig.shape[0]):
+        mesh_pred[i, mask[i], :] += shape_offsets[i, mask[i], :]
+    mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_pred, device=config['device']))[0].cpu().numpy()
+
+    gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+
+    mesh_posed_fixed = fix_collision(torch.tensor(mesh_posed, device=config['device']), batch['mesh_body_posed'], torch.tensor(mask, device=config['device'])).cpu().numpy()
+    body = batch['mesh_body_posed'].cpu().numpy()
+
+    garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
+    collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
+
+    return { **figs, **figs_coll, **{ type + '/images': imgs, } }
+
+
+def train():
+    valid_iterator = iter(valid_dataloader)
+
+    step = 0
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            if not args['debug']:
+                wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if step % 10000 == 0:
+            ckpt_saver_enc(model.encoder, step, optimizer, scaler)
+            ckpt_saver_dec(model.decoder, step, optimizer, scaler)
+
+        step += 1
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+    
+    use_amp = True
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
+    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model : Model = Model(config)
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
+
+    # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
+
+    if config['pretrained_model'] != '':
+        load_model(model.encoder, config['pretrained_model'] + '_enc')
+        load_model(model.decoder, config['pretrained_model'] + '_dec')
+
+    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}')
+
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+
+    train()
+
+
+

scripts/upscaler_train.py

@@ -0,0 +1,293 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+
+
+import time
+import argparse
+import wandb
+import random
+import numpy as np
+import torch
+# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str
+import plotly.graph_objects as go
+import plotly.express as ex
+import torch.optim as optim
+import torch.nn.functional as F
+import torchvision
+
+from data.dataset import Dataset
+# from torch.utils.data import DataLoader
+from custom_utils.custom_utils import plot_3d
+from models.upscaler import Upscaler
+from utils.checkpoints import CheckpointSaver, load_model
+from custom_utils.custom_utils import Accumulator
+from utils.utils import InfiniteDataLoader as DataLoader
+
+from pytorch_msssim import MS_SSIM
+
+
+low_res = 64
+is_skirt = False
+is_cape = True
+is_male = True
+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')
+smpl_res = 3638 if is_cape else 14475
+use_cache = False
+name=f'iccv_{low_res}' + skirt_modifier + cape_modifier + gender_modifier
+
+if not is_skirt:
+    train_dir = '/data-local/127/train'
+    valid_dir = '/data-local/127/valid'
+else:
+    train_dir = '/data-local/128/train_skirt'
+    valid_dir = '/data-local/127/valid_skirt'
+
+if is_cape:
+    if is_male:
+        train_dir = '/data-local/127/cape_processed/train/male'
+        valid_dir = '/data-local/127/cape_processed/test/male'
+    else:
+        train_dir = '/data-local/127/cape_processed/train/female'
+        valid_dir = '/data-local/127/cape_processed/test/female'   
+
+
+
+def init_wandb():
+    run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccv'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/upscaler_train.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256, is_skirt=is_skirt, is_smpl_orig=is_cape)
+
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces):
+    mesh_pred_unposed[:,2] += 1
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_gt[:,0],
+            y=mesh_gt[:,2],
+            z=mesh_gt[:,1],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2]
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_pred_unposed[:,0],
+            y=mesh_pred_unposed[:,2],
+            z=mesh_pred_unposed[:,1],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+        )
+    ])
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=2, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,3],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    return fig
+
+
+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', 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))
+    # 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)
+
+    # 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():
+        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 { 'train/loss': loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } # 'train/l1_loss': l1_loss,
+        
+
+def valid_step(batch):
+    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', 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))
+        # 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/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', align_corners=True, antialias=True)
+        pred = resizer(pred)
+    
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_unposed']
+
+    _F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if is_cape else "") + ".npy")
+    mask = batch['mesh_mask'].cpu().numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred.cpu().numpy(), _F, mask)
+    # gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    gt = [mesh_orig[i].cpu().numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i]) for i in range(4) }
+    # figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    orig_img = batch['uv_unposed']
+    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} }
+
+
+def train():
+    global step
+    valid_iterator = iter(valid_dataloader)
+
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 1000 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            if not args['debug']:
+                wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if step % 10000 == 0 and not args['debug']:
+            ckpt_saver(model, step, optimizer, None)
+
+        step += 1
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+    
+    use_amp = True
+    uv_mask_body = torch.ones((low_res, low_res//2), device='cuda')
+    low_resolution_modifier = '_' + str(low_res)
+    
+    signature = { 'id', 'frame_nr',
+                          'pose_shape',
+                          'uv_unposed' + low_resolution_modifier,
+                          'uv_unposed',
+                          'mesh_unposed',
+                          'mesh_posed',
+                         'mesh_mask',
+                         'uv_mask',
+                         'uv_mask' + low_resolution_modifier,
+                        }
+    
+    # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
+    dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape)
+    dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape)
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True)
+    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    if low_res == 32 or low_res == 64:
+        kernel_size = 3
+    elif low_res == 128 or low_res == 256:
+        kernel_size = 5
+    else:
+        kernel_size = 7
+        
+    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)
+    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.8, 0.2])
+    _MSSSIM_small = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[1.0])
+
+    step = 0
+    ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + name, top_n=5)
+
+    print(f'started training: {name}')
+
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+
+    train()
+
+
+

smpl/smpl_torch.py

@@ -8,7 +8,6 @@ import numpy as np
 import scipy.io as sio
 from time import time
 from torch.nn import Module
-from utils import plot_mesh
 
 class SMPL(Module):
   def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False):
@@ -53,7 +52,7 @@ class SMPL(Module):
     self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
 
     self.weights_all = self.weights
-    # self.weights = self.weights[:,self.to_keep,:]
+    self.weights = self.weights[:,self.to_keep,:]
     
     self.shape = (1,27578,3)
     self.device = device
@@ -135,7 +134,7 @@ class SMPL(Module):
     return ret
 
 
-  def forward(self, betas, pose, remove=True):
+  def forward(self, pose, betas, verts: torch.Tensor = None, remove: bool = True):
     
     """
           Construct a compute graph that takes in parameters and outputs a tensor as
@@ -164,14 +163,19 @@ 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])])
-    R_cube_big = self.rodrigues(pose.view(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
+    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)
 
-    R_cube = R_cube_big[:, 1:, :, :]
-    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])])
+    if verts is not None:
+      v_posed = verts
+    else:
+      R_cube = R_cube_big[:, 1:, :, :]
+      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, :, :, 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(
@@ -202,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.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_all[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
     )
@@ -216,7 +220,14 @@ class SMPL(Module):
     result = result - root_joint
     joints = joints - root_joint
 
-    if remove:
+    # # plut result as a plotly 3d mesh
+    # import plotly.graph_objects as go
+    # fig = go.Figure(data=[go.Mesh3d(x=result[0,:,0].cpu().numpy(), y=result[0,:,2].cpu().numpy(), z=result[0,:,1].cpu().numpy(), color='lightpink', opacity=0.50, i=self.f[:,0], j=self.f[:,1], k=self.f[:,2], name='body')])
+    # # plot the joints as spheres
+    # fig.add_trace(go.Scatter3d(x=joints[0,:,0].cpu().numpy(), y=joints[0,:,2].cpu().numpy(), z=joints[0,:,1].cpu().numpy(), mode='markers', marker=dict(size=5, color='red'), name='joints'))
+    # fig.show()
+
+    if remove and verts is None:
       result = result[:,self.to_keep,:]
 
     return result, joints

utils/checkpoints.py

@@ -16,9 +16,9 @@ class CheckpointSaver:
         self.top_n = top_n
         self.top_model_paths = []
         
-    def __call__(self, model, iteration, optimizer):
+    def __call__(self, model, iteration, optimizer, scaler):
         model_path = os.path.join(self.dirpath, f'iteration_{iteration}.pt')
-        torch.save({ 'model': model.state_dict(), 'optimizer': optimizer }, model_path)
+        torch.save({ 'model': model.state_dict(), 'optimizer': optimizer, 'scaler': scaler }, model_path)
         self.top_model_paths.append({'path': model_path})
         if len(self.top_model_paths)>self.top_n: 
             self.cleanup()
@@ -31,10 +31,13 @@ class CheckpointSaver:
         self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:]
 
 
-def load_model(model, checkpoint, optimizer=None, iteration=''):
+def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=None, scaler: torch.nn.Module=None, iteration=''):
     # list of checkpoint files in the folder checkpoints
     checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
     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 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
 

utils/postprocess.py

@@ -1,55 +1,53 @@
 import os
 
-# os.environ["CUDA_VISIBLE_DEVICES"] = "3"
-import tensorflow as tf
 
 import sys
 import shutil
 import tqdm
 import random
 import math
+import torch
 import numpy as np
 import scipy as sp
 import scipy.io as sio
 import plotly.graph_objects as go
 import plotly.express as ex
+from pytorch3d.structures import Meshes
 
 from custom_utils.custom_utils import _maskFaces
-
-from scipy.spatial import cKDTree
-from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
+from losses.losses import _nearest_neighbour
 
 
 # from hbertiche/PBNS
 """ COLLISION FUNCTIONS START """
-F = np.load("custom_utils/smpl_faces.npy")
+# F = np.load("smpl/smpl_faces.npy")
+to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
+F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
+smpl_to_keep_orig_post = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
+smpl_to_keep_orig_post = smpl_to_keep_orig_post[smpl_to_keep_orig_post < 6890]
+F_fix_coll_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')[:,smpl_to_keep_orig_post,:]
 
 
-def _nearest_neighbour(V, B):
-	tree = cKDTree(B.numpy())
-	return tree.query(V.numpy(), workers=-1)[1]
-
-@tf.function
-def _tf_nn_parallel(V, B):
-	return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0])
-
-
-def fix_collision(V, B, mask, thr=.004):
-    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+def fix_collision(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr=.004):
+    if V.shape[1] == 14475:
+        F_fix_coll_local = F_fix_coll
+    elif V.shape[1] == 3638:
+        F_fix_coll_local = F_fix_coll_orig
+    else:
+        raise ValueError("V must have 14475 or 3638 vertices")
+    mesh = Meshes(B, F_fix_coll_local)
+    B_vn = torch.stack(mesh.verts_normals_list())
    
     # compute correspondences
-    idx = _tf_nn_parallel(V, B)
-    N, M = V.shape[:2]
-    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
-    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
-    idx = tf.stack((_idx, idx), axis=-1)
+    idx = _nearest_neighbour(V, B)
     # compute loss
-    D = V - tf.gather_nd(B, idx)
-    b_vn = tf.gather_nd(B_vn, idx)
-    dot = tf.einsum('abc,abc->ab', D, b_vn)
-    loss = tf.minimum(dot - thr, 0)  * mask
-    B_offset_closest = tf.gather_nd(B, idx) + (b_vn / 2e2)
-    output = tf.where(loss[:,:,tf.newaxis] < 0, B_offset_closest, V)
+    B_idx = torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
+    D = V - B_idx
+    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)
+    loss = torch.minimum(dot - thr, torch.zeros_like(dot))  * mask
+    B_offset_closest = B_idx + (b_vn / 2e2)
+    output = torch.where(loss[:,:,None] < 0, B_offset_closest, V)
 
     return output