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

only load encoders of AEs

Browse Source
This commit is contained in:
luni
2022-09-26 11:35:08 +00:00
parent 55f30c0fa9
commit fdf36b27e7
3 changed files with 32 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_model.py
+9 -27
View File
@@ -8,19 +8,7 @@ import scipy as sp
# import plotly.express as ex
from custom_utils.custom_utils import load_model
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from models.base import Bottleneck, VariationalEncoder, Decoder
from models.cond_encoder import AutoEncoder
# from models.ae_body_model import AutoEncoder as BodyAE
# from models.ae_curv_model import AutoEncoder as CurvatureAE
# from models.ae_static_model import AutoEncoder as StaticAE
# from models.ae_body_model import AutoEncoder32 as BodyAE32
# from models.ae_curv_model import AutoEncoder32 as CurvatureAE32
# from models.ae_static_model import AutoEncoder32 as StaticAE32
# from models.ae_body_model import AutoEncoder128 as BodyAE128
# from models.ae_curv_model import AutoEncoder128 as CurvatureAE128
# from models.ae_static_model import AutoEncoder128 as StaticAE128
from models.base import AEEncoder, VariationalEncoder, Decoder
class CVAE(tf.keras.Model):
@@ -30,28 +18,23 @@ class CVAE(tf.keras.Model):
self.latent_dim = config['latent_dim']
self.size = config['resolution']
model : AutoEncoder = AutoEncoder(config['cond_ae'])
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(model, config['static_model'])
# self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
self.static_enc = model.encoder
load_model(self.static_enc, config['static_model'] + '_enc')
model : AutoEncoder = AutoEncoder(config['cond_ae'])
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(model, config['body_model'])
self.body_enc = model.encoder
load_model(self.body_enc, config['body_model'] + '_enc')
model : AutoEncoder = AutoEncoder(config['cond_ae'])
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(model, config['curvature_model'])
self.curv_enc = model.encoder
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.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'])
@@ -65,7 +48,6 @@ class CVAE(tf.keras.Model):
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(config['latent_dim']),])
# self.shuffle = tf.keras.layers.Lambda(lambda x: tf.stack([tf.gather(x[i], tf.random.shuffle(tf.range(x[i].shape[0]))) for i in range(x.shape[0])]))
# @tf.function
def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False):
@@ -153,7 +135,7 @@ class CVAE(tf.keras.Model):
return z_cond
def call(self, inputs, training=None, use_mask=False):
def call(self, inputs, training=None, use_mask=True):
if use_mask:
if self.size == 512: