diff --git a/configs/iccv/config_abl_pyr_32_64_128_seq.py b/configs/iccv/config_abl_pyr_32_64_128_seq.py new file mode 100644 index 0000000..8605270 --- /dev/null +++ b/configs/iccv/config_abl_pyr_32_64_128_seq.py @@ -0,0 +1,187 @@ +config = dict() + +config['gpu'] = '4' +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_127/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' \ No newline at end of file diff --git a/configs/iccv/config_abl_pyr_32_seq.py b/configs/iccv/config_abl_pyr_32_seq.py new file mode 100644 index 0000000..f8627bb --- /dev/null +++ b/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_127/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' \ No newline at end of file diff --git a/models/pyramid_level.py b/models/pyramid_level.py index d4508da..6a91197 100644 --- a/models/pyramid_level.py +++ b/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 \ No newline at end of file + return x, mean, logvar, z_latent \ No newline at end of file diff --git a/models/pyramid_model_optimized.py b/models/pyramid_model_optimized.py index 5019c7a..c4539f0 100644 --- a/models/pyramid_model_optimized.py +++ b/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)