update cvae

cvae.py

@@ -21,7 +21,7 @@ import wandb
 import random
 import numpy as np
 import scipy as sp
-from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
@@ -44,7 +44,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='cvae', name=config['name'], notes=config['notes'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
@@ -85,71 +85,66 @@ def train_step(batch):
         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_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 = edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])
+        # is the output already masked?
-        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'] - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
+        e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-        l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
+        c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], 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)
+        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.image.ssim_multiscale(pred + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
+        kl = kl_loss(mean, logvar)
-        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
 
         reconstr_loss = l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']
-        reconstr_loss = tf.reduce_mean(reconstr_loss)
+        loss = reconstr_loss + g_balancing_coeff * kl
-        kl = kl_loss(mean, logvar)
+
-        loss = reconstr_loss + balancing_coeff * kl
-        ms_ssim = tf.reduce_mean(ms_ssim)
-        l1_loss = tf.reduce_mean(l1_loss)
-        l2_loss = tf.reduce_mean(l2_loss)
-        e_loss = tf.reduce_mean(e_loss)
-        c_loss = tf.reduce_mean(c_loss)
         mean = tf.reduce_mean(mean)
         logvar = tf.reduce_mean(logvar)
 
+    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)))
     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'])
     gradients = tape.gradient(loss, model.trainable_variables)
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
              'train/reconstr_loss': reconstr_loss,
-             'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
+             'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
-             'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
+             'train/l2_loss': l2_loss }
         
 
 
 @tf.function
 def valid_step(batch):
-    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
-    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
 
     mesh_pred = uv2mesh(pred)
     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']))
+    shape = batch['pose_shape'][:, 72:]
-    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])
+    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-    b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask'])
+    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])])
-    # 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)
+    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-    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_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
-    # 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'])
+    # is the output already masked?
-    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'])
+
-  
+    e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    reconstr_loss = l1_loss
+    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]))
     kl = kl_loss(mean, logvar)
-    # loss = reconstr_loss / (2 * (gamma ** 2)) + kl
+
-    loss = reconstr_loss + balancing_coeff * kl
+    reconstr_loss = l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']
-    ms_ssim = tf.reduce_mean(ms_ssim)
+    loss = reconstr_loss + g_balancing_coeff * kl
-    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)
     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,
+    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)))
-             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss,
+    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'])
+
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
+             'train/reconstr_loss': reconstr_loss,
              'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar }
 
 
@@ -315,6 +310,57 @@ def sample(folder):
         print()
 
 
+
+
+
+def test():
+    load_model(model, 'checkpoints/' + config['name'].split('_')[0])
+
+    dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
+    static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
+   
+    step = 0
+    losses = []
+
+    for batch in dataset_test:
+        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+
+        test_loss = valid_step(batch)
+        losses.append(test_loss)
+
+        to_log = { **{ 'current_step': step }, **test_loss }
+        if step % 100 == 0:
+            figs = plot_everything(batch, 'test')
+            to_log = { **to_log, **figs }
+
+        wandb.log(to_log, step = step)
+
+        if step > 30:
+            break
+        
+        print(step)
+        step += 1
+   
+    # averages = dict()
+    path_to_plotly_html = "temp/plotly_figure.html"
+
+    histograms = dict()
+    for key in test_loss.keys():
+        accumulated = [losses[i][key].numpy() for i in range(len(losses))]
+        fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
+        fig.write_html(path_to_plotly_html, auto_play = False)
+
+        histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
+        # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
+
+    
+
+    wandb.log({ **histograms }, step = step)
+    wandb.finish()
+    pass
+
+
+
 if __name__ == '__main__':
     init_wandb()
     allow_memory_growth()
@@ -327,55 +373,55 @@ 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']])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
-    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'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
 
     model = CVAE(config)
 
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
-    balancing_coeff = config['starting_balancing_coeff']
-    reg_coeff = 1
-    valid_iterator = iter(dataset_valid)
-    gamma = 1
     balancing_activated_step = 0
+    g_balancing_coeff = config['starting_balancing_coeff']
 
     # sample('checkpoints/cvae_13')
+    test()
 
-    ckpt = tf.train.Checkpoint(model=model)
+    # ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
+    # manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
 
-    step = 0
+    # valid_iterator = iter(dataset_valid)
-    for batch in dataset_train:
+    # step = 0
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+    # for batch in dataset_train:
+    #     batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
-        train_loss = train_step(batch)
+    #     train_loss = train_step(batch)
 
-        # data_valid = valid_iterator.get_next()
+    #     data_valid = valid_iterator.get_next()
-        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+    #     data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
-        # valid_loss = valid_step(data_valid)
+    #     valid_loss = valid_step(data_valid)
-        valid_loss = {}
 
-        to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
+    #     to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
-        if step % 2500 == 0:
+    #     if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
-            train_figs = plot_everything(batch, 'train')
+    #         train_figs = plot_everything(batch, 'train')
-            # valid_figs = plot_everything(data_valid, 'valid')
+    #         valid_figs = plot_everything(data_valid, 'valid')
-            valid_figs = {}
+    #         to_log = { **to_log, **train_figs, **valid_figs }
-            to_log = { **to_log, **train_figs, **valid_figs }
 
-        if step % 10000 == 0:
+    #     if step % 10000 == 0:
-            manager.save(checkpoint_number=step)
+    #         manager.save(checkpoint_number=step)
+    #     if step == 0:
+    #         to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        if train_loss['train/euclid_error'] < 0.01:
+    #     if train_loss['train/euclid_error'] < 0.01:
-            if balancing_activated_step == 0:
+    #         if balancing_activated_step == 0:
-                balancing_activated_step = step
+    #             balancing_activated_step = step
-            balancing_coeff = balancing_coeff + 1 / 15000      
+    #         g_balancing_coeff = g_balancing_coeff + 1 / 15000      
 
-        to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
+    #     to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
         
-        wandb.log(to_log, step = step)
+    #     wandb.log(to_log, step = step)
 
-        print(step)
+    #     print(step)
 
-        step += 1
+    #     step += 1
+