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
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