From a5eff778df6a0331e417b628fc60f71fb2eedf51 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 27 Sep 2022 21:21:49 +0000 Subject: [PATCH] small updates, dataset, low dim cvae --- auto_encoder_body.py | 12 +++--------- cvae.py | 7 ++----- cvae_low.py | 16 +++++++++------- downscaling_test_curv.py | 2 +- losses.py | 13 +++++++++++++ 5 files changed, 28 insertions(+), 22 deletions(-) diff --git a/auto_encoder_body.py b/auto_encoder_body.py index b092832..b87d86a 100644 --- a/auto_encoder_body.py +++ b/auto_encoder_body.py @@ -21,13 +21,7 @@ from models.cond_encoder import AutoEncoder def init_wandb(): - run = wandb.init(project='vae', name=config['name'], - notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\ - flattened bottleneck with 128 latent_dim\ remove global orientation\ - fixed gender issue with deoriented bodies, again, aaand again\ - trained on 100 garments\ - remove multiplications with garment uv mask\ - ') + run = wandb.init(project='vae', name=config['name'], notes=config['notes']) artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_body.py') artifact.add_file('losses.py') @@ -158,9 +152,9 @@ if __name__ == '__main__': factory = DatasetFactory() - dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']]) + dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']]) - dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic']) + dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static']) model = AutoEncoder(config['cond_ae']) diff --git a/cvae.py b/cvae.py index 001bc8d..eaddc9d 100644 --- a/cvae.py +++ b/cvae.py @@ -29,7 +29,7 @@ from PIL import Image from utils import str_to_bool from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing from load_models import load_generator, load_discriminator -from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss +from losses import edge_loss, uv_mask_loss, bend_loss, kl_loss # from losses import collision_penalty as collision_penalty from losses import collision_loss as collision_penalty from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d @@ -39,8 +39,7 @@ from models.cvae_model import CVAE def init_wandb(): - run = wandb.init(project='vae', name=config['name'], - notes=config['notes']) + run = wandb.init(project='vae', name=config['name'], notes=config['notes']) artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae.py') artifact.add_file('models/cvae_model.py') @@ -322,8 +321,6 @@ if __name__ == '__main__': parser.add_argument('--name', type=str) args = vars(parser.parse_args()) - batch_size = config['batch_size'] - smpl = SMPL() factory = DatasetFactory() diff --git a/cvae_low.py b/cvae_low.py index 85e6405..7e18fbc 100644 --- a/cvae_low.py +++ b/cvae_low.py @@ -77,7 +77,7 @@ def train_step(batch): batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True) pred_enhanced = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False) - pred_full_res = tf.image.resize(pred_enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) + pred_full_res = tf.image.resize(pred_enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) * batch['uv_mask'][:,:,:,tf.newaxis] mesh_pred = uv2mesh(pred_full_res) mesh_orig = batch['mesh_unposed'] @@ -87,8 +87,8 @@ 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) - 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'])) + e_loss = edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']) + c_loss = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + res_modifier] - pred)), axis=-1), axis=[1, 2]) l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + res_modifier] - pred) ** 2), axis=-1), axis=[1, 2]) l2_loss_full = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full_res) ** 2), axis=-1), axis=[1, 2]) @@ -98,7 +98,7 @@ def train_step(batch): else: ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3, power_factors=(0.5, 0.3, 0.1, 0.1)) - reconstr_loss = l2_loss_full + reconstr_loss = l2_loss_full #+ e_loss + c_loss reconstr_loss = tf.reduce_mean(reconstr_loss) kl = kl_loss(mean, logvar) loss = reconstr_loss + balancing_coeff * kl @@ -107,6 +107,8 @@ def train_step(batch): l1_loss = tf.reduce_mean(l1_loss) l2_loss = tf.reduce_mean(l2_loss) l2_loss_full = tf.reduce_mean(l2_loss_full) + e_loss = tf.reduce_mean(e_loss) + c_loss = tf.reduce_mean(c_loss) mean = tf.reduce_mean(mean) logvar = tf.reduce_mean(logvar) @@ -215,7 +217,7 @@ def sample(folder): load_model(model,folder) count = 0 for batch in dataset_train: - batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) + batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) uv_curvature = tf.one_hot(batch['uv_curvature'], 7) nr_of_samples=2 @@ -371,9 +373,9 @@ if __name__ == '__main__': factory = DatasetFactory() - dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[512, config['resolution']]) + dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[512, config['resolution']], in_memory=config['dataset_in_memory']) static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[512, config['resolution']]) - dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic']) + dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[512, config['resolution']], in_memory=config['dataset_in_memory']) static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static']) model = CVAE(config) diff --git a/downscaling_test_curv.py b/downscaling_test_curv.py index 11e2fe2..f2bfb3a 100644 --- a/downscaling_test_curv.py +++ b/downscaling_test_curv.py @@ -104,7 +104,7 @@ if __name__ == '__main__': 'mesh_posed': (14475, 3,), # 'mesh_unposed': (14475, 3,), # 'uv_unposed': (3, 512, 256,) - 'uv_curvature': (512, 256), + 'uv_curvature': (512, 256, 7), 'mesh_curvature': (14475) } signature_static = { #'uv_static': (3, 512, 256,), diff --git a/losses.py b/losses.py index 020fc2f..8321fb9 100644 --- a/losses.py +++ b/losses.py @@ -258,6 +258,19 @@ def bend_loss(x, F, L, mask): # return loss +def log_normal_pdf(sample, mean, logvar, raxis=1): + log2pi = tf.math.log(2. * np.pi) + res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis) + return res + + +def kl_loss(mean, logvar): + kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar) + kl_loss = tf.reduce_mean(kl_loss) + kl_loss *= -0.5 + return kl_loss + + if __name__ == '__main__': V = tf.ones((4, 14475, 3)) B = tf.zeros((4, 14475, 3))