small updates, dataset, low dim cvae

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'],
+    run = wandb.init(project='vae', name=config['name'], notes=config['notes'])
-            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\
-    ')
     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'])

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'],
+    run = wandb.init(project='vae', name=config['name'], notes=config['notes'])
-            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()

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']))
+        e_loss = 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']))
+        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)

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

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