train with fixed reg losses

configs/abl_norm.py

@@ -1,192 +0,0 @@
-from utils.convnext_utils import CNBlockConfig
-
-config = dict()
-
-config['gpu'] = '0'
-config['name'] = 'abl_torch_norm'
-config['notes'] = ''
-config['tags'] = ['iccv', 'abl_norm', 'torch']
-config['script'] = 'scripts/cvae.py'
-
-config['pretrained_model'] = ''
-
-config['use_skinning'] =              True # and deshaping
-config['remove_global_orientation'] = True
-config['use_normals'] =             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_normals'], 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['normals_model'] = 'checkpoints/iccv/torch/norm_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_normals',
-                            'uv_static',
-                            'mesh_mask',
-                            'uv_mask',
-                        }
-
-if config['use_skinning']:
-  config['signature'].add('uv_unposed')
-  config['signature'].add('mesh_unposed')
-else:
-  config['signature'].add('uv_posed')
-
-config['upscale_enhancer'] = ''
-config['freeze_encoders'] = True
-config['use_one_hot'] = False
-config['bottleneck_scale'] = 2
-
-config['cvae'] = dict()
-config['cvae']['device'] = config['device']
-config['cvae']['is_label'] = False
-config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
-config['cvae']['kernel_size'] = 7
-config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
-config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
-
-block_setting = [
-    CNBlockConfig(32, 64, 3),
-    CNBlockConfig(64, 128, 3),
-    CNBlockConfig(128, 256, 3),
-    CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
-    CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
-]
-
-block_setting_reverse = [
-    CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
-    CNBlockConfig(256, 128, 9),
-    CNBlockConfig(128, 64, 9),
-    CNBlockConfig(64, 32, 3),
-    CNBlockConfig(32, 3, 3),
-]
-
-config['cvae']['encoder_block_settings'] = block_setting
-config['cvae']['decoder_block_settings'] = block_setting_reverse
-
-
-
-####################### cond_ae_norm ############################
-config['cond_ae_norm'] = dict()
-config['cond_ae_norm']['device'] = config['device']
-config['cond_ae_norm']['is_label'] = True
-config['cond_ae_norm']['latent_dim'] = config['latent_dim']
-config['cond_ae_norm']['kernel_size'] = 7
-config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
-
-block_setting = [
-    CNBlockConfig(16, 32, 3),
-    CNBlockConfig(32, 64, 3),
-    CNBlockConfig(64, 128, 9),
-    CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
-    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
-]
-
-block_setting_reverse = [
-    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
-    CNBlockConfig(128, 64, 9),
-    CNBlockConfig(64, 32, 9),
-    CNBlockConfig(32, 16, 3),
-    CNBlockConfig(16, 3, 3),
-]
-
-config['cond_ae_norm']['encoder_block_settings'] = block_setting
-config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
-config['cond_ae_norm']['activation'] = 'tanh'
-
-####################### cond_ae_body ############################
-config['cond_ae_body'] = dict()
-config['cond_ae_body']['device'] = config['device']
-config['cond_ae_body']['is_label'] = False
-config['cond_ae_body']['latent_dim'] = config['latent_dim']
-config['cond_ae_body']['kernel_size'] = 7
-config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
-
-block_setting = [
-    CNBlockConfig(16, 32, 3),
-    CNBlockConfig(32, 64, 3),
-    CNBlockConfig(64, 128, 9),
-    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
-    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
-]
-
-block_setting_reverse = [
-    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
-    CNBlockConfig(128, 64, 9),
-    CNBlockConfig(64, 32, 9),
-    CNBlockConfig(32, 16, 3),
-    CNBlockConfig(16, 3, 3),
-]
-
-config['cond_ae_body']['encoder_block_settings'] = block_setting
-config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
-config['cond_ae_body']['activation'] = 'tanh'
-
-####################### cond_ae_static ############################
-config['cond_ae_static'] = dict()
-config['cond_ae_static']['device'] = config['device']
-config['cond_ae_static']['is_label'] = False
-config['cond_ae_static']['latent_dim'] = config['latent_dim']
-config['cond_ae_static']['kernel_size'] = 7
-config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
-
-block_setting = [
-    CNBlockConfig(8, 16, 3),
-    CNBlockConfig(16, 32, 3),
-    CNBlockConfig(32, 64, 3),
-    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
-    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
-]
-
-block_setting_reverse = [
-    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
-    CNBlockConfig(64, 32, 9),
-    CNBlockConfig(32, 16, 3),
-    CNBlockConfig(16, 8, 3),
-    CNBlockConfig(8, 3, 3),
-]
-
-config['cond_ae_static']['encoder_block_settings'] = block_setting
-config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
-config['cond_ae_static']['activation'] = 'tanh'

configs/cvae_norm_1.py

@@ -1,172 +0,0 @@
-from utils.convnext_utils import CNBlockConfig
-
-config = dict()
-
-config['gpu'] = '0'
-config['name'] = 'cvae_norm_1'
-config['notes'] = ''
-config['tags'] = ['iccv', 'cvae_norm', 'torch']
-config['script'] = 'scripts/cvae_norm.py'
-
-config['pretrained_model'] = ''
-
-
-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['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',
-                            'uv_normals',
-                            'uv_static',
-                            'mesh_mask',
-                            'uv_mask',
-                            'mesh_normals',
-                        }
-
-
-
-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
-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(176, 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/direct_1.py

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

configs/direct_2.py

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

configs/direct_3.py

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

configs/direct_4.py

@@ -0,0 +1,82 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'direct_5'
+config['notes'] = 'with normals input but only single garment sequence' 
+config['tags'] = ['torch', 'direct'] 
+config['script'] = 'scripts/direct_train.py'
+
+config['pretrained_model'] = ''
+config['use_curvature'] = False
+config['use_normals'] = True
+
+# assert that only use curvature or normals
+assert not (config['use_curvature'] and config['use_normals'])
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/train'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = 'frames_00761.txt'
+config['valid_frames'] = 'frames_00761.txt'
+
+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['input_features'] = 9
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['is_label'] = False
+config['kernel_size'] = 7
+config['bottleneck_size'] = 64 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 128, 3),
+    CNBlockConfig(128, 256, 9),
+    CNBlockConfig(256, config['bottleneck_size'], 9),
+    CNBlockConfig(config['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['bottleneck_size'], 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['encoder_block_settings'] = block_setting
+config['decoder_block_settings'] = block_setting_reverse
+config['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['reg_loss_coeff'] = 0.1
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_normals',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        'mesh_posed',
+                        'uv_unposed',
+                        'mesh_unposed',
+                        'uv_static',
+                        'uv_mask',
+                        'uv_body_posed',
+                        }

configs/direct_6.py

@@ -0,0 +1,83 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'direct_6'
+config['notes'] = 'with normals input but only single garment sequence with msssim and l2 mesh loss only' 
+config['tags'] = ['torch', 'direct'] 
+config['script'] = 'scripts/direct_train.py'
+
+config['pretrained_model'] = ''
+config['use_curvature'] = False
+config['use_normals'] = True
+config['use_skinning'] = True
+
+# assert that only use curvature or normals
+assert not (config['use_curvature'] and config['use_normals'])
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/train'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = 'frames_00761.txt'
+config['valid_frames'] = 'frames_00761.txt'
+
+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['input_features'] = 9
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['is_label'] = False
+config['kernel_size'] = 7
+config['bottleneck_size'] = 64 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 128, 3),
+    CNBlockConfig(128, 256, 9),
+    CNBlockConfig(256, config['bottleneck_size'], 9),
+    CNBlockConfig(config['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['bottleneck_size'], 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['encoder_block_settings'] = block_setting
+config['decoder_block_settings'] = block_setting_reverse
+config['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['reg_loss_coeff'] = 0.1
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_normals',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        'mesh_posed',
+                        'uv_unposed',
+                        'mesh_unposed',
+                        'uv_static',
+                        'uv_mask',
+                        'uv_body_posed',
+                        }

configs/direct_7.py

@@ -0,0 +1,87 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'direct_9'
+config['notes'] = 'with normals input whole dataset with msssim and l2 mesh loss and reg losses and without skinning' 
+config['tags'] = ['torch', 'direct'] 
+config['script'] = 'scripts/direct_train.py'
+
+config['pretrained_model'] = ''
+config['use_curvature'] = False
+config['use_normals'] = True
+config['use_skinning'] = False
+
+# assert that only use curvature or normals
+assert not (config['use_curvature'] and config['use_normals'])
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/train'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+config['input_features'] = 9
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['is_label'] = False
+config['kernel_size'] = 7
+config['bottleneck_size'] = 64 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 128, 3),
+    CNBlockConfig(128, 256, 9),
+    CNBlockConfig(256, config['bottleneck_size'], 9),
+    CNBlockConfig(config['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['bottleneck_size'], 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['encoder_block_settings'] = block_setting
+config['decoder_block_settings'] = block_setting_reverse
+config['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['reg_loss_coeff'] = 0.1
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_normals',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        'mesh_posed',
+                        'uv_static',
+                        'uv_mask',
+                        'uv_body_posed',
+                        }
+
+if config['use_skinning']:
+  config['signature'].add('uv_unposed')
+  config['signature'].add('mesh_unposed')
+else:
+  config['signature'].add('uv_posed')

configs/direct_8.py

@@ -0,0 +1,87 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '6'
+config['name'] = 'direct_8'
+config['notes'] = 'with normals input but only single garment sequence with msssim and l2 mesh loss only and without skinning and remasked' 
+config['tags'] = ['torch', 'direct'] 
+config['script'] = 'scripts/direct_train.py'
+
+config['pretrained_model'] = ''
+config['use_curvature'] = False
+config['use_normals'] = True
+config['use_skinning'] = False
+
+# assert that only use curvature or normals
+assert not (config['use_curvature'] and config['use_normals'])
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-net/dataset/00761'
+config['valid_dir'] = '/data-net/dataset/00761'
+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['input_features'] = 9
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['is_label'] = False
+config['kernel_size'] = 7
+config['bottleneck_size'] = 64 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 128, 3),
+    CNBlockConfig(128, 256, 9),
+    CNBlockConfig(256, config['bottleneck_size'], 9),
+    CNBlockConfig(config['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['bottleneck_size'], 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['encoder_block_settings'] = block_setting
+config['decoder_block_settings'] = block_setting_reverse
+config['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['reg_loss_coeff'] = 0.1
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_normals',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        'mesh_posed',
+                        'uv_static',
+                        'uv_mask',
+                        'uv_body_posed',
+                        }
+
+if config['use_skinning']:
+  config['signature'].add('uv_unposed')
+  config['signature'].add('mesh_unposed')
+else:
+  config['signature'].add('uv_posed')

custom_utils/custom_utils.py

@@ -43,6 +43,7 @@ class UV2Mesh(nn.Module):
         if uv.shape[1] == 3 or uv.shape[1] == 7:
             uv = torch.permute(uv, [0,2,3,1])
         if uv.shape[1] != 512:
+            raise NotImplementedError
             uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR)
             
         # batch_size = uv.shape[0]

data/calculate_normals.py

@@ -62,8 +62,8 @@ def get_frames(folder):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "end of range", default='/data-net/dataset/00761')
+    parser.add_argument("-f", "--folder", help = "end of range", default='')
-    parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/data-net/dataset/cloth3d/train')
+    parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/')
     parser.add_argument("-s", "--start", help = "start of range", default=0)
     parser.add_argument("-e", "--end", help = "end of range", default=10)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')

data/cloth3d_io.py

@@ -342,6 +342,7 @@ import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
 import plotly.graph_objects as go
+from custom_utils.custom_utils import _maskFaces
 
 def plot_3d(mesh, faces):
     fig = go.Figure()
@@ -369,19 +370,7 @@ def plot_3d(mesh, faces):
     return fig
 
 
-def quads2tris(F):
-    out = []
-    for f in F:
-        if len(f) == 3: out += [f]
-        elif len(f) == 4: out += [[f[0],f[1],f[2]],
-                                [f[0],f[2],f[3]]]
-        else: print("This should not happen...")
-    return np.array(out, np.int32)
-
-
-
 if __name__ == '__main__':
-    from custom_utils.custom_utils import _maskFaces
 
     folder = '/home/workspace/data/dataset/cloth3d/train'
     F = np.load("custom_utils/smpl_faces.npy")#.astype(np.int32)

data/investigate.py

@@ -4,26 +4,11 @@ import random
 import numpy as np
 import plotly.graph_objects as go
 from pymeshfix import MeshFix
-import trimesh
 
 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
+from custom_utils.custom_utils import _maskFaces
 
-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
-
-    V_ind = np.zeros(V_mask.shape)
-    V_ind[V_mask] = np.arange(V_mask.sum())
-    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
 
 def get_frames(folder, frames_file):
     frames = []
@@ -148,7 +133,7 @@ def check(folder, frame):
     v = mesh_posed[mask,:]
     is_skirt = (garment == 'Dress' or garment == 'Skirt')
     skirt_mod = '_skirt' if is_skirt else ''
-    F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
+    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)
@@ -195,22 +180,7 @@ def check(folder, frame):
 
     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(v_2, faces_3)
-
-    # save stl
-    # mesh = trimesh.Trimesh(vertices=v_2, faces=faces_3, process=False, validate=False)
-    # mesh.export('temp/investigate_00761.stl')
-
-    # # read stl file
-    # mesh = trimesh.load('temp/fixed_00761.stl', process=False, validate=False)
-    # mesh = trimesh.load('temp/fixed_00761.ply', process=False, validate=False)
-
-    # plotly plot of mesh with v_2 vertices and faces from trimesh mesh
-    # faces_4 = faces_3 #np.asarray(mesh.faces)
-    # go.Figure(data = [go.Mesh3d(x=v_2[:, 0], y=v_2[:, 1], z=v_2[:, 2], i=faces_4[:, 0], j=faces_4[:, 1], k=faces_4[:, 2],)
-    #                      ]).show()
-    
-    
 
 
     # plot repeated
@@ -219,9 +189,9 @@ def check(folder, frame):
         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.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'),
+                      go.Scatter3d(x=y[:, 0], y=y[:, 1], z=y[:, 2], marker=dict(size=5, line=dict(width=0),), mode='markers'),
-    #                  ]).show()
+                     ]).show()
 
 
     # unq, idx, count = np.unique(v, axis=0, return_counts=True, return_index=True)
@@ -237,49 +207,30 @@ def check(folder, frame):
     # faces_4 = mask_faces(faces.astype(np.int64), mask_3, mask_2_raw, repeated, v)
     # plot(v[mask_3,:], faces_4)
 
-    return mask_2_raw, faces_3
+
     pass
     
 
-def remask(folder, frame, mask, out_folder):
-    id, frame_nr = frame.split('_')
-    frame_nr = int(frame_nr)
-
-    if frame_nr == 0:
-        mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static']
-        mesh_static_new = np.zeros((14475,3), dtype=np.float32)
-        mesh_static_new[mask,:] = mesh_static[mask,:]
-        np.savez_compressed(os.path.join(out_folder, 'mesh_static', f'{id}.npz'), mesh_static=mesh_static_new)
-
-    # mesh_posed = read_posed_garment_frame(folder, id, garment, frame_nr, mask=None, shape=None, pose=info['poses'][:,frame_nr], deorient=True, unshape=None)
-    mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{frame_nr}.npz'))['mesh_posed_no_orient']
-    mesh_posed_new = np.zeros((14475,3), dtype=np.float32)
-    mesh_posed_new[mask,:] = mesh_posed[mask,:]
-    np.savez_compressed(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{frame_nr}.npz'), mesh_posed_no_orient=mesh_posed)
-
-
 
 if __name__ == "__main__":
-    folder = '/data-local/127/train'
+    folder = '/home/workspace/data/dataset/cloth3d/train'
-    frames = [f'00761_{i}' for i in range(300)]
+    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 = 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 ''
 
-    mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
-    faces = np.load('/data-net/dataset/00761/faces/00761.npy')
 
-    for frame in frames:
+    while True:
-        # frame = random.choice(frames)
+        frame = random.choice(frames)
-        # frame = '00761 0 Tshirt'
+        frame = '05936 0 Dress'
 
-        remask(folder, frame, mask, '/data-net/dataset/00761')
+        check(folder, frame)
 
         pass

data/registration.py

@@ -1,461 +0,0 @@
-import os
-import sys
-import random
-import numpy as np
-import plotly.graph_objects as go
-import torch
-from torch.nn.functional import mse_loss
-from pytorch3d.structures import Meshes
-from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
-from pymeshfix import MeshFix
-
-# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
-
-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,
-    readOBJ,
-    quads2tris
-)
-from smpl.smpl_torch import SMPL
-# from custom_utils.custom_utils import _maskFaces
-
-device = 'cuda'
-smpl = SMPL(device, 'smpl')
-
-
-import torch.nn.functional as F
-from pytorch3d.structures import Meshes
-
-def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
-    """
-    Given two PyTorch3D meshes and a set of vertex indices for mesh1,
-    find the closest vertices from mesh2 to these vertices.
-
-    Args:
-        mesh1: A PyTorch3D Meshes object representing the first mesh.
-        mesh2: A PyTorch3D Meshes object representing the second mesh.
-        vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
-
-    Returns:
-        A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
-    """
-    # Extract the vertices of mesh1 and mesh2.
-    vertices1 = mesh1.verts_packed()
-    vertices2 = mesh2.verts_packed()
-
-    # Extract the vertices of mesh1 that correspond to the given indices.
-    vertices1_selected = vertices1[vertex_indices]
-
-    # Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
-    distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
-
-    # Find the indices of the closest vertices in mesh2.
-    closest_indices = torch.argmin(distances, dim=1)
-
-    return closest_indices
-
-
-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
-
-    V_ind = np.zeros(V_mask.shape)
-    V_ind[V_mask] = np.arange(V_mask.sum())
-    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
-
-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
-
-
-from data.cloth3d_io import readPC2Frame
-import scipy.io as sio
-import scipy
-import scipy.spatial as spatial
-
-def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
-    pc16_path = os.path.join(folder, id, garment + '.pc16')
-    verts = readPC2Frame(pc16_path, frame, True)
-
-    map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
-
-    if deorient:
-        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
-        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
-        verts = R.apply(verts)
-
-    return verts, map
-
-def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
-    pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
-    verts = readPC2Frame(pc16_path, frame, True)
-
-    if deorient:
-        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
-        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
-        verts = R.apply(verts)
-
-    return verts
-
-def edge_loss_custom(x, y, edges):
-    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
-    x_e = torch.sum(x_e ** 2, -1)
-    x_e = torch.sqrt(x_e + 1e-9)
-    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
-    y_e = torch.sum(y_e ** 2, -1)
-    y_e = torch.sqrt(y_e + 1e-9)
-    d_e = y_e - x_e
-    loss = torch.mean(d_e ** 2, dim=-1)[0]
-    return loss
-
-
-def crotch(mesh, body, crotch_idx, nr=75):
-    crotch_middle_on_body = crotch_idx
-    # create kdtree for body and mesh
-    body_tree = spatial.cKDTree(body)
-    mesh_tree = spatial.cKDTree(mesh)
-    # find closest point on mesh from crotch middle on body along the z axis
-
-    crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
-
-    crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
-    crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
-
-    body_verts = body[crotch_body]
-    mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
-    mesh_verts = spatial.cKDTree(mesh_verts)
-    closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
-    crotch_mesh = crotch_mesh[closest_mesh_verts]
-
-    return crotch_body, crotch_mesh
-
-
-def find_closest_point_along_z(point, verts, eps=0.01):
-    min_dist = np.inf
-    closest_point = None
-    for i, v in enumerate(verts):
-        dist = np.linalg.norm(point - v)
-        if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
-            min_dist = dist
-            closest_point = i
-    return closest_point
-
-
-def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
-    """
-    Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
-
-    Args:
-        mesh: A PyTorch3D Meshes object representing the mesh.
-
-    Returns:
-        A tensor of shape (E,) representing the lengths of all edges in the mesh,
-        where E is the total number of edges in the mesh.
-    """
-    # Extract the vertices and faces of the mesh.
-    vertices = mesh.verts_packed()
-    faces = mesh.faces_packed()
-
-    # Compute the pairwise distances between all pairs of vertices that form an edge.
-    edge_indices = faces.view(-1, 3)
-    edge_vertices = vertices[edge_indices]
-    edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
-
-    return edge_lengths
-
-
-def register(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)
-    # convert it to indices
-    mask = raw_mask > 0
-    mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
-    body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
-    template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
-    faces = quads2tris(faces)
-    
-    # boundaries = raw_mask == 2
-    # # invert map and use it to invert boudaries
-    # map_inv = np.zeros(map.max()+1, dtype=np.int32)
-    # map_inv[map] = np.arange(len(map))
-    # boundaries = boundaries[map_inv]
-    # # get the boundary indices
-    
-    # mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
-
-
-    
-    # pose = torch.zeros(1, 72).to('cuda')
-    shape = torch.zeros(1, 11).to(device)
-    shape[:, :10] = torch.tensor(info['shape']).to(device)
-    shape[:, -1] = torch.tensor(info['gender']).to(device)
-    # for star pose
-    # pose[:, 5] = 3.1415/4
-    # pose[:,8] = -3.1415/4
-    STAR_POSE = np.zeros((24, 3), np.float32)
-    STAR_POSE[1,2] = .15
-    STAR_POSE[2,2] = -.15
-    pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
-    star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
-    # rotate star body with 90 degrees around x axis
-    R = spatial.transform.Rotation.from_euler('x', np.pi/2)
-    star_body = R.apply(star_body).astype(np.float32)
-
-    is_skirt = (garment == 'Dress' or garment == 'Skirt')
-    skirt_mod = '_skirt' if is_skirt else ''
-    F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
-
-    ##################### mask body #####################
-    new_crotch = 4729
-
-    # star_body = star_body[mask]
-    # new_crotch = np.zeros((14475), bool)
-    # new_crotch[4729] = True
-    # new_crotch = new_crotch[mask]
-    # new_crotch = np.where(new_crotch)[0][0]
-    # _, F = _maskFaces(np.zeros((14475, 3)), F, mask)
-
-    # plot(template, faces)
-    # v = mesh_posed[mask,:]
-    
-    ##################### boundaries #####################
-    # meshfix = MeshFix(template, faces)
-    # holes = meshfix.extract_holes()
-    # new_border = []
-    # for border in np.array(holes.points):
-    #     # find the index of the vertex in the original mesh
-    #     ind = np.argwhere(np.all(template == border, axis=1))
-    #     new_border.extend(ind.ravel())
-    # mesh1_boundaries = torch.tensor(new_border).to(device)
-    # np.save('temp/body_boundaries.npy', new_border)
-    mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
-
-    # meshfix = MeshFix(template, faces)
-    # holes = meshfix.extract_holes()
-    # new_border = []
-    # for border in np.array(holes.points):
-    #     # find the index of the vertex in the original mesh
-    #     ind = np.argwhere(np.all(template == border, axis=1))
-    #     new_border.extend(ind.ravel())
-    # mesh1_boundaries = torch.tensor(new_border).to(device)
-    # np.save('temp/body_boundaries.npy', new_border)
-    # mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
-
-    ##################### crotch #####################
-    crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
-
-
-
-
-    # export star_body as obj
-    # writeOBJ('temp/star_body.obj', star_body, F)
-
-
-    template = torch.tensor(template).to(device)
-    template_orig = template.clone()
-    faces = torch.tensor(faces).to(device)
-    star_body = torch.tensor(star_body).to(device)
-    F = torch.tensor(F).to(device)
-
-    # verts = template.cpu().numpy()
-    # star_body = star_body.cpu().numpy()
-    # faces = faces.detach().cpu().numpy()
-    # F = F.detach().cpu().numpy()
-    # fig = go.Figure(data=[
-    #     go.Mesh3d(
-    #         x=verts[:,0], y=verts[:,1], z=verts[:,2],
-    #         i=faces[:,0], j=faces[:,1], k=faces[:,2],
-    #         color='lightpink', opacity=0.50),
-    #     go.Mesh3d(
-    #         x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
-    #         i=F[:,0], j=F[:,1], k=F[:,2],
-    #         color='lightblue', opacity=0.50),
-    #     go.Scatter3d(
-    #         x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
-    #         mode='markers', marker=dict(size=2, color='red')),
-    #     go.Scatter3d(
-    #         x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
-    #         mode='markers', marker=dict(size=2, color='blue')),
-    # ])
-    # # make it cuve and limit to -1 and -1
-    # fig.update_layout(scene = dict(
-    #     xaxis = dict(range=[-1,1],),
-    #     yaxis = dict(range=[-1,1],),
-    #     zaxis = dict(range=[-1,1],),),)
-    # fig.show()
-
-    # Define the meshes to optimize
-    mesh1 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
-    mesh2 = Meshes(verts=star_body[None,:,:], faces=F[None,:,:])
-
-    mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
-
-    avg_edge_length = mesh_edge_lengths(mesh1).mean()
-
-    # Define the optimizer and learning rate
-    mesh1.verts_packed().requires_grad = True
-    optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
-    # optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
-
-    # Define the loss functions
-    def laplacian_smoothing_loss(mesh):
-        return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
-
-    def edge_loss(mesh):
-        return mesh_edge_loss(mesh, target_length=avg_edge_length)
-
-    def normal_consistency_loss(mesh):
-        return mesh_normal_consistency(mesh)
-
-    # Optimization loop
-    for i in range(10000):
-        optimizer.zero_grad()
-
-        # find boundaries with the find holes function?
-
-        # Forward pass
-        mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
-        ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
-        bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
-        # crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
-        crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
-        loss: torch.Tensor = ch_loss * 10 +  (bound_ch_loss + crotch_ch_loss) * 5
-        # loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
-        loss += laplacian_smoothing_loss(mesh1_out) * 0.05
-        loss += edge_loss(mesh1_out) * 100
-        # loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
-        loss += normal_consistency_loss(mesh1_out) * 0.1
-
-        # Backward pass
-        loss.backward()
-        optimizer.step()
-
-        # Print the loss
-        if i % 10 == 0:
-            print(f"Loss: {loss.item()} chamfer: {ch_loss.item()} ")
-
-    # create a plotly plot with the body and the mesh1_out
-    verts = mesh1_out.verts_packed().detach().cpu().numpy()
-    star_body = star_body.detach().cpu().numpy()
-    faces = faces.detach().cpu().numpy()
-    F = F.detach().cpu().numpy()
-    mesh1_boundaries = mesh1_boundaries.detach().cpu().numpy()
-    mesh2_boundaries = mesh2_boundaries.detach().cpu().numpy()
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            x=verts[:,0], y=verts[:,1], z=verts[:,2],
-            i=faces[:,0], j=faces[:,1], k=faces[:,2],
-            color='lightpink', opacity=0.80),
-        go.Mesh3d(
-            x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
-            i=F[:,0], j=F[:,1], k=F[:,2],
-            color='lightblue', opacity=0.50),
-        go.Scatter3d(
-            x=verts[mesh1_boundaries,0], y=verts[mesh1_boundaries,1], z=verts[mesh1_boundaries,2],
-            mode='markers', marker=dict(size=2, color='red')),
-        go.Scatter3d(
-            x=star_body[mesh2_boundaries,0], y=star_body[mesh2_boundaries,1], z=star_body[mesh2_boundaries,2],
-            mode='markers', marker=dict(size=2, color='blue')),   
-        go.Scatter3d(
-            x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
-            mode='markers', marker=dict(size=2, color='red')),
-        go.Scatter3d(
-            x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
-            mode='markers', marker=dict(size=2, color='blue')),
-    ])
-    # make it cuve and limit to -1 and -1
-    fig.update_layout(scene = dict(
-        xaxis = dict(range=[-1,1],),
-        yaxis = dict(range=[-1,1],),
-        zaxis = dict(range=[-1,1],),),)
-    fig.update_layout(scene_aspectmode='cube')
-    fig.show()
-
-    pass
-
-if __name__ == "__main__":
-    folder = '/data-net/code/cloth3d/train'
-    # folder = r'D:\Work\code\cloth3d\train'
-    # frames = [f'00761_{i}' for i in range(300)]
-    # folder = r'Z:\Work\Huni\Barcelona\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 = False
-
-    skirt_modifier = '_skirt' if is_skirt else ''
-    orient_mod = '_no_orient' if use_no_orient else ''
-    shape_mod = '_no_shape' if no_shape else ''
-
-    # mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
-    # faces = np.load('/data-net/dataset/00761/faces/00761.npy')
-
-    for frame in frames:
-        frame = random.choice(frames)
-        # frame = '00761 0 Tshirt'
-        frame = '01893 256 Jumpsuit'
-        
-
-        register(folder, frame)
-        break
-        pass

data/registration_post.py

@@ -1,482 +0,0 @@
-import os
-import sys
-import random
-import numpy as np
-import plotly.graph_objects as go
-import torch
-from torch.nn.functional import mse_loss
-from pytorch3d.structures import Meshes
-from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
-from pymeshfix import MeshFix
-
-# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
-
-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,
-    readOBJ,
-    quads2tris
-)
-from smpl.smpl_torch import SMPL
-# from custom_utils.custom_utils import _maskFaces
-
-device = 'cuda'
-smpl = SMPL(device, 'smpl')
-
-
-import torch.nn.functional as F
-from pytorch3d.structures import Meshes
-
-def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
-    """
-    Given two PyTorch3D meshes and a set of vertex indices for mesh1,
-    find the closest vertices from mesh2 to these vertices.
-
-    Args:
-        mesh1: A PyTorch3D Meshes object representing the first mesh.
-        mesh2: A PyTorch3D Meshes object representing the second mesh.
-        vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
-
-    Returns:
-        A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
-    """
-    # Extract the vertices of mesh1 and mesh2.
-    vertices1 = mesh1.verts_packed()
-    vertices2 = mesh2.verts_packed()
-
-    # Extract the vertices of mesh1 that correspond to the given indices.
-    vertices1_selected = vertices1[vertex_indices]
-
-    # Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
-    distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
-
-    # Find the indices of the closest vertices in mesh2.
-    closest_indices = torch.argmin(distances, dim=1)
-
-    return closest_indices
-
-
-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
-
-    V_ind = np.zeros(V_mask.shape)
-    V_ind[V_mask] = np.arange(V_mask.sum())
-    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
-
-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
-
-
-from data.cloth3d_io import readPC2Frame
-import scipy.io as sio
-import scipy
-import scipy.spatial as spatial
-
-def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
-    pc16_path = os.path.join(folder, id, garment + '.pc16')
-    verts = readPC2Frame(pc16_path, frame, True)
-
-    map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
-
-    if deorient:
-        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
-        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
-        verts = R.apply(verts)
-
-    return verts, map
-
-def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
-    pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
-    verts = readPC2Frame(pc16_path, frame, True)
-
-    if deorient:
-        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
-        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
-        verts = R.apply(verts)
-
-    return verts
-
-def edge_loss_custom(x, y, edges):
-    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
-    x_e = torch.sum(x_e ** 2, -1)
-    x_e = torch.sqrt(x_e + 1e-9)
-    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
-    y_e = torch.sum(y_e ** 2, -1)
-    y_e = torch.sqrt(y_e + 1e-9)
-    d_e = y_e - x_e
-    loss = torch.mean(d_e ** 2, dim=-1)[0]
-    return loss
-
-
-def crotch(mesh, body, crotch_idx, nr=75):
-    crotch_middle_on_body = crotch_idx
-    # create kdtree for body and mesh
-    body_tree = spatial.cKDTree(body)
-    mesh_tree = spatial.cKDTree(mesh)
-    # find closest point on mesh from crotch middle on body along the z axis
-
-    crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
-
-    crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
-    crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
-
-    body_verts = body[crotch_body]
-    mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
-    mesh_verts = spatial.cKDTree(mesh_verts)
-    closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
-    crotch_mesh = crotch_mesh[closest_mesh_verts]
-
-    return crotch_body, crotch_mesh
-
-
-def find_closest_point_along_z(point, verts, eps=0.01):
-    min_dist = np.inf
-    closest_point = None
-    for i, v in enumerate(verts):
-        dist = np.linalg.norm(point - v)
-        if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
-            min_dist = dist
-            closest_point = i
-    return closest_point
-
-
-def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
-    """
-    Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
-
-    Args:
-        mesh: A PyTorch3D Meshes object representing the mesh.
-
-    Returns:
-        A tensor of shape (E,) representing the lengths of all edges in the mesh,
-        where E is the total number of edges in the mesh.
-    """
-    # Extract the vertices and faces of the mesh.
-    vertices = mesh.verts_packed()
-    faces = mesh.faces_packed()
-
-    # Compute the pairwise distances between all pairs of vertices that form an edge.
-    edge_indices = faces.view(-1, 3)
-    edge_vertices = vertices[edge_indices]
-    edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
-
-    return edge_lengths
-
-# try to optimize the already registered garment ??????????
-
-def register(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)
-    # convert it to indices
-    mask = raw_mask > 0
-    mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
-    body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
-    template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
-    faces = quads2tris(faces)
-    
-    boundaries = raw_mask == 2
-    registered_boundaries = torch.tensor(np.where(boundaries[mask])[0]).to(device)
-    # invert map and use it to invert boudaries
-    # map_inv = np.zeros(map.max()+1, dtype=np.int32)
-    # map_inv[map] = np.arange(len(map))
-    # boundaries = boundaries[map_inv]
-    # # get the boundary indices
-    # mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
-
-
-    
-    # # pose = torch.zeros(1, 72).to('cuda')
-    # shape = torch.zeros(1, 11).to(device)
-    # shape[:, :10] = torch.tensor(info['shape']).to(device)
-    # shape[:, -1] = torch.tensor(info['gender']).to(device)
-    # # for star pose
-    # # pose[:, 5] = 3.1415/4
-    # # pose[:,8] = -3.1415/4
-    # STAR_POSE = np.zeros((24, 3), np.float32)
-    # STAR_POSE[1,2] = .15
-    # STAR_POSE[2,2] = -.15
-    # pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
-    # star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
-    # # rotate star body with 90 degrees around x axis
-    # R = spatial.transform.Rotation.from_euler('x', np.pi/2)
-    # star_body = R.apply(star_body).astype(np.float32)
-
-    is_skirt = (garment == 'Dress' or garment == 'Skirt')
-    skirt_mod = '_skirt' if is_skirt else ''
-    F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
-
-    ##################### mask body #####################
-    new_crotch = 4729
-
-    # star_body = star_body[mask]
-    new_crotch = np.zeros((14475), bool)
-    new_crotch[4729] = True
-    new_crotch = new_crotch[mask]
-    new_crotch = np.where(new_crotch)[0][0]
-    _, F = _maskFaces(np.zeros((14475, 3)), F, mask)
-    registered = template[map][mask]
-
-    # plot(template, faces)
-    # v = mesh_posed[mask,:]
-    
-    ##################### boundaries #####################
-    # meshfix = MeshFix(template, faces)
-    # holes = meshfix.extract_holes()
-    # new_border = []
-    # for border in np.array(holes.points):
-    #     # find the index of the vertex in the original mesh
-    #     ind = np.argwhere(np.all(template == border, axis=1))
-    #     new_border.extend(ind.ravel())
-    # mesh1_boundaries = torch.tensor(new_border).to(device)
-    # np.save('temp/body_boundaries.npy', new_border)
-    template_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
-
-    # meshfix = MeshFix(template, faces)
-    # holes = meshfix.extract_holes()
-    # new_border = []
-    # for border in np.array(holes.points):
-    #     # find the index of the vertex in the original mesh
-    #     ind = np.argwhere(np.all(template == border, axis=1))
-    #     new_border.extend(ind.ravel())
-    # mesh1_boundaries = torch.tensor(new_border).to(device)
-    # np.save('temp/body_boundaries.npy', new_border)
-    # mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
-
-    ##################### crotch #####################
-    # crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
-
-
-
-
-    # export star_body as obj
-    # writeOBJ('temp/star_body.obj', star_body, F)
-
-
-    template = torch.tensor(template).to(device)
-    template_orig = template.clone()
-    faces = torch.tensor(faces).to(device)
-    registered = torch.tensor(registered).to(device)
-    F = torch.tensor(F).to(device)
-
-    # verts = template.cpu().numpy()
-    # star_body = star_body.cpu().numpy()
-    # faces = faces.detach().cpu().numpy()
-    # F = F.detach().cpu().numpy()
-    # fig = go.Figure(data=[
-    #     go.Mesh3d(
-    #         x=verts[:,0], y=verts[:,1], z=verts[:,2],
-    #         i=faces[:,0], j=faces[:,1], k=faces[:,2],
-    #         color='lightpink', opacity=0.50),
-    #     go.Mesh3d(
-    #         x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
-    #         i=F[:,0], j=F[:,1], k=F[:,2],
-    #         color='lightblue', opacity=0.50),
-    #     go.Scatter3d(
-    #         x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
-    #         mode='markers', marker=dict(size=2, color='red')),
-    #     go.Scatter3d(
-    #         x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
-    #         mode='markers', marker=dict(size=2, color='blue')),
-    # ])
-    # # make it cuve and limit to -1 and -1
-    # fig.update_layout(scene = dict(
-    #     xaxis = dict(range=[-1,1],),
-    #     yaxis = dict(range=[-1,1],),
-    #     zaxis = dict(range=[-1,1],),),)
-    # fig.show()
-
-    # Define the meshes to optimize
-    mesh1 = Meshes(verts=registered[None,:,:], faces=F[None,:,:])
-    mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
-    mesh1_boundaries = registered_boundaries
-    mesh2_boundaries = template_boundaries
-
-    # mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
-
-    avg_edge_length = mesh_edge_lengths(mesh1).mean()
-
-    # Define the optimizer and learning rate
-    mesh1.verts_packed().requires_grad = True
-    optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
-    # optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
-
-    # Define the loss functions
-    def laplacian_smoothing_loss(mesh):
-        return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
-
-    def edge_loss(mesh):
-        return mesh_edge_loss(mesh, target_length=avg_edge_length)
-
-    def normal_consistency_loss(mesh):
-        return mesh_normal_consistency(mesh)
-
-    # Optimization loop
-    for i in range(10000):
-        optimizer.zero_grad()
-
-        # find boundaries with the find holes function?
-
-        # Forward pass
-        mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
-        ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
-        bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
-        # crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
-        # crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
-        loss: torch.Tensor = ch_loss * 3 +  (bound_ch_loss) * 20
-        # loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
-        lap_loss = laplacian_smoothing_loss(mesh1_out)
-        e_loss = edge_loss(mesh1_out)
-        # loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
-        norm_loss = normal_consistency_loss(mesh1_out)
-        loss += lap_loss * 0.05 + e_loss * 100 + norm_loss * 0.1
-
-        # Backward pass
-        loss.backward()
-        optimizer.step()
-
-        # Print the loss
-        if i % 10 == 0:
-            print(f"it: {i} : loss: {loss:.5f} chamfer: {ch_loss:.6f} bound_chamfer: {bound_ch_loss:.6f} laplacian: {lap_loss:.6f} edge: {e_loss:.6f} normal: {norm_loss:.5f}")
-    
-    # create a plotly plot with the body and the mesh1_out
-    registered = mesh1_out.verts_packed().detach().cpu().numpy()
-    template = mesh2.verts_packed().detach().cpu().numpy()
-    faces = faces.detach().cpu().numpy()
-    F = F.detach().cpu().numpy()
-    # faces, F = F, faces
-    # verts, star_body = star_body, verts
-    mesh1_boundaries, mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy(), mesh2_boundaries.detach().cpu().numpy()
-    # mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy()
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            x=template[:,0], y=template[:,1], z=template[:,2],
-            i=faces[:,0], j=faces[:,1], k=faces[:,2],
-            color='lightblue', opacity=0.80),
-        go.Mesh3d(
-            x=registered[:,0], y=registered[:,1], z=registered[:,2],
-            i=F[:,0], j=F[:,1], k=F[:,2],
-            color='lightpink', opacity=0.50),
-        go.Scatter3d(
-            x=registered[mesh1_boundaries,0], y=registered[mesh1_boundaries,1], z=registered[mesh1_boundaries,2],
-            mode='markers', marker=dict(size=2, color='red')),
-        go.Scatter3d(
-            x=template[mesh2_boundaries,0], y=template[mesh2_boundaries,1], z=template[mesh2_boundaries,2],
-            mode='markers', marker=dict(size=2, color='blue')),   
-        # go.Scatter3d(
-        #     x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
-        #     mode='markers', marker=dict(size=2, color='red')),
-        # go.Scatter3d(
-        #     x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
-        #     mode='markers', marker=dict(size=2, color='blue')),
-    ])
-    # make it cuve and limit to -1 and -1
-    fig.update_layout(scene = dict(
-        xaxis = dict(range=[-1,1],),
-        yaxis = dict(range=[-1,1],),
-        zaxis = dict(range=[-1,1],),),)
-    fig.update_layout(scene_aspectmode='cube')
-    fig.show()
-
-    # create a plotly wireframe plot of registered
-    fig = go.Figure()
-    for edge in mesh1_out.edges_packed().detach().cpu().numpy():
-        fig.add_trace(go.Scatter3d(
-            x=registered[edge,0], y=registered[edge,1], z=registered[edge,2],
-            mode='lines', line=dict(color='red', width=1)))
-    fig.update_layout(scene = dict(
-        xaxis = dict(range=[-1,1],),
-        yaxis = dict(range=[-1,1],),
-        zaxis = dict(range=[-1,1],),),)
-    fig.update_layout(scene_aspectmode='cube')
-    fig.show()
-
-
-    pass
-
-if __name__ == "__main__":
-    folder = '/data-net/code/cloth3d/train'
-    # folder = r'D:\Work\code\cloth3d\train'
-    # frames = [f'00761_{i}' for i in range(300)]
-    # folder = r'Z:\Work\Huni\Barcelona\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 = False
-
-    skirt_modifier = '_skirt' if is_skirt else ''
-    orient_mod = '_no_orient' if use_no_orient else ''
-    shape_mod = '_no_shape' if no_shape else ''
-
-    # mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
-    # faces = np.load('/data-net/dataset/00761/faces/00761.npy')
-
-    for frame in frames:
-        frame = random.choice(frames)
-        # frame = '00761 0 Tshirt'
-        frame = '01893 256 Jumpsuit'
-        
-
-        register(folder, frame)
-        break
-        pass

models/base.py

@@ -35,10 +35,10 @@ class VariationalEncoder(nn.Module):
 
 class AEEncoder(nn.Module):
     
-    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda'):
+    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda', input_features=3):
         super(AEEncoder, self).__init__()
 
-        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device)
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device, input_features=input_features)
         if activation == 'tanh':
             self.act = nn.Tanh()
         elif activation == 'sigmoid':

models/cond_encoder.py

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

models/cvae_model.py

@@ -27,16 +27,16 @@ class CVAE(nn.Module):
             if config['body_model'] != '':
                 load_model(self.body_enc, config['body_model'] + '_enc')
 
-        if config['use_normals']:
+        if config['use_curvature']:
-            self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device'])
+            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['normals_model'] != '':
+            if config['curvature_model'] != '':
-                load_model(self.norm_enc, config['normals_model'] + '_enc')
+                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
-            if config['use_normals']:
+            if config['use_curvature']:
-                for param in self.norm_enc.parameters(): param.requires_grad = False
+                for param in self.curv_enc.parameters(): param.requires_grad = False
 
         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'])
@@ -48,7 +48,7 @@ class CVAE(nn.Module):
         else:
             self.avg_garment = None
 
-        self.use_normals = config['use_normals']
+        self.use_curvature = config['use_curvature']
         self.use_offsets_over_body = not config['use_skinning']
 
         if self.use_offsets_over_body:
@@ -98,9 +98,9 @@ class CVAE(nn.Module):
     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_normals:
+        if self.use_curvature:
-            z_norm = self.norm_enc(uv_curv)
+            z_curv = self.curv_enc(uv_curv)
-            z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
+            z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
         else:
             z_cond = torch.concat([z_static, z_body], dim=1)
 
@@ -108,13 +108,13 @@ class CVAE(nn.Module):
 
 
     def forward(self, inputs):
-        uv_unposed, uv_template, uv_body_posed, uv_norm, uv_mask = inputs
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
-        x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_norm], dim=1)
+        x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
         
         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_norm)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
 
         z = torch.concat([z, z_cond], dim=1)
         x = self.decoder(z)

models/cvae_norm_model.py

@@ -1,97 +0,0 @@
-import os
-import time
-import random
-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 utils.checkpoints import load_model
-from models.base import AEEncoder, VariationalEncoder, Decoder
-
-
-class CVAENorm(nn.Module):
-    
-    def __init__(self, config):
-        super(CVAENorm, self).__init__()
-        self.size = config['resolution']
-
-        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']['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['freeze_encoders']:
-            for param in self.static_enc.parameters(): param.requires_grad = False
-            for param in self.body_enc.parameters(): param.requires_grad = False
-
-        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.latent_shape = (config['cvae']['latent_dim'], 8, 4)
-
-
-    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], 8, 4, 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: torch.Tensor, logvar: torch.Tensor, training: bool = True, fix_noise: bool = False):
-        if training == True:
-            if not fix_noise:
-                eps = torch.randn(size=mean.shape).to(mean.device)
-            else:
-                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 = torch.exp(logvar * 0.5) + mean
-        return out
-
-
-    def encode_conds(self, uv_template, uv_body_posed):
-        z_body = self.body_enc(uv_body_posed)
-        z_static = self.static_enc(uv_template)
-        z_cond = torch.concat([z_static, z_body], dim=1)
-        return z_cond
-
-
-    def forward(self, inputs):
-        uv_unposed, uv_template, uv_body_posed, uv_mask = inputs
-        x = torch.concat([uv_unposed, uv_template, uv_body_posed], dim=1)
-        
-        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)
-
-        z = torch.concat([z, z_cond], dim=1)
-        x = self.decoder(z)
-        x = x * uv_mask[:,None,:,:]
-
-        return x, mean, logvar

models/direct.py

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

run.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.cvae_norm_1 import config
+from configs.direct_6 import config
 import argparse
 
 # parser = argparse.ArgumentParser(description='')

scripts/cvae.py

@@ -13,8 +13,7 @@ args = vars(parser.parse_args())
 config_module = importlib.import_module(args['config'])
 config = config_module.config
 if args['debug']:
-    from configs.abl_norm import config
+    pass
-
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -51,7 +50,7 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
 
 
 def init_wandb():
-    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
+    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/cvae.py')
     artifact.add_file('losses/losses.py')
@@ -84,7 +83,7 @@ def train_step(batch):
         mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
 
     model.train()
-    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
+    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']:
@@ -140,7 +139,7 @@ def valid_step(batch):
 
     model.eval()
     with torch.inference_mode():
-        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
+        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']:
@@ -189,16 +188,15 @@ def plot_everything(batch: dict, type: str):
 
     model.eval()
     with torch.inference_mode():
-        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
+        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
     
     mesh_pred = uv2mesh(pred).detach().cpu().numpy()
     
     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])
-    uv_normals_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
     pred_img = torch.permute(pred, [0,2,3,1])
-    img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals_img, orig_img, pred_img), axis=2)
+    img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
     
     F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -264,10 +262,10 @@ def sample(folder):
     count = 0
     for batch in dataset_train:
         batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
-        uv_normals = tf.one_hot(batch['uv_normals'], 7)
+        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
         nr_of_samples=2
 
-        pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
+        pred = model.sample_single(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])
 
@@ -283,7 +281,7 @@ def sample(folder):
         mask = batch['mesh_mask'].numpy().astype(bool)[0]
         _, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
 
-        pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
+        pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
         mesh_pred_zero = uv2mesh(pred_zero)
         mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
 
@@ -459,7 +457,7 @@ def train():
             g_balancing_coeff += 1 / 10000 
 
         if step % 10000 == 0:
-            ckpt_saver(model, step, optimizer, None)
+            ckpt_saver(model, step, optimizer)
 
         step += 1
 
@@ -493,6 +491,6 @@ if __name__ == '__main__':
 
     if not args['debug']:
         # model = torch.compile(model)
-        wandb.watch(model, log='all', log_freq=10000, log_graph=True)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
 
     train()

scripts/cvae_norm.py

@@ -1,400 +0,0 @@
-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.cvae_norm_1 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 tqdm
-import random
-import numpy as np
-import plotly.graph_objects as go
-import plotly.express as ex
-from PIL import Image
-
-if config['is_skirt']:
-    from losses.losses import edge_loss_skirt as edge_loss
-else:
-    from losses.losses import edge_loss
-
-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_norm_model import CVAENorm
-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='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
-    artifact = wandb.Artifact('code', type='code')
-    artifact.add_file('scripts/cvae_norm.py')
-    artifact.add_file('losses/losses.py')
-    artifact.add_file('models/base.py')
-    artifact.add_file('models/cond_encoder.py')
-    artifact.add_file('models/convnext.py')
-    artifact.add_file('models/cvae_norm_model.py')
-    artifact.add_file('data/dataset.py')
-    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
-    for config_file_name in config_file_names:
-        file_name = '/'.join(config_file_name.split('.')) + '.py'
-        if os.path.exists(file_name):
-            artifact.add_file(file_name)
-    run.log_artifact(artifact)
-
-uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-
-
-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
-
-
-def train_step(batch):
-    model.train()
-    pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
-    pred = pred * 1.1
-    mesh_pred = uv2mesh(pred)
-
-
-    l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
-    l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred) ** 2, axis=-1), axis=-1))
-    ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
-    kl = kl_loss(mean, logvar)
-
-    reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
-    loss = reconstr_loss + g_balancing_coeff * kl
-
-    optimizer.zero_grad()
-    loss.backward()
-    optimizer.step()
-
-    model.eval()
-    with torch.inference_mode():
-        mean = torch.mean(mean)
-        logvar = torch.mean(logvar)
-        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
-    
-    return { 'train/loss': loss, 'train/euclid_error': metric,
-             'train/reconstr_loss': reconstr_loss,
-             '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/ms_ssim': ms_ssim, 
-            }
- 
-
-def valid_step(batch):
-    model.eval()
-    with torch.inference_mode():
-        pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
-        pred = pred * 1.1
-        mesh_pred = uv2mesh(pred)
-
-        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
-        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_pred - batch['mesh_normals']) ** 2, axis=-1), axis=-1))
-        ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_static'] + 1.1))
-        kl = kl_loss(mean, logvar)
-
-        reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
-        loss = reconstr_loss + g_balancing_coeff * kl
-
-        mean = torch.mean(mean)
-        logvar = torch.mean(logvar)
-        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
-    
-
-    return { 'valid/loss': loss, 'valid/euclid_error': metric, 
-             'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
-             'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
-             'valid/ms_ssim': ms_ssim, 
-              }
-
-
-def plot_everything(batch: dict, type: str):
-    model.eval()
-    with torch.inference_mode():
-        pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
-    
-    pred = pred * 1.1
-    # mesh_pred = uv2mesh(pred).detach().cpu().numpy()
-    
-    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])
-    orig_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
-    pred_img = torch.permute(pred, [0,2,3,1])
-    img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
-    imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
-    
-    # F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
-    # mask = batch['mesh_mask'].cpu().numpy().astype(bool)
-
-    mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0))
-    logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0))
-
-    return { **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
-
-
-def sample(folder): 
-    # run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
-    load_model(model,folder)
-    count = 0
-    for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
-        uv_normals = tf.one_hot(batch['uv_normals'], 7)
-        nr_of_samples=2
-
-        pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
-        pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
-        pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
-
-        mesh_pred = uv2mesh(pred)
-        mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
-        mesh_orig = batch['mesh_posed'][0].numpy()
-        mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
-
-        body = batch['mesh_body_posed'][0].numpy()
-        body_color = np.ones_like(body, dtype=np.int32)*100
-
-        F = np.load("custom_utils/smpl_faces.npy")
-        mask = batch['mesh_mask'].numpy().astype(bool)[0]
-        _, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
-
-        pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
-        mesh_pred_zero = uv2mesh(pred_zero)
-        mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
-
-
-        meshes = []
-        meshes.append(go.Mesh3d(
-                    # 8 vertices of a cube
-                    x=mesh_orig_reposed[mask,0],
-                    y=mesh_orig_reposed[mask,2],
-                    z=mesh_orig_reposed[mask,1],
-                    # i, j and k give the vertices of triangles
-                    i = faces[:,0],
-                    j = faces[:,1],
-                    k = faces[:,2]
-                    ))
-
-        for i in range(nr_of_samples):
-            x = mesh_posed[i][mask, :].copy()
-            # x[:, 2] += 1 * (i + 1)
-            x[:, 2] += 1
-            meshes.append(go.Mesh3d(
-                        # 8 vertices of a cube
-                        x=x[:,0],
-                        y=x[:,2],
-                        z=x[:,1],
-                        # i, j and k give the vertices of triangles
-                        i = faces[:,0],
-                        j = faces[:,1],
-                        k = faces[:,2]
-                        ))
-
-        
-        for i in range(nr_of_samples):
-            x = mesh_posed_zero[i][mask, :].copy()
-            # x[:, 2] += 1 * (i + 1)
-            x[:, 2] += 2
-            meshes.append(go.Mesh3d(
-                        # 8 vertices of a cube
-                        x=x[:,0],
-                        y=x[:,2],
-                        z=x[:,1],
-                        # i, j and k give the vertices of triangles
-                        i = faces[:,0],
-                        j = faces[:,1],
-                        k = faces[:,2]
-                        ))
-
-        # for i in range(5):
-        #     x = body.copy()
-        #     x[:, 2] += 1 * i
-        #     meshes.append(go.Mesh3d(
-        #                 # 8 vertices of a cube
-        #                 x=x[:,0],
-        #                 y=x[:,2],
-        #                 z=x[:,1],
-        #                 # i, j and k give the vertices of triangles
-        #                 i = F[:,0],
-        #                 j = F[:,1],
-        #                 k = F[:,2],
-        #                 vertexcolor=body_color
-        #                 ))
-
-
-        fig = go.Figure(data=meshes)
-
-        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
-        # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
-        fig.update_layout(
-            scene = dict(
-                xaxis = dict(nticks=4, range=[-1,1],),
-                            # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
-                            yaxis = dict(nticks=4, range=[-1,3],),
-                            zaxis = dict(nticks=4, range=[-1,1],),)
-                            )
-
-        # wandb.log({'plot_' + str(count): fig})
-        fig.show()
-
-        pred = pred.numpy()
-        frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
-        frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
-               save_all=True, duration=500, loop=0)
-
-
-        if count > 10:
-            break
-
-        print(count)
-        count += 1
-
-        print()
-
-
-
-def test():
-    str_split = config['name'].split('_')
-    test_idx = str_split.index('test')
-    checkpoint_name = '_'.join(str_split[:test_idx])
-    load_model(model, 'checkpoints/' + checkpoint_name)
-
-    dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
-    static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
-   
-    step = 0
-    losses = []
-
-    for batch in tqdm.tqdm(dataset_test):
-        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
-
-        test_loss = test_step(batch)
-        losses.append(test_loss)
-
-        to_log = { **{ 'current_step': step }, **test_loss }
-        if step % 1000 == 0:
-            figs = plot_everything(batch, 'test')
-            to_log = { **to_log, **figs }
-
-        wandb.log(to_log, step = step)
-        
-        # print(step)
-        step += 1
-   
-    # averages = dict()
-    path_to_plotly_html = "temp/plotly_figure.html"
-
-    histograms = dict()
-    for key in test_loss.keys():
-        accumulated = [losses[i][key].numpy() for i in range(len(losses))]
-        fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
-        fig.write_html(path_to_plotly_html, auto_play = False)
-
-        histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
-        # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
-
-    
-
-    wandb.log({ **histograms }, step = step)
-    wandb.finish()
-    pass
-
-
-
-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'] 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'])
-    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
-    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
-    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
-    train_accumulator = Accumulator()
-    valid_accumulator = Accumulator()
-
-    model = CVAENorm(config)
-
-    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
-    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
-    g_balancing_coeff = config['starting_balancing_coeff']
-
-    ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
-
-    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=10000, log_graph=True)
-
-    train()

scripts/direct_train.py

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