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

working pyramid update

Browse Source
This commit is contained in:
luni
2022-09-20 22:56:44 +00:00
parent 456a8b2c37
commit 494ceaa276
5 changed files with 557 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_pyramid_model.py
+26 -3
View File
@@ -8,6 +8,7 @@ import numpy as np
from custom_utils.custom_utils import load_model
from models.cvae_model import CVAE, CVAE128, CVAE32
from models.upscaler import UpscalerBasic as Upscaler
class Pyramid(tf.keras.Model):
@@ -17,7 +18,8 @@ class Pyramid(tf.keras.Model):
avg_garment_path='',
static_model='',
body_model='',
curvature_model='',
curvature_model='',
upscaler_model='',
freeze_encoders=False,
use_full_res_encoders=True):
super(Pyramid, self).__init__()
@@ -42,6 +44,18 @@ class Pyramid(tf.keras.Model):
body_model=body_model,
curvature_model=curvature_model,
freeze_encoders=freeze_encoders)
self.upscale_enhancer = None
if upscaler_model != '':
self.upscale_enhancer = Upscaler(512, low_res)
load_model(self.upscale_enhancer, upscaler_model)
self.upscale_enhancer = tf.keras.Sequential([
self.upscale_enhancer.flatten,
self.upscale_enhancer.seq,
self.upscale_enhancer.reshape
])
self.upscale_enhancer.trainable = False
self.use_full_res_encoders = use_full_res_encoders
@@ -51,18 +65,27 @@ class Pyramid(tf.keras.Model):
self.avg_garment = None
def call(self, inputs, training=None):
uv_unposed, uv_static, uv_body_posed, uv_curv, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low = inputs
uv_unposed, uv_static, uv_body_posed, uv_curv, uv_mask, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_mask_low = inputs
if self.use_full_res_encoders:
pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_static, uv_body_posed, uv_curv], training=training)
else:
pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training)
x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred_low = pred_low * uv_mask_low[:,:,:,tf.newaxis]
if self.upscale_enhancer is not None:
x = self.upscale_enhancer(pred_low, training=training)
x = x * (1 - uv_mask_low[:,:,:,tf.newaxis])
pred_low = x + pred_low
x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = x * uv_mask[:,:,:,tf.newaxis]
pred_full, mean_full, logvar_full, _ = self.model_full([uv_unposed, uv_static, uv_body_posed, uv_curv], training=training)
x = x + pred_full
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean_full, logvar_full, pred_low, mean_low, logvar_low,