59 lines
2.2 KiB
Python
59 lines
2.2 KiB
Python
import tensorflow as tf
|
|
|
|
from stylegan2.layers.embedding import PoseShapeEmbedding, GarmentEmbedding
|
|
from stylegan2.layers.dense import Dense
|
|
from stylegan2.layers.bias_act import BiasAct
|
|
|
|
|
|
class Mapping(tf.keras.layers.Layer):
|
|
def __init__(self, w_dim, labels_dim, n_mapping, name, **kwargs):
|
|
super(Mapping, self).__init__(name=name, **kwargs)
|
|
self.w_dim = w_dim
|
|
self.labels_dim = labels_dim
|
|
self.n_mapping = n_mapping
|
|
self.gain = 1.0
|
|
self.lrmul = 0.01
|
|
|
|
self.pose_shape_embedding = PoseShapeEmbedding()
|
|
self.garment_embedding = GarmentEmbedding((self.w_dim * 2) - (216 + 11))
|
|
# self.garment_embedding = GarmentEmbedding(self.w_dim)
|
|
|
|
self.normalize = tf.keras.layers.Lambda(lambda x: x * tf.math.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8))
|
|
|
|
self.dense_layers = list()
|
|
self.bias_act_layers = list()
|
|
for ii in range(self.n_mapping):
|
|
self.dense_layers.append(Dense(w_dim, gain=self.gain, lrmul=self.lrmul, name='dense_{:d}'.format(ii)))
|
|
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, conditions = inputs
|
|
x = latents
|
|
|
|
# embed label if any
|
|
pose_shape_latent = self.pose_shape_embedding([conditions['pose_shape'][:, :72], conditions['pose_shape'][:, 72:]])
|
|
garment_latent = self.garment_embedding(conditions['uv_static'])
|
|
x = tf.concat([x, pose_shape_latent, garment_latent], axis=1)
|
|
# x = tf.concat([x, garment_latent], axis=1)
|
|
|
|
# normalize inputs
|
|
x = self.normalize(x)
|
|
|
|
# apply mapping blocks
|
|
for dense, apply_bias_act in zip(self.dense_layers, self.bias_act_layers):
|
|
x = dense(x)
|
|
x = apply_bias_act(x)
|
|
|
|
return x
|
|
|
|
def get_config(self):
|
|
config = super(Mapping, self).get_config()
|
|
config.update({
|
|
'w_dim': self.w_dim,
|
|
'labels_dim': self.labels_dim,
|
|
'n_mapping': self.n_mapping,
|
|
'gain': self.gain,
|
|
'lrmul': self.lrmul,
|
|
})
|
|
return config
|