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

also use vae latent for second level

Browse Source
This commit is contained in:
luni
2023-01-30 19:51:14 +00:00
parent 2b8abd776b
commit 6e24c1f41b
5 changed files with 21 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_level.py
+1 -1
View File
@@ -92,7 +92,7 @@ class PyramidLevel(tf.keras.Model):
# x = x * uv_mask_high[:,:,:,tf.newaxis]
return x, mean, logvar
return x, z, mean, logvar
models/pyramid_model_optimized.py
+15 -3
View File
@@ -78,6 +78,11 @@ class Pyramid(tf.keras.Model):
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
self.upscale_enhancer.trainable = False
self.cond_encoder_2 = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(config['pyramid']['latent_dim'], activation='tanh'),])
def enhance(self, low, mask, training):
x = self.upscale_enhancer([low, mask], training=training)
return x
@@ -90,19 +95,25 @@ class Pyramid(tf.keras.Model):
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond, training=training)
return z_cond, z_static, z_body, z_curv
def encode_conds_2(self, z, z_static, z_body, z_curv, training):
z_cond = tf.concat([z, z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder_2(z_cond, training=training)
return z_cond
def call(self, inputs, training):
levels_input, conds = inputs
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
z_cond, z_static, z_body, z_curv = self.encode_conds(uv_static, uv_body, uv_curv, training)
if self.add_noise == True and training == True:
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
out_levels = []
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
pred, z, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
if self.use_offsets:
pred = pred + levels_input[0][2] # uv_body_posed_32
if 1 != len(self.levels):
@@ -119,8 +130,9 @@ class Pyramid(tf.keras.Model):
for i in range(1, len(self.levels)):
input = levels_input[i]
input[0] = pred # use previous level output as input for next level
z_cond_2 = self.encode_conds_2(z, z_static, z_body, z_curv, training=training)
pred_level, mean, logvar = self.levels[i](input, z_cond, training=training)
pred_level, z, mean, logvar = self.levels[i](input, z_cond_2, training=training)
pred = pred + pred_level
if i != len(self.levels) - 1:
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)