progress save

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:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        # if self.avg_garment is not None:
+        #     x = x + self.avg_garment[tf.newaxis,:,:,:]
     
-        if self.use_offsets_over_body:
-            x = x + uv_body_posed
+        # if self.use_offsets_over_body:
+        #     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