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