diff --git a/configs/config_abl_pyr_32_skirt.py b/configs/config_abl_pyr_32_skirt.py index 7fec07d..3e129a9 100644 --- a/configs/config_abl_pyr_32_skirt.py +++ b/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' diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index 77508c1..f509156 100644 --- a/custom_utils/custom_utils.py +++ b/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 diff --git a/cvae_pyramid.py b/cvae_pyramid.py index 7594a89..af88b98 100644 --- a/cvae_pyramid.py +++ b/cvae_pyramid.py @@ -803,12 +803,14 @@ if __name__ == '__main__': if len(config['pyramid']['resolutions']) == 1: model.static_enc.trainable = False model.body_enc.trainable = False - model.curv_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 - model.curv_enc.trainable = False + if config['use_curvature']: + model.curv_enc.trainable = False model.cond_encoder.trainable = False # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer']) diff --git a/models/pyramid_model_optimized.py b/models/pyramid_model_optimized.py index 5301872..f540c43 100644 --- a/models/pyramid_model_optimized.py +++ b/models/pyramid_model_optimized.py @@ -24,15 +24,17 @@ class Pyramid(tf.keras.Model): 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['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') if config['pyramid']['freeze_encoders']: self.static_enc.trainable = False self.body_enc.trainable = False - self.curv_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'), tf.keras.layers.Dense(256, 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) - - z_cond = tf.concat([z_static, z_body, z_curv], axis=-1) + 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,7 +299,8 @@ 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'] != '': 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)) @@ -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) - - z_cond = tf.concat([z_static, z_body, z_curv], axis=-1) + 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