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

backup

Browse Source
This commit is contained in:
luni
2022-10-04 21:50:58 +00:00
parent 5b82041acb
commit 3c798b473d
8 changed files with 270 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_model.py
+36 -6
View File
@@ -10,6 +10,7 @@ from custom_utils.custom_utils import load_model
from models.base import AEEncoder
from models.cvae_model import CVAE
from models.pyramid_level import PyramidLevel
from models.upscaler import UpscalerBasic, UpscalerBasicNew
class Pyramid(tf.keras.Model):
@@ -41,6 +42,16 @@ class Pyramid(tf.keras.Model):
self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
self.upscale_enhancer = None
if config['pyramid']['final_upscale_enhancer'] != '':
self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
self.upscale_enhancer.trainable = False
def enhance(self, low, mask, training=None):
x = self.upscale_enhancer([low, mask], training=training)
return x
def encode_conds(self, uv_template, uv_body_posed, uv_curv):
z_static = self.static_enc(uv_template)
@@ -53,18 +64,37 @@ class Pyramid(tf.keras.Model):
return z_cond
def call(self, inputs, training=None):
uv_static, uv_body, uv_curv = inputs[-1][1], inputs[-1][2], inputs[-1][3]
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)
out_levels = []
pred, mean, logvar = self.levels[0](inputs[0], z_cond, training=training)
out_levels.append([pred, mean, logvar])
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
if 1 != len(self.levels):
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 = pred * levels_input[1][-2][:,:,:,tf.newaxis]
else:
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
out_levels.append([pred, mean, logvar])
for i in range(1, len(self.levels)):
pred_level, mean, logvar = self.levels[i](inputs[i], z_cond, training=training)
pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
pred = pred + pred_level
out_levels.append([pred, mean, logvar])
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 = 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)
out_levels.append([pred, mean, logvar])
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
return out_levels