try smoothing with gnn

cvae.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_all_post import config
+from configs.iccv.config_abl_all import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -42,6 +42,7 @@ 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
 
 
 def init_wandb():
@@ -53,6 +54,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')
@@ -75,7 +77,7 @@ def kl_loss(mean, logvar):
 
 # 3d loss on posed !!!!!!!!!!!!
 
-# @tf.function
+@tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -87,6 +89,7 @@ def train_step(batch):
         pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
 
         mesh_pred = uv2mesh(pred)
+        mesh_pred = smoother(mesh_pred, training=True)
 
         if config['use_skinning']:
             shape = batch['pose_shape'][:, 72:]
@@ -145,6 +148,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:]
@@ -254,7 +258,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']
@@ -478,7 +483,7 @@ def test():
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
@@ -519,6 +524,7 @@ if __name__ == '__main__':
                                         is_cape=config['is_cape'])
 
     model = CVAE(config)
+    smoother = Smoother()
 
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
     balancing_activated_step = 0