fix

aaa_helpers.txt

@@ -11,8 +11,11 @@ nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1" > nohup/iccv/encod
 nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_skirt" > nohup/iccv/encoders/curv_iccv_1_skirt.out 2>&1 &
 nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_cape_f" > nohup/iccv/encoders/curv_iccv_1_cape_f.out 2>&1 &
 nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_cape_m" > nohup/iccv/encoders/curv_iccv_1_cape_m.out 2>&1 &
+nohup python cvae.py --name="lhunor_cvae_iccv_gnn_1" > nohup/iccv/cvae_iccv_gnn_1.out 2>&1 &
 
 nohup python cvae.py --name="lhunor_abl_all" > nohup/iccv/ablation/abl_all.out 2>&1 &
+nohup python cvae.py --name="lhunor_abl_all_gnn" > nohup/iccv/ablation/abl_all_gnn.out 2>&1 &
+nohup python cvae.py --name="lhunor_abl_all_gnn_3" > nohup/iccv/ablation/abl_all_gnn_3.out 2>&1 &
 
 nohup python upscale_test.py --name="iccv_32_2" > nohup/iccv/upscale/iccv_32_2.out 2>&1 &
 nohup python upscale_test.py --name="iccv_64_2" > nohup/iccv/upscale/iccv_64_2.out 2>&1 &

configs/iccv/config_abl_all.py

@@ -7,23 +7,24 @@ config = dict()
 # DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
 # DON'T FORGET TO CHECK AVG GARMENT USED
 
-config['use_skinning'] =              True # and deshaping
+config['use_skinning'] =              False # and deshaping
 config['remove_global_orientation'] = True
 config['use_curvature'] =             True
 config['use_extra_losses'] =          True # edge, collision, 3d mesh
 config['use_offsets'] =               True
 
 config['use_norm_loss'] =             True
+config['only_curvatures'] =           False
 
 ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
 assert ablation_count == 1 or ablation_count == 0
 
-config['gpu'] = '3'
-config['name'] = 'abl_all'
+config['gpu'] = '2'
+config['name'] = 'abl_all_gnn_3'
 config['notes'] = ''
-config['tags'] = ['iccv', 'abl_all']
+config['tags'] = ['iccv', 'abl_all_gnn']
 
-config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+config['avg_garment_path'] = '' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
 if config['use_skinning']:
   assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
 
@@ -43,35 +44,38 @@ config['static_model'] = 'checkpoints/iccv/static_iccv_1'
 config['body_model'] = 'checkpoints/iccv/body_iccv_1'
 config['curvature_model'] = 'checkpoints/iccv/curv_iccv_1'
 
+# config['static_model'] = 'checkpoints/post/static_post_2'
+# config['body_model'] = 'checkpoints/post/body_post_2'
+# config['curvature_model'] = 'checkpoints/post/curv_post_1'
+
 
 config['smpl_res'] = 14475
 if config['is_cape']:
   config['smpl_res'] = 3638
 
-config['latent_dim'] = 128
 
-config['pretrained_model'] = ''
+config['pretrained_model'] = 'checkpoints/iccv/abl_skin'
 
 ###############    
 config['lr'] = 0.0001#                 0.0001
 ###############
-config['starting_balancing_coeff'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
 
 config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                           'uv_body_posed': (512, 256, 3),
-                          'mesh_body_posed': (14475, 3,),
-                          'mesh_posed': (14475, 3,),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
                           'uv_curvature': (512, 256, 7),
                         }
 config['signature_static'] = { 'uv_static': (512, 256, 3),
-                              'mesh_mask': (14475,),
+                              'mesh_mask': (config['smpl_res'],),
                               'uv_mask': (512, 256,),
                               }
 
 if config['use_skinning']:
   config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
-  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+  config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,)
 else:
   config['signature_dynamic']['uv_posed'] = (512, 256, 3)
 

cvae.py

@@ -42,11 +42,12 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mes
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
 from models.cvae_model import CVAENew as CVAE
+from models.gnn import Smoother
 from custom_utils.custom_utils import Accumulator
 
 
 def init_wandb():
-    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
+    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
@@ -54,6 +55,7 @@ def init_wandb():
     # artifact.add_file('configs/config_base.py')
     # artifact.add_file('configs/config_cvae.py')
     artifact.add_file('models/base.py')
+    artifact.add_file('models/gnn.py')
     artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
@@ -85,9 +87,10 @@ def train_step(batch):
         else:
             mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
 
-        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
+        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
 
         mesh_pred = uv2mesh(pred)
+        mesh_pred = smoother(mesh_pred, training=True)
 
         if config['use_skinning']:
             shape = batch['pose_shape'][:, 72:]
@@ -104,34 +107,29 @@ def train_step(batch):
         e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
         c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
         c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
-        l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
         l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
         n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-        # conc_loss = concept_loss(pred, batch['uv_unposed'])
-        kl = kl_loss(mean, logvar)
 
-        reconstr_loss = l2_loss
+        reconstr_loss = l2_mesh_loss
         if config['use_extra_losses']:
             if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True:
-                reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff']
+                reconstr_loss = reconstr_loss + (e_loss) * config['reg_loss_coeff']
             else:
-                reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+                reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
         if config['use_norm_loss'] and step > 2000:
             reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
 
-        loss = reconstr_loss + g_balancing_coeff * kl
+        loss = reconstr_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))
+    gradients = tape.gradient(loss, smoother.trainable_variables)
+    optimizer.apply_gradients(zip(gradients, smoother.trainable_variables))
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
              'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
-             'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
-             'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
+             'train/mean': mean, 'train/logvar': logvar, 'train/var': np.e ** logvar,
+             'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
              # 'train/ms_ssim': ms_ssim, 
             }
 
@@ -146,6 +144,7 @@ def valid_step(batch):
     pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
 
     mesh_pred = uv2mesh(pred)
+    mesh_pred = smoother(mesh_pred, training=False)
 
     if config['use_skinning']:
         shape = batch['pose_shape'][:, 72:]
@@ -162,22 +161,16 @@ def valid_step(batch):
     e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
     c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
     c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
-    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
     l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
     n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-    # conc_loss = concept_loss(pred, batch['uv_unposed'])
-    kl = kl_loss(mean, logvar)
 
-    reconstr_loss = l2_loss
+    reconstr_loss = l2_mesh_loss
     if config['use_extra_losses']:
-        reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+        reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
     if config['use_norm_loss']:
         reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
 
-    loss = reconstr_loss + g_balancing_coeff * kl
-
-    mean = tf.reduce_mean(mean)
-    logvar = tf.reduce_mean(logvar)
+    loss = reconstr_loss
 
     # 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'])
@@ -185,7 +178,7 @@ def valid_step(batch):
 
     return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric, 
              'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/collision_ratio': c_ratio,
-             'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
+             'valid/mean': mean, 'valid/logvar': logvar,  'valid/var': np.e ** logvar,
               'valid/normal_loss': n_loss,# 'valid/ms_ssim': ms_ssim, 
               }
 
@@ -255,7 +248,8 @@ def plot_everything(batch: dict, type: str):
 
     pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
     
-    mesh_pred = uv2mesh(pred).numpy()
+    mesh_pred = uv2mesh(pred)
+    mesh_pred = smoother(mesh_pred, training=False).numpy()
     
     orig_img = uv_orig
     uv_static_img = batch['uv_static']
@@ -286,7 +280,6 @@ def plot_everything(batch: dict, type: str):
 
     gt_posed = [batch['mesh_posed'].numpy()[i, mask[i], :] for i in range(mask.shape[0])]
 
-
     body = batch['mesh_body_posed'].numpy()
     mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
 
@@ -520,6 +513,8 @@ if __name__ == '__main__':
                                         is_cape=config['is_cape'])
 
     model = CVAE(config)
+    model.trainable = False
+    smoother = Smoother()
 
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
     balancing_activated_step = 0
@@ -552,17 +547,6 @@ if __name__ == '__main__':
             valid_loss = valid_step(data_valid)
             valid_accumulator.update(valid_loss)
 
-        # 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')
-        #     to_log = { **to_log, **train_figs, **valid_figs }
-
-        # if step % 5000 == 0:
-        #     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'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
             if balancing_activated_step == 0:
                 balancing_activated_step = step
@@ -584,6 +568,9 @@ if __name__ == '__main__':
             train_accumulator.reset()
             valid_accumulator.reset()
 
+        if step % 5000 == 0:
+            manager.save(checkpoint_number=step)
+
         print(step)
 
         step += 1

models/gnn.py

@@ -0,0 +1,40 @@
+import os
+import sys
+
+import scipy.io as sio
+import tensorflow as tf
+from spektral.layers import GCNConv, GeneralConv
+from spektral.utils import sp_matrix_to_sp_tensor
+from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
+                                     GlobalAveragePooling2D, Layer,
+                                     LayerNormalization, Flatten, Reshape, UpSampling2D)
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+class Smoother(tf.keras.Model):
+    def __init__(self, ):
+        super().__init__()
+        self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo())
+        self.conv1 = GCNConv(16, activation='tanh', use_bias=False)
+        self.conv2 = GCNConv(3, activation='tanh', use_bias=False)
+        # self.conv3 = GCNConv(512, activation='tanh', use_bias=False)
+        # self.conv4 = GCNConv(512, activation='tanh', use_bias=False)
+        # self.conv5 = GCNConv(64, activation='tanh', use_bias=False)
+        self.conv6 = GCNConv(3, activation='tanh', use_bias=False)
+
+    def call(self, inputs, training=None):
+        x = inputs
+        x = self.conv1([x, self.a], training=training)
+        x = self.conv2([x, self.a], training=training)
+        # x = self.conv3([x, self.a], training=training)
+        # x = self.conv4([x, self.a], training=training)
+        # x = self.conv5([x, self.a], training=training)
+        x = self.conv6([x, self.a], training=training)
+        return x
+
+
+
+
+
+