add no curvature option

configs/config_abl_pyr_32_skirt.py

@@ -1,8 +1,8 @@
 config = dict()
 
-config['gpu'] = '3'
-config['name'] = 'pyr_abl_32_skirt_r1'
-config['notes'] = ''
+config['gpu'] = '6'
+config['name'] = 'pyr_abl_32_skirt_no_curv_r1'
+config['notes'] = 'review_1'
 
 # DON'T FORGET TO RESET SSIM LOSS USAGE
 # DON'T FORGET TO CHECK MODEL LOADING
@@ -18,6 +18,7 @@ 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_curvature'] =             False
 
 config['pyramid'] = dict()
 config['pyramid']['resolutions'] = [32]
@@ -33,6 +34,8 @@ config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2'
 config['batch_size'] = 4
 config['is_skirt'] = True
 config['is_cape'] = False
+config['frames_file'] = 'frames_skirt_10p.txt'
+config['frames_file_valid'] = 'frames_skirt_10p.txt'
 
 config['pyramid']['static_model'] = 'checkpoints/static_1_skirt'
 config['pyramid']['body_model'] = 'checkpoints/body_1_skirt_2_2'

custom_utils/custom_utils.py

@@ -311,7 +311,7 @@ def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False
 
 def get_summary_str(model):
     lines = []
-    model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True)
+    model.summary(print_fn=lines.append, line_length=150, expand_nested=True, show_trainable=True)
     # Add initial spaces to avoid markdown formatting in TensorBoard
     return lines
 

cvae_pyramid.py

@@ -803,11 +803,13 @@ if __name__ == '__main__':
     if len(config['pyramid']['resolutions']) == 1:
         model.static_enc.trainable = False
         model.body_enc.trainable = False
+        if config['use_curvature']:
             model.curv_enc.trainable = False
     else:
         load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else ''))
         model.static_enc.trainable = False
         model.body_enc.trainable = False
+        if config['use_curvature']:
             model.curv_enc.trainable = False
         model.cond_encoder.trainable = False
 

models/pyramid_model_optimized.py

@@ -25,6 +25,7 @@ class Pyramid(tf.keras.Model):
         if config['pyramid']['body_model'] != '':
             load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
         
+        if config['use_curvature']:
             self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
             if config['pyramid']['curvature_model'] != '':
                 load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
@@ -32,6 +33,7 @@ class Pyramid(tf.keras.Model):
         if config['pyramid']['freeze_encoders']:
             self.static_enc.trainable = False
             self.body_enc.trainable = False
+            if config['use_curvature']:
                 self.curv_enc.trainable = False
 
         self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
@@ -70,11 +72,12 @@ class Pyramid(tf.keras.Model):
         self.add_noise = config['pyramid']['add_noise']
         self.noise_amount = config['pyramid']['noise_amount']
         self.use_offsets = config['use_offsets']
+        self.use_curvature = config['use_curvature']
  
         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)
+            self.upscale_enhancer : tf.keras.Model = 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
 
@@ -85,9 +88,11 @@ class Pyramid(tf.keras.Model):
     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)
-
+        if self.use_curvature:
+            z_curv = self.curv_enc(uv_curv, training=training)
             z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        else:
+            z_cond = tf.concat([z_static, z_body], axis=-1)
         z_cond = self.cond_encoder(z_cond, training=training)
 
         return z_cond
@@ -294,6 +299,7 @@ class PyramidNew(tf.keras.Model):
         self.add_noise = config['pyramid']['add_noise']
         self.noise_amount = config['pyramid']['noise_amount']
         self.use_offsets = config['use_offsets']
+        self.use_curvature = config['use_curvature']
 
         self.upscale_enhancer = None
         if config['pyramid']['final_upscale_enhancer'] != '':
@@ -309,9 +315,11 @@ class PyramidNew(tf.keras.Model):
     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)
-
+        if self.use_curvature:
+            z_curv = self.curv_enc(uv_curv, training=training)
             z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        else:
+            z_cond = tf.concat([z_static, z_body], axis=-1)
 
         return z_cond