40 lines
1.2 KiB
Python
40 lines
1.2 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, CVAE32
|
|
|
|
|
|
class Pyramid(tf.keras.Model):
|
|
|
|
def __init__(self, latent_dim=128, avg_garment_path=''):
|
|
super(Pyramid, self).__init__()
|
|
|
|
self.model_32 = CVAE32(latent_dim=latent_dim)
|
|
self.model_full = CVAE(latent_dim=latent_dim)
|
|
|
|
if avg_garment_path != '':
|
|
self.avg_garment = np.load(avg_garment_path)
|
|
else:
|
|
self.avg_garment = None
|
|
|
|
def call(self, inputs, training=None):
|
|
uv_unposed, uv_static, uv_body_posed, uv_curv, uv_unposed_32, uv_static_32, uv_body_posed_32, uv_curv_32 = inputs
|
|
|
|
pred_32, mean_32, logvar_32, _ = self.model_32([uv_unposed_32, uv_static_32, uv_body_posed_32, uv_curv_32], training=training)
|
|
|
|
x = tf.image.resize(pred_32, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
|
|
|
pred_full, mean_full, logvar_full, _ = self.model_full([uv_unposed, uv_static, uv_body_posed, uv_curv], training=training)
|
|
|
|
x = x + pred_full
|
|
|
|
return x, mean_full, logvar_full, pred_32, mean_32, logvar_32,
|
|
|
|
|