136 lines
4.5 KiB
Python
136 lines
4.5 KiB
Python
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
|