93 lines
3.5 KiB
Python
93 lines
3.5 KiB
Python
import os
|
|
import time
|
|
import random
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
# import plotly.graph_objects as go
|
|
# import plotly.express as ex
|
|
|
|
from custom_utils.custom_utils import load_model
|
|
from models.cvae_model import CVAE
|
|
from models.upscaler import UpscalerBasic as Upscaler
|
|
|
|
|
|
class Pyramid(tf.keras.Model):
|
|
|
|
def __init__(self, low_res=32,
|
|
latent_dim=128,
|
|
avg_garment_path='',
|
|
static_model='',
|
|
body_model='',
|
|
curvature_model='',
|
|
upscaler_model='',
|
|
freeze_encoders=False,
|
|
use_full_res_encoders=True):
|
|
super(Pyramid, self).__init__()
|
|
|
|
|
|
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,
|
|
encoder_res=512)
|
|
|
|
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
|
|
|
|
if avg_garment_path != '':
|
|
self.avg_garment = np.load(avg_garment_path)
|
|
else:
|
|
self.avg_garment = None
|
|
|
|
def enhance(self, low, mask, training):
|
|
x = self.upscale_enhancer(low, training=training)
|
|
x = x * (1 - mask[:,:,:,tf.newaxis])
|
|
low = x + low
|
|
return low
|
|
|
|
def call(self, inputs, training):
|
|
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)
|
|
|
|
pred_low = pred_low * uv_mask_low[:,:,:,tf.newaxis]
|
|
|
|
if self.upscale_enhancer is not None:
|
|
pred_low = self.enhance(pred_low, uv_mask_low, training=training)
|
|
|
|
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)
|
|
|
|
x = x + pred_full
|
|
x = x * uv_mask[:,:,:,tf.newaxis]
|
|
|
|
return x, mean_full, logvar_full, pred_low, mean_low, logvar_low,
|
|
|
|
|