try less gnn

aaa_helpers.txt

@@ -3,14 +3,12 @@ nohup python calculate_error.py --name error_al_all_150k > nohup/error/post_al_a
 nohup python auto_encoder_static.py --name="lhunor_static_iccv_1" > nohup/iccv/encoders/static_iccv_1.out 2>&1 &
 nohup python auto_encoder_static.py --name="lhunor_static_iccv_1_skirt" > nohup/iccv/encoders/static_iccv_1_skirt.out 2>&1 &
 nohup python auto_encoder_static.py --name="lhunor_static_iccv_1_cape_f" > nohup/iccv/encoders/static_iccv_1_cape_f.out 2>&1 &
-nohup python auto_encoder_static.py --name="lhunor_static_iccv_1_cape_m" > nohup/iccv/encoders/static_iccv_1_cape_m.out 2>&1 &
 nohup python auto_encoder_body.py --name="lhunor_body_iccv_1" > nohup/iccv/encoders/body_iccv_1.out 2>&1 &
 nohup python auto_encoder_body.py --name="lhunor_body_iccv_1_skirt" > nohup/iccv/encoders/body_iccv_1_skirt.out 2>&1 &
-nohup python auto_encoder_body.py --name="lhunor_body_iccv_1_skirt" > nohup/iccv/encoders/body_iccv_1_skirt.out 2>&1 &
 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_decoder_1" > nohup/iccv/cvae_iccv_gnn_decoder_1.out 2>&1 &
+nohup python cvae.py --name="lhunor_cvae_iccv_gnn_5" > nohup/iccv/cvae_iccv_gnn_5.out 2>&1 &
 
 skirts!!!!
 pyramid next week

auto_encoder_static.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.iccv.config_ae_static_cape_m import config
+from configs.iccv.config_ae_static_cape_f import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
 import tensorflow as tf

configs/iccv/config_abl_all.py

@@ -19,14 +19,14 @@ 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'] = '4'
+config['gpu'] = '0'
-config['name'] = 'iccv_gnn_decoder_1'
+config['name'] = 'iccv_gnn_5'
 config['notes'] = ''
-config['tags'] = ['iccv', 'gnn', 'ggn_decoder']
+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
+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']:
+if config['use_skinning']:
-  # assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
 
 config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
 

configs/iccv/config_ae_static_cape_m.py

@@ -1,51 +0,0 @@
-config = dict()
-
-
-config['batch_size'] = 4
-config['train_dir'] = '/home/workspace/local-data/cape_processed/train/male'
-config['valid_dir'] = '/home/workspace/local-data/cape_processed/test/male'
-config['test_dir'] = '/home/workspace/local-data/cape_processed/test/male'
-config['use_filled'] = False
-config['resolution'] = 512
-config['is_skirt'] = False
-config['dataset_in_memory'] = True
-config['is_cape'] = True
-config['remove_global_orientation'] = True
-
-config['smpl_res'] = 14475
-if config['is_cape']:
-  config['smpl_res'] = 3638
-
-config['latent_dim'] = 128
-
-config['bottleneck_scale'] = 2
-
-config['cond_ae'] = dict()
-config['cond_ae']['is_label'] = False
-config['cond_ae']['upscale_output'] = False
-config['cond_ae']['latent_dim'] = config['latent_dim']
-config['cond_ae']['kernel_size'] = 7
-config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
-config['cond_ae']['encoder_depths'] = [3, 3, 3, 27, 3]
-config['cond_ae']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae']['bottleneck_size']]
-config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths']
-config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims']
-config['cond_ae']['activation'] = 'tanh'
-
-
-config['gpu'] = '5'
-config['lr'] = 0.0005
-config['name'] = 'static_iccv_1_cape_m'
-config['notes'] = ' ' 
-config['tags'] = ['iccv', 'static', 'cape_m'] 
-
-config['cond_ae']['upscale_output'] = False
-config['pretrained_model'] = ''
-
-
-config['signature_dynamic'] = { 'id': (), 'frame_nr': (), }
-config['signature_static'] = { 'uv_static': (512, 256, 3),
-                         'mesh_static': (config['smpl_res'], 3,),
-                         'mesh_mask': (config['smpl_res'],),
-                         'uv_mask': (512, 256,),
-                        }

cvae.py

@@ -88,7 +88,8 @@ 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 = pred
+        mesh_pred = uv2mesh(pred)
+        mesh_pred = smoother(mesh_pred, training=True)
 
         if config['use_skinning']:
             shape = batch['pose_shape'][:, 72:]
@@ -105,18 +106,18 @@ 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_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_mesh_loss
+        reconstr_loss = l2_loss
         if config['use_extra_losses']:
             if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True:
-                reconstr_loss = reconstr_loss + (e_loss) * config['reg_loss_coeff']
+                reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff']
             else:
-                reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+                reconstr_loss = reconstr_loss + l2_mesh_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']
 
@@ -127,12 +128,12 @@ def train_step(batch):
 
     # 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)
+    gradients = tape.gradient(loss, model.trainable_variables + smoother.trainable_variables)
-    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables + 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_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
+             'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
              # 'train/ms_ssim': ms_ssim, 
             }
 
@@ -146,7 +147,8 @@ 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 = pred
+    mesh_pred = uv2mesh(pred)
+    mesh_pred = smoother(mesh_pred, training=False)
 
     if config['use_skinning']:
         shape = batch['pose_shape'][:, 72:]
@@ -163,15 +165,15 @@ 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_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_mesh_loss
+    reconstr_loss = l2_loss
     if config['use_extra_losses']:
-        reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+        reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
     if config['use_norm_loss']:
         reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
 
@@ -256,12 +258,14 @@ 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 = pred.numpy()
+    mesh_pred = uv2mesh(pred)
+    mesh_pred = smoother(mesh_pred, training=False).numpy()
     
     orig_img = uv_orig
     uv_static_img = batch['uv_static']
     uv_body_posed_img = batch['uv_body_posed']
-    img = tf.concat((uv_static_img, uv_body_posed_img, orig_img), axis=2)
+    pred_img = pred
+    img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
     
     F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -428,6 +432,8 @@ def sample(folder):
 
 
 
+
+
 def test():
     str_split = config['name'].split('_')
     test_idx = str_split.index('test')
@@ -518,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/cvae_model.py

@@ -10,7 +10,7 @@ import scipy as sp
 from custom_utils.custom_utils import load_model
 from models.base import AEEncoder, VariationalEncoder, Decoder
 from models.base import AEEncoderNew, VariationalEncoderNew, DecoderNew
-from models.gnn import GNNDecoder
+
 
 class CVAE(tf.keras.Model):
     
@@ -150,6 +150,8 @@ class CVAE(tf.keras.Model):
         return x, mean, logvar
 
 
+
+
 class CVAENew(tf.keras.Model):
     
     def __init__(self, config):
@@ -178,8 +180,7 @@ class CVAENew(tf.keras.Model):
                 self.curv_enc.trainable = False
 
         self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
-        self.decoder = GNNDecoder()
+        self.decoder = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
-
 
         self.reshape = tf.keras.layers.Reshape((8, 4, 64))
 
@@ -272,13 +273,13 @@ class CVAENew(tf.keras.Model):
         z = tf.concat([z, z_cond], axis=-1)
         x = self.decoder(z, training=training)
 
-        # if self.avg_garment is not None:
+        if self.avg_garment is not None:
-        #     x = x + self.avg_garment[tf.newaxis,:,:,:]
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
     
-        # if self.use_offsets_over_body:
+        if self.use_offsets_over_body:
-        #     x = x + uv_body_posed
+            x = x + uv_body_posed
 
-        # x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,:,:,tf.newaxis]
 
         return x, mean, logvar
 

models/gnn.py

@@ -3,7 +3,7 @@ import sys
 
 import scipy.io as sio
 import tensorflow as tf
-from spektral.layers import GCNConv, GeneralConv, GATConv
+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,
@@ -16,61 +16,25 @@ 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.conv1 = GCNConv(3, activation='tanh', use_bias=False)
-        self.conv2 = GeneralConv(channels=128, batch_norm=True, dropout=0.1, aggregate='mean', activation='prelu', use_bias=True)
+        self.conv2 = GCNConv(3, activation='tanh', use_bias=False)
-        self.conv3 = GeneralConv(channels=64, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True)
+        # self.conv3 = GCNConv(3, activation='tanh', use_bias=False)
-        self.conv4 = GeneralConv(channels=16, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True)
+        # self.conv4 = GCNConv(64, activation='tanh', use_bias=False)
-
+        # self.conv5 = GCNConv(64, activation='tanh', use_bias=False)
-        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.conv3([x, self.a], training=training)
-        x = self.conv4([x, self.a], training=training)
+        # x = self.conv4([x, self.a], training=training)
-        x = self.conv5([x, self.a], training=training)
+        # x = self.conv5([x, self.a], training=training)
-        x = self.conv6([x, self.a], training=training)
-        return x
-
-import numpy as np
-import scipy.sparse as sp
-
-class GNNDecoder(tf.keras.Model):
-    def __init__(self, ):
-        super().__init__()
-        self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo())
-        # self.a_dense = sio.loadmat("custom_utils/ssmpl0.mat")['L'].todense()
-        # self.a_dense = np.ones((100,100), dtype=np.float32)
-        # self.a = sp_matrix_to_sp_tensor(sp.coo_matrix(self.a_dense))
-        
-        self.flatten = Flatten()
-        # self.conv1 = GATConv(channels=4, attn_heads=1, concat_heads=True, dropout_rate=0.0, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True)
-        # self.conv2 = GATConv(channels=4, attn_heads=1, concat_heads=True, dropout_rate=0.1, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True)
-        # self.conv3 = GATConv(channels=4, attn_heads=1, concat_heads=True, dropout_rate=0.1, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True)
-        self.dense1 = Dense(512)
-        self.dense2 = Dense(256)
-        self.conv1 = GCNConv(256, use_bias=True)
-        self.conv2 = GCNConv(256, use_bias=True)
-        self.conv3 = GeneralConv(channels=128, batch_norm=True, dropout=0.1, aggregate='sum', activation='prelu', use_bias=True)
-        self.conv4 = GeneralConv(channels=128, batch_norm=True, dropout=0.1, aggregate='sum', activation='prelu', use_bias=True)
-        self.conv5 = GeneralConv(channels=64, batch_norm=True, dropout=0.0, aggregate='sum', activation='prelu', use_bias=True)
-        self.conv6 = GCNConv(3, activation='tanh', use_bias=False)
-
-    def call(self, inputs, training=None):
-        x = inputs
-        x = self.flatten(x)
-        x = tf.tile(x[:,None,:], [1,14475,1])
-        x = self.dense1(x)
-        x = self.dense2(x)
-        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
 
 
 
+
+
+