try smoothing with gnn

aaa_helpers.txt

@@ -8,6 +8,7 @@ nohup python auto_encoder_body.py --name="lhunor_body_iccv_1_skirt" > nohup/iccv
 nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1 " > nohup/iccv/encoders/curv_iccv_1.out 2>&1 &
 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 cvae.py --name="lhunor_cvae_iccv_gnn_1" > nohup/iccv/cvae_iccv_gnn_1.out 2>&1 &
 
 skirts!!!!
 pyramid next week

configs/iccv/config_abl_all.py

@@ -14,13 +14,15 @@ 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'] = '0'
+config['gpu'] = '3'
-config['name'] = 'abl_all'
+config['name'] = 'iccv_gnn_2'
 config['notes'] = ''
+config['tags'] = ['iccv', '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
 if config['use_skinning']:
@@ -38,39 +40,42 @@ config['train_dir'] = '/home/workspace/local-data/train'
 config['valid_dir'] = '/home/workspace/local-data/valid'
 config['test_dir'] = '/home/workspace/data/dataset/test_ready'
 
-config['static_model'] = 'checkpoints/static_1'
+# config['frames_file'] = 'frames_dress_30k.txt'
-config['body_model'] = 'checkpoints/body_1'
+# config['frames_file_valid'] = 'frames_dress_30k.txt'
-config['curvature_model'] = 'checkpoints/curv_1'
+# config['frames_file_test'] = 'frames_dress_30k.txt'
+
+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['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_body_posed': (config['smpl_res'], 3,),
-                          'mesh_posed': (14475, 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)
 
@@ -81,37 +86,42 @@ config['use_one_hot'] = False
 
 config['cvae'] = dict()
 config['cvae']['is_label'] = False
-config['cvae']['latent_dim'] = config['latent_dim']
 config['cvae']['add_noise'] = True
-config['cvae']['noise_amount'] = 0.025
+config['cvae']['noise_amount'] = 0.001
 
 config['cvae']['kernel_size'] = 7
-config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3]
-config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256, 64]
-config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = config['cvae']['encoder_depths']
-config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = config['cvae']['encoder_dims']
-config['cvae']['decoder_dims'] = [256, 128, 64, 16]
 
 
-config['cond_ae'] = dict()
+config['cond_ae_static'] = dict()
-config['cond_ae']['is_label'] = False
+config['cond_ae_static']['is_label'] = False
-config['cond_ae']['upscale_output'] = False
+config['cond_ae_static']['kernel_size'] = 7
-config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['bottleneck_size'] = 8
-config['cond_ae']['kernel_size'] = 7
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
-config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
-config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
-config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
-config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae_static']['activation'] = 'tanh'
-config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+
-config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
 
 config['cond_ae_body'] = dict()
-config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['is_label'] = False
 config['cond_ae_body']['kernel_size'] = 7
-config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 16
-config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
-config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
-config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
-config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
-config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['activation'] = 'tanh'
-config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

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

models/gnn.py

@@ -0,0 +1,41 @@
+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 = GeneralConv(channels=32, batch_norm=True, dropout=0.1, aggregate='mean', activation='prelu', use_bias=True)
+        self.conv2 = GeneralConv(channels=128, batch_norm=True, dropout=0.1, aggregate='mean', activation='prelu', use_bias=True)
+        self.conv3 = GeneralConv(channels=64, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True)
+        self.conv4 = GeneralConv(channels=16, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True)
+
+        self.conv5 = GCNConv(3, 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
+
+
+
+
+
+