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

backup

Browse Source
This commit is contained in:
luni
2022-11-04 16:17:58 +00:00
parent ebe8522905
commit 1c6b805519
18 changed files with 1413 additions and 486 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_model.py
+23 -15
View File
@@ -26,14 +26,16 @@ class CVAE(tf.keras.Model):
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['use_curvature']:
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
if config['use_curvature']:
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'),
@@ -50,7 +52,13 @@ class CVAE(tf.keras.Model):
self.add_noise = config['cvae']['add_noise']
self.noise_amount = config['cvae']['noise_amount']
self.use_curvature = config['use_curvature']
self.use_offsets_over_body = not config['use_skinning']
if self.use_offsets_over_body:
assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
if self.avg_garment is not None:
assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
@tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None):
@@ -84,27 +92,24 @@ class CVAE(tf.keras.Model):
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training=None):
z_body = self.body_enc(uv_body_posed, training=training)
z_static = self.static_enc(uv_template, training=training)
z_curv = self.curv_enc(uv_curv, training=training)
if self.use_curvature:
z_curv = self.curv_enc(uv_curv, training=training)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
else:
z_cond = tf.concat([z_static, z_body], axis=-1)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond, training=training)
return z_cond
def call(self, inputs, training=None):
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
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = 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, training=training)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
if self.add_noise and training == True:
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
@@ -115,7 +120,10 @@ class CVAE(tf.keras.Model):
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if self.use_offsets_over_body:
x = x + uv_body_posed
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar