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

@@ -0,0 +1,72 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '1'
+config['name'] = 'static_torch_1'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'static'] 
+
+config['pretrained_model'] = 'checkpoints/iccv/static_torch_1'
+
+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'] = True
+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'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0005
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

configs/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

@@ -0,0 +1,72 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '2'
+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'
+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'] = True
+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'] = 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'], 27),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 27),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0005
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

configs/ae_static_orig_2.py

@@ -0,0 +1,73 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '3'
+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'
+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'] = True
+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'] = 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/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

@@ -1,12 +1,16 @@
 import json
 import numpy as np
-import tensorflow as tf
+# import tensorflow as tf
+import torch
+import torch.nn as nn
 import scipy as sp
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
 
-from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
-from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
+# from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
+# from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
 
 
 def m2uv_array(m2uv, is_smpl_orig=False):
@@ -23,8 +27,8 @@ def m2uv_array(m2uv, is_smpl_orig=False):
     return result
 
 
-class UV2Mesh(tf.keras.Model):
-    def __init__(self, res=256, is_skirt=False, is_smpl_orig=False, is_adaptive=False):
+class UV2Mesh(nn.Module):
+    def __init__(self, res=256, is_skirt=False, is_smpl_orig=False):
         super(UV2Mesh, self).__init__()
         self.res = res
         uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy')
@@ -33,48 +37,41 @@ class UV2Mesh(tf.keras.Model):
         m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig)
         self.uv_pixels = self.uv_map_discreet[m2uv]
         self.smpl_res = 14475 if not is_smpl_orig else 3638
-        self.is_adaptive = is_adaptive
-
-        if self.is_adaptive:
-            laplacian = sp.sparse.load_npz("custom_utils/laplacian/laplacian0" + ("_skirt" if is_skirt else "") + ".npz")
-            self.gc = GraphConvolution((3, 3), L=laplacian, batchSize=4) 
 
 
-    @tf.function
     def uv_to_mesh(self, uv):
-        if uv.shape[1] == 3:
-            uv = tf.transpose(uv, [0,2,3,1])
+        if uv.shape[1] == 3 or uv.shape[1] == 7:
+            uv = torch.permute(uv, [0,2,3,1])
         if uv.shape[1] != 512:
-            uv = tf.image.resize(uv, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+            raise NotImplementedError
+            uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR)
             
-        batch_size = uv.shape[0]
+        # batch_size = uv.shape[0]
 
-        x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
-        y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
-        z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
-        x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
-        y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
-        z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
+        # x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
+        # y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
+        # z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
+        # x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
+        # y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
+        # z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
 
 
-        x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
-        y1 = tf.gather_nd(uv, y1_ind, batch_dims=1)
-        z1 = tf.gather_nd(uv, z1_ind, batch_dims=1)
-        x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
-        y2 = tf.gather_nd(uv, y2_ind, batch_dims=1)
-        z2 = tf.gather_nd(uv, z2_ind, batch_dims=1)
+        # x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
+        # y1 = tf.gather_nd(uv, y1_ind, batch_dims=1)
+        # z1 = tf.gather_nd(uv, z1_ind, batch_dims=1)
+        # x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
+        # y2 = tf.gather_nd(uv, y2_ind, batch_dims=1)
+        # z2 = tf.gather_nd(uv, z2_ind, batch_dims=1)
 
-        p1 = tf.stack((x1, y1, z1), axis=-1)
-        p2 = tf.stack((x2, y2, z2), axis=-1)
+        # p1 = tf.stack((x1, y1, z1), axis=-1)
+        # p2 = tf.stack((x2, y2, z2), axis=-1)
 
-        mesh = (p1 + p2) / 2
+        mesh = torch.mean(uv[:, self.uv_pixels[:,:,0], self.uv_pixels[:,:,1]].float(), dim=2)
         return mesh
 
 
     def __call__(self, uv):
         x = self.uv_to_mesh(uv)
-        if self.is_adaptive:
-            x = self.gc(x)
         return x
 
 
@@ -107,7 +104,7 @@ class UV2MeshLabel():
         return self.uv_to_mesh(uv)
 
 
-def _maskFaces(V, F, mask):
+def _maskFaces(V, F, mask, return_unfixed=False):
     V_mask = mask
     F_mask = np.take(V_mask, F.reshape(-1), axis=0)
     F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
@@ -117,7 +114,10 @@ def _maskFaces(V, F, mask):
     F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
     V_new = V[V_mask, :]
 
-    return V_new, F_new
+    if return_unfixed:
+        return V_new, F_new, F[F_mask, :]
+    else:
+        return V_new, F_new
 
 
 # def _maskFaces_tf(V, F, mask):
@@ -292,28 +292,18 @@ def sparse_to_tensor(L):
     return tf.SparseTensor(idx, data, shape)
 
 
-def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration=''):
-    ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
-    if not disable_assert:
-        assert manager.latest_checkpoint is not None, checkpoint
-        assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0', checkpoint
-    if assert_consumed_on:
-        ckpt.restore(manager.latest_checkpoint).assert_consumed()
-    else:
-        if iteration != '':
-            name_base = manager.latest_checkpoint.split('-')[0]
-            ckpt_ind = [i for i in range(len(manager.checkpoints)) if manager.checkpoints[i].endswith(name_base + '-' + iteration)][0]
-            ckpt.restore(manager.checkpoints[ckpt_ind]).assert_existing_objects_matched()#.assert_consumed()
-        else:
-            ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
-
 
 def get_summary_str(model):
-    lines = []
-    model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True)
+    # from torchview import draw_graph
+    # draw_graph(model, input_size=((4,3,512,256), (4,512,256)), expand_nested=True)
+
+    # lines = []
+    # model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True)
     # Add initial spaces to avoid markdown formatting in TensorBoard
-    return lines
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params = sum(param.numel() for param in model.parameters())
+
+    return [f'all_params: {total_params}, trainable: {trainable_params}']
 
 def get_summary_html(model):
     lines = get_summary_str(model)
@@ -338,7 +328,7 @@ def merge_verts(Vs, batchDim):
     return V
 
 
-class GraphConvolution(tf.keras.layers.Layer):
+class GraphConvolution(nn.Module):
     def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
         super(GraphConvolution, self).__init__()
         self.shape = shape
@@ -403,3 +393,45 @@ def group_inputs_for_levels(batch, config):
         levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
     conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
     return levels_input, conds
+
+
+class Accumulator():
+    def __init__(self):
+        self.dict = dict()
+
+    def update(self, dict):
+        for key in dict.keys():
+            if key in self.dict.keys():
+                self.dict[key].append(dict[key])
+            else:
+                self.dict[key] = [dict[key]]
+
+    def get(self):
+        result = dict()
+        for key in self.dict.keys():
+            result[key] = np.mean(self.dict[key])
+        return result
+
+    def reset(self):
+        self.dict = dict()
+
+
+from stl import mesh as mesh_stl
+
+def save_stl(file, mesh, faces):
+    obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype))
+    for i, f in enumerate(faces):
+        for j in range(3):
+            obj.vectors[i][j] = mesh[f[j],:]    
+
+    obj.save(file)
+
+
+def write_obj(file, vertices, faces, colors):
+    with open(file, 'w') as file:
+        for v, c in zip(vertices, colors):
+            line = 'v ' + ' '.join([str(_) for _ in v]) + ' ' + ' '.join([str(_) for _ in c]) + '\n'
+            file.write(line)
+        for f in faces:
+            line = 'f ' + ' '.join([str(_) for _ in f]) + '\n'
+            file.write(line)

data/cloth3d_io.py

@@ -260,9 +260,12 @@ def read_posed_garment(folder, id, garment, map=None, mask=None, shape=None, sha
     return verts
 
 
-def read_garment_mask(folder, id, garment):
+def read_garment_mask(folder, id, garment, return_raw=False):
     raw_mask = sio.loadmat(os.path.join(folder, id, garment + '_mask.mat'))['mask'].reshape((-1))
-    mask = raw_mask > 0
+    if not return_raw:
+        mask = raw_mask > 0
+    else:
+        mask = raw_mask
     return mask
 
 

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,79 +1,104 @@
 import os
 import sys
+
+# os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+
+import math
 import time
-import cv2
+import torch
 import random
-import pickle
+import pytorch3d
+import torchvision as tv
 import numpy as np
-import tensorflow as tf
 import scipy.io as sio
+import plotly.express as ex
+import torchvision
+
+from torch.utils.data import Dataset
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
-from custom_utils.custom_utils import nest_conditions, batch_to_dict
-from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame
+from data.utils import Frame
 
 
-class StaticDataset():
-    def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False, use_no_orient=True) -> None:
+def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
+    if frames_file is None:
+        uv_static_folder = os.path.join(folder, 'pose_shape') 
+        frames = sorted(os.listdir(uv_static_folder))
+
+        if is_cape:
+            frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
+            frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames]
+        else:
+            frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+            frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
+
+        # if testing:
+        #     frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
+        #     frames = [frame.split('.')[0] + '_4' for frame in frames]
+
+
+    if frames_file is not None:
+        frames = []
+        with open(os.path.join(folder, frames_file), 'r') as f:
+            lines = f.readlines()
+            lines = [line.strip() for line in lines]
+        frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines]
+
+    if shuffle:
+        random.shuffle(frames)
+
+    return frames
+
+
+class StaticCache():
+    def __init__(self, folder, signature, resolutions=[512], is_cape=False, device='cuda') -> None:
         self.store = dict()
         self.signature = signature
         self.folder = folder
-        self.in_memory = in_memory
         assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
-        self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256',
-                              'uv_mask', 'uv_mask_32', 'uv_mask_64', 'uv_mask_128', 'uv_mask_256',
-                              'mesh_static', 'mesh_mask', 'uv_shaped_body']
-        self.keys = [value for value in signature.keys() if value in self.possible_keys]
-        if 'uv_static_128' in signature.keys() and 'uv_static_128' not in self.keys:
-            self.keys.append('uv_static_128')
         self.resolutions = resolutions
         self.filled = ''
         self.is_cape = is_cape
-        self.use_no_orient = use_no_orient
+        self.device = device
         
         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 = tf.image.resize(uv_mask[:,:,tf.newaxis], size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
-                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, garment):
+    def read(self, id):
         data = dict()
 
-        info = loadInfo(os.path.join(self.folder, id))
-
-        data['shape'] = info['shape']
-        data['poses'] = info['poses']
-
-        data['map'] = sio.loadmat(os.path.join(self.folder, id, garment + '_map.mat'))['map'].reshape((-1))
-
-
-        if 'mesh_mask' in self.signature.keys():
-            if self.is_cape:
-                data['mesh_mask'] = np.ones((3638), np.float32)
-            else:
-                data['mesh_mask'] = read_garment_mask(self.folder, id, garment).astype(np.float32)
-
         for res in self.resolutions:
             res_string = ''
             if res != 512:
                 res_string = '_' + str(res)
 
-            if 'uv_static' + res_string in self.signature.keys():
+            if 'uv_static' + res_string in self.signature:
                 uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id +  '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
-            if 'uv_mask' + res_string in self.signature.keys():
+            if 'uv_mask' + res_string in self.signature:
                 if self.is_cape:
                     data['uv_mask' + res_string] = self.uv_masks[res]
                 else:
                     uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
 
-        if 'mesh_static' in self.signature.keys():
-            data['mesh_static'] = read_garment_template(self.folder, id, garment)
+        if 'mesh_static' in self.signature:
+            mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
+        if 'mesh_mask' in self.signature:
+            if self.is_cape:
+                data['mesh_mask'] = np.ones((3638), np.float32)
+            else:
+                mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
+        # if 'uv_shaped_body' in self.signature:
+        #     uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
 
         return data
 
@@ -81,471 +106,195 @@ class StaticDataset():
     def __getitem__(self, id) -> dict:
         if id not in self.store:
             item = self.read(id)
-            if not self.in_memory:
-                return item
-            else:
-                self.store[id] = item
+            stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
+            for key, value in stored_data.items():
+                if key.startswith('uv_') and value.shape[-1] == 3:
+                    stored_data[key] = torch.permute(value, [2,0,1])
+            self.store[id] = stored_data
         
-        return self.store[id]
+        data = { key: value.to(self.device) for key, value in self.store[id].items()}
+        return data
 
 
-    def expand_batch(self, batch):
-        for key in self.keys:
-                batch[key] = tf.zeros(((0,) + self.signature[key]))
-        for i in range(batch['id'].shape[0]):
-            item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8'), batch['garment'][i].numpy().decode('UTF-8'))
-            for key in self.keys:
-                if key == 'pose':
-                    frame_nr = int(batch['frame_nr'][i].numpy().decode('UTF-8'))
-                    pose = item[key][frame_nr, :]
-                    if self.use_no_orient:
-                        pose[:3] = 0
-                    batch[key] = tf.concat((batch[key], tf.expand_dims(pose, axis=0)), axis=0)
-                else:
-                    batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0)
-        return batch
+class Dataset(Dataset):
+    def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
+        self.folder = folder
+        self.device = device
+        self.is_cape = is_cape
+        self.signature = signature
+        self.resolutions = resolutions
+
+        self.transforms = tv.transforms.ToTensor()
+        self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
+
+        self.body_faces = np.load('smpl/smpl_faces.npy')
+        self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device)
+
+        self.cache = dict() if use_cache else None
 
 
-
-class DatasetIdx(tf.data.Dataset):
-
-    def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
-        # if frames_file is None:
-        #     uv_static_folder = os.path.join(folder, 'pose_shape') 
-        #     frames = sorted(os.listdir(uv_static_folder))
-
-        #     if is_cape:
-        #         frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
-        #     else:
-        #         frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
-
-        #     if testing:
-        #         frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
-        #         frames = [frame.split('.')[0] + '_4' for frame in frames]
-
-        frames = []
-        with open(os.path.join(folder, frames_file), 'r') as f:
-            lines = f.readlines()
-            lines = [line.strip() for line in lines]
-        frames = lines
-
-        if shuffle:
-            random.shuffle(frames)
-
-        return tf.data.Dataset.from_tensor_slices(frames)
+    def __len__(self):
+        return len(self.frames)
 
 
-class DatasetFactory():
-
-    def __init__(self):
-        super(DatasetFactory, self).__init__()
-
-    def _build_tf_signature(self, signature):
-        tf_signature = []
-        for key in signature.keys():
-            if key == 'id':
-                tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
-            elif key == 'frame_nr':
-                tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
-            elif key == 'garment':
-                tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
-            elif key == 'mesh_curvature':
-                tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
-            else:
-                tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
-        return tf_signature
-
-    
-    def get_dataset(self,
-                    folder, 
-                    batch_size, 
-                    signature, 
-                    shuffle=True, 
-                    resolutions=[512], 
-                    in_memory=False, 
-                    do_repeat=True, 
-                    is_cape=False,
-                    use_no_orient=True,
-                    testing=False,
-                    frames_file=None):
-        assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
-
-        dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, frames_file=frames_file)
+    def __getitem__(self, idx):
+        frame : Frame = self.frames[idx]
+        if self.cache is not None:
+            if frame in self.cache.keys():
+                data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()}
+                data = { **data, **self.static_cache[frame.id] }
+                return data
 
         separator = '_'
-        if is_cape:
+        if self.is_cape:
             separator = '.'
 
-        uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
+        uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2'
+        if self.is_cape:
+            uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
 
-        if use_no_orient == False:
-            assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!'
+        item = dict()
 
-        to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
-        shapedirs = []
-        with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
-            female = pickle.load(f, encoding="latin1")
-            shapedirs.append(female['shapedirs'][to_keep, :])
-        with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
-            male = pickle.load(f, encoding="latin1")
-            shapedirs.append(male['shapedirs'][to_keep, :])
+        if 'id' in self.signature:
+            item['id'] = frame.id
+        if 'frame_nr' in self.signature:
+            item['frame_nr'] = frame.frame_nr
+        if 'pose_shape' in self.signature:
+            pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
+            pose_shape[:3] = 0
+            item['pose_shape'] = pose_shape
+        if 'mesh_body_posed' in self.signature:
+            mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
+            item['mesh_body_posed'] = mesh_posed_body_no_orient
+        if 'mesh_posed' in self.signature:
+            mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
+            item['mesh_posed'] = mesh_posed
+        if 'mesh_unposed' in self.signature:
+            mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
+            item['mesh_unposed'] = mesh_unposed_no_shape
+        if 'mesh_curvature' in self.signature:
+            mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
+            item['mesh_curvature'] = mesh_curvature.astype(np.int32)
+
+        for res in self.resolutions:
+            res_string = ''
+            if res != 512:
+                res_string = '_' + str(res)
+
+            if 'uv_curvature' + res_string in self.signature:
+                uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz')); 
+                if res_string != '':
+                    uv_curvature = uv_curvature['uv_curvature'] 
+                else:
+                    uv_curvature = uv_curvature['curvature_uv']
+                uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device)
+                item['uv_curvature' + res_string] = uv_curvature
+
+            if 'uv_unposed' + res_string in self.signature:
+                uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
+                item['uv_unposed' + res_string] = uv_unposed_no_shape
+
+            if 'uv_body_posed' + res_string in self.signature:
+                uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz'))
+                uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
+                item['uv_body_posed' + res_string] = uv_posed_body_no_orient
+
+            if 'uv_posed' + res_string in self.signature:
+                uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
+                item['uv_posed' + res_string] = uv_posed
 
 
+        data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
 
-        def read_item(frame):
-            frame_split = tf.strings.split(frame, sep=separator)
-            id, frame_nr, garment = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8"), frame_split[2].numpy().decode("utf-8")
-            item = dict()
+        for key, value in data.items():
+            if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7):
+                data[key] = torch.permute(value, [2,0,1])
 
-            # redunda ntly loading info again
-            pose = loadInfo(os.path.join(folder, id))['poses'][:,frame_nr]
-            # !!!!!!!!!!!!!!!!!!!
-            map = None
-
-            if 'id' in signature.keys():
-                item['id'] = id
-            if 'frame_nr' in signature.keys():
-                item['frame_nr'] = frame_nr
-            if 'garment' in signature.keys():
-                item['garment'] = garment
-
-            if 'mesh_body_posed' in signature.keys():
-                if use_no_orient:
-                    item['mesh_body_posed'] = read_posed_body_frame(folder, id, frame_nr, pose, deorient=use_no_orient)
-            if 'mesh_posed' in signature.keys():
-                if use_no_orient:
-                    item['mesh_posed'] = read_posed_garment_frame(folder, id, garment, frame, map, mask=None, shape=None, shapedirs=None, pose=None, unshape=False, deorient=False)
-            # if 'mesh_unposed' in signature.keys():
-            #     mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
-            #     item['mesh_unposed'] = mesh_unposed_no_shape
-
-            if 'mesh_curvature' in signature.keys():
-                mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
-                item['mesh_curvature'] = mesh_curvature.astype(np.int32)
-
-
-            for res in resolutions:
-                res_string = ''
-                if res != 512:
-                    res_string = '_' + str(res)
-
-                if 'uv_curvature' + res_string in signature.keys():
-                    uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz')); 
-                    if res_string != '':
-                        uv_curvature = uv_curvature['uv_curvature'] 
-                    else:
-                        uv_curvature = uv_curvature['curvature_uv']
-                    uv_curvature = tf.one_hot(uv_curvature, 7)
-                    item['uv_curvature' + res_string] = uv_curvature
-
-                if 'uv_unposed' + res_string in signature.keys():
-                    uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
-                    item['uv_unposed' + res_string] = uv_unposed_no_shape
-
-                if 'uv_body_posed' + res_string in signature.keys():
-                    if use_no_orient:
-                        uv_posed_body_no_orient = np.load(os.path.join(folder, uv_posed_body_no_orient_folder_name + res_string, id + separator + frame_nr + '.npz'))
-                        uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
-                        item['uv_body_posed' + res_string] = uv_posed_body_no_orient
-                    else:
-                        uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']
-                        item['uv_body_posed' + res_string] = uv_body_posed
-
-                if 'uv_posed' + res_string in signature.keys():
-                    if use_no_orient:
-                        uv_posed = np.load(os.path.join(folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
-                        item['uv_posed' + res_string] = uv_posed
-                    else:
-                        uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
-                        item['uv_posed' + res_string] = uv_posed
-
-            as_list = []
-            for key in signature.keys():
-                if key in item:
-                    as_list.append(item[key])
-
-            return as_list
-
-
-        if in_memory:
-            dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, 
-                                                inp=[i], 
-                                                Tout=self._build_tf_signature(signature)
-                                                ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True)
-        else:
-            dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, 
-                                                inp=[i], 
-                                                Tout=self._build_tf_signature(signature)
-                                                ), num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size=batch_size, drop_remainder=True)
-
-        if do_repeat:
-            dataset = dataset.repeat()
-        
-        return dataset
+        if self.cache is not None:
+            self.cache[frame] = data
         
+        data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
+        data = { **data, **self.static_cache[frame.id] }
+        return data
 
 
 if __name__ == '__main__':
-    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
-    from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
-    import plotly.graph_objects as go
-    import plotly.express as ex
-    # import cmapy as cm
-    uv2mesh = UV2Mesh(is_skirt=True)
+    from torch.utils.data import DataLoader
     
-    F = np.load("custom_utils/smpl_faces.npy")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
+    signature = { 'id', 
+                  'frame_nr',
+                  'pose_shape',
+                  'uv_posed',
+                  'uv_body_posed',
+                  'mesh_body_posed',
+                  'mesh_posed',
+                  'mesh_unposed',
+                  'uv_unposed',
+                  'uv_curvature',
+                  'mesh_curvature',
+                  'uv_static',
+                  'mesh_static',
+                  'mesh_mask',
+                  'uv_mask',
+                  'uv_shaped_body',
+                  }
     
-    # dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
-    # for d in dataset:
-    #     print(d)
+    dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
+    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
 
-    factory = DatasetFactory()
-    signature_dynamic = { 'id': (), 'frame_nr': (), 'garment': (),
-                        #   'pose_shape': (72 + 11,),
-                        #   'uv_posed': (3, 512, 256,),
-                        #   'uv_body_posed': (512, 256, 3),
-                        #   'uv_body_posed_32': (32, 16, 3),
-                        #   'mesh_body_posed': (14475, 3,),
-                        #   'mesh_posed': (14475, 3,),
-                        #   'mesh_unposed': (14475, 3,),
-                          'uv_unposed': (512, 256, 3),
-                        #   'uv_unposed_32': (3, 32, 16,),
-                          'uv_curvature': (512, 256, 7),
-                        #   'mesh_curvature':  (14475)
-                        }
-    signature_static = { #'uv_static': (3, 512, 256,),
-                        #  'uv_static_32': (3, 32, 16,),
-                        #  'mesh_static': (14475, 3,),
-                        #  'mesh_mask': (14475,),
-                         'uv_mask': (512, 256,),
-                        #  'uv_shaped_body': (3, 512, 256,),
-                          }
+    for batch in train_dataloader:
+        batch
+        pass
 
-    dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
-                                  batch_size=1,
-                                  signature=signature_dynamic,
-                                  resolutions=[512],
-                                  do_repeat=False,
-                                  use_one_hot=False,
-                                  is_cape=False,
-                                  frames_file='frames_dress_30k.txt')
-    static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
-                                   signature=signature_static,
-                                   resolutions=[512],
-                                  is_cape=False)
-
-
-    # for i, batch in enumerate(dataset):
-    #     print(i)
-
-
-    sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
-    sum_mask = np.zeros((512, 256), dtype=np.float32)
-    for i, batch in enumerate(dataset):
-        print(i)
-        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
-        ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
-        sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
-        sum_mask += batch['uv_mask'].numpy().sum(axis=0)
-
-    sum_mask[sum_mask == 0] = 1
-    avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
-
-    # x = np.transpose(avg_garment, [1,2,0])
-    x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0]
-    go.Figure(data=[
-            go.Mesh3d(
-                # 8 vertices of a cube
-                x=x[:,0],
-                y=x[:,1],
-                z=x[:,2],
-                # i, j and k give the vertices of triangles
-                i = F[:,0],
-                j = F[:,1],
-                k = F[:,2]
-            )]).show()
-
-    np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment)
-
-    print()
-
-    # pass
-
-
-    # uv_res = 256
-
-    # def m2uv_array(m2uv):
-    #     result = np.zeros((14475, 2), dtype=np.int32)
-    #     for k, v in m2uv.items():
-    #         if len(v) == 2:
-    #             result[k] = list(v)
-    #         elif len(v) == 1:
-    #             x = v.pop()
-    #             result[k] = [x, x]
-    #         else:
-    #             assert False
-
-    #     return result
-
-    # def unique_uv_map_mask(uv_map_discreet):
-    #     mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
-    #     temp = np.zeros((uv_res * 2, uv_res))
-    #     for p in uv_map_discreet:
-    #         temp[p[0], p[1]] += 1
-    #     # temp = temp == 1
-    #     for i, p in enumerate(uv_map_discreet):
-    #         if temp[p[0], p[1]] == 1:
-    #             mask[i] = True
-    #         elif temp[p[0], p[1]] == 2:
-    #             mask[i] = True
-    #             temp[p[0], p[1]] = 0
-    #     return mask
-
-    # def uv_to_mesh(uv, m2uv):
-    #     res = uv.shape[0]
-    #     mesh = np.zeros((14475, 3))
-    #     uv_pixels = uv_map_discreet[m2uv]
-    #     # counter = np.ones((14475))
-    #     # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
-    #     mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
-    #     return mesh
-
-
-    # uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
-    # m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
-    # m2uv = m2uv_array(m2uv)
-    # uv_mask = np.load('custom_utils/mask_256x512.npy')
-    # curvature_label_type = '6_25'
-
-    # uv_res = 256
-    # uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
-    # uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
-
-    # uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
-    # uv_map_mask = unique_uv_map_mask(uv_map_discreet)
-    # uv_counter = np.zeros((uv_map.shape[0]))
-    # uv_counter[m2uv[:, 0]] = 1
-    # uv_counter[m2uv[:, 1]] += 1
-    # uv_counter = uv_counter[uv_map_mask]
-
-
-
-    # for i, batch in enumerate(dataset):
-    #     print(i)
-    #     batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
-
-    #     uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
-    #     mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
-    #     mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
-
-
-    #     values = mesh_unposed
-    #     uv_values = np.zeros((uv_map.shape[0], 3))
-    #     uv_values[m2uv[:, 0]] = values
-    #     uv_values[m2uv[:, 1]] += values
-    #     uv_counter = np.zeros((uv_map.shape[0]))
-    #     uv_counter[m2uv[:, 0]] = 1
-    #     uv_counter[m2uv[:, 1]] += 1
-    #     uv_values /= uv_counter[:, np.newaxis]
-    #     x = uv_values
-    #     # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
-    #     im = ((x - x.min()) / (x.max() - x.min())) * 255
-
-
-    #     x = np.transpose(uv_unposed, [1,2,0])
-    #     x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
-    #     _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
-    #     go.Figure(data=[
-    #             go.Mesh3d(
-    #                 # 8 vertices of a cube
-    #                 x=x[:,0],
-    #                 y=x[:,1],
-    #                 z=x[:,2],
-    #                 # i, j and k give the vertices of triangles
-    #                 i = faces[:,0],
-    #                 j = faces[:,1],
-    #                 k = faces[:,2]
-    #             )]).show()
-
-
-    #     print()
+    pass 
 
 
 
 
-    # # @tf.function
-    # def euclid_error(mask, gt, pred):
-    #     # return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)), axis=-1) / tf.reduce_sum(mask, axis=-1)
-    #     return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))
-
-    # all_error = []
-    # for batch in dataset:
-    #     batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
-    #     mesh_reconstruct = uv2mesh(batch['uv_posed'])
-    #     err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed']))
-
-    #     i = 296
-    #     err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min())
-    #     colors = np.array([cm.color('jet', 1 - i) for i in err_normed])
-
-    #     _, faces = _maskFaces(batch['mesh_posed'][i].numpy(), F, batch['mesh_mask'][i].numpy().astype(bool))
-    #     x = batch['mesh_posed'][i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
-    #     y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
-    #     colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :]
-    #     fig = go.Figure(data=[
-    #         go.Mesh3d(
-    #             # 8 vertices of a cube
-    #             x=x[:,0],
-    #             y=x[:,1],
-    #             z=x[:,2],
-    #             # i, j and k give the vertices of triangles
-    #             i = faces[:,0],
-    #             j = faces[:,1],
-    #             k = faces[:,2]
-    #         ),
-    #         go.Mesh3d(
-    #             # 8 vertices of a cube
-    #             x=y[:,0],
-    #             y=y[:,1],
-    #             z=y[:,2],
-    #             # i, j and k give the vertices of triangles
-    #             i = faces[:,0],
-    #             j = faces[:,1],
-    #             k = faces[:,2],
-    #         )])
-    #     fig.show()
-
-    #     fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2])
-    #     fig.update_traces(marker={'color': colors, 'size': 5})
-    #     fig.show()
-
-    #     all_error.append(err.numpy().sum())
-    #     # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic)))
-    #     print(len(all_error))
 
 
-    # dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2)
-
-    # from losses import collision_loss as collision_penalty
-    # from losses import collision_penalty as collision_penalty_2
 
 
-    # for d in dataset:
-    #     d
 
-    #     v = uv2mesh(d[0])
-    #     b = d[3]
-    #     m = d[4]
 
-    #     b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0)
+# if __name__ == '__main__':
+#     os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+#     from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
+#     import plotly.graph_objects as go
+#     import plotly.express as ex
+#     # import cmapy as cm
+#     uv2mesh = UV2Mesh(is_skirt=True)
 
-    #     c = collision_penalty(v, b, m)
-    #     c2 = collision_penalty(v, b2, m)
+#     F = np.load("custom_utils/smpl_faces.npy")
 
-    #     c_ = collision_penalty_2(v, b, m)
-    #     c2_ = collision_penalty_2(v, b2, m)
+#     sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
+#     sum_mask = np.zeros((512, 256), dtype=np.float32)
+#     for i, batch in enumerate(dataset):
+#         print(i)
+#         batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+#         ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
+#         sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
+#         sum_mask += batch['uv_mask'].numpy().sum(axis=0)
 
-    #     print()
-    #     pass
+#     sum_mask[sum_mask == 0] = 1
+#     avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
 
+#     # x = np.transpose(avg_garment, [1,2,0])
+#     x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0]
+#     go.Figure(data=[
+#             go.Mesh3d(
+#                 # 8 vertices of a cube
+#                 x=x[:,0],
+#                 y=x[:,1],
+#                 z=x[:,2],
+#                 # i, j and k give the vertices of triangles
+#                 i = F[:,0],
+#                 j = F[:,1],
+#                 k = F[:,2]
+#             )]).show()
+
+#     np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment)
+
+#     print()

data/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

data/investigate.py

@@ -3,6 +3,7 @@ import sys
 import random
 import numpy as np
 import plotly.graph_objects as go
+from pymeshfix import MeshFix
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures, read_curvature_frame, writeOBJ
@@ -51,7 +52,7 @@ def build_lookup_table(repeated):
             lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
     return lookup
 
-def mask_faces(all_faces, mask, repeated, verts):
+def mask_faces(all_faces, mask, raw_mask, repeated, verts):
     # if any of the vertices of a face is masked, mask the face
     lookup = build_lookup_table(repeated)
     faces = []
@@ -79,6 +80,30 @@ def mask_faces(all_faces, mask, repeated, verts):
                                     if mask[v1[i]] and mask[v2[j]] and mask[v3[k]]:
                                         # if the euclidean distance between any two of the vertices of the face is less than 0.05 then append the face
                                         # if np.linalg.norm(verts[v1[i]] - verts[v2[j]]) < 0.05 and np.linalg.norm(verts[v1[i]] - verts[v3[k]]) < 0.05 and np.linalg.norm(verts[v2[j]] - verts[v3[k]]) < 0.05:
+                                        
+                                        #only add the face if this face originally existed among the faces of the garment
+                                        # if not np.any([np.array([v1[i], v2[j], v3[k]], dtype=np.int32) == face for face in all_faces]):
+                                        #     continue
+                                        
+                                        # if all vertices of the face have a value of 2 in the raw mask then don't add the face
+                                        # if np.all(raw_mask[[v1[i], v2[j], v3[k]]] == 2):
+                                        #     continue
+
+                                        # skip face if it's area is more than 0.001 and at least one of the vertices has value of 2 in the raw mask
+                                        # if np.linalg.norm(np.cross(verts[v1[i]] - verts[v2[j]], verts[v1[i]] - verts[v3[k]])) > 0.001 and np.any(raw_mask[[v1[i], v2[j], v3[k]]] == 2):
+                                        #     continue
+                                        
+                                        # if np.linalg.norm(np.cross(verts[v1[i]] - verts[v2[j]], verts[v1[i]] - verts[v3[k]])) > 0.001:
+                                        #     continue
+
+                                        # if any edge of the face is longer than 0.1 and both vertices of the edge have a value of 2 in the raw mask then skip the face
+                                        # if np.linalg.norm(verts[v1[i]] - verts[v2[j]]) > 0.05 and (raw_mask[v1[i]] == 2 or raw_mask[v2[j]] == 2):
+                                        #     continue
+                                        # if np.linalg.norm(verts[v1[i]] - verts[v3[k]]) > 0.05 and (raw_mask[v1[i]] == 2 or raw_mask[v3[k]] == 2):
+                                        #     continue
+                                        # if np.linalg.norm(verts[v2[j]] - verts[v3[k]]) > 0.05 and (raw_mask[v2[j]] == 2 or raw_mask[v3[k]] == 2):
+                                        #     continue
+
                                         faces.append(np.array([v1[i], v2[j], v3[k]], dtype=np.int32))
                                         assert np.all(verts[face] == verts[faces[-1]])
                                         found_replacement = True
@@ -100,22 +125,38 @@ def check(folder, frame):
     id, frame_nr, garment = frame.split()
     frame_nr = int(frame_nr)
     info = loadInfo(os.path.join(folder, id, 'info.mat'))
-    mask = read_garment_mask(folder, id, garment)
+    raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
+    mask = raw_mask > 0
     mesh_posed = read_posed_garment_frame(folder, id, garment, frame_nr, mask=mask, shape=info['shape'], pose=info['poses'][:,frame_nr], deorient=use_no_orient, unshape=no_shape)
 
+
     v = mesh_posed[mask,:]
     is_skirt = (garment == 'Dress' or garment == 'Skirt')
     skirt_mod = '_skirt' if is_skirt else ''
     F = np.load(f'custom_utils/smpl_faces{skirt_mod}.npy').astype(np.int32)
-    _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+    _, faces, unfixed = _maskFaces(np.zeros((14475,3)), F, mask, True)
 
+    meshfix = MeshFix(v, faces)
+    holes = meshfix.extract_holes()
     # plot(v, faces)
 
     unq, idx, count = np.unique(mesh_posed, axis=0, return_counts=True, return_index=True)
     repeated_groups = unq[count > 1]
 
+    new_border = []
+    for border in np.array(holes.points):
+        # find the index of the vertex in the original mesh
+        ind = np.argwhere(np.all(mesh_posed == border, axis=1))
+        new_border.extend(ind.ravel())
+
+    raw_mask[new_border] = 2    
     mask_2 = np.zeros((14475), dtype=bool)
+    mask_2_raw = np.zeros((14475), dtype=np.uint8)
     mask_2[idx] = True
+    mask_2_raw = mask_2.astype(np.uint8)
+    mask_3 = np.zeros((14475), dtype=bool)
+    mask_3[mask_2] = raw_mask[mask_2] == 2
+    mask_2_raw[mask_3] = 2
     v_2 = mesh_posed[mask_2,:]
     _, faces_2 = _maskFaces(np.zeros((14475,3)), F, mask_2)
     # plot(v_2, faces_2)
@@ -127,11 +168,46 @@ def check(folder, frame):
         repeated_idx = np.argwhere(np.all(mesh_posed == repeated_group, axis=1))
         repeated.append(repeated_idx.ravel())
 
-    faces_3 = mask_faces(F, mask_2, repeated, mesh_posed)
-    writeOBJ('temp/investigate_3.obj', v_2, faces_3)
+    # # if a vertex from v_2 is not part of v then set the corresponding value in mask_0 to 0
+
+    # mask_0 = np.zeros((v_2.shape[0]), dtype=bool)
+    # for i, _v in enumerate(v_2):
+    #     if not np.any([np.all(_v == _w) for _w in v]):
+    #         mask_0[i] = True
+    # mask_0_0 = np.zeros((14475), dtype=bool)
+    # mask_0_0[mask_2] = 1 - mask_0
+    # mask_4 = np.logical_not(mask_0)
+
+    faces_3 = mask_faces(F, mask_2, mask_2_raw, repeated, mesh_posed)
+    # writeOBJ('temp/investigate_3.obj', v_2, faces_3)
     plot(v_2, faces_3)
 
 
+    # plot repeated
+    r = []
+    for rr in repeated:
+        r.extend(rr)
+    x = mesh_posed
+    y = mesh_posed[r,:]
+    go.Figure(data = [go.Scatter3d(x=x[:, 0], y=x[:, 1], z=x[:, 2], marker=dict(size=3, line=dict(width=0),), mode='markers'),
+                      go.Scatter3d(x=y[:, 0], y=y[:, 1], z=y[:, 2], marker=dict(size=5, line=dict(width=0),), mode='markers'),
+                     ]).show()
+
+
+    # unq, idx, count = np.unique(v, axis=0, return_counts=True, return_index=True)
+    # repeated_groups = unq[count > 1]
+    # repeated = []
+    # for repeated_group in repeated_groups:
+    #     repeated_idx = np.argwhere(np.all(v == repeated_group, axis=1))
+    #     repeated.append(repeated_idx.ravel())
+
+    # mask_3 = np.zeros((v.shape[0]), dtype=bool)
+    # mask_3[idx] = True
+
+    # faces_4 = mask_faces(faces.astype(np.int64), mask_3, mask_2_raw, repeated, v)
+    # plot(v[mask_3,:], faces_4)
+
+
     pass
     
 

data/utils.py

@@ -1,16 +1,16 @@
 import tables as tb
 
 class Frame:
-    def __init__(self, id, frame_nr, garment):
+    def __init__(self, id, frame_nr, garment=''):
         self.id = id
         self.frame_nr = frame_nr
         self.garment = garment
 
     def __hash__(self):
-        return hash((self.id, self.frame_nr))
+        return hash((self.id, self.frame_nr, self.garment))
 
     def __eq__(self, other):
-        return (self.id, self.frame_nr) == (other.id, other.frame_nr)
+        return (self.id, self.frame_nr, self.garment) == (other.id, other.frame_nr, self.garment)
 
 
 def save_to_h5file(h5file, id, frame_nr, garment, data, name, type):

docker/Dockerfile

@@ -1,7 +1,5 @@
-FROM tensorflow/tensorflow:2.9.0-gpu
-
-RUN apt-get update
-RUN apt-get upgrade -y
+# Use an existing image as the base
+FROM nvcr.io/nvidia/pytorch:22.06-py3
 
 # Add the current user to the image
 ARG USER_NAME=lhunor
@@ -10,7 +8,7 @@ ARG GROUP_NAME=guest
 ARG GROUP_ID=11000
 
 # Update the system and install required packages
-RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git
+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
@@ -28,10 +26,12 @@ RUN usermod -aG sudo $USER_NAME
 # Set the default user to the new user
 USER $USER_NAME
 
-WORKDIR //home-net
+WORKDIR /workspace
 
-RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip htop nano -y
-RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python cmapy
+ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
+RUN pip install --upgrade pip
+COPY requirements.txt /tmp/
+RUN pip install --requirement /tmp/requirements.txt
 
 # defining the volumes that will be mounted
 VOLUME [ "/data-net" ]

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/Dockerfile_old

@@ -0,0 +1,42 @@
+FROM tensorflow/tensorflow:2.9.0-gpu
+
+RUN apt-get update
+RUN apt-get upgrade -y
+
+# 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
+
+# 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 //home-net
+
+RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip htop nano -y
+RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python cmapy
+
+# 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/requierements.txt

@@ -1,12 +0,0 @@
-numpy
-matplotlib
-opencv-python
-wandb
-OpenEXR
-tensorflow-graphics
-scipy
-plotly
-pandas
-tensorboardX
-opencv-python
-cmapy

docker/requirements.txt

@@ -0,0 +1,15 @@
+torch==2.0.0    
+numpy
+matplotlib
+opencv-python
+wandb
+scipy
+plotly
+pandas
+cmapy
+pytorch3d
+tables
+numpy-stl
+torchmetrics
+pytorch_msssim
+importlib

losses/losses.py

@@ -0,0 +1,215 @@
+import torch
+
+import numpy as np
+# import tensorflow_graphics as tfg
+import scipy as sp
+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
+
+
+
+# from hbertiche/PBNS
+""" COLLISION FUNCTIONS START """
+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
+
+
+# @tf.function
+def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004):
+    if B.shape[1] == 14475:
+        F_norm_local = F_norm
+    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 = _nearest_neighbour(V, B)
+    # compute loss
+    D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
+    b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
+    dot = torch.einsum('abc,abc->ab', D, b_vn)
+    thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))  * mask
+    loss = torch.sum(thresholded ** 2, axis=-1)
+    ratio = torch.sum((thresholded < 0).to(torch.float32)) / torch.sum(mask)
+
+    return loss, ratio
+
+
+
+""" COLLISION FUNCTIONS END """
+
+
+laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
+edges: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda')
+
+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(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):
+    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_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
+    d_e = d_e * e_mask
+    loss = torch.sum(d_e ** 2, dim=-1)# / tf.reduce_sum(e_mask, axis=-1)
+    return loss
+
+def get_edge_mask_skirt(masks: torch.Tensor):
+    if edges.get_device() != masks.get_device():
+        edges.to(masks.get_device())
+    return torch.min(masks[:,edges_skirt], dim=-1)[0]
+
+def edge_loss_skirt(x, y, mask):
+    e_mask = get_edge_mask_skirt(mask)
+    x_e = x[:,edges_skirt[:,0],:] - x[:,edges_skirt[:,1],:]
+    x_e = torch.sum(x_e ** 2, -1)
+    x_e = torch.sqrt(x_e + 1e-9)
+    y_e = y[:,edges_skirt[:,0],:] - y[:,edges_skirt[:,1],:]
+    y_e = torch.sum(y_e ** 2, -1)
+    y_e = torch.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    d_e = d_e * e_mask
+    loss = torch.sum(d_e ** 2, dim=-1)# / tf.reduce_sum(e_mask, axis=-1)
+    return loss
+
+
+
+whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy')
+
+def uv_mask_loss(pred, uv_mask):
+    inv_mask = (1 - uv_mask)  * whole_garment_uv_mask
+    pred = tf.reduce_sum(pred, axis=1)
+    invalid = pred * inv_mask
+    invalid = invalid > 0.01
+    invalid_count = tf.math.count_nonzero(invalid, dtype=tf.float32)
+    total_count = tf.math.count_nonzero(inv_mask, dtype=tf.float32)
+    return invalid_count / total_count
+
+
+
+
+def merge_verts(Vs, batchDim):
+    V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
+    return V
+    
+
+def _bend_loss_single(x, F, L, mask):
+    verts = tf.boolean_mask(x, mask, axis=0)
+    VN = vertex_normals(verts, F)
+    bend = edge_convolution_template(
+        VN, L, sizes=None,
+        edge_function=lambda x,y: (x - y)**2,
+        reduction='weighted',
+        edge_function_kwargs={}
+    )
+    loss = tf.reduce_sum(bend)
+    return loss
+
+def bend_loss(x, F, L, mask):
+    # x: predicted outfits
+    # F: faces
+    # L: laplacian
+    loss = tf.stack([_bend_loss_single(x[i], F[i], L[i], mask[i]) for i in range(x.shape[0])])
+    return loss
+
+
+# def bend_loss(x, F, L, mask):
+#     # x: predicted outfits
+#     # F: faces
+#     # L: laplacian
+#     # x = x * mask[:,:,tf.newaxis]
+#     batch_size = x.shape[0]
+#     verts, faces = maskFacesBatch(x, F, mask)
+#     nr_valid_verts = [tf.reduce_sum(mask[i]) for i in range(batch_size)]
+#     x = tf.concat(verts, axis=0)
+#     F = tf.concat([faces[i] + tf.reduce_sum(nr_valid_verts[:i]) for i in range(1, batch_size)])
+#     F = tf.concat((faces[0], F), axis=0)
+#     VN = vertex_normals(x, F)
+#     bend = edge_convolution_template(
+#         VN, L, sizes=None,
+#         edge_function=lambda x,y: (x - y)**2,
+#         reduction='weighted',
+#         edge_function_kwargs={}
+#     )
+#     loss = tf.reduce_sum(bend, axis=[1])
+#     loss = tf.reshape(loss, (batch_size, 14475))
+#     loss = tf.reduce_mean(loss, axis=1)
+#     return loss
+
+
+def log_normal_pdf(sample, mean, logvar, raxis=1):
+    log2pi = tf.math.log(2. * np.pi)
+    res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
+    return res
+
+
+def kl_loss(mean, logvar):
+    kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar)
+    kl_loss = torch.mean(kl_loss)
+    kl_loss *= -0.5
+    return kl_loss
+
+
+def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
+    f_norm = None
+    if not is_skirt and not is_smpl_orig:
+        f_norm = F_norm
+    elif is_skirt:
+        f_norm = F_norm_skirt
+    elif is_smpl_orig:
+        f_norm = F_norm_orig
+    else:
+        raise NotImplementedError
+
+    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,140 +1,29 @@
-import tensorflow as tf
-from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
-
-class Bottleneck(tf.keras.Model):
-    
-    def __init__(self, dim, latent_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)]
-            )
-            
-        self.flatten = tf.keras.layers.Flatten()
-        self.dense = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training):
-        x = inputs
-
-        x = self.bottleneck(x, training=training)
-        x = self.flatten(x)
-        x = self.dense(x, training=training)
-
-        return x
+import torch
+from torch import nn
+from models.convnext_orig import ConvNeXt, CNBlock, CNBlockConfig
+from typing import Any, Callable, List, Optional, Sequence
 
 
-class VariationalEncoder(tf.keras.Model):
+class VariationalEncoder(nn.Module):
     
-    def __init__(self, latent_dim, depths, dims, kernel_size=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(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
 
-        self.flatten = tf.keras.layers.Flatten()
-        self.mean = tf.keras.layers.Dense(latent_dim)
-        self.log_var = tf.keras.layers.Dense(latent_dim)
+        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)
 
-    def call(self, inputs, training):
+        self.flatten_mean = nn.Flatten()
+        self.flatten_var = nn.Flatten()
+
+
+    def forward(self, inputs):
         x = inputs
 
-        x = self.encoder(x, training=training)
-        x = self.flatten(x)
-        mean = self.mean(x)
-        log_var = self.log_var(x)
-
-        return mean, log_var
-
-
-class AEEncoder(tf.keras.Model):
-    
-    def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7):
-        super(AEEncoder, self).__init__()
-
-        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)
-
-
-    def call(self, inputs, training):
-        if len(inputs) == 2:
-            x, _ = inputs
-        else:
-            x = inputs
-
-        x = self.encoder(x, training=training)
-        x = self.bottleneck(x, training=training)
-
-        if len(inputs) == 2:
-            return x, inputs[1]
-        else:
-            return x
-
-
-class Decoder(tf.keras.Model):
-    
-    def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
-        super(Decoder, self).__init__()
-
-        self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
-        self.reshape = tf.keras.layers.Reshape(input_size)
-        self.upscale = upscale
-        
-        if not is_label:
-            self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
-        else:
-            self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
-            
-
-    def call(self, inputs, training):
-        if len(inputs) == 2:
-            x, mask = inputs
-        else:
-            x = inputs
-
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x, training=training)
-
-        if self.upscale:
-            x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
-
-
-        if len(inputs) == 2:
-            if len(mask.shape) == 2:
-                x = x * mask[tf.newaxis,:,:,tf.newaxis]
-            else:
-                x = x * mask[:,:,:,tf.newaxis]
-
-        return x
-
-
-
-class VariationalEncoderNew(tf.keras.Model):
-    
-    def __init__(self, depths, dims, kernel_size=7):
-        super(VariationalEncoderNew, self).__init__()
-
-        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
-
-        self.mean_branch = tf.keras.Sequential(
-                [Block(dim=64, drop_path=0, kernel_size=kernel_size,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-
-        self.var_branch = tf.keras.Sequential(
-                [Block(dim=64, drop_path=0, kernel_size=kernel_size,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-
-        self.flatten_mean = tf.keras.layers.Flatten()
-        self.flatten_var = tf.keras.layers.Flatten()
-
-
-
-    def call(self, inputs, training):
-        x = inputs
-
-        x = self.encoder(x, training=training)
+        x = self.encoder(x)
         mean = self.mean_branch(x)
         log_var = self.var_branch(x)
 
@@ -144,23 +33,27 @@ class VariationalEncoderNew(tf.keras.Model):
         return mean, log_var
         
 
+class AEEncoder(nn.Module):
     
-class AEEncoderNew(tf.keras.Model):
+    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__()
 
-    def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
-        super(AEEncoderNew, self).__init__()
-
-        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
-        self.act = tf.keras.layers.Activation(activation)
+        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':
+            self.act = nn.Sigmoid()
+        else:
+            raise ValueError('Activation function not supported')
 
 
-    def call(self, inputs, training):
+    def forward(self, inputs):
         if len(inputs) == 2:
             x, _ = inputs
         else:
             x = inputs
 
-        x = self.encoder(x, training=training)
+        x = self.encoder(x)
         x = self.act(x)
 
         if len(inputs) == 2:
@@ -169,39 +62,37 @@ class AEEncoderNew(tf.keras.Model):
             return x
 
 
-class DecoderNew(tf.keras.Model):
+class Decoder(nn.Module):
     
-    def __init__(self, depths, dims, kernel_size=7, is_label=False):
-        super(DecoderNew, self).__init__()
+    def __init__(self, block_setting: CNBlockConfig, kernel_size=7, is_label=False, nr_of_classes=None, device='cuda'):
+        super(Decoder, self).__init__()
 
         self.is_label = is_label
 
-        if not is_label:
-            self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
-        else:
-            self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
+        self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_classes, device=device, reversed=True)
             
+        if self.is_label:
+            self.softmax = nn.Softmax(dim=1)
 
-    def call(self, inputs, training):
+    def forward(self, inputs):
         if len(inputs) == 2:
             x, mask = inputs
         else:
             x = inputs
 
-        x = self.decoder(x, training=training)
+        x = self.decoder(x)
 
         if self.is_label:
-            x = x * mask[:,:,:,tf.newaxis]
+            x = x * mask[:,None,:,:]
 
-            background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis])
-            x = tf.concat((background, x), axis=-1)
-
-            x = tf.math.softmax(x)
+            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:
             if len(mask.shape) == 2:
-                x = x * mask[tf.newaxis,:,:,tf.newaxis]
+                x = x * mask[None,None,:,:]
             else:
-                x = x * mask[:,:,:,tf.newaxis]
+                x = x * mask[:,None,:,:]
 
         return x

models/cond_encoder.py

@@ -1,38 +1,21 @@
-import tensorflow as tf
-from models.base import AEEncoder, Decoder, AEEncoderNew, DecoderNew
+import torch
+from torch import nn
+from models.base import AEEncoder, Decoder
 
 
-class AutoEncoder(tf.keras.Model):
+class AutoEncoder(nn.Module):
     
     def __init__(self, config):
         super(AutoEncoder, self).__init__()
 
-        self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
-        self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
+        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 call(self, inputs, training):
-        x = inputs
+    def forward(self, input):
+        x = input
 
-        x = self.encoder(x, training=training)
-        x = self.decoder(x, training=training)
-
-        return x
-
-
-class AutoEncoderNew(tf.keras.Model):
-    
-    def __init__(self, config):
-        super(AutoEncoderNew, self).__init__()
-
-        self.encoder : AEEncoderNew = AEEncoderNew(config['encoder_depths'], config['encoder_dims'], config['kernel_size'], config['activation'])
-        self.decoder : DecoderNew = DecoderNew(config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
-
-
-    def call(self, inputs, training):
-        x = inputs
-
-        x = self.encoder(x, training=training)
-        x = self.decoder(x, training=training)
+        x = self.encoder(x)
+        x = self.decoder(x)
 
         return x

models/convnext.py

@@ -1,306 +1,259 @@
-from unittest import skip
-import numpy as np
-import tensorflow as tf
-from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
-                                     GlobalAveragePooling2D, Layer,
-                                     LayerNormalization, Flatten, Reshape, UpSampling2D)
+# taken from torchvision.models.convnext
 
-def drop_path(inputs, drop_prob, is_training):
-    # borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
-    if (not is_training) or (drop_prob == 0.):
-        return inputs
+from functools import partial
+from typing import Any, Callable, List, Optional, Sequence
 
-    # Compute keep_prob
-    keep_prob = 1.0 - drop_prob
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
 
-    # Compute drop_connect tensor
-    random_tensor = keep_prob
-    shape = (tf.shape(inputs)[0],) + (1,) * \
-        (len(tf.shape(inputs)) - 1)
-    random_tensor += tf.random.uniform(shape, dtype=inputs.dtype)
-    binary_tensor = tf.floor(random_tensor)
-    output = tf.math.divide(inputs, keep_prob) * binary_tensor
-    return output
+from torchvision.ops.misc import Conv2dNormActivation, Permute
+from torchvision.ops.stochastic_depth import StochasticDepth
+from utils.convnext_utils import CNBlockConfig
 
 
-class DropPath(tf.keras.layers.Layer):
-    # borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
-    def __init__(self, drop_prob=None):
-        super().__init__()
-        self.drop_prob = drop_prob
-
-    def call(self, x, training):
-        return drop_path(x, self.drop_prob, training)
-
-
-class Block(Layer):
-    r""" ConvNeXt Block. There are two equivalent implementations:
-    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
-    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
-    We use (2) as we find it slightly faster in PyTorch
-
-    Args:
-        dim (int): Number of input channels.
-        drop_path (float): Stochastic depth rate. Default: 0.0
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-    """
-
-    def __init__(self, dim, kernel_size=7, drop_path=0., layer_scale_init_value=1e-6, prefix=''):
-        super().__init__()
-        self.dwconv = DepthwiseConv2D(
-            kernel_size=kernel_size, padding='same')  # depthwise conv
-        self.norm = LayerNormalization(epsilon=1e-6)
-        # pointwise/1x1 convs, implemented with linear layers
-        self.pwconv1 = Dense(4 * dim)
-        self.act = tf.keras.activations.gelu
-        self.pwconv2 = Dense(dim)
-        self.drop_path = DropPath(drop_path)
-        self.dim = dim
-        self.layer_scale_init_value = layer_scale_init_value
-        self.prefix = prefix
-
-    def build(self, input_shape):
-        self.gamma = tf.Variable(
-            initial_value=self.layer_scale_init_value * tf.ones((self.dim)),
-            trainable=True,
-            name=f'{self.prefix}/gamma')
-        self.built = True
-
-    def call(self, x, training):
-        input = x
-        x = self.dwconv(x, training=training)
-        # x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
-        x = self.norm(x, training=training)
-        x = self.pwconv1(x, training=training)
-        x = self.act(x)
-        x = self.pwconv2(x, training=training)
-        if self.gamma is not None:
-            x = self.gamma * x
-        # x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
-
-        x = input + self.drop_path(x, training=training)
+class LayerNorm2d(nn.LayerNorm):
+    def forward(self, x: Tensor) -> Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.permute(0, 3, 1, 2)
         return x
 
 
-class ConvNeXt(tf.keras.Model):
-    r""" ConvNeXt
-        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
-          https://arxiv.org/pdf/2201.03545.pdf
-
-    Args:
-        num_classes (int): Number of classes for classification head. Default: 1000
-        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
-        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
-        include_top (bool): whether to add head or just use it as feature extractor. Default: True
-        drop_path_rate (float): Stochastic depth rate. Default: 0.
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
-    """
-
-    def __init__(self, num_classes=1000, kernel_size=7,
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True,
-                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
-                 ):
+class CNBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        layer_scale: float = 1e-6,
+        stochastic_depth_prob: float = 0.0,
+        kernel_size: int = 7,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
         super().__init__()
-        self.include_top = include_top
-        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
-        # stem = tf.keras.Sequential([
-        #     Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
-        #     LayerNormalization(epsilon=1e-6)]
-        # )
-        stem = [Conv2D(dims[0], kernel_size=4, strides=4, padding='same'), LayerNormalization(epsilon=1e-6)]
-        self.downsample_layers.append(stem)
-        for i in range(len(depths) - 1):
-            # downsample_layer = tf.keras.Sequential([
-            #     LayerNormalization(epsilon=1e-6),
-            #     Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
-            # )
-            downsample_layer = [LayerNormalization(epsilon=1e-6), Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
-            self.downsample_layers.append(downsample_layer)
+        if norm_layer is None:
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
 
-        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
-        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
-        cur = 0
-        for i in range(len(depths)):
-            stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
-                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
-            )
-            self.stages.append(stage)
-            cur += depths[i]
-
-        if self.include_top:
-            self.avg = GlobalAveragePooling2D()
-            self.norm = LayerNormalization(epsilon=1e-6)  # final norm layer
-            self.head = Dense(num_classes)
-        else:
-            self.avg = None
-            self.norm = None
-            self.head = None
-
-    def forward_features(self, x, training):
-        for i in range(len(self.downsample_layers)):
-            # x = self.downsample_layers[i](x)
-            x = self.downsample_layers[i][0](x, training=training)
-            x = self.downsample_layers[i][1](x, training=training)
-            x = self.stages[i](x, training=training)
-        return x
-
-    def call(self, x, training):
-        x = self.forward_features(x, training=training)
-        return x
-        
-
-class ConvNeXtReverse(tf.keras.Model):
-    r""" ConvNeXt
-        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
-          https://arxiv.org/pdf/2201.03545.pdf
-
-    Args:
-        num_classes (int): Number of classes for classification head. Default: 1000
-        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
-        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
-        include_top (bool): whether to add head or just use it as feature extractor. Default: True
-        drop_path_rate (float): Stochastic depth rate. Default: 0.
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
-    """
-
-    def __init__(self, 
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7,
-                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
-                 ):
-        super().__init__()
-        self.upsample_layers = []  # stem and 3 intermediate downsampling conv layers
-        stem = tf.keras.Sequential([
-            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
-            # UpSampling2D((4, 4), interpolation='nearest'),
-            # Conv2D(dims[0], kernel_size=4, strides=1, padding='same'),
-            # Dense(dims[0]),
-            LayerNormalization(epsilon=1e-6)]
+        self.block = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=kernel_size, padding='same', groups=dim, bias=True),
+            Permute([0, 2, 3, 1]),
+            norm_layer(dim),
+            nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
+            nn.GELU(),
+            nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
+            Permute([0, 3, 1, 2]),
         )
-        self.upsample_layers.append(stem)
-        for i in range(len(depths) - 1):
-            downsample_layer = tf.keras.Sequential([
-                LayerNormalization(epsilon=1e-6),
-                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')
-                # UpSampling2D((2,2), interpolation='nearest'),
-                # Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'),
-                # Dense(dims[i+1]),
-            ])
-            self.upsample_layers.append(downsample_layer)
+        self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
 
-        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
-        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
-        cur = 0
-        for i in range(len(depths)):
-            stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
-                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
-            )
-            self.stages.append(stage)
-            cur += depths[i]
-
-        self.last = Dense(3, activation='tanh')
-
-    def call(self, x, training):
-        for i in reversed(range(len(self.upsample_layers))):
-            x = self.upsample_layers[i](x, training=training)
-            x = self.stages[i](x, training=training)
-        x = self.last(x, training=training)
-        return x
+    def forward(self, input: Tensor) -> Tensor:
+        result = self.layer_scale * self.block(input)
+        result = self.stochastic_depth(result)
+        result += input
+        return result
 
 
-class ConvNeXtReverseLabels(tf.keras.Model):
-    r""" ConvNeXt
-        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
-          https://arxiv.org/pdf/2201.03545.pdf
-
-    Args:
-        num_classes (int): Number of classes for classification head. Default: 1000
-        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
-        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
-        include_top (bool): whether to add head or just use it as feature extractor. Default: True
-        drop_path_rate (float): Stochastic depth rate. Default: 0.
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
-    """
-
-    def __init__(self, 
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
-                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
-                 nr_of_classes=7):
+class ConvNeXtBase(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        kernel_size: int = 7,
+    ) -> None:
         super().__init__()
-        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
-        stem = tf.keras.Sequential([
-            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
-            LayerNormalization(epsilon=1e-6)]
-        )
-        self.downsample_layers.append(stem)
-        for i in range(len(depths) - 1):
-            downsample_layer = tf.keras.Sequential([
-                LayerNormalization(epsilon=1e-6),
-                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
-            )
-            self.downsample_layers.append(downsample_layer)
 
-        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
-        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
-        cur = 0
-        for i in range(len(depths)):
-            stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
-                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
-            )
-            self.stages.append(stage)
-            cur += depths[i]
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
 
-        # self.last = Dense(nr_of_classes, activation='softmax')
-        self.last = Dense(nr_of_classes - 1)
+        if block is None:
+            block = CNBlock
 
-    def call(self, x, training):
-        for i in reversed(range(len(self.downsample_layers))):
-            x = self.downsample_layers[i](x, training=training)
-            x = self.stages[i](x, training=training)
-        x = self.last(x, training=training)
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        layers_list: List[nn.Module] = []
+
+        # Stem
+        firstconv_output_channels = block_setting[0].input_channels
+        self.stem = Conv2dNormActivation(
+                3,
+                firstconv_output_channels,
+                kernel_size=4,
+                stride=4,
+                padding=0,
+                norm_layer=norm_layer,
+                activation_layer=None,
+                bias=True,
+            ).to(device)
+        
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        for cnf in block_setting:
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                # adjust stochastic depth probability based on the depth of the stage block
+                sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
+                stage_block_id += 1
+            layers_list.append(nn.Sequential(*stage))
+            if cnf.out_channels is not None:
+                # Downsampling
+                layers_list.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    )
+                )
+
+        self.layers = nn.Sequential(*layers_list).to(device)
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        x = self.stem(x)
+        x = self.layers(x)
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+
+class ConvNeXtReverseBase(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        use_skip: bool = False,
+        is_label: bool = False,
+        nr_of_labels: Optional[int] = None,
+        kernel_size: int = 7,
+    ) -> None:
+        super().__init__()
+
+        self.use_skip = use_skip
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        self.first = nn.Sequential(
+                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
+                        norm_layer(block_setting[0].out_channels),
+                    ).to(device)
+
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        layers_list: List[nn.Module] = []
+        
+        for cnf in block_setting[1:]:
+            if cnf.input_channels is not None:
+                # Bottlenecks
+                stage: List[nn.Module] = []
+                for _ in range(cnf.num_layers):
+                    # adjust stochastic depth probability based on the depth of the stage block
+                    sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                    stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
+                    stage_block_id += 1
+                layers_list.append(nn.Sequential(*stage))
+            else:
+                layers_list.append(nn.Identity())
+
+            if cnf.out_channels is not None:
+                # Upsampling
+                layers_list.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    )
+                )
+
+        self.layers = nn.Sequential(*layers_list).to(device)
+
+        self.last = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    3 if not is_label else nr_of_labels,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
+        
+        self.activation = nn.Tanh() if not is_label else nn.Softmax(dim=1)
+
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+
+    def _forward_impl(self, x) -> Tensor:
+        x = self.first(x)
+        x = self.layers(x)
+        x = self.last(x)
+        x = self.activation(x)
+        return x
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
+def _convnext(
+    block_setting: List[CNBlockConfig],
+    stochastic_depth_prob: float,
+    reversed: bool,
+    is_label: bool,
+    nr_of_labels: Optional[int],
+    kernel_size: int,
+    device='cuda',
+    **kwargs: Any,
+) -> ConvNeXtBase:
+
+    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)
+
+
+    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:
+        super().__init__()
+
+        self.model = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, device=device, **kwargs)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.model(x)
         return x
         
 
 
+if __name__ == '__main__':
+    model = ConvNeXt().to('cuda')
+    test = model(torch.zeros((4, 3, 512, 256)).to('cuda'))
 
-
-model_urls = {
-    "convnext_tiny_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_tiny_1k_224_ema.h5",
-    "convnext_small_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_small_1k_224_ema.h5",
-    "convnext_base_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_224.h5",
-    "convnext_base_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_384.h5",
-    "convnext_large_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_224.h5",
-    "convnext_large_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_384.h5",
-    "convnext_xlarge_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_224_ema.h5",
-    "convnext_xlarge_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_384_ema.h5",
-}
-
-
-model_configs = dict(
-    convnext_tiny=dict(
-        depths=[3, 3, 9, 3],
-        dims=[8, 16, 32, 512]
-    ),
-    convnext_small=dict(
-        depths=[3, 3, 27, 3],
-        dims=[96, 192, 384, 768]
-    ),
-    convnext_base=dict(
-        depths=[3, 3, 27, 3],
-        dims=[128, 256, 512, 1024]
-    ),
-    convnext_large=dict(
-        depths=[3, 3, 27, 3],
-        dims=[192, 384, 768, 1536]
-    ),
-    convnext_xlarge=dict(
-        depths=[3, 3, 27, 3],
-        dims=[256, 512, 1024, 2048]
-    ),
-)
-
+    pass

models/convnext_orig.py

@@ -0,0 +1,263 @@
+# taken from torchvision.models.convnext
+
+from functools import partial
+from typing import Any, Callable, List, Optional, Sequence
+
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+
+from torchvision.ops.misc import Conv2dNormActivation, Permute
+# from torchvision.ops.misc import ConvNormActivation as Conv2dNormActivation
+# from torchvision.ops import Permute
+from torchvision.ops.stochastic_depth import StochasticDepth
+from utils.convnext_utils import CNBlockConfig
+
+
+class LayerNorm2d(nn.LayerNorm):
+    def forward(self, x: Tensor) -> Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.permute(0, 3, 1, 2)
+        return x
+
+
+class CNBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        layer_scale: float = 1e-6,
+        stochastic_depth_prob: float = 0.0,
+        kernel_size: int = 7,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.block = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=kernel_size, padding='same', groups=dim, bias=True),
+            Permute([0, 2, 3, 1]),
+            norm_layer(dim),
+            nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
+            nn.GELU(),
+            nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
+            Permute([0, 3, 1, 2]),
+        )
+        self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = self.layer_scale * self.block(input)
+        result = self.stochastic_depth(result)
+        result += input
+        return result
+
+
+class ConvNeXtBase(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        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__()
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        layers_list: List[nn.Module] = []
+
+        # Stem
+        firstconv_output_channels = block_setting[0].input_channels
+        self.stem = Conv2dNormActivation(
+                input_features if not is_label else nr_of_labels,
+                firstconv_output_channels,
+                kernel_size=4,
+                stride=4,
+                padding=0,
+                norm_layer=norm_layer,
+                activation_layer=None,
+                bias=True,
+            ).to(device)
+        layers_list.append(self.stem)
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        for cnf in block_setting:
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                # adjust stochastic depth probability based on the depth of the stage block
+                sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
+                stage_block_id += 1
+            layers_list.append(nn.Sequential(*stage).to(device))
+            if cnf.out_channels is not None:
+                # Downsampling
+                layers_list.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    ).to(device)
+                )
+
+        self.layers = nn.Sequential(*layers_list).to(device)
+        self.act = nn.Tanh().to(device)
+
+        # for m in self.modules():
+        #     if isinstance(m, (nn.Conv2d, nn.Linear)):
+        #         nn.init.trunc_normal_(m.weight, std=0.02)
+        #         if m.bias is not None:
+        #             nn.init.zeros_(m.bias)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        x = self.layers(x)
+        x = self.act(x)
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+
+class ConvNeXtReverseBase(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        use_skip: bool = False,
+        is_label: bool = False,
+        nr_of_labels: Optional[int] = None,
+        kernel_size: int = 7,
+    ) -> None:
+        super().__init__()
+
+        self.use_skip = use_skip
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        self.stem = nn.Sequential(
+                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=4, stride=4),
+                        norm_layer(block_setting[0].out_channels),
+                    ).to(device)
+        
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        layers_list: List[nn.Module] = []
+
+
+        for i, cnf in enumerate(block_setting):
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                # adjust stochastic depth probability based on the depth of the stage block
+                sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
+                stage_block_id += 1
+            layers_list.append(nn.Sequential(*stage))
+
+            if i == 0:
+                layers_list.append(self.stem)
+            else:
+                # Upsampling
+                layers_list.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    )
+                )
+
+        self.last = nn.Sequential(
+                        Permute([0, 2, 3, 1]),
+                        nn.Linear(in_features=block_setting[-1].out_channels, out_features=3 if not is_label else nr_of_labels, bias=True),
+                        Permute([0, 3, 1, 2]),
+                    ).to(device)
+        
+        self.layers = nn.Sequential(*layers_list).to(device)
+        self.activation = nn.Tanh().to(device) if not is_label else nn.Identity().to(device)
+
+
+        # for m in self.modules():
+        #     if isinstance(m, (nn.Conv2d, nn.Linear)):
+        #         nn.init.trunc_normal_(m.weight, std=0.02)
+        #         if m.bias is not None:
+        #             nn.init.zeros_(m.bias)
+
+
+    def _forward_impl(self, x) -> Tensor:
+        x = self.layers(x)
+        x = self.last(x)
+        x = self.activation(x)
+        return x
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
+
+def _convnext(
+    block_setting: List[CNBlockConfig],
+    stochastic_depth_prob: float,
+    reversed: bool,
+    is_label: bool,
+    nr_of_labels: Optional[int],
+    kernel_size: int,
+    input_features: int,
+    device='cuda',
+    **kwargs: Any,
+) -> ConvNeXtBase:
+
+    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, 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, input_features=3, device='cuda', **kwargs) -> None:
+        super().__init__()
+
+        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)
+        return x
+        
+
+
+if __name__ == '__main__':
+    model = ConvNeXt().to('cuda')
+    test = model(torch.zeros((4, 3, 512, 256)).to('cuda'))
+
+    pass

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

@@ -0,0 +1,25 @@
+import os
+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:])
+    
+    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

@@ -0,0 +1,224 @@
+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.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
+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
+
+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_static_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):
+    pred = model([batch['uv_static'], batch['uv_mask']])
+
+    ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
+    loss = 1e3 * (1 - ms_ssim)
+    # loss = l1_loss
+
+    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():
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_static']
+        l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1))
+        e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+    
+    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_static'], batch['uv_mask']])
+
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_static']
+
+        e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1))
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
+        loss = 1e3 * (1 - ms_ssim)
+
+        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/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
+       
+
+def plot_everything(batch: dict, type: str):
+    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)
+    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)
+    garment_static, faces = maskFacesBatch(mesh_pred.detach().cpu().numpy(), F, mask)
+    gt = [mesh_orig[i].cpu().numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    body = gt
+    garment_posed = garment_static
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((faces[i], faces[i] + np.sum(mask[i])), 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_static[i], faces[i], 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)
+            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 : 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)
+
+    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/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

@@ -0,0 +1,233 @@
+# based on https://github.com/CalciferZh/SMPL/blob/master/smpl_torch_batch.py
+
+import os
+import sys
+import pickle
+import torch
+import numpy as np
+import scipy.io as sio
+from time import time
+from torch.nn import Module
+
+class SMPL(Module):
+  def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False):
+    
+    super(SMPL, self).__init__()
+    self.J_regressor = []
+    self.v_template = []
+    self.shapedirs = []
+    self.posedirs = []
+    self.weights = []
+    self.is_skirt = is_skirt
+    # Load SMPL models
+    # path = os.path.abspath(os.path.dirname(__file__))
+    with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
+        female = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(female['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(female['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(female['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(female['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(female['weights'], dtype=torch.float32))
+    with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f:
+        male = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(male['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(male['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(male['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(male['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(male['weights'], dtype=torch.float32))
+        kintree_table = male['kintree_table']
+        id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
+        self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
+    self.J_regressor = torch.stack(self.J_regressor, axis=0).to(device)
+    self.v_template = torch.stack(self.v_template, axis=0).to(device)
+    self.shapedirs = torch.stack(self.shapedirs, axis=0).to(device)
+    self.posedirs = torch.stack(self.posedirs, axis=0).to(device)
+    self.weights = torch.stack(self.weights, axis=0).to(device)
+
+    # Topology no head/hands/feet
+    F = sio.loadmat(os.path.join(path, 'ssmpl_faces.mat'))
+    self.f = np.array(F['faces_default'])
+    self.f_skirt = np.array(F['faces_skirt'])
+    self.nn = np.array(sio.loadmat(os.path.join(path, 'NN.mat')))
+    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.shape = (1,27578,3)
+    self.device = device
+
+    if orig_smpl:
+        self.to_keep = self.to_keep[self.to_keep < 6890]
+        self.J_regressor = self.J_regressor[:,:,:6890]
+        self.v_template = self.v_template[:,:6890,:]
+        self.shapedirs = self.shapedirs[:,:6890,:,:]
+        self.posedirs = self.posedirs[:,:6890,:,:]
+        self.weights_all = self.weights_all[:,:6890,:]
+        self.weights = self.weights_all[:,self.to_keep,:]
+        self.shape = (1, 6890, 3)
+
+  @staticmethod
+  def rodrigues(r):
+    """
+    Rodrigues' rotation formula that turns axis-angle tensor into rotation
+    matrix in a batch-ed manner.
+    Parameter:
+    ----------
+    r: Axis-angle rotation tensor of shape [batch_size * angle_num, 1, 3].
+    Return:
+    -------
+    Rotation matrix of shape [batch_size * angle_num, 3, 3].
+    """
+    # eps = r.clone().normal_(std=1e-8)
+    # theta = torch.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
+    theta = torch.norm(r, dim=(1, 2), keepdim=True)
+    theta = torch.maximum(theta, torch.ones_like(theta) * np.finfo(np.float32).eps)
+
+    theta_dim = theta.shape[0]
+    r_hat = r / theta
+    # cos = torch.cos(theta)
+    z_stick = torch.zeros(theta_dim, dtype=torch.float32).to(r.device)
+    m = torch.stack(
+      (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+       -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1)
+    m = torch.reshape(m, (-1, 3, 3))
+    i_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) \
+             + torch.zeros((theta_dim, 3, 3), dtype=torch.float32)).to(r.device)
+    A = r_hat.permute(0, 2, 1)
+    # dot = torch.matmul(A, r_hat)
+    # R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
+    R = i_cube + m * torch.sin(theta) + torch.matmul(m, m) * (1 - torch.cos(theta))
+    return R
+
+  @staticmethod
+  def with_zeros(x):
+    """
+    Append a [0, 0, 0, 1] tensor to a [3, 4] tensor.
+    Parameter:
+    ---------
+    x: Tensor to be appended.
+    Return:
+    ------
+    Tensor after appending of shape [4,4]
+    """
+    ones = torch.tensor(
+      [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
+    ).expand(x.shape[0],-1,-1).to(x.device)
+    ret = torch.cat((x, ones), dim=1)
+    return ret
+
+  @staticmethod
+  def pack(x):
+    """
+    Append zero tensors of shape [4, 3] to a batch of [4, 1] shape tensor.
+    Parameter:
+    ----------
+    x: A tensor of shape [batch_size, 4, 1]
+    Return:
+    ------
+    A tensor of shape [batch_size, 4, 4] after appending.
+    """
+    zeros43 = torch.zeros(
+      (x.shape[0], x.shape[1], 4, 3), dtype=torch.float32).to(x.device)
+    ret = torch.cat((zeros43, x), dim=3)
+    return ret
+
+
+  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
+          model vertices. Face indices are also returned as a numpy ndarray.
+          
+          20190128: Add batch support.
+          Parameters:
+          ---------
+          pose: Also known as 'theta', an [N, 24, 3] tensor indicating child joint rotation
+          relative to parent joint. For root joint it's global orientation.
+          Represented in a axis-angle format.
+          betas: Parameter for model shape. A tensor of shape [N, 10] as coefficients of
+          PCA components. Only 10 components were released by SMPL author.
+          trans: Global translation tensor of shape [N, 3].
+          Return:
+          ------
+          A 3-D tensor of [N * 6890 * 3] for vertices,
+          and the corresponding [N * 19 * 3] joint positions.
+    """
+    batch_num = betas.shape[0]
+
+    # v_shaped = torch.tensordot(betas, self.shapedirs, dims=([1], [2])) + self.v_template
+    v_shaped = []
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + torch.reshape(self.v_template[0], shape=self.shape)) # female
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male
+    v_shaped = torch.stack(v_shaped, axis=-1)
+    
+    # J = torch.matmul(self.J_regressor, v_shaped)
+    J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
+    J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
+    J = J - J0
+    R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
+
+    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(
+      self.with_zeros(torch.cat((R_cube_big[:, 0], torch.reshape(J[:, 0, :], (-1, 3, 1))), dim=2))
+    )
+    for i in range(1, 24):
+      results.append(
+        torch.matmul(
+          results[self.parent[i]],
+          self.with_zeros(
+            torch.cat(
+              (R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))),
+              dim=2
+            )
+          )
+        )
+      )
+    
+    stacked = torch.stack(results, dim=1)
+    results = stacked - \
+      self.pack(
+        torch.matmul(
+          stacked,
+          torch.reshape(
+            torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2),
+            (batch_num, 24, 4, 1)
+          )
+        )
+      )
+    
+    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_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
+    )
+    v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1)))
+    v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
+    result = v
+
+    root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
+    result = result - root_joint
+    joints = joints - root_joint
+
+    # # 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

@@ -0,0 +1,43 @@
+import torch
+import numpy as np 
+import os
+import logging
+
+
+class CheckpointSaver:
+    def __init__(self, dirpath, top_n=5):
+        """
+        dirpath: Directory path where to store all model weights 
+        decreasing: If decreasing is `True`, then lower metric is better
+        top_n: Total number of models to track based on validation metric value
+        """
+        if not os.path.exists(dirpath): os.makedirs(dirpath)
+        self.dirpath = dirpath
+        self.top_n = top_n
+        self.top_model_paths = []
+        
+    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, 'scaler': scaler }, model_path)
+        self.top_model_paths.append({'path': model_path})
+        if len(self.top_model_paths)>self.top_n: 
+            self.cleanup()
+    
+    def cleanup(self):
+        to_remove = self.top_model_paths[:len(self.top_model_paths)-self.top_n]
+        logging.info(f"Removing extra models.. {to_remove}")
+        for o in to_remove:
+            os.remove(o['path'])
+        self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:]
+
+
+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'].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/convnext_utils.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+class CNBlockConfig:
+    # Stores information listed at Section 3 of the ConvNeXt paper
+    def __init__(
+        self,
+        input_channels: int,
+        out_channels: Optional[int],
+        num_layers: int,
+    ) -> None:
+        self.input_channels = input_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "input_channels={input_channels}"
+        s += ", out_channels={out_channels}"
+        s += ", num_layers={num_layers}"
+        s += ")"
+        return s.format(**self.__dict__)

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
 

utils/utils.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import numpy as np
+from torch.utils.data import DataLoader
 
 
 
@@ -25,3 +26,20 @@ class Accumulator():
         self.dict = dict()
 
 
+class InfiniteDataLoader(DataLoader):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Initialize an iterator over the dataset.
+        self.dataset_iterator = super().__iter__()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            batch = next(self.dataset_iterator)
+        except StopIteration:
+            # Dataset exhausted, use a new fresh iterator.
+            self.dataset_iterator = super().__iter__()
+            batch = next(self.dataset_iterator)
+        return batch