import os import time import random import tensorflow as tf import numpy as np import scipy as sp # import plotly.graph_objects as go # import plotly.express as ex from custom_utils.custom_utils import load_model from models.base import AEEncoder, VariationalEncoder, Decoder class CVAE(tf.keras.Model): def __init__(self, config): super(CVAE, self).__init__() self.latent_dim = config['cvae']['latent_dim'] self.size = config['resolution'] self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size']) if config['static_model'] != '': load_model(self.static_enc, config['static_model'] + '_enc') self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size']) if config['body_model'] != '': load_model(self.body_enc, config['body_model'] + '_enc') self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size']) if config['curvature_model'] != '': load_model(self.curv_enc, config['curvature_model'] + '_enc') if config['freeze_encoders']: self.static_enc.trainable = False self.body_enc.trainable = False self.curv_enc.trainable = False self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'), tf.keras.layers.Dense(256, activation='relu'), tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dense(config['latent_dim']),]) self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size']) self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size']) if config['avg_garment_path'] != '': self.avg_garment = np.load(config['avg_garment_path']) else: self.avg_garment = None # @tf.function def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False): if not fix_noise: eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim)) else: eps = tf.random.normal(shape=(1, self.latent_dim)) eps = tf.tile(eps, [nr_of_samples,1]) uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1]) uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1]) if len(uv_curv.shape) == 3: uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1]) uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1]) # uv_template, uv_body_posed = 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_cond = tf.concat([z_static, z_body, z_curv], axis=-1) z_cond = self.cond_encoder(z_cond) z = tf.concat([eps, z_cond], axis=-1) x = self.decoder(z) x = x + self.avg_garment[tf.newaxis,:,:,:] x = x * uv_mask[:,:,:,tf.newaxis] return x def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None): if noise is None: eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim)) else: eps = noise z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv) z = tf.concat([eps, z_cond], axis=-1) x = self.decoder(z) if self.avg_garment is not None: x = x + self.avg_garment[tf.newaxis,:,:,:] if uv_mask is not None: x = x * uv_mask[:,:,:,tf.newaxis] return x def reparameterize(self, mean, logvar, fix_noise=False): if not fix_noise: 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]) return eps * tf.exp(logvar * .5) + mean def decode(self, z, apply_sigmoid=False): logits = self.decoder(z) if apply_sigmoid: probs = tf.sigmoid(logits) return probs return logits def encode_conds(self, uv_template, uv_body_posed, uv_curv): if uv_template.shape[1] == 3: uv_template = tf.transpose(uv_template, [0,2,3,1]) if uv_body_posed.shape[1] == 3: uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1]) if uv_curv.shape[1] == 3: uv_curv = tf.transpose(uv_curv, [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_cond = tf.concat([z_static, z_body, z_curv], axis=-1) z_cond = self.cond_encoder(z_cond) return z_cond def call(self, inputs, training=None, use_mask=True): if use_mask: if self.size == 512: uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv else: uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_template_high, uv_body_posed_high, uv_curv_high, uv_mask_high = inputs else: if self.size == 512: uv_unposed, uv_template, uv_body_posed, uv_curv = inputs uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv else: uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1) mean, logvar = self.encoder(x, training=training) z = self.reparameterize(mean, logvar) z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high) z = tf.concat([z, z_cond], axis=-1) x = self.decoder(z, training=training) if self.avg_garment is not None: x = x + self.avg_garment[tf.newaxis,:,:,:] if use_mask: x = x * uv_mask[:,:,:,tf.newaxis] return x, mean, logvar, z