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_model.py
+5 -5
View File
@@ -161,7 +161,7 @@ class CVAE(tf.keras.Model):
def reparameterize(self, mean, logvar, fix_noise=False):
if not fix_noise:
eps = tf.random.normal(shape=mean.shape)
eps = tf.random.normal(shape=mean.shape)
else:
eps = tf.random.normal(shape=(1, mean.shape[1]))
eps = tf.tile(eps, [mean.shape[0],1])
@@ -358,11 +358,11 @@ class CVAE32(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
+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,