luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update pyramid

Browse Source
This commit is contained in:
luni
2022-10-07 21:58:18 +00:00
parent 9d34c540d9
commit 01c8c3dc27
3 changed files with 179 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/config_pyramid.py
+30 -12
View File
@@ -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:
cvae_pyramid.py
+51 -5
View File
@@ -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
models/pyramid_model.py
+98 -2
View File
@@ -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