pyramid with common encoders per level
This commit is contained in:
@@ -12,15 +12,39 @@ from models.cvae_model import CVAE, CVAE128, CVAE32
|
||||
|
||||
class Pyramid(tf.keras.Model):
|
||||
|
||||
def __init__(self, low_res=32, latent_dim=128, avg_garment_path=''):
|
||||
def __init__(self, low_res=32,
|
||||
latent_dim=128,
|
||||
avg_garment_path='',
|
||||
static_model='',
|
||||
body_model='',
|
||||
curvature_model='',
|
||||
freeze_encoders=False,
|
||||
use_full_res_encoders=True):
|
||||
super(Pyramid, self).__init__()
|
||||
|
||||
if low_res == 32:
|
||||
self.model_low = CVAE32(latent_dim=latent_dim)
|
||||
self.model_low = CVAE32(latent_dim=latent_dim,
|
||||
static_model=static_model,
|
||||
body_model=body_model,
|
||||
curvature_model=curvature_model,
|
||||
freeze_encoders=freeze_encoders,
|
||||
use_full_res_encoders=use_full_res_encoders)
|
||||
if low_res == 128:
|
||||
self.model_low = CVAE128(latent_dim=latent_dim)
|
||||
self.model_full = CVAE(latent_dim=latent_dim)
|
||||
self.model_low = CVAE128(latent_dim=latent_dim,
|
||||
static_model=static_model,
|
||||
body_model=body_model,
|
||||
curvature_model=curvature_model,
|
||||
freeze_encoders=freeze_encoders,
|
||||
use_full_res_encoders=use_full_res_encoders)
|
||||
|
||||
self.model_full = CVAE(latent_dim=latent_dim,
|
||||
static_model=static_model,
|
||||
body_model=body_model,
|
||||
curvature_model=curvature_model,
|
||||
freeze_encoders=freeze_encoders)
|
||||
|
||||
self.use_full_res_encoders = use_full_res_encoders
|
||||
|
||||
if avg_garment_path != '':
|
||||
self.avg_garment = np.load(avg_garment_path)
|
||||
else:
|
||||
@@ -29,7 +53,10 @@ class Pyramid(tf.keras.Model):
|
||||
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
|
||||
|
||||
pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training)
|
||||
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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user