progress save

configs/iccv/config_abl_all.py

@@ -20,11 +20,11 @@ ablation_count = [config['use_skinning'], config['remove_global_orientation'], c
 assert ablation_count == 1 or ablation_count == 0
 
 config['gpu'] = '3'
-config['name'] = 'iccv_gnn_2'
+config['name'] = 'iccv_gnn_3'
 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
+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']:
   assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
 

cvae.py

@@ -89,7 +89,6 @@ 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:]
@@ -148,7 +147,6 @@ 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:]
@@ -259,7 +257,6 @@ 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)
-    mesh_pred = smoother(mesh_pred, training=False).numpy()
     
     orig_img = uv_orig
     uv_static_img = batch['uv_static']
@@ -432,8 +429,6 @@ def sample(folder):
 
 
 
-
-
 def test():
     str_split = config['name'].split('_')
     test_idx = str_split.index('test')
@@ -524,7 +519,6 @@ 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,8 +150,6 @@ class CVAE(tf.keras.Model):
         return x, mean, logvar
 
 
-
-
 class CVAENew(tf.keras.Model):
     
     def __init__(self, config):
@@ -180,7 +178,8 @@ 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 = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
+        self.decoder = GNNDecoder()
+
 
         self.reshape = tf.keras.layers.Reshape((8, 4, 64))
 
@@ -273,13 +272,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
+from spektral.layers import GCNConv, GeneralConv, GATConv
 from spektral.utils import sp_matrix_to_sp_tensor
 from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
                                      GlobalAveragePooling2D, Layer,
@@ -35,7 +35,32 @@ class Smoother(tf.keras.Model):
         return x
 
 
+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.flatten = Flatten()
+        self.conv1 = GATConv(channels=1024, attn_heads=4, concat_heads=True, dropout_rate=0.0, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True)
+        self.conv2 = GATConv(channels=512, attn_heads=4, concat_heads=True, dropout_rate=0.1, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True)
+        self.conv3 = GATConv(channels=256, attn_heads=4, concat_heads=True, dropout_rate=0.1, return_attn_coef=False, add_self_loops=True, activation=None, 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 = GeneralConv(channels=3, batch_norm=False, dropout=0.0, aggregate='sum', activation='tanh', use_bias=True)
+
+    def call(self, inputs, training=None):
+        x = inputs
+        x = self.flatten(x)
+        x = tf.tile(x[:,None,:], [1,14475,1])
+
+        x = self.conv1([x, self.a_dense], training=training)
+        x = self.conv2([x, self.a_dense], training=training)
+        x = self.conv3([x, self.a_dense], 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