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

add no curvature option

Browse Source
This commit is contained in:
luni
2023-01-26 12:03:50 +00:00
parent 2ac1a812fc
commit 44e1d02e5a
4 changed files with 32 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_model_optimized.py
+21 -13
View File
@@ -24,15 +24,17 @@ class Pyramid(tf.keras.Model):
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'])
if config['pyramid']['body_model'] != '':
load_model(self.body_enc, config['pyramid']['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['pyramid']['curvature_model'] != '':
load_model(self.curv_enc, config['pyramid']['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['pyramid']['curvature_model'] != '':
load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
if config['pyramid']['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'),
@@ -70,11 +72,12 @@ class Pyramid(tf.keras.Model):
self.add_noise = config['pyramid']['add_noise']
self.noise_amount = config['pyramid']['noise_amount']
self.use_offsets = config['use_offsets']
self.use_curvature = config['use_curvature']
self.upscale_enhancer = None
if config['pyramid']['final_upscale_enhancer'] != '':
k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
self.upscale_enhancer : tf.keras.Model = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
self.upscale_enhancer.trainable = False
@@ -85,9 +88,11 @@ class Pyramid(tf.keras.Model):
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
z_static = self.static_enc(uv_template, training)
z_body = self.body_enc(uv_body_posed, training)
z_curv = self.curv_enc(uv_curv, training)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
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 = self.cond_encoder(z_cond, training=training)
return z_cond
@@ -294,7 +299,8 @@ class PyramidNew(tf.keras.Model):
self.add_noise = config['pyramid']['add_noise']
self.noise_amount = config['pyramid']['noise_amount']
self.use_offsets = config['use_offsets']
self.use_curvature = config['use_curvature']
self.upscale_enhancer = None
if config['pyramid']['final_upscale_enhancer'] != '':
k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
@@ -309,9 +315,11 @@ class PyramidNew(tf.keras.Model):
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
z_static = self.static_enc(uv_template, training)
z_body = self.body_enc(uv_body_posed, training)
z_curv = self.curv_enc(uv_curv, training)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
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)
return z_cond