extrapolation train

.gitignore

@@ -172,3 +172,4 @@ ids/
 examples/
 custom_data/dijkstra*
 .vscode/*
+.vscode/launch.json

.vscode/launch.json

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

compression_test.py

@@ -0,0 +1,32 @@
+import compress_pickle as cp
+import numpy as np
+import os
+
+def convert_bytes(num):
+    """
+    this function will convert bytes to MB.... GB... etc
+    """
+    for x in ['bytes', 'KB', 'MB', 'GB', 'TB']:
+        if num < 1024.0:
+            return "%3.1f %s" % (num, x)
+        num /= 1024.0
+
+
+def file_size(file_path):
+    """
+    this function will return the file size
+    """
+    if os.path.isfile(file_path):
+        file_info = os.stat(file_path)
+        return convert_bytes(file_info.st_size)
+    
+
+if __name__ == '__main__':
+    # Create a random array
+    uv_mask = np.load('custom_utils/mask_256x512.npy')
+    arr = np.stack([np.random.rand(512, 256, 3) * uv_mask[:,:,None] for _ in range(5)])
+    cp.dump(arr, 'temp/test_pickle', 'bz2')
+    np.savez_compressed('temp/test_numpy.npz', arr)
+
+    print('Pickle file size: ', file_size('temp/test_pickle.' + 'bz'))
+    print('Numpy file size: ', file_size('temp/test_numpy.npz'))

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

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

configs/cape_03223_SL.py

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

configs/cape_03223_SS.py

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

configs/pyr_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,40 @@ 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)
+            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 +103,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,6 +113,9 @@ 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, :]
 
+    if return_unfixed:
+        return V_new, F_new, F[F_mask, :]
+    else:
         return V_new, F_new
 
 
@@ -292,28 +291,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 +327,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 +392,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

@@ -231,10 +231,6 @@ def read_posed_garment_frame(folder, id, garment, frame, map=None, mask=None, sh
         map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
     verts = verts[map, :]
 
-    if unshape:
-        shape_offsets = np.tensordot(shape, shapedirs, axes=[[1], [2]])
-        verts[mask, :] = verts[mask, :] - np.squeeze(shape_offsets)[mask, :]
-
     if deorient:
         R = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
         # R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
@@ -244,22 +240,17 @@ def read_posed_garment_frame(folder, id, garment, frame, map=None, mask=None, sh
 
 
 # can i cache map?
-# NOT WORKING WELL
 def read_posed_garment(folder, id, garment, map=None, mask=None, shape=None, shapedirs=None, poses=None, unshape=False, deorient=False):
     # body_posed = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V']
     # offsets = readPC2(os.path.join(folder, id, garment + '_offsets_dynamic.pc16'), True)['V']
     # verts = (body_posed + offsets) * mask[:, None]
     pc16_path = os.path.join(folder, id, garment + '.pc16')
-    info = loadInfo(os.path.join(folder, id, 'info.mat'))
-    verts = np.stack([readPC2Frame(pc16_path, frame, True) for frame in range(info['poses'].shape[1])])
+    # info = loadInfo(os.path.join(folder, id, 'info.mat'))
+    verts = readPC2(pc16_path, True)['V']
     if map is None:
         map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
     verts = verts[:, map, :]
 
-    if unshape:
-        shape_offsets = np.tensordot(shape, shapedirs, axes=[[1], [2]])
-        verts[:, mask, :] = verts[:, mask, :] - np.squeeze(shape_offsets)[None, mask, :]
-
     if deorient:
         for i in range(verts.shape[0]):
             R = scipy.spatial.transform.Rotation.from_euler('xyz', poses[:3, i]).inv()
@@ -269,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))
+    if not return_raw:
         mask = raw_mask > 0
+    else:
+        mask = raw_mask
     return mask
 
 
@@ -298,7 +292,7 @@ def read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
 
 def read_posed_body(folder, id, poses=None, deorient=False):
     pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
-    verts = readPC2(pc16_path, True)
+    verts = readPC2(pc16_path, True)['V']
 
     if deorient:
         for i in range(verts.shape[0]):
@@ -309,17 +303,40 @@ def read_posed_body(folder, id, poses=None, deorient=False):
     return verts
 
 
-def read_curvatures(h5file, id, frames, neighbourhood=None):
-    frames_filtered = [frame for frame in frames if frame.split()[0] == id]
-    garment = frames_filtered[0].split()[2]
-    curvatures_all = []
-    for frame_nr in range(len(frames_filtered)):
-        curv = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw').read()
-        curvatures_all.append(curv)
-    return np.array(curvatures_all)
+# def read_curvatures(h5file, id, garment, frames, neighbourhood=None):
+#     frames_filtered = [frame for frame in frames if (frame.split()[0] == id and frame.split()[2] == garment)]
+#     if len(frames_filtered) == 0:
+#         print(f'No frames found for {id} {garment}')
+#     if len(frames_filtered[0].split()) < 3:
+#         print(f'No garment specified for {id} {garment}')
+#         print(frames_filtered[0].split())
+#     garment = frames_filtered[0].split()[2]
+#     curvatures_all = []
+#     for frame_nr in range(len(frames_filtered)):
+#         curv = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw').read()
+#         curvatures_all.append(curv)
+#     return np.array(curvatures_all)
 
+def read_curvatures(folder, id, garment, frames=None, count=None, neighbourhood=None, per_garment=False):
+    if neighbourhood is not None:
+        raise NotImplementedError('Neighbourhood not implemented for reading curvatures')
+    if count is None:
+        frames_filtered = [frame for frame in frames if (frame.split()[0] == id and frame.split()[2] == garment)]
+        count = len(frames_filtered)
 
+    all = []
+    for i in range(count):
+        curv = np.load(os.path.join(folder, 'curvature_mesh_global' if not per_garment else 'curvature_mesh_garment', f'{id}_{i}_{garment}.npz'))['mesh_curvature_raw'] # not actually raw
+        all.append(curv)
 
+    return all
+
+def read_curvature_frame(folder, id, garment, frame, neighbourhood=None, per_garment=False):
+    if neighbourhood is not None:
+        raise NotImplementedError('Neighbourhood not implemented for reading curvatures')
+    curv = np.load(os.path.join(folder, 'curvature_mesh_global' if not per_garment else 'curvature_mesh_garment', f'{id}_{frame}_{garment}.npz'))['mesh_curvature_raw'] # not actually raw
+
+    return curv
 
 import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')

data/convert_to_uv.py

@@ -96,6 +96,73 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
         return result 
 
 
+def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
+    batch = verts.shape[0]
+    verts_uv = np.zeros((batch, res * 2, res, 3))
+    uv_values = np.zeros((batch, uv_map.shape[0], 3))
+    uv_values[:, m2uv[:, 0]] = verts
+    uv_values[:, m2uv[:, 1]] += verts
+    uv_values = uv_values[:, uv_map_mask, :]
+    uv_values /= uv_counter[np.newaxis, :, np.newaxis]
+    return_garment_uv_mask = False
+    # test
+    verts_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+
+    if mask is not None:
+        garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
+        garment_mask[m2uv[:, 0]] = mask
+        garment_mask[m2uv[:, 1]] += mask
+        garment_mask = garment_mask[uv_map_mask] #   add garment mask
+    else:
+        garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0
+
+    unique_uv_values = uv_values[:, garment_mask, :]
+    uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
+    rbf_dest_current = rbf_dest
+
+    # add back mesh vertices that are not part of the garment
+    uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :]
+    rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
+    rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
+    
+    if garment_uv_mask is None and mask is not None:
+        garment_uv_mask = np.zeros((res * 2, res))
+        garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
+        garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4)
+        garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
+        return_garment_uv_mask = True
+
+    if garment_uv_mask is not None:
+        rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
+
+    rbf_dest_current = np.nonzero(rbf_dest_current)
+    rbf_dest_current = np.array(rbf_dest_current).T
+
+    final = []
+    for i in range(batch):
+        print('processing frame: ' + str(i))
+        if method == 'rbf':
+            interpolator =  spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[i,:,:], neighbors=neighbours, kernel='linear')
+        elif method == 'nearest':
+            interpolator =  spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:,:])
+        else:
+            raise NotImplementedError
+
+        result = np.zeros((res * 2, res, 3), np.float32)
+        result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current)
+        result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[i,:,:]
+
+        if garment_uv_mask is not None:
+            result *= garment_uv_mask[:,:,np.newaxis]
+        
+        final.append(result)
+
+    if return_garment_uv_mask:
+        return final, garment_uv_mask
+    else:
+        return final 
+    
+
 def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None, id=None):
     batch = len(verts)
 
@@ -333,8 +400,8 @@ def get_shapedirs():
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/cloth3d/train')
-    parser.add_argument("-s", "--start", help = "start of range", default=0)
-    parser.add_argument("-e", "--end", help = "end of range", default=1000)
+    parser.add_argument("-s", "--start", help = "start of range", default=70500)
+    parser.add_argument("-e", "--end", help = "end of range", default=90000)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
     parser.add_argument("-sk", "--skirt", help = "end of range", default=0)
     args = parser.parse_args()
@@ -388,15 +455,25 @@ if __name__ == '__main__':
     with open(os.path.join(folder, frames_file), 'r') as f:
         lines = f.readlines()
         lines = [line.strip() for line in lines]
-    frames = lines
+    frames = lines[int(args.start):int(args.end)]
     ids = [frame.split()[0] for frame in frames]
-    ids = sorted(list(set(ids)))[int(args.start):int(args.end)]
+    ids = sorted(list(set(ids)))
 
     shapedirs = get_shapedirs()
 
     root = '/'.join(folder.split('/')[:-1])
     type = folder.split('/')[-1]
-    h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a')
+    # h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')
+    # h5file = tb.open_file(os.path.join(root, f'{type}_extras_uv.h5'), mode='a')
+
+    extras_folder = os.path.join(root, f'{type}_extras')
+    os.makedirs(extras_folder, exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_posed' + orient_mod), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_static'), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_curvature'), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_unposed'), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_body_posed' + orient_mod), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_mask'), exist_ok=True)
 
 
     for id in tqdm.tqdm(ids):
@@ -411,121 +488,73 @@ if __name__ == '__main__':
         # also masking is not the same as just filling the empty values with neighbouring
         # values, at least the plots are different
 
-        garment = garments[0]
+        for garment in garments:
+        # garment = garments[0]
+
+            # if all the uv maps have been generated, skip this id
+            # if all([os.path.exists(os.path.join(extras_folder, 'uv_body_posed' + orient_mod,  f'{id}_{i}_{garment}.npz')) for i in range(info['poses'].shape[1])]):
+            #     continue
 
             mask = read_garment_mask(folder, id, garment)
 
             if 'mesh_posed' in data_to_convert:
-            mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
+                # mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
                 mesh_posed = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], poses=info['poses'], deorient=use_no_orient, unshape=no_shape)
             if 'mesh_unposed' in data_to_convert:
+                raise NotImplementedError()
                 mesh_unposed = read_unposed_garment(folder, id, garment, mask=mask, shape=info['shape'], unshape=no_shape)
             if 'mesh_static' in data_to_convert:
                 mesh_static = read_garment_template(folder, id, garment, mask)
             if 'mesh_curvature' in data_to_convert:
-            mesh_curvature = read_curvatures(h5file, id, frames, neighbourhood=None)
+                mesh_curvature = read_curvatures(folder + '_extras', id, garment, count=info['poses'].shape[1])
             if 'body_posed' in data_to_convert:
                 body_posed = read_posed_body(folder, id, info['poses'], use_no_orient)
 
 
-        # verts_static = read_garment_template(folder, id, garment, mask)
-        # body = read_posed_body(folder, id, info['poses'], use_no_orient)
-        # gt = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], ...)
-        # gt = readPC2(os.path.join(folder, id, garment + '_registered.pc16'),  True)['V']
-        # curvature = np.stack([np.load(os.path.join(folder, f'mesh_curvature_{neighbourhood}_2',  f'{id}_{i}.npz'))['mesh_curvature'] for i in range(frames_len)])
-        # mesh_unposed = [np.load(os.path.join(folder, id, 'mesh_unposed',  id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
-        # mesh_unposed = np.stack([mesh_unposed[i]['mesh_unposed'] for i in range(gt.shape[0])])
-        # mesh_unposed_no_shape = [np.load(os.path.join(folder, id, 'mesh_unposed_no_shape',  id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
-        # mesh_unposed_no_shape = np.stack([mesh_unposed_no_shape[i]['mesh_unposed_no_shape'] for i in range(gt.shape[0])])
-        # body_no_orient = [np.load(os.path.join(folder, id, 'mesh_posed_body_no_orient',  id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
-        # body_no_orient = np.stack([body_no_orient[i]['mesh_posed_body_no_orient'] for i in range(gt.shape[0])])
-        # mesh_unposed_no_shape_one_hot = [np.load(os.path.join(local_folder, 'mesh_unposed_no_shape_one_hot',  id + '_' + str(i) + '.npz')) for i in range(frames)]
-        # mesh_unposed_no_shape_one_hot = [mesh_unposed_no_shape_one_hot[i]['mesh_unposed_no_shape'] for i in range(frames)]
-        # mesh_posed_no_orient = [np.load(os.path.join(local_folder, 'mesh_posed_no_orient',  id + '_' + str(i) + '.npz')) for i in range(frames)]
-        # mesh_posed_no_orient = [mesh_posed_no_orient[i]['mesh_posed_no_orient'] for i in range(frames)]
-
-
-        # verts_dynamic = gt
-
-        # TODO: check if have to mask out non garment vertices or not? are they the same vertices erpeated many times? if remove, have to interpolate from almost nothing?
-
-        # test(data.f.Vt, data.f.Ft, verts)
-        # test(uv_map, data.f.Ft, verts)
-
             start = time.perf_counter()
 
-        uv_static, uv_mask, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
+            uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
+            uv_static = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_static[None, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
             uv_static = uv_static.astype(np.float32)
 
             if 'mesh_posed' in data_to_convert:
-            uv_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
+                uv_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
             if 'mesh_unposed' in data_to_convert:
-            uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
+                raise NotImplementedError()
+                uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
             if 'mesh_curvature' in data_to_convert:
-            uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
+                uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
             if 'body_posed' in data_to_convert:
-            uv_body_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
+                uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv, uv_map, body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
+
+
 
             print(time.perf_counter() - start)
 
             # mesh_static
-        save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
-        save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
+            # save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
+            # save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
+            np.savez_compressed(os.path.join(extras_folder, 'uv_static',  f'{id}_{garment}.npz'), uv_static=uv_static)
+            np.savez_compressed(os.path.join(extras_folder, 'uv_mask',  f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
+
             if 'mesh_posed' in data_to_convert:
-            for i in range(uv_posed.shape[0]):
-                save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
+                for i in range(len(uv_posed)):
+                    np.savez_compressed(os.path.join(extras_folder, 'uv_posed' + orient_mod,  f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
+                    # save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
             if 'mesh_unposed' in data_to_convert:
+                raise NotImplementedError()
                 for i in range(uv_unposed.shape[0]):
                     save_to_h5file(h5file, id, i, garment, uv_unposed[i, :, :, :], 'uv_unposed' + orient_mod + shape_mod, 'float32')
             if 'mesh_curvature' in data_to_convert:
-            for i in range(uv_curvature.shape[0]):
-                save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
+                for i in range(len(uv_curvature)):
+                    np.savez_compressed(os.path.join(extras_folder, 'uv_curvature',  f'{id}_{i}_{garment}.npz'), uv_curvature=uv_curvature[i].astype(np.uint8))
+                    # save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
             if 'body_posed' in data_to_convert:
-            for i in range(uv_body_posed.shape[0]):
-                save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
-        
-
-        # start = time.perf_counter()
-        # np.save(os.path.join(folder, id, 'body_static_uv.npy'), uv_body_static)
-
-        # os.makedirs(os.path.join(folder, id, 'body_posed_uv'), exist_ok=True)
-        # for i in range(uv_body.shape[0]):
-        #     np.savez_compressed(os.path.join(folder, id, 'body_posed_uv', garment + '_' + str(i) + '.npz'), uv_body=uv_body[i, :, :, :])
-        
-        # os.makedirs(os.path.join(folder, id, 'uv_posed_body_no_orient'), exist_ok=True)
-        # for i in range(uv_posed_body_no_orient.shape[0]):
-        #     np.savez_compressed(os.path.join(folder, id, 'uv_posed_body_no_orient', garment + '_' + str(i) + '.npz'), uv_posed_body_no_orient=uv_posed_body_no_orient[i, :, :, :])
-
-        # os.makedirs(os.path.join(folder, id, 'gt_uv'), exist_ok=True)
-        # for i in range(uv_gt.shape[0]):
-        #     np.savez_compressed(os.path.join(folder, id, 'gt_uv', garment + '_' + str(i) + '.npz'), uv_gt=uv_gt[i, :, :, :])
-
-        os.makedirs(os.path.join(folder, f'uv_curvature_{neighbourhood}_2'), exist_ok=True)
-        for i in range(uv_curvature.shape[0]):
-            np.savez_compressed(os.path.join(folder, f'uv_curvature_{neighbourhood}_2', id + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :])
-
-        # os.makedirs(os.path.join(folder, id, 'uv_unposed'), exist_ok=True)
-        # for i in range(uv_unposed.shape[0]):
-        #     np.savez_compressed(os.path.join(folder, id, 'uv_unposed', garment + '_' + str(i) + '.npz'), uv_unposed=uv_unposed[i, :, :, :])
-
-        # os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape'), exist_ok=True)
-        # for i in range(uv_unposed_no_shape.shape[0]):
-        #     np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape[i, :, :, :])
-
-        # os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot'), exist_ok=True)
-        # for i in range(uv_unposed_no_shape_one_hot.shape[0]):
-        #     np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape_one_hot[i, :, :, :])
-
-        # os.makedirs(os.path.join(folder, id, 'uv_shaped_body'), exist_ok=True)
-        # np.savez_compressed(os.path.join(folder, id, 'uv_shaped_body.npz'), uv_shaped_body=uv_shaped_body)
-
-        # print(time.perf_counter() - start)
-        
-        # mesh = uv_to_mesh(uv_verts_static, m2uv)
-        # verify_reconstruction(verts_static * mask[:, np.newaxis], mesh * mask[:, np.newaxis])
-
-        # mesh = uv_to_mesh(uv_body[0,:,:], m2uv)
-        # verify_reconstruction(body[0,:,:] * mask[:, np.newaxis], mesh * mask[:, np.newaxis])
+                for i in range(len(uv_body_posed)):
+                    np.savez_compressed(os.path.join(extras_folder, 'uv_body_posed' + orient_mod,  f'{id}_{i}_{garment}.npz'), uv_body_posed=uv_body_posed[i])
+                    # save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
             
             pass
 
+
+    h5file_curv.close()

data/curvatures.py

@@ -103,37 +103,43 @@ class Cache:
 
 
     def read(self, main_id):
-        frames = [frame for frame in self.frames if frame.split()[0] == main_id]
-        for frame in frames:
-            id = frame.split()[0]
-            garment = frame.split()[2]
-            if id in self.cache.keys():
-                continue
+        # frames = [frame for frame in self.frames if (frame.split()[0] == main_id.split()[0] and frame.split()[2] == main_id.split()[1])]
+        # for frame in frames:
+        id = main_id.split()[0]
+        garment = main_id.split()[1]
 
-            self.cache[id] = Data()
+        self.cache[f'{id} {garment}'] = Data()
+
+        is_skirt = (garment == 'Dress' or garment == 'Skirt')
 
         mask = read_garment_mask(self.folder, id, garment)
         
-            _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+        _, faces = _maskFaces(np.zeros((14475,3)), F if not is_skirt else F_skirt, mask)
         v = read_posed_garment_frame(folder, id, garment, 0)[mask,:]
         m = tm.Trimesh(v, faces, process=False, validate=False)
         outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
         outliers = np.nonzero(outliers)[0]
         outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
         outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
+        # sorted_faces = [sorted(f) for f in faces.astype(np.int64)]
+        # for outlier_face in outliers_faces:
+        #     assert sorted(outlier_face) not in sorted_faces
         faces = np.concatenate((faces, outliers_faces), axis=0)
         m = tm.Trimesh(v, faces, process=False, validate=False)
 
         if self.include_neighbours:
+            if not is_skirt:
                 k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :self.neighbourhood]
-                self.cache[id].k_neighbors = k_neighbors
+            else:
+                k_neighbors = np.argsort(dist_skirt[mask, :][:, mask], axis=-1)[:, :self.neighbourhood]
+            self.cache[f'{id} {garment}'].k_neighbors = k_neighbors
             G = tm.graph.vertex_adjacency_graph(m)
             for i in range(v.shape[0]):
                 for k in range(25):
                     assert nx.has_path(G, i, k_neighbors[i, k]), f'No path between {i} and {k_neighbors[i, k]}'
 
-            self.cache[id].mask = mask
-            self.cache[id].faces = faces
+        self.cache[f'{id} {garment}'].mask = mask
+        self.cache[f'{id} {garment}'].faces = faces
 
     def __getitem__(self, id):
         if id not in self.cache.keys():
@@ -164,10 +170,10 @@ def calculate_dist(folder, neighbourhood, frames):
 def generate_curvatures(folder, neighbourhood, frames):
     # ids = os.listdir(folder)[int(args.start):int(args.end)]
     # frames = frames[int(args.start):int(args.end)]
-    os.makedirs(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'mesh_curvature_raw'), exist_ok=True)
+    os.makedirs(os.path.join(folder + '_extras', f'mesh_curvature_raw'), exist_ok=True)
 
     cache = Cache(folder, neighbourhood, frames)
-    filters = tb.Filters(complevel=9, complib='blosc2')
+    # filters = tb.Filters(complevel=9, complib='blosc2')
 
     for frame in tqdm.tqdm(frames):
 
@@ -175,29 +181,29 @@ def generate_curvatures(folder, neighbourhood, frames):
         frame_nr = int(frame.split()[1])
         garment = frame.split()[2]
     
-        mask = cache[id].mask
-        faces = cache[id].faces
+        mask = cache[f'{id} {garment}'].mask
+        faces = cache[f'{id} {garment}'].faces
         v = read_posed_garment_frame(folder, id, garment, frame_nr)[mask,:]
         
         m = tm.Trimesh(v, faces, process=False, validate=False)
         export_to_ply(m, temp_file)
         cloud = PyntCloud.from_file(temp_file)
 
-        k_neighbors = cache[id].k_neighbors
+        k_neighbors = cache[f'{id} {garment}'].k_neighbors
         ev = cloud.add_scalar_field("eigen_values", k_neighbors=k_neighbors)
         cloud.add_scalar_field("curvature", ev=ev)
         curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])
         out = np.zeros((14475), dtype=np.float32)
-        out[mask] = curvature
+        out[mask] = curvature.real if str(curvature.dtype) == 'complex64' else curvature
 
-        np.savez_compressed(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'mesh_curvature_raw', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
-        group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
-        if f'curvature_raw' in group:
-            array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw')
-            array[:] = out
-        else:
-            carray = h5file.create_carray(group, "curvature_raw", atom=tb.Float32Atom(), shape=out.shape, filters=filters)
-            carray[:] = out
+        np.savez_compressed(os.path.join(folder + '_extras', f'mesh_curvature_raw', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
+        # group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
+        # if f'curvature_raw' in group:
+        #     array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw')
+        #     array[:] = out
+        # else:
+        #     carray = h5file.create_carray(group, "curvature_raw", atom=tb.Float32Atom(), shape=out.shape, filters=filters)
+        #     carray[:] = out
 
     return None
 
@@ -214,34 +220,34 @@ def get_bins(values, nr_of_bins):
 def generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood):
     cache = Cache(folder, neighbourhood, frames, include_neighbours=False)
 
-    os.makedirs(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_global'), exist_ok=True)
-    os.makedirs(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_garment'), exist_ok=True)
+    os.makedirs(os.path.join(folder + '_extras', f'curvature_mesh_global'), exist_ok=True)
+    os.makedirs(os.path.join(folder + '_extras', f'curvature_mesh_garment'), exist_ok=True)
     labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
-    filters = tb.Filters(complevel=9, complib='blosc2')
+    # filters = tb.Filters(complevel=9, complib='blosc2')
 
     for i, frame in enumerate(tqdm.tqdm(frames)):
         id = frame.split()[0]
         frame_nr = int(frame.split()[1])
         garment = frame.split()[2]
 
-        curvature = curvatures[i][cache[id].mask]
+        curvature = curvatures[i][cache[f'{id} {garment}'].mask]
 
         frame_labels = np.zeros((curvature.shape[0]), dtype=np.uint8)
         for j in range(bins.shape[0] - 2):
             frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
         frame_labels[curvature >= bins[-2]] = labels[-1]
         out = np.zeros((14475), dtype=np.uint8)
-        out[cache[id].mask] = frame_labels + 1
+        out[cache[f'{id} {garment}'].mask] = frame_labels + 1
 
-        group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
-        if f'curvature_mesh_global' in group:
-            array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_global')
-            array[:] = out
-        else:
-            carray = h5file.create_carray(group, "curvature_mesh_global", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
-            carray[:] = out
+        # group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
+        # if f'curvature_mesh_global' in group:
+        #     array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_global')
+        #     array[:] = out
+        # else:
+        #     carray = h5file.create_carray(group, "curvature_mesh_global", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
+        #     carray[:] = out
 
-        np.savez_compressed(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_global', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
+        np.savez_compressed(os.path.join(folder + '_extras', f'curvature_mesh_global', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
         # writePC2(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".pc16"), unmasked_label, True)
 
         # redo it with local bins
@@ -251,17 +257,17 @@ def generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood):
             frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
         frame_labels[curvature >= bins[-2]] = labels[-1]
         out = np.zeros((14475), dtype=np.uint8)
-        out[cache[id].mask] = frame_labels + 1
+        out[cache[f'{id} {garment}'].mask] = frame_labels + 1
 
-        group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
-        if f'curvature_mesh_garment' in group:
-            array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_garment')
-            array[:] = out
-        else:
-            carray = h5file.create_carray(group, "curvature_mesh_garment", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
-            carray[:] = out
+        # group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
+        # if f'curvature_mesh_garment' in group:
+        #     array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_garment')
+        #     array[:] = out
+        # else:
+        #     carray = h5file.create_carray(group, "curvature_mesh_garment", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
+        #     carray[:] = out
 
-        np.savez_compressed(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_garment', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
+        np.savez_compressed(os.path.join(folder + '_extras', f'curvature_mesh_garment', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
        
 
 def visualize_curvature_labels(folder, frames):
@@ -319,6 +325,7 @@ def read_raw_curvatures(folder, frames, neighbourhood):
     
     return curvatures_all
 
+
 def read_raw_curvatures_2(folder, frames, neighbourhood):
     curvatures_all = []
     for frame in tqdm.tqdm(frames):
@@ -331,17 +338,17 @@ def read_raw_curvatures_2(folder, frames, neighbourhood):
     
     return curvatures_all
 
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/data/dataset/cloth3d/train')
-    parser.add_argument("-s", "--start", help = "start of range", default=0)
-    parser.add_argument("-e", "--end", help = "end of range", default=1000)
+    parser.add_argument("-s", "--start", help = "start of range", default=300)
+    parser.add_argument("-e", "--end", help = "end of range", default=600)
     parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
     args = parser.parse_args()
 
-    is_skirt = False
-    skirt_modifier = '_skirt' if is_skirt else ''
-    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+    F = np.load("custom_utils/smpl_faces.npy").astype(np.int32)
+    F_skirt = np.load("custom_utils/smpl_faces_skirt.npy").astype(np.int32)
     # k_neighbours(k=25)
 
 
@@ -352,13 +359,10 @@ if __name__ == '__main__':
     neighbourhood = 25
 
     temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
-    if not is_skirt:
     dist = np.load('custom_data/dijkstra.npy')
-    else:
-        dist = np.load('custom_data/dijkstra_skirt.npy')
+    dist_skirt = np.load('custom_data/dijkstra_skirt.npy')
 
-    print('SKIRT' if is_skirt else 'NOT SKIRT')
-    print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -skirt {is_skirt} -neigh {neighbourhood}')
+    print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -neigh {neighbourhood}')
 
     # frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)]
     # frames = [frame.split('.')[0] for frame in frames]
@@ -371,21 +375,21 @@ if __name__ == '__main__':
 
     root = '/'.join(folder.split('/')[:-1])
     type = folder.split('/')[-1]
-    h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a')
+    # h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a')
 
     start = time.perf_counter()
+    # frames = [f'05235 {i} Tshirt' for i in range(300)]
 
-    # generate_curvatures(folder, neighbourhood, frames)
-    curvatures = read_raw_curvatures(folder, frames, neighbourhood)
+    generate_curvatures(folder, neighbourhood, frames)
+    # curvatures = read_raw_curvatures_2(folder, frames, neighbourhood)
     # bins = get_bins(curvatures, nr_of_labels)
     # np.save(f'custom_data/{type}_bins{skirt_modifier}_{neighbourhood}.npy', bins)
     # print(bins)
-    bins = np.load(f'custom_data/{type}_bins{skirt_modifier}_{neighbourhood}.npy')
-    generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood)
+    # bins = np.load(f'custom_data/{type}_bins_{neighbourhood}.npy')
+    # generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood)
     # visualize_curvature_labels(folder, frames)
 
     print(time.perf_counter() - start)
 
-    h5file.close()
 
     pass

data/dataset.py

@@ -1,79 +1,107 @@
 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]
+        if is_cape:
+            frames = [Frame(line.split('.')[0], line.split('.')[1]) for line in lines]
+        else:
+            frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines]
+
+    if shuffle:
+        random.shuffle(frames)
+
+    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,381 +109,183 @@ 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_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!'
-
-        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, :])
-
-
-
-        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()
 
-            # 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']
+        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 resolutions:
+        for res in self.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 '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 = tf.one_hot(uv_curvature, 7)
+                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 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']
+            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 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'))
+            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
-                    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']
+            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
 
-            as_list = []
-            for key in signature.keys():
-                if key in item:
-                    as_list.append(item[key])
 
-            return as_list
+        data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
 
+        for key, value in data.items():
+            if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7):
+                data[key] = torch.permute(value, [2,0,1])
 
-        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"
 
-
-    # dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
-    # for d in dataset:
-    #     print(d)
-
-    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,),
+    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 = 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)
+    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)
 
+    for batch in train_dataloader:
+        batch
+        pass
 
-    # 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]
+    pass 
 
 
 
+
+
+
+
+
+
+
+# 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)
+
+#     F = np.load("custom_utils/smpl_faces.npy")
+
+#     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)
 
-    #     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)
+#     sum_mask[sum_mask == 0] = 1
+#     avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
 
-
-    #     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)
+#     # 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
@@ -463,89 +293,11 @@ if __name__ == '__main__':
 #                 y=x[:,1],
 #                 z=x[:,2],
 #                 # i, j and k give the vertices of triangles
-    #                 i = faces[:,0],
-    #                 j = faces[:,1],
-    #                 k = faces[:,2]
+#                 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()
-
-
-
-
-    # # @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)
-
-    #     c = collision_penalty(v, b, m)
-    #     c2 = collision_penalty(v, b2, m)
-
-    #     c_ = collision_penalty_2(v, b, m)
-    #     c2_ = collision_penalty_2(v, b2, m)
-
-    #     print()
-    #     pass
-    

data/investigate.py

@@ -0,0 +1,236 @@
+import os
+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
+from custom_utils.custom_utils import _maskFaces
+
+
+def get_frames(folder, frames_file):
+    frames = []
+    with open(os.path.join(folder, frames_file), 'r') as f:
+        for line in f:
+            frames.append(line.strip())
+    return frames
+
+
+def plot(mesh, faces):
+    fig = go.Figure()
+
+    fig.add_trace(go.Mesh3d(
+                x=mesh[:,0],
+                y=mesh[:,1],
+                z=mesh[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),)
+    
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                    )
+    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                yaxis = dict(nticks=4, range=[-1,1],),
+                zaxis = dict(nticks=4, range=[-1,1],),)
+    fig.update_layout(scene = helper,
+                        )
+
+    fig.show()
+
+
+def build_lookup_table(repeated):
+    # iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
+
+    lookup = {}
+    for i in range(len(repeated)):
+        for j in range(len(repeated[i])):
+            if len(repeated[i]) == 1:
+                continue
+            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, raw_mask, repeated, verts):
+    # if any of the vertices of a face is masked, mask the face
+    lookup = build_lookup_table(repeated)
+    faces = []
+    for face in all_faces:
+        # if any two or all three vertices of the face have the same coordinates then continue
+        if np.allclose(verts[face[0]], verts[face[1]]) or np.allclose(verts[face[0]], verts[face[2]]) or np.allclose(verts[face[1]], verts[face[2]]):
+            continue
+
+        if mask[face[0]] and mask[face[1]] and mask[face[2]]:
+            faces.append(face)
+        else:
+            # create three arrays of vertices from the lookup table for each vertex of the face
+            # but only for those vertices that are in the lookup table, otherwise createa a single element array with the vertex
+            v1 = lookup[face[0]] if face[0] in lookup else [face[0]]
+            v2 = lookup[face[1]] if face[1] in lookup else [face[1]]
+            v3 = lookup[face[2]] if face[2] in lookup else [face[2]]
+            # try each combination of vertices from the three arrays and check if the face is masked
+            found_replacement = False
+            for i in range(len(v1)):
+                if not found_replacement:
+                    for j in range(len(v2)):
+                        if not found_replacement:
+                            for k in range(len(v3)):
+                                if not found_replacement:
+                                    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
+                                        break
+                # for j in range(len(v2)):
+                #     for k in range(len(v3)):
+                #         if mask[v1[i]] and mask[v2[j]] and mask[v3[k]]:
+                #             faces.append([v1[i], v2[j], v3[k]])
+                #             break
+
+    faces = np.array(faces)
+    V_ind = np.zeros(mask.shape, dtype=np.int32)
+    V_ind[mask] = np.arange(mask.sum())
+    F_new = np.take(V_ind, faces.reshape(-1), axis=0).reshape((-1, 3))
+    return F_new
+
+
+def check(folder, frame):
+    id, frame_nr, garment = frame.split()
+    frame_nr = int(frame_nr)
+    info = loadInfo(os.path.join(folder, id, 'info.mat'))
+    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, 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)
+
+    # writeOBJ('temp/investigate.obj', v, faces)
+
+    repeated = []
+    for repeated_group in repeated_groups:
+        repeated_idx = np.argwhere(np.all(mesh_posed == repeated_group, axis=1))
+        repeated.append(repeated_idx.ravel())
+
+    # # 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
+    
+
+
+if __name__ == "__main__":
+    folder = '/home/workspace/data/dataset/cloth3d/train'
+    frames = get_frames(folder, 'frames_all.txt')
+    root = '/'.join(folder.split('/')[:-1])
+    type = folder.split('/')[-1]
+
+    use_no_orient = False
+    no_shape = False
+    is_skirt = True
+
+    skirt_modifier = '_skirt' if is_skirt else ''
+    orient_mod = '_no_orient' if use_no_orient else ''
+    shape_mod = '_no_shape' if no_shape else ''
+
+
+    while True:
+        frame = random.choice(frames)
+        frame = '05936 0 Dress'
+
+        check(folder, frame)
+
+        pass

data/preprocessing.py

@@ -111,7 +111,7 @@ def create_extras_file(folder):
     # create a new pytables file
     root = '/'.join(folder.split('/')[:-1])
     type = folder.split('/')[-1]
-    h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='w', title='Extras')
+    h5file = tb.open_file(os.path.join(root, f'{type}_extras_test.h5'), mode='w', title='Extras')
 
     frames = []
     with open(os.path.join(folder, 'frames_all.txt'), 'r') as f:
@@ -119,7 +119,7 @@ def create_extras_file(folder):
         lines = [line.strip() for line in lines]
     frames = lines
 
-    for frame in tqdm(frames[:3000]):
+    for frame in tqdm(frames[:1000]):
         # add a new empty group to the root group for each frame
         id, frame_nr, garment = frame.split()
         if not id in h5file.root:

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):

data/verify.py

@@ -0,0 +1,138 @@
+import os
+import sys
+import random
+import numpy as np
+import tables as tb
+import plotly.graph_objects as go
+import plotly.express as ex
+from plotly.subplots import make_subplots
+
+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
+from custom_utils.custom_utils import _maskFaces
+
+
+def get_frames(folder, frames_file):
+    frames = []
+    with open(os.path.join(folder, frames_file), 'r') as f:
+        for line in f:
+            frames.append(line.strip())
+    return frames
+
+
+def verify_uv(folder, frame):
+    id, frame_nr, garment = frame.split()
+    uv_static = np.load(os.path.join(folder, 'uv_static', f'{id}_{garment}.npz'))['uv_static']
+    uv_mask = np.load(os.path.join(folder, 'uv_mask', f'{id}_{garment}.npz'))['uv_mask']
+    uv_posed = np.load(os.path.join(folder, 'uv_posed' + orient_mod, f'{id}_{frame_nr}_{garment}.npz'))['uv_posed']
+    uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed' + orient_mod, f'{id}_{frame_nr}_{garment}.npz'))['uv_body_posed']
+    uv_curvature = np.load(os.path.join(folder, 'uv_curvature', f'{id}_{frame_nr}_{garment}.npz'))['uv_curvature']
+
+    all_1 = np.concatenate([uv_static, uv_posed, uv_body_posed], axis=1)
+    all_1 = (all_1 - all_1.min()) / (all_1.max() - all_1.min())
+
+    ex.imshow(all_1).show()
+    ex.imshow(uv_mask).show()
+    ex.imshow(uv_curvature).show()
+
+
+def verify_mesh(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)
+    mesh_static = read_garment_template(folder, id, garment, mask)[mask,:]
+    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)[mask,:]
+    body_posed = read_posed_body_frame(folder, id, frame_nr, info['poses'], use_no_orient)
+    mesh_curvature = read_curvature_frame(folder + '_extras', id, garment, frame_nr)[mask]
+
+    fig = make_subplots(rows=1, cols=4, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(4)] ])
+    F = np.load(f'custom_utils/smpl_faces{skirt_modifier}.npy').astype(np.int32)
+    _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+
+
+    fig.add_trace(go.Mesh3d(
+                x=mesh_static[:,0],
+                y=mesh_static[:,1],
+                z=mesh_static[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),
+            row=1, col=1)
+    
+    fig.add_trace(go.Mesh3d(
+                x=mesh_posed[:,0],
+                y=mesh_posed[:,1],
+                z=mesh_posed[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),
+            row=1, col=2)
+
+    fig.add_trace(go.Mesh3d(
+                x=body_posed[:,0],
+                y=body_posed[:,1],
+                z=body_posed[:,2],
+                i = F[:,0],
+                j = F[:,1],
+                k = F[:,2],
+            ),
+            row=1, col=3)
+
+    fig.add_trace(go.Mesh3d(
+                x=mesh_posed[:,0],
+                y=mesh_posed[:,1],
+                z=mesh_posed[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=np.tile(mesh_curvature[:,None], (1,3)) * 40,
+            ),
+            row=1, col=4)
+
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                    scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
+                    scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
+                    scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
+                    )
+    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                yaxis = dict(nticks=4, range=[-1,1],),
+                zaxis = dict(nticks=4, range=[-1,1],),)
+    fig.update_layout(scene = helper,
+                        scene2 = helper,
+                        scene3 = helper,
+                        scene4 = helper,
+                        )
+
+    fig.show()
+
+    pass
+
+
+if __name__ == '__main__':
+    folder = '/home/workspace/data/dataset/cloth3d/train'
+    frames = get_frames(folder, 'frames_all.txt')
+    root = '/'.join(folder.split('/')[:-1])
+    type = folder.split('/')[-1]
+    # h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')
+
+    use_no_orient = False
+    no_shape = False
+    is_skirt = True
+
+    skirt_modifier = '_skirt' if is_skirt else ''
+    orient_mod = '_no_orient' if use_no_orient else ''
+    shape_mod = '_no_shape' if no_shape else ''
+
+
+    while True:
+        frame = random.choice(frames)
+        frame = '05936 0 Dress'
+
+        verify_mesh(folder, frame)
+        verify_uv(folder + '_extras', frame)
+
+        pass
+

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'):
+        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)
+        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)
+        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.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.stages.append(stage)
-            cur += depths[i]
+        self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
 
-        if self.include_top:
-            self.avg = GlobalAveragePooling2D()
-            self.norm = LayerNormalization(epsilon=1e-6)  # final norm layer
-            self.head = Dense(num_classes)
+    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,
+    ) -> 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(
+                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:
-            self.avg = None
-            self.norm = None
-            self.head = None
+                layers_list.append(nn.Identity())
 
-    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)
+            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 call(self, x, training):
-        x = self.forward_features(x, training=training)
-        return x
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(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
+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:
 
-    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.
-    """
+    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)
 
-    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.,
-                 ):
+
+    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.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.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.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.model = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, device=device, **kwargs)
 
-        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
-
-
-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):
-        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]
-
-        # self.last = Dense(nr_of_classes, activation='softmax')
-        self.last = Dense(nr_of_classes - 1)
-
-    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)
+    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)
-        else:
-            z_body = self.body_enc(uv_body_posed, training=training)
-            z_static = self.static_enc(uv_template, 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, training=training)
-                z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+            z_curv = self.curv_enc(uv_curv)
+            z_cond = torch.concat([z_static, z_body, z_curv], dim=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/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.cape_03223_SS 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 &')

run_parallel.py

@@ -0,0 +1,29 @@
+import os
+
+
+
+if __name__ == '__main__':
+
+    folder = '/home/workspace/data/dataset/cloth3d/train'
+    # folder = '/home/workspace/local-data/cape_processed/train/male'
+    # folder = '/home/workspace/local-data/valid_skirt'
+    # folder = '/home/workspace/local-data/valid'
+    # ids = os.listdir(folder)
+    # ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:]
+    skirt = 0
+
+    count = int(500 * 1e3)
+    step = 10000
+
+    # !!!!!!! change skirt parameter in script
+
+    while count < 510 * 1e3:
+        # os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('nohup python convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('nohup python convert_to_uv_local.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
+        # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_valid_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_valid_skirt_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('python Data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + '')
+        count += step

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,25 @@
 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())
+
+if args['debug']:
+    from configs.cape_00032_SL import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
-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,36 +30,37 @@ import numpy as np
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
-from postprocessing import fix_collision
 
 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 collision_ratio, kl_loss, normal_loss #, concept_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, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
-from dataset import DatasetFactory, StaticDataset
-from custom_utils.smpl import SMPL
-from models.cvae_model import CVAENew as CVAE
+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='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
-    artifact.add_file('cvae.py')
-    artifact.add_file('models/cvae_model.py')
-    artifact.add_file('losses.py')
-    # artifact.add_file('configs/config_base.py')
-    # artifact.add_file('configs/config_cvae.py')
+    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('dataset.py')
-    artifact.add_file('convnext/models/convnext_tf.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'
@@ -63,231 +68,177 @@ def init_wandb():
             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):
-
-    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']
 
-        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
-
+    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)
 
     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])])
+        # 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_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]))
-        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))
+    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'])
-        # conc_loss = concept_loss(pred, batch['uv_unposed'])
     kl = kl_loss(mean, logvar)
 
     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']
+        reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
 
     loss = reconstr_loss + g_balancing_coeff * kl
 
-        mean = tf.reduce_mean(mean)
-        logvar = tf.reduce_mean(logvar)
+    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'])
     
-    # 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)
 
         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])])
+            # 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_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]))
-    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))
+        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'])
-    # 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']
-    if config['use_norm_loss']:
-        reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
+            reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
 
         loss = reconstr_loss + g_balancing_coeff * kl
 
-    mean = tf.reduce_mean(mean)
-    logvar = tf.reduce_mean(logvar)
+        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'])
     
-    # 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_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]))
-    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))]
@@ -301,8 +252,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 } }
 
@@ -427,8 +378,6 @@ def sample(folder):
 
 
 
-
-
 def test():
     str_split = config['name'].split('_')
     test_idx = str_split.index('test')
@@ -477,97 +426,72 @@ def test():
 
 
 
-if __name__ == '__main__':
-    # init_wandb()
-    parser = argparse.ArgumentParser(description='')
-    parser.add_argument('--name', type=str)
-    args = vars(parser.parse_args())
-
-    smpl = SMPL(orig_smpl=config['is_cape'])
-    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'])
-
-    model = CVAE(config)
-
-    optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
-    balancing_activated_step = 0
-    g_balancing_coeff = config['starting_balancing_coeff']
-
-    # sample('checkpoints/cvae_13')
-    # test()
-
-    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']))
+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)
 
-        data_valid = valid_iterator.get_next()
-        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
             valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
 
-        to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
-        if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
+        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 % 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'] < (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      
-
-        to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
-        
+            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, None)
+
         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 ''
+
+    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 = CVAE(config)
+
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
+    g_balancing_coeff = config['starting_balancing_coeff']
+
+    ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + config['name'], top_n=10)
+
+    step = 0
+    if config['pretrained_model'] != '':
+        step = load_model(model, config['pretrained_model'])
+
+    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,12 +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)
+    out_levels = model([levels_input, conds])
 
     if config['use_skinning']:
         mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
@@ -236,24 +239,24 @@ def train_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?
+    # 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 = 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'])
+    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']:
@@ -263,33 +266,44 @@ def train_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,
+                                '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):
+    model.eval()
+    with torch.inference_mode():
         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 +324,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 +351,35 @@ 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)
+        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'])
+        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,
+                                    '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']
     
+    model.eval()
+    with torch.inference_mode():
         levels_input = group_inputs_for_levels(batch, config)
-    output_levels = model(levels_input, training=False)
+        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/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.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
+            
+    T = torch.stack([torch.tensordot(results[i], self.weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
+    rest_shape_h = torch.cat(
+      (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
+    )
+    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, v_shaped, v_posed

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

@@ -0,0 +1,270 @@
+import tensorflow as tf
+
+import numpy as np
+import tensorflow_graphics as tfg
+import scipy as sp
+import scipy.io as sio
+
+from scipy.spatial import cKDTree
+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 """
+F_coll = np.load("custom_utils/smpl_faces.npy")
+F_coll_skirt = np.load("custom_utils/smpl_faces_skirt.npy")
+F_coll_orig = np.load("custom_utils/smpl_faces_orig.npy")
+F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1])
+F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1])
+F_norm_orig = np.tile(F_coll_orig[np.newaxis,:,:], [4,1,1])
+
+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])
+
+# @tf.function
+def collision_loss(V, B, mask, thr=.004):
+    B_vn = vertex_normals(B, tf.tile(F_coll[None], [B.shape[0], 1, 1]))
+   
+    # compute correspondences
+    idx = _tf_nn_parallel(V, B)
+    N, M = V.shape[:2]
+    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
+    idx = tf.stack((_idx, idx), axis=-1)
+    # compute loss
+    D = V - tf.gather_nd(B, idx)
+    b_vn = tf.gather_nd(B_vn, idx)
+    dot = tf.einsum('abc,abc->ab', D, b_vn)
+    thresholded = tf.minimum(dot - thr, 0)  * mask
+    loss = tf.reduce_sum(thresholded ** 2, axis=-1)
+    ratio = tf.reduce_sum(tf.cast(thresholded < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
+
+    return loss, ratio
+
+
+def collision_ratio(V, B, mask, thr=.004):
+    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+   
+    # compute correspondences
+    idx = _tf_nn_parallel(V, B)
+    N, M = V.shape[:2]
+    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
+    idx = tf.stack((_idx, idx), axis=-1)
+    # compute loss
+    D = V - tf.gather_nd(B, idx)
+    b_vn = tf.gather_nd(B_vn, idx)
+    dot = tf.einsum('abc,abc->ab', D, b_vn)
+    loss = tf.minimum(dot - thr, 0)  * mask # / tf.reduce_sum(mask, axis=-1)
+    ratio = tf.reduce_sum(tf.cast(loss < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
+
+    return ratio
+
+""" COLLISION FUNCTIONS END """
+
+
+laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
+edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
+
+laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo()
+edges_skirt = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
+
+
+def get_edge_mask(masks):
+    return tf.reduce_min(tf.gather(masks, edges, axis=1), axis=-1)
+
+def edge_loss(x, y, mask):
+    e_mask = get_edge_mask(mask)
+    x_e = tf.gather(x, edges[:,0], axis=1) - tf.gather(x, edges[:,1], axis=1)
+    x_e = tf.reduce_sum(x_e ** 2, -1)
+    x_e = tf.sqrt(x_e + 1e-9)
+    y_e = tf.gather(y, edges[:,0], axis=1) - tf.gather(y, edges[:,1], axis=1)
+    y_e = tf.reduce_sum(y_e ** 2, -1)
+    y_e = tf.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    d_e = d_e * e_mask
+    loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1)
+    return loss
+
+def get_edge_mask_skirt(masks):
+    return tf.reduce_min(tf.gather(masks, edges_skirt, axis=1), axis=-1)
+
+def edge_loss_skirt(x, y, mask):
+    e_mask = get_edge_mask_skirt(mask)
+    x_e = tf.gather(x, edges_skirt[:,0], axis=1) - tf.gather(x, edges_skirt[:,1], axis=1)
+    x_e = tf.reduce_sum(x_e ** 2, -1)
+    x_e = tf.sqrt(x_e + 1e-9)
+    y_e = tf.gather(y, edges_skirt[:,0], axis=1) - tf.gather(y, edges_skirt[:,1], axis=1)
+    y_e = tf.reduce_sum(y_e ** 2, -1)
+    y_e = tf.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    d_e = d_e * e_mask
+    loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1)
+    return loss
+
+
+class PairwiseDistance(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(PairwiseDistance, self).__init__(**kwargs)
+
+    def call(self, A, B):
+        rA = tf.reduce_sum(tf.square(A), axis=-1)
+        rB = tf.reduce_sum(tf.square(B), axis=-1)
+        transpose_axes = [0, 2, 1] 
+        distances = - 2*tf.matmul(A, tf.transpose(B, transpose_axes)) + rA[:, :, tf.newaxis] + rB[:, tf.newaxis, :]
+        return distances
+
+
+class NearestNeighbour(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(NearestNeighbour, self).__init__(**kwargs)
+
+    def call(self, A, B):
+        distances = PairwiseDistance(dtype=self.dtype)(A, B)
+        nearest_neighbour = tf.argmin(distances, axis=-1)
+        return tf.cast(nearest_neighbour, dtype=tf.int32)
+
+#
+# @tf.function
+def collision_penalty(va, vb, mask, eps=1e-3):
+    nb = vertex_normals(vb, tf.tile(F[None], [vb.shape[0], 1, 1]))
+    batch_size = tf.cast(tf.shape(va)[0], va.dtype)
+
+    closest_vertices = NearestNeighbour(dtype=va.dtype)(va, vb)
+    vb = tf.gather(vb, closest_vertices, batch_dims=1)
+    nb = tf.gather(nb, closest_vertices, batch_dims=1)
+
+    distance = tf.reduce_sum(nb*(va - vb), axis=-1) 
+    interpenetration = tf.maximum(eps - distance * mask, 0)
+
+    return tf.reduce_sum(interpenetration**3) / batch_size
+
+whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy')
+
+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 - tf.square(mean) - tf.exp(logvar)
+    kl_loss = tf.reduce_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
+
+    pred_norm = vertex_normals(pred, f_norm)
+    gt_norm = vertex_normals(gt, f_norm)
+    return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))
+
+# vgg = tf.keras.applications.vgg16.VGG16(include_top=False, input_shape=(512,256,3))
+# vgg_block_1 = tf.keras.Sequential(vgg.layers[:4])
+# vgg_block_2 = tf.keras.Sequential(vgg.layers[4:7])
+# vgg_block_3 = tf.keras.Sequential(vgg.layers[7:11])
+
+# def concept_loss(pred, gt):
+#     pred = tf.keras.applications.vgg16.preprocess_input((pred + 1) * 128)
+#     gt = tf.keras.applications.vgg16.preprocess_input((gt + 1) * 128)
+
+#     pred_features_1 = vgg_block_1(pred, training=False)
+#     pred_features_2 = vgg_block_2(pred_features_1, training=False)
+#     pred_features_3 = vgg_block_3(pred_features_2, training=False)
+#     gt_features_1 = vgg_block_1(gt, training=False)
+#     gt_features_2 = vgg_block_2(gt_features_1, training=False)
+#     gt_features_3 = vgg_block_3(gt_features_2, training=False)
+
+#     l1_1 = tf.reduce_mean(tf.math.abs(pred_features_1 - gt_features_1))
+#     l1_2 = tf.reduce_mean(tf.math.abs(pred_features_2 - gt_features_2))
+#     l1_3 = tf.reduce_mean(tf.math.abs(pred_features_3 - gt_features_3))
+
+#     return l1_1 + l1_2 + l1_3
+
+
+if __name__ == '__main__':
+
+
+    V = tf.ones((4, 512, 256, 3))
+    B = tf.ones((4, 512, 256, 3))
+    m = tf.ones((4, 14475))
+    loss_1 = concept_loss(V, B)
+    print()

utils/postprocess.py

@@ -0,0 +1,135 @@
+import os
+
+
+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 losses.losses import _nearest_neighbour
+
+
+# from hbertiche/PBNS
+""" COLLISION FUNCTIONS START """
+# 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 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 = _nearest_neighbour(V, B)
+    # compute loss
+    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
+
+
+""" COLLISION FUNCTIONS END """
+
+
+
+if __name__ == '__main__':
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    posed = np.load('/home/workspace/local-data/train/mesh_posed/00016_10.npz')['mesh_posed'][np.newaxis,:,:]
+    body = np.load('/home/workspace/local-data/train/mesh_body_posed/00016_10.npy')[np.newaxis,:,:]
+    mask = np.load('/home/workspace/local-data/train/mesh_mask/00016.npy')
+    posed_2 = np.load('/home/workspace/local-data/train/mesh_posed/00016_15.npz')['mesh_posed'][np.newaxis,:,:]
+
+    fixed = fix_collision(posed_2, body, mask).numpy()
+
+    x = posed_2[0, mask,:]
+    y = body[0]
+    _, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
+    fig = go.Figure(data=[
+            go.Mesh3d(
+                x=x[:,0],
+                y=x[:,1],
+                z=x[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+            ),
+            go.Mesh3d(
+                x=y[:,0],
+                y=y[:,1],
+                z=y[:,2],
+                i = F[:,0],
+                j = F[:,1],
+                k = F[:,2],
+            ),
+        ])
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=1, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,1],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    fig.show()
+
+    x = fixed[0, mask,:]
+    y = body[0]
+    _, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
+    fig = go.Figure(data=[
+            go.Mesh3d(
+                x=x[:,0],
+                y=x[:,1],
+                z=x[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+            ),
+            go.Mesh3d(
+                x=y[:,0],
+                y=y[:,1],
+                z=y[:,2],
+                i = F[:,0],
+                j = F[:,1],
+                k = F[:,2],
+            ),
+        ])
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=1, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,1],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    fig.show()
+
+    print()

utils/utils.py

@@ -0,0 +1,45 @@
+import os
+import sys
+import numpy as np
+from torch.utils.data import DataLoader
+
+
+
+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()
+
+
+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