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

large refactor

Browse Source
This commit is contained in:
luni
2022-09-25 08:31:13 +00:00
parent 84ec5a056e
commit 92891f5b26
33 changed files with 1661 additions and 1635 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_pyramid_model.py
+20 -20
View File
@@ -7,7 +7,7 @@ import numpy as np
# import plotly.express as ex
from custom_utils.custom_utils import load_model
from models.cvae_model import CVAE, CVAE128, CVAE32
from models.cvae_model import CVAE
from models.upscaler import UpscalerBasic as Upscaler
@@ -24,26 +24,22 @@ class Pyramid(tf.keras.Model):
use_full_res_encoders=True):
super(Pyramid, self).__init__()
if low_res == 32:
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,
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_low = CVAE(latent_dim=latent_dim,
size=low_res,
static_model=static_model,
body_model=body_model,
curvature_model=curvature_model,
freeze_encoders=freeze_encoders,
encoder_res=512 if use_full_res_encoders else 32)
self.model_full = CVAE(latent_dim=latent_dim,
size=512,
static_model=static_model,
body_model=body_model,
curvature_model=curvature_model,
freeze_encoders=freeze_encoders)
freeze_encoders=freeze_encoders,
encoder_res=512)
self.upscale_enhancer = None
if upscaler_model != '':
@@ -64,6 +60,12 @@ class Pyramid(tf.keras.Model):
else:
self.avg_garment = None
def enhance(self, low, mask, training=None):
x = self.upscale_enhancer(low, training=training)
x = x * (1 - mask[:,:,:,tf.newaxis])
low = x + low
return low
def call(self, inputs, training=None):
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
@@ -75,11 +77,9 @@ class Pyramid(tf.keras.Model):
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
pred_low = self.enhance(pred_low, uv_mask_low, training=training)
x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = tf.image.resize(pred_low, (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)