prototype conditional pipeline

stylegan2/layers/mapping_block.py

@@ -1,6 +1,6 @@
 import tensorflow as tf
 
-from stylegan2.layers.embedding import LabelEmbedding
+from stylegan2.layers.embedding import PoseShapeEmbedding, GarmentEmbedding
 from stylegan2.layers.dense import Dense
 from stylegan2.layers.bias_act import BiasAct
 
@@ -14,10 +14,8 @@ class Mapping(tf.keras.layers.Layer):
         self.gain = 1.0
         self.lrmul = 0.01
 
-        if self.labels_dim != self.w_dim:
-            self.labels_embedding = LabelEmbedding(embed_dim=self.w_dim, name='labels_embedding')
-        else:
-            self.labels_embedding = tf.keras.layers.Lambda(lambda x: x)
+        self.pose_shape_embedding = PoseShapeEmbedding()
+        self.garment_embedding = GarmentEmbedding((self.w_dim * 2) - (216 + 11))
 
         self.normalize = tf.keras.layers.Lambda(lambda x: x * tf.math.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8))
 
@@ -28,13 +26,13 @@ class Mapping(tf.keras.layers.Layer):
             self.bias_act_layers.append(BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_{:d}'.format(ii)))
 
     def call(self, inputs, training=None, mask=None):
-        latents, labels = inputs
+        latents, conditions = inputs
         x = latents
 
         # embed label if any
-        if self.labels_dim > 0:
-            y = self.labels_embedding(labels)
-            x = tf.concat([x, y], axis=1)
+        pose_shape_latent = self.pose_shape_embedding([conditions[0], conditions[1]])
+        garment_latent = self.garment_embedding(conditions[2])
+        x = tf.concat([x, pose_shape_latent, garment_latent], axis=1)
 
         # normalize inputs
         x = self.normalize(x)