pyramid available with 128 res

This commit is contained in:
2022-09-15 10:14:35 +00:00
parent c3a2f558de
commit c6f6f8710d
6 changed files with 337 additions and 58 deletions
+10 -7
View File
@@ -7,15 +7,18 @@ import numpy as np
# import plotly.express as ex
from custom_utils.custom_utils import load_model
from models.cvae_model import CVAE, CVAE32
from models.cvae_model import CVAE, CVAE128, CVAE32
class Pyramid(tf.keras.Model):
def __init__(self, latent_dim=128, avg_garment_path=''):
def __init__(self, low_res=32, latent_dim=128, avg_garment_path=''):
super(Pyramid, self).__init__()
self.model_32 = CVAE32(latent_dim=latent_dim)
if low_res == 32:
self.model_low = CVAE32(latent_dim=latent_dim)
if low_res == 128:
self.model_low = CVAE128(latent_dim=latent_dim)
self.model_full = CVAE(latent_dim=latent_dim)
if avg_garment_path != '':
@@ -24,16 +27,16 @@ 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_32, uv_static_32, uv_body_posed_32, uv_curv_32 = inputs
uv_unposed, uv_static, uv_body_posed, uv_curv, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low = inputs
pred_32, mean_32, logvar_32, _ = self.model_32([uv_unposed_32, uv_static_32, uv_body_posed_32, uv_curv_32], training=training)
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_32, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred_full, mean_full, logvar_full, _ = self.model_full([uv_unposed, uv_static, uv_body_posed, uv_curv], training=training)
x = x + pred_full
return x, mean_full, logvar_full, pred_32, mean_32, logvar_32,
return x, mean_full, logvar_full, pred_low, mean_low, logvar_low,