update curvature CVAE

models/cvae_curv_model.py

@@ -122,21 +122,33 @@ class CVAE(tf.keras.Model):
         return x
 
 
-    def sample(self, uv_template, uv_body_posed, uv_mask, eps=None):
-        eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
-
-        uv_template = tf.transpose(uv_template, [0,2,3,1])
-        uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
+    def encode_conds(self, uv_template, uv_body_posed):
+        if uv_template.shape[1] == 3:
+            uv_template = tf.transpose(uv_template, [0,2,3,1])
+        if uv_body_posed.shape[1] == 3:
+            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
 
         z_static = self.static_enc(uv_template)
         z_body = self.body_enc(uv_body_posed)
 
         z_cond = tf.concat([z_static, z_body], axis=-1)
         z_cond = self.cond_encoder(z_cond)
+
+        return z_cond
+
+    def sample(self, uv_template, uv_body_posed, uv_mask=None, noise=None):
+        if noise is None:
+            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+        else:
+            eps = noise
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed)
         z = tf.concat([eps, z_cond], axis=-1)
 
         x = self.decoder(z)
-        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        if uv_mask is not None:
+            x = x * uv_mask[:,:,:,tf.newaxis]
 
         return x
 
@@ -168,17 +180,20 @@ class CVAE(tf.keras.Model):
 
 
     def call(self, inputs, training=None):
-        uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs
-        uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
-        z_static = self.static_enc(uv_template)
-        z_body = self.body_enc(uv_body_posed)
+        uv_template, uv_body_posed, uv_curv, uv_mask = inputs
+        if uv_template.shape[1] == 3:
+            uv_template = tf.transpose(uv_template, [0,2,3,1])
+        if uv_body_posed.shape[1] == 3:
+            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
+        if uv_curv.shape[1] == 3:
+            uv_curv = tf.transpose(uv_curv, [0,2,3,1])
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed)
 
         x = tf.concat([uv_template, uv_body_posed, uv_curv], axis=-1)
-
         mean, logvar = self.encoder(x)
         z = self.reparameterize(mean, logvar)
-        z_cond = tf.concat([z_static, z_body], axis=-1)
-        z_cond = self.cond_encoder(z_cond)
+
         z = tf.concat([z, z_cond], axis=-1)
         x = self.decoder(z)