luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

prototype conditional pipeline

Browse Source
This commit is contained in:
luni
2022-06-10 06:59:58 +00:00
parent 03c22c1ba6
commit 31607cc754
24 changed files with 1093 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
stylegan2/discriminator.py
+2 -12
View File
@@ -130,14 +130,9 @@ class Discriminator(tf.keras.Model):
self.last_dense = Dense(max(self.labels_dim, 1), gain=1.0, lrmul=1.0, name='last_dense')
self.last_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
# conditional output
self.last_block_condition = DiscriminatorLastBlock(n_f0, n_f1, res, name='{:d}x{:d}'.format(res, res))
self.condition_dense = Dense(self.labels_dim, gain=1.0, lrmul=1.0, name='last_dense')
self.condition_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
def call(self, inputs, training=None, mask=None):
images, _ = inputs
images = inputs
x = self.initial_fromrgb(images)
for block in self.blocks:
@@ -150,13 +145,8 @@ class Discriminator(tf.keras.Model):
x = self.last_dense(x)
x = self.last_bias(x)
# condition output
cond = self.last_block_condition(features)
cond = self.condition_dense(cond)
cond = self.condition_bias(cond)
scores_out = x
return scores_out, cond
return scores_out
def compute_output_shape(self, input_shape):
return input_shape[0][0], max(self.labels_dim, 1)
stylegan2/generator.py
+5 -5
View File
@@ -53,10 +53,10 @@ class Generator(tf.keras.Model):
self.w_avg.assign(lerp(batch_avg, self.w_avg, self.w_ema_decay))
return
def style_mixing_regularization(self, latents1, labels, w_broadcasted1):
def style_mixing_regularization(self, latents1, conditions, w_broadcasted1):
# get another w and broadcast it
latents2 = tf.random.normal(shape=tf.shape(latents1), dtype=tf.dtypes.float32)
dlatents2 = self.g_mapping([latents2, labels])
dlatents2 = self.g_mapping([latents2, conditions])
w_broadcasted2 = self.broadcast(dlatents2)
# find mixing limit index
@@ -89,14 +89,14 @@ class Generator(tf.keras.Model):
return truncated_w_broadcasted
def call(self, inputs, ret_w_broadcasted=False, truncation_psi=1.0, truncation_cutoff=None, training=None, mask=None):
latents, labels = inputs
latents, conditions = inputs
dlatents = self.g_mapping([latents, labels])
dlatents = self.g_mapping([latents, conditions])
w_broadcasted = self.broadcast(dlatents)
if training:
self.update_moving_average_of_w(w_broadcasted)
w_broadcasted = self.style_mixing_regularization(latents, labels, w_broadcasted)
w_broadcasted = self.style_mixing_regularization(latents, conditions, w_broadcasted)
if not training:
w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff)
stylegan2/layers/embedding.py
+120 -9
View File
@@ -1,24 +1,135 @@
import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt
class LabelEmbedding(tf.keras.layers.Layer):
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(LabelEmbedding, self).__init__(**kwargs)
super(ParameterEmbedding, 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)
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 = tf.matmul(inputs, self.w)
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(LabelEmbedding, self).get_config()
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
stylegan2/layers/mapping_block.py
+7 -9
View File
@@ -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)