train with fixed reg losses

configs/ae_norm_1.py

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

configs/cape_00032_SL.py

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

configs/cape_03223_SL.py

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

configs/cape_03223_SS.py

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

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

@@ -0,0 +1,112 @@
+import os
+import sys
+import numpy as np
+import trimesh as tm
+import scipy.io as sio
+import plotly.graph_objects as go
+import plotly.express as px
+from tqdm import tqdm
+import torch
+from scipy.spatial.transform import Rotation as R
+import argparse
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+from data.cloth3d_io import readOBJ, readPC2
+from data.utils import Frame
+from smpl.smpl_torch import SMPL
+
+
+def calculate_normals(vertices, mesh_file, poses):
+    _, faces, _, _ = readOBJ(mesh_file)
+    rodr = smpl.rodrigues(torch.tensor(poses.T.reshape(-1, 1, 3), dtype=torch.float32)).numpy().reshape(vertices.shape[0], -1, 3, 3)
+
+    normals = []
+    for i in range(vertices.shape[0]):
+        r = R.from_matrix(rodr[i, 0, :, :])
+        # assert r.as_matrix() == rodr[frame, 0, :, :]
+        out = r.apply(vertices[i], inverse=True)
+
+        # create a plotly mesh plot of out and vertices [i]
+        # fig = go.Figure(data=[
+        #     go.Mesh3d(
+        #         x=out[:, 0], y=out[:, 2], z=out[:, 1],
+        #         # i, j and k give the vertices of triangles
+        #         i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
+        #         name='y',
+        #         showscale=False
+        #     ),
+        #     go.Mesh3d(
+        #         x=vertices[i][:, 0], y=vertices[i][:, 1], z=vertices[i][:, 2],
+        #         # i, j and k give the vertices of triangles
+        #         i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
+        #         name='x',
+        #         showscale=False
+        #     )
+        # ])
+        # fig.show()
+    
+        mesh = tm.Trimesh(vertices=out, faces=faces, process=False, validate=False)
+        normals.append(mesh.vertex_normals)
+    return np.stack(normals)
+
+
+def get_frames(folder):
+    uv_static_folder = os.path.join(folder, 'pose_shape') 
+    frames = sorted(os.listdir(uv_static_folder))
+
+    frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+    frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
+    return frames
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--folder", help = "end of range", default='')
+    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')
+    args = parser.parse_args()
+
+    obj_folder = args.folder_object # '/data-net/dataset/cloth3d/train'
+    folder = args.folder #'/data-local/127/train'
+
+    os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True)
+    frames = get_frames(folder)
+    smpl = SMPL('cpu', 'smpl')
+
+    ids = {frame.id for frame in frames}
+    ids = sorted(list(ids))[int(args.start):int(args.end)]
+
+    for id in tqdm(ids):
+        # find obj files in the folder
+        garment_obj_folder = os.path.join(obj_folder, id[:5])
+        garments = [file.split('.')[0] for file in os.listdir(garment_obj_folder) if file.endswith('.obj')]
+
+        if len(id) == 5:
+            garment = garments[0]
+        elif len(id) == 6:
+            garment = garments[1]
+        else:
+            raise ValueError('Invalid frame id')
+
+        mesh_file = os.path.join(garment_obj_folder, f'{garment}.obj')
+        vertices =readPC2(os.path.join(garment_obj_folder, f'{garment}.pc16'), True)['V']
+        poses = sio.loadmat(os.path.join(garment_obj_folder, f'info.mat'))['poses']
+        normals = calculate_normals(vertices, mesh_file, poses)
+
+        map = sio.loadmat(os.path.join(garment_obj_folder, f'{garment}_map.mat'))['map'][0]
+        mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy'))
+
+        mesh_normals = normals[:,map,:]
+        mesh_normals = mesh_normals * mask[None,:, None]
+
+        for i in range(vertices.shape[0]):
+            # np.savez(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals)
+            np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals[i])
+
+        # pass
+
+        # print(frame.id, frame.frame_nr)
+

data/dataset.py

@@ -45,9 +45,6 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
         with open(os.path.join(folder, frames_file), 'r') as f:
             lines = f.readlines()
             lines = [line.strip() for line in lines]
-        if is_cape:
-            frames = [Frame(line.split('.')[0], line.split('.')[1]) for line in lines]
-        else:
         frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines]
 
     if shuffle:
@@ -178,6 +175,9 @@ class Dataset(Dataset):
         if 'mesh_curvature' in self.signature:
             mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
             item['mesh_curvature'] = mesh_curvature.astype(np.int32)
+        if 'mesh_normals' in self.signature:
+            mesh_normals = np.load(os.path.join(self.folder, 'mesh_normals', frame.id + separator + frame.frame_nr + '.npz')); mesh_normals = mesh_normals['mesh_normals']
+            item['mesh_normals'] = mesh_normals
 
         for res in self.resolutions:
             res_string = ''
@@ -206,6 +206,10 @@ class Dataset(Dataset):
                 uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
                 item['uv_posed' + res_string] = uv_posed
 
+            if 'uv_normals' + res_string in self.signature:
+                uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals']
+                item['uv_normals' + res_string] = uv_normals
+
 
         data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
 

docker/docker-compose.yml

@@ -34,6 +34,7 @@ services:
       - '/data-net/guest/lhunor/data:/data-net'
       - '/data-fast/126-data1/lhunor/data:/data-local/126'
       - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/127-data2/lhunor/data:/data-local/127_2'
       - '/data-fast/128-data1/lhunor/data:/data-local/128'
     deploy:
       resources:

losses/losses.py

@@ -18,12 +18,12 @@ from torch_cluster import knn
 smpl_to_keep_orig = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
 smpl_to_keep_orig = smpl_to_keep_orig[smpl_to_keep_orig < 6890]
 
-to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
-F_coll = torch.tensor(np.load("smpl/smpl_faces.npy")[to_keep,:], device='cuda')
-F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:], device='cuda')
+# to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
+F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda')
+F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda')
 F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device='cuda')
-F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
-F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
+F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [4,1,1]), device='cuda')
+F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [4,1,1]), device='cuda')
 F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')
 F_norm_orig = F_norm_orig[:,smpl_to_keep_orig,:]
 

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/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.cape_03223_SS import config
+from configs.direct_6 import config
 import argparse
 
 # parser = argparse.ArgumentParser(description='')
@@ -10,9 +10,9 @@ import argparse
 
 
 if __name__ == '__main__':
-    # import wandb
+    import wandb
     # wandb.login(key='******************************************')
-    # wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838')
+    wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838')
     name = config['name']
     script = config['script']
     config = [key for key in sys.modules.keys()  if key.startswith('configs.')][-1]

run_parallel.py

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

scripts/ae_norm_train.py

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

scripts/cvae.py

@@ -10,12 +10,10 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
-if args['debug']:
-    from configs.cape_00032_SL import config
-else:
 config_module = importlib.import_module(args['config'])
 config = config_module.config
-
+if args['debug']:
+    pass
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -52,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')
@@ -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
 
@@ -471,8 +469,8 @@ if __name__ == '__main__':
     smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
-    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], 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'])
+    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()
@@ -480,11 +478,12 @@ if __name__ == '__main__':
 
     model = CVAE(config)
 
+
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
     g_balancing_coeff = config['starting_balancing_coeff']
 
-    ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + config['name'], top_n=10)
+    ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
 
     step = 0
     if config['pretrained_model'] != '':

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()
+
+
+

utils/postprocess.py

@@ -22,7 +22,7 @@ from losses.losses import _nearest_neighbour
 """ COLLISION FUNCTIONS START """
 # F = np.load("smpl/smpl_faces.npy")
 to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
-F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
+F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [4,1,1]), device='cuda')
 smpl_to_keep_orig_post = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
 smpl_to_keep_orig_post = smpl_to_keep_orig_post[smpl_to_keep_orig_post < 6890]
 F_fix_coll_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')[:,smpl_to_keep_orig_post,:]