diff --git a/cvae_pyramid.py b/cvae_pyramid.py index 024742b..81b8717 100644 --- a/cvae_pyramid.py +++ b/cvae_pyramid.py @@ -43,7 +43,7 @@ from models.pyramid_model import Pyramid def init_wandb(): - run = wandb.init(project='vae', name=config['name'], notes=config['notes']) + run = wandb.init(project='cvae', name=config['name'], notes=config['notes']) artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae_pyramid.py') @@ -111,6 +111,36 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col ), row=1, col=1+len(mesh_pred_unposed_levels)+1+1) + if len(mesh_pred_unposed_levels) == 1: + 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), + ) + 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, + ) + if len(mesh_pred_unposed_levels) == 2: + 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), + ) + 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, + ) if len(mesh_pred_unposed_levels) == 3: 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), @@ -129,22 +159,6 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col scene5 = helper, scene6 = helper, ) - if len(mesh_pred_unposed_levels) == 2: - 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), - ) - 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, - ) 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), @@ -230,8 +244,8 @@ def train_step(batch): 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 + 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)) + # 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) reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']) @@ -257,41 +271,54 @@ def train_step(batch): @tf.function def valid_step(batch): - uv_curvature = tf.one_hot(batch['uv_curvature'], 7) - pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']]) + levels_input, conds = group_inputs_for_levels(batch) + out_levels = model([levels_input, conds], training=False) - mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) + shape = batch['pose_shape'][:, 72:] + gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) + shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) mesh_orig = batch['mesh_unposed'] - e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])) - l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2]) - b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask']) - # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2) - ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3) - fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3]) - # mesh_orig_2 = uv2mesh(batch['uv_unposed']) - metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) - metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) - - reconstr_loss = l1_loss - # logpx_z = -reconstr_loss - # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.) - # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar) - # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x)) + + output = dict() + reconstr_loss_total = 0 + kl_total = 0 + + #for i, res in enumerate(config['pyramid']['resolutions']): + i = -1 + res = config['pyramid']['resolutions'][-1] + res_modifier = '_' + str(res) if res != 512 else '' + pred, mean, logvar = out_levels[i] + + mesh_pred = uv2mesh(pred) + 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 + 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) - # loss = reconstr_loss / (2 * (gamma ** 2)) + kl - loss = reconstr_loss + balancing_coeff * kl - ms_ssim = tf.reduce_mean(ms_ssim) - l1_loss = tf.reduce_mean(l1_loss) - fft_loss = tf.reduce_mean(fft_loss) - b_loss = tf.reduce_mean(b_loss) - reconstr_loss = tf.reduce_mean(reconstr_loss) + + reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']) + # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']) + kl_total += kl + mean = tf.reduce_mean(mean) logvar = tf.reduce_mean(logvar) - return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, - 'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss, - 'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar } + metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + + output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, + 'valid/ms_ssim' + res_modifier: ms_ssim, 'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl, + 'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss, + 'valid/euclid_error': metric } } + + loss = reconstr_loss_total + balancing_coeff * kl_total + + return { **{ 'valid/loss': loss, 'valid/reconstr_loss_total': reconstr_loss_total, 'valid/kl_total': kl_total }, **output } def plot_everything(batch: dict, type: str): @@ -484,7 +511,7 @@ if __name__ == '__main__': dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory']) static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions) dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory']) - static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static']) + static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions) model: Pyramid = Pyramid(config) @@ -511,12 +538,12 @@ if __name__ == '__main__': # 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 + # might not 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.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 @@ -532,16 +559,14 @@ if __name__ == '__main__': train_loss = train_step(batch) - # data_valid = valid_iterator.get_next() - # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic)) - # valid_loss = valid_step(data_valid) - valid_loss = {} + data_valid = valid_iterator.get_next() + data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic'])) + valid_loss = valid_step(data_valid) to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } if step % 5000 == 0 or (step < 5000 and step % 500 == 0): train_figs = plot_everything(batch, 'train') - # valid_figs = plot_everything(data_valid, 'valid') - valid_figs = {} + valid_figs = plot_everything(data_valid, 'valid') to_log = { **to_log, **train_figs, **valid_figs } if step % 10000 == 0: