add ae sub folder

stylegan2_ae/layers/embedding.py

@@ -0,0 +1,135 @@
+import tensorflow as tf
+
+from convnext.models.convnext_tf import ConvNeXt
+
+def rodrigues(r):
+    theta = tf.norm(r, axis=(1, 2), keepdims=True)
+    theta = tf.maximum(theta, 1e-7)
+    r_hat = r / theta
+    z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32)
+    m = tf.stack(
+        (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+        -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1)
+    m = tf.reshape(m, (-1, 3, 3))
+    i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros(
+        (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32)
+    # R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m
+    R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta))
+    return R
+
+
+def run_rodrigues(pose):
+    batch_size = pose.shape[0]
+    R_out = rodrigues(tf.reshape(pose, (-1, 1, 3)))
+    R_out = tf.reshape(R_out, (batch_size, -1, 3, 3))
+    return R_out
+
+
+# if __name__ == '__main__':
+#     x = tf.ones((4, 72))
+#     y = run_rodrigues(x)
+#     print()
+#     pass
+
+
+class ParameterEmbedding(tf.keras.layers.Layer):
+    def __init__(self, embed_dim, **kwargs):
+        super(ParameterEmbedding, self).__init__(**kwargs)
+        self.embed_dim = embed_dim
+
+        self.dense_1 = tf.keras.layers.Dense(128, activation='tanh', name='dense_1')
+        self.bn_1 = tf.keras.layers.BatchNormalization()
+        self.dense_2 = tf.keras.layers.Dense(512, activation='tanh', name='dense_2')
+        self.bn_2 = tf.keras.layers.BatchNormalization()
+        self.dense_3 = tf.keras.layers.Dense(self.embed_dim, activation='tanh', name='dense_3')
+
+    def call(self, inputs, training=None, mask=None):
+        x = inputs
+        x = self.dense_1(x)
+        x = self.bn_1(x)
+        x = self.dense_2(x)
+        x = self.bn_2(x)
+        x = self.dense_3(x)
+        return x
+
+    def get_config(self):
+        config = super(ParameterEmbedding, self).get_config()
+        config.update({
+            'embed_dim': self.embed_dim,
+        })
+        return config
+
+
+class PoseShapeEmbedding(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(PoseShapeEmbedding, self).__init__(**kwargs)
+
+        # self.pose_embedding = ParameterEmbedding(int(self.embed_dim/2))
+        # self.shape_embedding = ParameterEmbedding(int(self.embed_dim/2))
+        self.flatten = tf.keras.layers.Flatten()
+
+    def call(self, inputs, training=None, mask=None):
+        pose, shape = inputs
+        rotations = run_rodrigues(pose)
+        rotations = self.flatten(rotations)
+        # pose_latent = self.pose_embedding(rotations, training)
+        # shape_latent = self.shape_embedding(shape, training)
+        return tf.concat((rotations, shape), axis=-1)
+
+    def get_config(self):
+        config = super(PoseShapeEmbedding, self).get_config()
+        config.update({
+            'embed_dim': 216 + 11,
+        })
+        return config
+
+
+class GarmentEmbedding(tf.keras.layers.Layer):
+    def __init__(self, embed_dim, **kwargs):
+        super(GarmentEmbedding, self).__init__(**kwargs)
+        self.embed_dim = embed_dim
+
+        # TODO: channels_first?
+
+        self.encoder = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 256, self.embed_dim], include_top=False)
+        # self.encoder(tf.keras.Input(shape=(*input_shape, 3)))
+        self.pooling = tf.keras.layers.GlobalAveragePooling2D()
+
+    def call(self, inputs, training=None, mask=None):
+        x = inputs
+
+        x = tf.transpose(x, [0,2,3,1])
+        x = self.encoder(x, training=training)
+        x = self.pooling(x)
+
+        return x
+
+    def get_config(self):
+        config = super(GarmentEmbedding, self).get_config()
+        config.update({
+            'embed_dim': self.embed_dim,
+        })
+        return config
+
+
+# class LabelEmbedding(tf.keras.layers.Layer):
+#     def __init__(self, embed_dim, **kwargs):
+#         super(LabelEmbedding, self).__init__(**kwargs)
+#         self.embed_dim = embed_dim
+
+#     def build(self, input_shape):
+#         weight_shape = [input_shape[1], self.embed_dim]
+#         # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
+#         w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=1.0)
+#         self.w = tf.Variable(w_init, name='w', trainable=True)
+
+#     def call(self, inputs, training=None, mask=None):
+#         x = tf.matmul(inputs, self.w)
+#         return x
+
+#     def get_config(self):
+#         config = super(LabelEmbedding, self).get_config()
+#         config.update({
+#             'embed_dim': self.embed_dim,
+#         })
+#         return config