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

started training pyramid

Browse Source
This commit is contained in:
luni
2023-02-26 22:25:50 +00:00
parent a678724711
commit 64ecb900f5
8 changed files with 314 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_level.py
+3 -4
View File
@@ -102,13 +102,12 @@ class PyramidLevelNew(tf.keras.Model):
super(PyramidLevelNew, self).__init__()
self.size = config['resolution']
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'], latent_dim=config['latent_dim'])
self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
self.upscale_enhancer = None
if config['upscale_enhancer'] != '':
k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=config['kernel_size'])
load_model(self.upscale_enhancer, config['upscale_enhancer'])
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
@@ -117,7 +116,7 @@ class PyramidLevelNew(tf.keras.Model):
else:
self.avg_garment = None
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
self.reshape = tf.keras.layers.Reshape((8, 4, config['latent_dim']))
def enhance(self, low, mask, training=None):
models/pyramid_model_optimized.py
-4
View File
@@ -343,11 +343,7 @@ class PyramidNew(tf.keras.Model):
for i in range(1, len(self.levels)):
input = levels_input[i]
input[0] = input[0] - pred
input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3])
input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None]
input[0] = input[0] * input[-2][:,:,:,None] # uv_mask_res
pred_level, mean, logvar = self.levels[i](input, z_cond, training=training)
pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res
pred = pred + pred_level
if i != len(self.levels) - 1:
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)