diff --git a/configs/config_pyramid.py b/configs/config_pyramid.py index 10ea014..1e35434 100644 --- a/configs/config_pyramid.py +++ b/configs/config_pyramid.py @@ -1,7 +1,7 @@ from configs.config_base import config -config['gpu'] = '2' -config['name'] = 'cvae_pyr_6_32_64_128_3_x' +config['gpu'] = '1' +config['name'] = 'cvae_pyr_7_32_64_128_256_512' config['notes'] = '\ offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. | \ retrained body encoder with deoriented bodies | \ @@ -14,20 +14,24 @@ common condition encoders for all levels, only run once | \ full res loss on each level | \ all higher levels are added to the output of first level | \ using custom MLP upscale enhancement | \ -layers 32, 64, 128 | \ +layers 32 | \ testing plotting | \ only train highest level, all other use from pretrained pyramids with one less level | \ maybe upscale enhancer was not frozen? --> fixed pretrained approach | \ upscale from level to level, not to final res, only at the end to final res | \ updated upscale enhancers | \ -try an old trained model as pretraining for new level | \ +try with l2 loss | \ \ ' # 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 -config['dataset_in_memory'] = True + +config['dataset_in_memory'] = False config['resolution'] = None config['lr'] = 0.0001 @@ -52,9 +56,9 @@ config['pyramid'] = dict() config['pyramid']['is_label'] = False config['pyramid']['latent_dim'] = config['latent_dim'] -config['pyramid']['resolutions'] = [32, 64, 128] -config['pyramid']['loss_weights'] = [1,1,1] -config['pyramid']['pretrained_pyramid'] = 'checkpoints/cvae_pyr_6_32_64' +config['pyramid']['resolutions'] = [32, 64, 128, 256, 512] +config['pyramid']['loss_weights'] = [1,1,1,1,1] +config['pyramid']['pretrained_pyramid'] = 'checkpoints/cvae_pyr_7_32_64_128_256_50k' for res in config['pyramid']['resolutions']: if res != 512: @@ -75,7 +79,7 @@ config['pyramid']['curvature_model'] = '' config['pyramid']['freeze_encoders'] = True config['avg_garment_path'] = '' -config['pyramid']['final_upscale_enhancer'] = 'checkpoints/up_3_d128_256' +config['pyramid']['final_upscale_enhancer'] = '' config['pyramid']['levels'] = [] @@ -93,6 +97,20 @@ for res in config['pyramid']['resolutions']: level['avg_garment_path'] = '' level['upscale_enhancer'] = '' config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['latent_dim'] = config['latent_dim'] + level['kernel_size'] = 7 + level['encoder_depths'] = [9, 9, 9, 9] + level['encoder_dims'] = [32, 64, 128, 256] + level['decoder_input'] = (8, 4, 128) + level['decoder_depths'] = [9, 9, 9, 9] + level['decoder_dims'] = [256, 128, 64, 16] + level['avg_garment_path'] = '' + level['upscale_enhancer'] = 'checkpoints/up_3_d256_512_2' + level['upscale_enhancer_freeze'] = True + config['pyramid']['levels'].append(level) elif res == 128: level = dict() level['resolution'] = 128 @@ -104,7 +122,7 @@ for res in config['pyramid']['resolutions']: level['decoder_depths'] = [16, 16, 9, 9] level['decoder_dims'] = [256, 256, 128, 64] level['avg_garment_path'] = '' - level['upscale_enhancer'] = '' + level['upscale_enhancer'] = 'checkpoints/up_3_d128_256_2' level['upscale_enhancer_freeze'] = True config['pyramid']['levels'].append(level) elif res == 64: @@ -118,7 +136,7 @@ for res in config['pyramid']['resolutions']: level['decoder_depths'] = [27, 27, 16] level['decoder_dims'] = [256, 256, 256] level['avg_garment_path'] = '' - level['upscale_enhancer'] = 'checkpoints/up_3_d64_128' + level['upscale_enhancer'] = 'checkpoints/up_3_d64_128_2' level['upscale_enhancer_freeze'] = False config['pyramid']['levels'].append(level) elif res == 32: @@ -132,7 +150,7 @@ for res in config['pyramid']['resolutions']: level['decoder_depths'] = [27, 27] level['decoder_dims'] = [256, 256] level['avg_garment_path'] = '' - level['upscale_enhancer'] = 'checkpoints/up_3_d32_64' + level['upscale_enhancer'] = 'checkpoints/up_3_d32_64_2' level['upscale_enhancer_freeze'] = True config['pyramid']['levels'].append(level) else: diff --git a/cvae_pyramid.py b/cvae_pyramid.py index 5dfa1fd..a7fc91c 100644 --- a/cvae_pyramid.py +++ b/cvae_pyramid.py @@ -145,6 +145,48 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col scene4 = helper, scene5 = helper, ) + if len(mesh_pred_unposed_levels) == 4: + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), + scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), + scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1), + scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), + scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1), + scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,), + scene7_aspectmode='manual', scene7_aspectratio=dict(x=1, y=1, z=1,), + ) + helper = dict(xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,1],), + zaxis = dict(nticks=4, range=[-1,1],),) + fig.update_layout(scene = helper, + scene2 = helper, + scene3 = helper, + scene4 = helper, + scene5 = helper, + scene6 = helper, + scene7 = helper, + ) + if len(mesh_pred_unposed_levels) == 5: + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), + scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), + scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1), + scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), + scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1), + scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,), + scene7_aspectmode='manual', scene7_aspectratio=dict(x=1, y=1, z=1,), + scene8_aspectmode='manual', scene8_aspectratio=dict(x=1, y=1, z=1,), + ) + helper = dict(xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,1],), + zaxis = dict(nticks=4, range=[-1,1],),) + fig.update_layout(scene = helper, + scene2 = helper, + scene3 = helper, + scene4 = helper, + scene5 = helper, + scene6 = helper, + scene7 = helper, + scene8 = helper, + ) return fig @@ -183,10 +225,12 @@ def train_step(batch): mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped) + # is the output already masked? + e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])) l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2])) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred - tf.reduce_min(pred), batch['uv_unposed'] - tf.reduce_min(pred), max_val=tf.reduce_max(pred) - tf.reduce_min(pred), power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) #ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3)) kl = kl_loss(mean, logvar) @@ -462,15 +506,17 @@ if __name__ == '__main__': model.curv_enc.trainable = False model.cond_encoder.trainable = False + + # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer']) + # model.levels[0].upscale_enhancer.trainable = False + # the last level doesn't have an upscale enhancer as it's enhanced for the final resolutin directly # might now be necessary, as enhancing to final or a bit higher resolution seems to be the same for i in range(len(model.levels)-1): load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer']) model.levels[i].upscale_enhancer.trainable = False - load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer']) - model.upscale_enhancer.trainable = False - - load_model(model.levels[-1], 'checkpoints/cvae_low_2_d128_cont') + # load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer']) + # model.upscale_enhancer.trainable = False for i in range(len(model.levels)-1): model.levels[i].trainable = False diff --git a/models/pyramid_model.py b/models/pyramid_model.py index bfa95a9..0885f03 100644 --- a/models/pyramid_model.py +++ b/models/pyramid_model.py @@ -70,13 +70,15 @@ class Pyramid(tf.keras.Model): out_levels = [] pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training) - pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res 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)): @@ -91,10 +93,104 @@ class Pyramid(tf.keras.Model): 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.append([pred, mean, logvar]) return out_levels +# 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=None): +# 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=None): +# 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 \ No newline at end of file