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

pyramid with common encoders per level

Browse Source
This commit is contained in:
luni
2022-09-16 20:51:31 +00:00
parent eac8aad8ab
commit 37df8468e2
3 changed files with 81 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_model.py
+33 -13
View File
@@ -258,21 +258,31 @@ class CVAE32(tf.keras.Model):
static_model='',
body_model='',
curvature_model='',
freeze_encoders=False):
freeze_encoders=False,
use_full_res_encoders=True):
super(CVAE32, self).__init__()
self.latent_dim = latent_dim
model : StaticAE32 = StaticAE32()
if use_full_res_encoders:
model : StaticAE = StaticAE()
else:
model : StaticAE32 = StaticAE32()
if static_model != '':
load_model(model, static_model)
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
model : BodyAE32 = BodyAE32()
if use_full_res_encoders:
model : BodyAE = BodyAE()
else:
model : BodyAE32 = BodyAE32()
if body_model != '':
load_model(model, body_model)
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
model : CurvatureAE32 = CurvatureAE32()
if use_full_res_encoders:
model : CurvatureAE = CurvatureAE()
else:
model : CurvatureAE32 = CurvatureAE32()
if curvature_model != '':
load_model(model, curvature_model)
self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
@@ -429,21 +439,31 @@ class CVAE128(tf.keras.Model):
static_model='',
body_model='',
curvature_model='',
freeze_encoders=False):
freeze_encoders=False,
use_full_res_encoders=True):
super(CVAE128, self).__init__()
self.latent_dim = latent_dim
model : StaticAE128 = StaticAE128()
if use_full_res_encoders:
model : StaticAE = StaticAE()
else:
model : StaticAE128 = StaticAE128()
if static_model != '':
load_model(model, static_model)
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
model : BodyAE128 = BodyAE128()
if use_full_res_encoders:
model : BodyAE = BodyAE()
else:
model : BodyAE128 = BodyAE128()
if body_model != '':
load_model(model, body_model)
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
model : CurvatureAE128 = CurvatureAE128()
if use_full_res_encoders:
model : CurvatureAE = CurvatureAE()
else:
model : CurvatureAE128 = CurvatureAE128()
if curvature_model != '':
load_model(model, curvature_model)
self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
@@ -515,11 +535,11 @@ class CVAE128(tf.keras.Model):
def call(self, inputs, training=None, zero_latent=False):
uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
z_static = self.static_enc(uv_template_high)
z_body = self.body_enc(uv_body_posed_high)
z_curv = self.curv_enc(uv_curv_high)
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
models/cvae_pyramid_model.py
+32 -5
View File
@@ -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)