create end2end laplacian pyramid

This commit is contained in:
2023-01-26 20:39:54 +00:00
parent 8c0436cea8
commit 0dfe7f7659
5 changed files with 133 additions and 116 deletions
+8 -5
View File
@@ -107,27 +107,30 @@ class Pyramid(tf.keras.Model):
pred = pred + levels_input[0][2] # uv_body_posed_32
if 1 != len(self.levels):
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
out_levels.append([pred, mean, logvar])
pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR)
pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
else:
pred = self.enhance(pred, levels_input[0][-2], training=training)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
out_levels.append([pred, mean, logvar])
for i in range(1, len(self.levels)):
pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
input = levels_input[i]
input[0] = pred # use previous level output as input for next level
pred_level, mean, logvar = self.levels[i](input, z_cond, training=training)
pred = pred + pred_level
if i != len(self.levels) - 1:
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
out_levels.append([pred, mean, logvar])
pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.BILINEAR)
pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
else:
if self.upscale_enhancer is not None:
pred = self.enhance(pred, levels_input[-1][-2], training=training)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
out_levels.append([pred, mean, logvar])