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

prepare cvae_curv

Browse Source
This commit is contained in:
luni
2023-03-02 16:13:18 +00:00
parent dcb4ade0ca
commit d93678ee1a
4 changed files with 181 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_curv_model.py
+28 -18
View File
@@ -8,7 +8,7 @@ import scipy as sp
# import plotly.express as ex
from custom_utils.custom_utils import load_model
from models.base import AEEncoder, VariationalEncoder, Decoder
from models.base import AEEncoderNew, VariationalEncoderNew, DecoderNew
@@ -23,11 +23,11 @@ class CVAE(tf.keras.Model):
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'])
self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'])
if config['static_model'] != '':
load_model(self.static_enc, config['static_model'] + '_enc')
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size'])
if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc')
@@ -40,9 +40,10 @@ class CVAE(tf.keras.Model):
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(self.latent_dim, activation='tanh'),])
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'], is_label=config['cvae']['is_label'])
self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'], latent_dim=config['cvae']['latent_dim'])
self.decoder = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'], is_label=config['cvae']['is_label'])
self.reshape = tf.keras.layers.Reshape((8, 4, config['cvae']['latent_dim']))
self.add_noise = config['cvae']['add_noise']
self.noise_amount = config['cvae']['noise_amount']
@@ -89,12 +90,12 @@ class CVAE(tf.keras.Model):
x = self.decoder(z, training=False)
x = x * uv_mask[:,:,:,tf.newaxis]
# x = x * uv_mask[:,:,:,tf.newaxis]
background = tf.ones((uv_template.shape[0],512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
x = tf.concat((background, x), axis=-1)
# background = tf.ones((uv_template.shape[0],512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
# x = tf.concat((background, x), axis=-1)
x = tf.math.softmax(x)
# x = tf.math.softmax(x)
return x
@@ -112,9 +113,17 @@ class CVAE(tf.keras.Model):
return x
def reparameterize(self, mean, logvar):
eps = tf.random.normal(shape=mean.shape)
return eps * tf.exp(logvar * .5) + mean
def reparameterize(self, mean, logvar, training, fix_noise=False):
if training == True:
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])
out = eps * tf.exp(logvar * .5) + mean
else:
out = tf.exp(logvar * .5) + mean
return out
def call(self, inputs, training):
@@ -123,18 +132,19 @@ class CVAE(tf.keras.Model):
x = tf.concat([uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
z = self.reparameterize(mean, logvar, training=training)
z = self.reshape(z)
z_cond = self.encode_conds(uv_template, uv_body_posed, training=training)
z = tf.concat([z, z_cond], axis=-1)
x = self.decoder(z, training=training)
x = self.decoder([z, uv_mask], training=training)
x = x * uv_mask[:,:,:,tf.newaxis]
# x = x * uv_mask[:,:,:,tf.newaxis]
background = tf.ones((4,512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
x = tf.concat((background, x), axis=-1)
# background = tf.ones((4,512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
# x = tf.concat((background, x), axis=-1)
x = tf.math.softmax(x)
# x = tf.math.softmax(x)
return x, mean, logvar