save

aaa_helpers.txt

@@ -127,6 +127,7 @@ nohup python cvae_pyramid.py --name="lhunor_iccv_pyr_32_64_128_256_skin" > nohup
 nohup python cvae_pyramid.py --name="lhunor_iccv_pyr_32_64_128_256_curv" > nohup/iccv/pyr/iccv_pyr_32_64_128_256_curv.out 2>&1 &
 nohup python cvae_pyramid.py --name="lhunor_iccv_pyr_32_64_128_256_curv_skirt" > nohup/iccv/pyr/iccv_pyr_32_64_128_256_curv_skirt.out 2>&1 &
 nohup python cvae_pyramid.py --name="lhunor_iccv_pyr_32_64_128_256_reg" > nohup/iccv/pyr/iccv_pyr_32_64_128_256_reg.out 2>&1 &
+nohup python cvae_pyramid.py --name="lhunor_iccv_pyr_32_64_128_256_seq" > nohup/iccv/pyr/iccv_pyr_32_64_128_256_seq.out 2>&1 &
 
 nohup python cvae_pyramid.py --name="lhunor_iccv_pyr_32_64_128_256_512" > nohup/iccv/pyr/iccv_pyr_32_64_128_256_512.out 2>&1 &
 nohup python cvae_pyramid.py --name="lhunor_iccv_pyr_32_64_128_256_512_skirt" > nohup/iccv/pyr/iccv_pyr_32_64_128_256_512_skirt.out 2>&1 &
@@ -164,6 +165,7 @@ nohup python evaluate_ae_pyramid.py --name evaluate_ae_iccv_pyr_32_64_128_curv >
 nohup python evaluate_ae_pyramid.py --name evaluate_ae_iccv_pyr_32_64_128_curv_skirt > nohup/iccv/evaluate/iccv_pyr_32_64_128_curv_skirt.out 2>&1 &
 
 nohup python evaluate_ae_pyramid.py --name evaluate_ae_iccv_pyr_32_64_128_256 > nohup/iccv/evaluate/iccv_pyr_32_64_128_256.out 2>&1 &
+nohup python evaluate_ae_pyramid.py --name evaluate_ae_iccv_pyr_32_64_128_256_seq > nohup/iccv/evaluate/iccv_pyr_32_64_128_256_seq.out 2>&1 &
 nohup python evaluate_ae_pyramid.py --name evaluate_ae_iccv_pyr_32_64_128_256_skirt > nohup/iccv/evaluate/iccv_pyr_32_64_128_256_skirt.out 2>&1 &
 nohup python evaluate_ae_pyramid.py --name evaluate_ae_iccv_pyr_32_64_128_256_skin > nohup/iccv/evaluate/iccv_pyr_32_64_128_256_skin.out 2>&1 &
 nohup python evaluate_ae_pyramid.py --name evaluate_ae_iccv_pyr_32_64_128_256_curv > nohup/iccv/evaluate/iccv_pyr_32_64_128_256_curv.out 2>&1 &

calculate_error.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.iccv.config_abl_pyr_32_64_128_skin_skirt import config
+from configs.iccv.config_abl_pyr_32_64_128_seq import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -58,7 +58,7 @@ if __name__ == '__main__':
     # smpl = SMPL(use_one_hot=config['use_one_hot'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
-    error_type = 'collision' # 'euclidean' 'collision' edge normal
+    error_type = 'euclidean' # 'euclidean' 'collision' edge normal
 
     folder = config['test_dir']
 
@@ -92,9 +92,9 @@ if __name__ == '__main__':
     frames = [frame for frame in frames if not frame.endswith('_fixed_coll.npy')]
     # frames = frames[:100]
 
-    print('Reading bodies...')
-    bodies_folder = '/home/workspace/local-data' + ('' if config['is_skirt'] else '_127') + '/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/mesh_posed_body_no_orient'
-    bodies = {frame : np.load(bodies_folder + '/' + frame[:-1] + 'z')['mesh_posed_body_no_orient'] for frame in tqdm(frames) }
+    # print('Reading bodies...')
+    # bodies_folder = '/home/workspace/local-data' + ('' if config['is_skirt'] else '_127') + '/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/mesh_posed_body_no_orient'
+    # bodies = {frame : np.load(bodies_folder + '/' + frame[:-1] + 'z')['mesh_posed_body_no_orient'] for frame in tqdm(frames) }
 
     if config['remove_global_orientation']:
         gt_name = 'mesh_posed_no_orient'

configs/iccv/config_abl_pyr_32_64_128_256_seq.py

@@ -0,0 +1,187 @@
+config = dict()
+
+config['gpu'] = '3'
+config['name'] = 'iccv_pyr_32_64_128_256_seq_2'
+config['notes'] = ''
+config['tags'] = ['pyr', 'iccv']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+config['use_curvature'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64, 128, 256]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/iccv/iccv_pyr_32_64_128_seq_2'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/local-data/test_ready'
+
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+config['bottleneck_scale'] = 2
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = False
+config['pyramid']['noise_amount'] = 0.0
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+config['pyramid']['static_model'] = 'checkpoints/iccv/static_iccv_1'
+config['pyramid']['body_model'] = 'checkpoints/iccv/body_iccv_1'
+config['pyramid']['curvature_model'] = 'checkpoints/iccv/curv_iccv_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 3, 9, 3]
+    level['encoder_dims'] = [16, 32, 64, 128, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [64, 128, 128, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_256_2'#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [128, 128, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_128_2'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_64_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [16]
+    level['encoder_dims'] = [level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_32_2'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 27, 3]
+config['cond_ae_static']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/iccv/config_abl_pyr_32_64_128_seq.py

@@ -0,0 +1,187 @@
+config = dict()
+
+config['gpu'] = '1'
+config['name'] = 'iccv_pyr_32_64_128_seq_2'
+config['notes'] = ''
+config['tags'] = ['pyr', 'iccv']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+config['use_curvature'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64, 128]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/iccv/iccv_pyr_32_seq_2'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/local-data/test_ready'
+
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+config['bottleneck_scale'] = 2
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = False
+config['pyramid']['noise_amount'] = 0.0
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+config['pyramid']['static_model'] = 'checkpoints/iccv/static_iccv_1'
+config['pyramid']['body_model'] = 'checkpoints/iccv/body_iccv_1'
+config['pyramid']['curvature_model'] = 'checkpoints/iccv/curv_iccv_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 3, 9, 3]
+    level['encoder_dims'] = [16, 32, 64, 128, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [64, 128, 128, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_256_2'#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [128, 128, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_128_2'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_64_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [16]
+    level['encoder_dims'] = [level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_32_2'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 27, 3]
+config['cond_ae_static']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/iccv/config_abl_pyr_32_seq.py

@@ -0,0 +1,187 @@
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'iccv_pyr_32_seq_2'
+config['notes'] = ''
+config['tags'] = ['pyr', 'iccv']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+config['use_curvature'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/iccv/iccv_pyr_32'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/local-data/test_ready'
+
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+config['bottleneck_scale'] = 2
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = False
+config['pyramid']['noise_amount'] = 0.0
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+config['pyramid']['static_model'] = 'checkpoints/iccv/static_iccv_1'
+config['pyramid']['body_model'] = 'checkpoints/iccv/body_iccv_1'
+config['pyramid']['curvature_model'] = 'checkpoints/iccv/curv_iccv_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 3, 9, 3]
+    level['encoder_dims'] = [16, 32, 64, 128, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [64, 128, 128, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_256_2'#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [128, 128, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_128_2'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_64_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = 64 * config['bottleneck_scale']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [16]
+    level['encoder_dims'] = [level['latent_dim']]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/iccv/iccv_32_2'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 27, 3]
+config['cond_ae_static']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

evaluate_ae_pyramid.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.iccv.config_abl_pyr_32_64_128_skin_skirt import config
+from configs.iccv.config_abl_pyr_32_64_128_seq import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf

get_error.py

@@ -8,11 +8,11 @@ if __name__ == '__main__':
     garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
     folder = '/home/workspace/data/results/errors/iccv'
 
-    checkpoint_name = 'iccv_pyr_32_64_128_skin'
+    checkpoint_name = 'iccv_pyr_32_64_128_256_seq_2'
     checkpoint_name_skirt = 'iccv_pyr_32_64_128_skin_skirt'
-    iteration = '75' + '000'
+    iteration = '185' + '000'
     iteration_skirt : str = '155' + '000'
-    error_type = 'collision' # 'euclidean' 'collision' edge normal
+    error_type = 'euclidean' # 'euclidean' 'collision' edge normal
 
 
     errors = []
@@ -21,31 +21,31 @@ if __name__ == '__main__':
     per_garment_errors_skirt  = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
 
 
-    # with open(os.path.join(folder, f'{checkpoint_name}-{iteration}-{error_type}.csv'), 'r') as f:
-    #     csv_reader = csv.reader(f, delimiter=',')
-    #     for row in csv_reader:
-    #         per_garment_errors[garment_type_table[row[0].split('_')[0]]].append(float(row[1]))
-
-    # print('\nNon-skirt:')
-    # for key in per_garment_errors.keys():
-    #     avg = np.mean(per_garment_errors[key])
-    #     errors.append(avg)
-    #     print(key + ': ' + str(avg))
-
-    # print('\nAvg non-skirt: ' + str(np.mean(errors)))
-
-    with open(os.path.join(folder, f'{checkpoint_name_skirt}-{iteration_skirt}-{error_type}.csv'), 'r') as f:
+    with open(os.path.join(folder, f'{checkpoint_name}-{iteration}-{error_type}.csv'), 'r') as f:
         csv_reader = csv.reader(f, delimiter=',')
         for row in csv_reader:
-            per_garment_errors_skirt[garment_type_table[row[0].split('_')[0]]].append(float(row[1]))
+            per_garment_errors[garment_type_table[row[0].split('_')[0]]].append(float(row[1]))
 
-    print('\nSkirt:')
-    for key in per_garment_errors_skirt.keys():
-        avg = np.mean(per_garment_errors_skirt[key])
-        errors_skirt.append(avg)
+    print('\nNon-skirt:')
+    for key in per_garment_errors.keys():
+        avg = np.mean(per_garment_errors[key])
+        errors.append(avg)
         print(key + ': ' + str(avg))
 
-    print('Avg skirt: ' + str(np.mean(errors_skirt)))
+    print('\nAvg non-skirt: ' + str(np.mean(errors)))
+
+    # with open(os.path.join(folder, f'{checkpoint_name_skirt}-{iteration_skirt}-{error_type}.csv'), 'r') as f:
+    #     csv_reader = csv.reader(f, delimiter=',')
+    #     for row in csv_reader:
+    #         per_garment_errors_skirt[garment_type_table[row[0].split('_')[0]]].append(float(row[1]))
+
+    # print('\nSkirt:')
+    # for key in per_garment_errors_skirt.keys():
+    #     avg = np.mean(per_garment_errors_skirt[key])
+    #     errors_skirt.append(avg)
+    #     print(key + ': ' + str(avg))
+
+    # print('Avg skirt: ' + str(np.mean(errors_skirt)))
 
     # cum_avg_shirt = np.mean(per_garment_errors['Tshirt'] + per_garment_errors_skirt['Tshirt'])
     # cum_avg_top = np.mean(per_garment_errors['Top'] + per_garment_errors_skirt['Top'])

models/pyramid_level.py

@@ -155,15 +155,18 @@ class PyramidLevelNew(tf.keras.Model):
         return out
 
 
-    def call(self, inputs, z_cond, training=None):
+    def call(self, inputs, z_cond, prev_latent=None, training=None):
         uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
 
         x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
         mean, logvar = self.encoder(x, training=training)
         z = self.reparameterize(mean, logvar, training)
-        z = self.reshape(z)
-
-        z = tf.concat([z, z_cond], axis=-1)
+        z_latent = self.reshape(z)
+    
+        if prev_latent is not None:
+            z_latent = tf.concat([z_latent, prev_latent], axis=-1)
+            
+        z = tf.concat([z_latent, z_cond], axis=-1)
         x = self.decoder(z, training=training)
 
         if self.avg_garment is not None:
@@ -177,4 +180,4 @@ class PyramidLevelNew(tf.keras.Model):
         
         # x = x * uv_mask_high[:,:,:,tf.newaxis]
 
-        return x, mean, logvar
+        return x, mean, logvar, z_latent

models/pyramid_model_optimized.py

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