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

sequential pyramid

Browse Source
This commit is contained in:
luni
2023-03-11 22:53:30 +00:00
parent 4b1cb0b06d
commit eb3ad3b887
4 changed files with 384 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/iccv/config_abl_pyr_32_64_128_seq.py
+187
View File
@@ -0,0 +1,187 @@
config = dict()
config['gpu'] = '4'
config['name'] = 'iccv_pyr_32_64_128_seq_2'
config['notes'] = ''
config['tags'] = ['pyr', 'iccv']
# DON'T FORGET TO RESET SSIM LOSS USAGE
# DON'T FORGET TO CHECK MODEL LOADING
# DON'T FORGET TO CHECK IF STORE IN MEMORY
# DON'T FORGET TO CHECK REG LOSS COEFF
# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
# DON'T FORGET TO CHECK IF SKIRT
# DON'T FORGET TO CHECK IF CAPE
# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
config['use_skinning'] = True # and deshaping
config['remove_global_orientation'] = True
config['use_extra_losses'] = True # edge, collision, 3d mesh
config['use_offsets'] = False
config['use_norm_loss'] = True
config['use_curvature'] = True
config['pretrained_model'] = ''
config['pretrained_model_iteration'] = ''
config['pyramid'] = dict()
config['pyramid']['resolutions'] = [32, 64, 128]
config['pyramid']['pretrained_pyramid'] = 'checkpoints/iccv/iccv_pyr_32_seq_2'
config['pyramid']['final_upscale_enhancer'] = ''
config['train_dir'] = '/home/workspace/local-data/train'
config['valid_dir'] = '/home/workspace/local-data/valid'
config['test_dir'] = '/home/workspace/local-data_127/test_ready'
config['batch_size'] = 4
config['is_skirt'] = False
config['is_cape'] = False
config['dataset_in_memory'] = False
config['resolution'] = None
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['lr'] = 0.0001
config['starting_balancing_coeff'] = 0.001
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['avg_garment_path'] = ''
config['bottleneck_scale'] = 2
config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
'pose_shape': (72 + 11,),
'uv_body_posed': (512, 256, 3),
'mesh_body_posed': (config['smpl_res'], 3,),
'mesh_posed': (config['smpl_res'], 3,),
'mesh_unposed': (config['smpl_res'], 3,),
'uv_unposed': (512, 256, 3),
'uv_curvature': (512, 256, 7),
}
config['signature_static'] = { 'uv_static': (512, 256, 3),
'mesh_mask': (config['smpl_res'],),
'uv_mask': (512, 256,),
}
config['pyramid']['is_label'] = False
config['pyramid']['add_noise'] = False
config['pyramid']['noise_amount'] = 0.0
config['pyramid']['loss_weights'] = [1,1,1,1,1]
for res in config['pyramid']['resolutions']:
if res != 512:
res_modifier = '_' + str(res)
config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
config['pyramid']['static_model'] = 'checkpoints/iccv/static_iccv_1'
config['pyramid']['body_model'] = 'checkpoints/iccv/body_iccv_1'
config['pyramid']['curvature_model'] = 'checkpoints/iccv/curv_iccv_1'
config['pyramid']['freeze_encoders'] = True
config['pyramid']['levels'] = []
for res in config['pyramid']['resolutions']:
if res == 512:
level = dict()
level['resolution'] = 512
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 7
level['encoder_depths'] = [3, 3, 3, 9, 3]
level['encoder_dims'] = [16, 32, 64, 128, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = ''
config['pyramid']['levels'].append(level)
elif res == 256:
level = dict()
level['resolution'] = 256
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 5
level['encoder_depths'] = [3, 3, 9, 3]
level['encoder_dims'] = [64, 128, 128, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_256_2'#'checkpoints/up_3_d256_512_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 128:
level = dict()
level['resolution'] = 128
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 5
level['encoder_depths'] = [9, 9, 9]
level['encoder_dims'] = [128, 128, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_128_2'#'checkpoints/up_3_d128_256_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 64:
level = dict()
level['resolution'] = 64
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 3
level['encoder_depths'] = [9, 9]
level['encoder_dims'] = [64, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_64_2'#checkpoints/up_3_d64_128_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 32:
level = dict()
level['resolution'] = 32
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 3
level['encoder_depths'] = [16]
level['encoder_dims'] = [level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_32_2'#'checkpoints/up_3_d32_64_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
else:
raise NotImplementedError
config['cond_ae_static'] = dict()
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 27, 3]
config['cond_ae_static']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae_static']['bottleneck_size']]
config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
config['cond_ae_static']['activation'] = 'tanh'
config['cond_ae_curv'] = dict()
config['cond_ae_curv']['is_label'] = True
config['cond_ae_curv']['kernel_size'] = 7
config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
config['cond_ae_curv']['activation'] = 'tanh'
config['cond_ae_body'] = dict()
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
config['cond_ae_body']['activation'] = 'tanh'
configs/iccv/config_abl_pyr_32_seq.py
+187
View File
@@ -0,0 +1,187 @@
config = dict()
config['gpu'] = '4'
config['name'] = 'iccv_pyr_32_seq_2'
config['notes'] = ''
config['tags'] = ['pyr', 'iccv']
# DON'T FORGET TO RESET SSIM LOSS USAGE
# DON'T FORGET TO CHECK MODEL LOADING
# DON'T FORGET TO CHECK IF STORE IN MEMORY
# DON'T FORGET TO CHECK REG LOSS COEFF
# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
# DON'T FORGET TO CHECK IF SKIRT
# DON'T FORGET TO CHECK IF CAPE
# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
config['use_skinning'] = True # and deshaping
config['remove_global_orientation'] = True
config['use_extra_losses'] = True # edge, collision, 3d mesh
config['use_offsets'] = False
config['use_norm_loss'] = True
config['use_curvature'] = True
config['pretrained_model'] = ''
config['pretrained_model_iteration'] = ''
config['pyramid'] = dict()
config['pyramid']['resolutions'] = [32, 64]
config['pyramid']['pretrained_pyramid'] = 'checkpoints/iccv/iccv_pyr_32'
config['pyramid']['final_upscale_enhancer'] = ''
config['train_dir'] = '/home/workspace/local-data/train'
config['valid_dir'] = '/home/workspace/local-data/valid'
config['test_dir'] = '/home/workspace/local-data_127/test_ready'
config['batch_size'] = 4
config['is_skirt'] = False
config['is_cape'] = False
config['dataset_in_memory'] = False
config['resolution'] = None
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['lr'] = 0.0001
config['starting_balancing_coeff'] = 0.001
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['avg_garment_path'] = ''
config['bottleneck_scale'] = 2
config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
'pose_shape': (72 + 11,),
'uv_body_posed': (512, 256, 3),
'mesh_body_posed': (config['smpl_res'], 3,),
'mesh_posed': (config['smpl_res'], 3,),
'mesh_unposed': (config['smpl_res'], 3,),
'uv_unposed': (512, 256, 3),
'uv_curvature': (512, 256, 7),
}
config['signature_static'] = { 'uv_static': (512, 256, 3),
'mesh_mask': (config['smpl_res'],),
'uv_mask': (512, 256,),
}
config['pyramid']['is_label'] = False
config['pyramid']['add_noise'] = False
config['pyramid']['noise_amount'] = 0.0
config['pyramid']['loss_weights'] = [1,1,1,1,1]
for res in config['pyramid']['resolutions']:
if res != 512:
res_modifier = '_' + str(res)
config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
config['pyramid']['static_model'] = 'checkpoints/iccv/static_iccv_1'
config['pyramid']['body_model'] = 'checkpoints/iccv/body_iccv_1'
config['pyramid']['curvature_model'] = 'checkpoints/iccv/curv_iccv_1'
config['pyramid']['freeze_encoders'] = True
config['pyramid']['levels'] = []
for res in config['pyramid']['resolutions']:
if res == 512:
level = dict()
level['resolution'] = 512
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 7
level['encoder_depths'] = [3, 3, 3, 9, 3]
level['encoder_dims'] = [16, 32, 64, 128, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = ''
config['pyramid']['levels'].append(level)
elif res == 256:
level = dict()
level['resolution'] = 256
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 5
level['encoder_depths'] = [3, 3, 9, 3]
level['encoder_dims'] = [64, 128, 128, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_256_2'#'checkpoints/up_3_d256_512_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 128:
level = dict()
level['resolution'] = 128
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 5
level['encoder_depths'] = [9, 9, 9]
level['encoder_dims'] = [128, 128, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_128_2'#'checkpoints/up_3_d128_256_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 64:
level = dict()
level['resolution'] = 64
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 3
level['encoder_depths'] = [9, 9]
level['encoder_dims'] = [64, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_64_2'#checkpoints/up_3_d64_128_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 32:
level = dict()
level['resolution'] = 32
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 3
level['encoder_depths'] = [16]
level['encoder_dims'] = [level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_32_2'#'checkpoints/up_3_d32_64_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
else:
raise NotImplementedError
config['cond_ae_static'] = dict()
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 27, 3]
config['cond_ae_static']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae_static']['bottleneck_size']]
config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
config['cond_ae_static']['activation'] = 'tanh'
config['cond_ae_curv'] = dict()
config['cond_ae_curv']['is_label'] = True
config['cond_ae_curv']['kernel_size'] = 7
config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
config['cond_ae_curv']['activation'] = 'tanh'
config['cond_ae_body'] = dict()
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
config['cond_ae_body']['activation'] = 'tanh'
models/pyramid_level.py
+7 -4
View File
@@ -155,15 +155,18 @@ class PyramidLevelNew(tf.keras.Model):
return out
def call(self, inputs, z_cond, training=None):
def call(self, inputs, z_cond, prev_latent=None, training=None):
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = 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, training)
z = self.reshape(z)
z_latent = self.reshape(z)
z = tf.concat([z, z_cond], axis=-1)
if prev_latent is not None:
z_latent = tf.concat([z_latent, prev_latent], axis=-1)
z = tf.concat([z_latent, z_cond], axis=-1)
x = self.decoder(z, training=training)
if self.avg_garment is not None:
@@ -177,4 +180,4 @@ class PyramidLevelNew(tf.keras.Model):
# x = x * uv_mask_high[:,:,:,tf.newaxis]
return x, mean, logvar
return x, mean, logvar, z_latent
models/pyramid_model_optimized.py
+2 -255
View File
@@ -12,244 +12,6 @@ from models.pyramid_level import PyramidLevel, PyramidLevelNew
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
class Pyramid(tf.keras.Model):
def __init__(self, config):
super(Pyramid, self).__init__()
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['pyramid']['static_model'] != '':
load_model(self.static_enc, config['pyramid']['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'])
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['pyramid']['freeze_encoders']:
self.static_enc.trainable = False
self.body_enc.trainable = False
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'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(config['pyramid']['latent_dim'], activation='tanh'),])
self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
# if config['is_cape']:
# if len(config['pyramid']['resolutions']) == 2:
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# if len(config['pyramid']['resolutions']) == 3:
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# if len(config['pyramid']['resolutions']) == 4:
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# else:
if len(config['pyramid']['resolutions']) == 2:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
if len(config['pyramid']['resolutions']) == 3:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
if len(config['pyramid']['resolutions']) == 4:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
if len(config['pyramid']['resolutions']) == 5:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
self.add_noise = config['pyramid']['add_noise']
self.noise_amount = config['pyramid']['noise_amount']
self.use_offsets = config['use_offsets']
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)
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
self.upscale_enhancer.trainable = False
def enhance(self, low, mask, training):
x = self.upscale_enhancer([low, mask], training=training)
return x
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)
z_cond = self.cond_encoder(z_cond, training=training)
return z_cond
def call(self, inputs, training):
levels_input, conds = inputs
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
if self.add_noise == True and training == True:
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
out_levels = []
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
if self.use_offsets:
pred = pred + levels_input[0][2] # uv_body_posed_32
if 1 != len(self.levels):
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
out_levels.append([pred, mean, logvar])
pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
else:
pred = self.enhance(pred, levels_input[0][-2], training=training)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
out_levels.append([pred, mean, logvar])
for i in range(1, len(self.levels)):
pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
pred = pred + pred_level
if i != len(self.levels) - 1:
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
out_levels.append([pred, mean, logvar])
pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
else:
if self.upscale_enhancer is not None:
pred = self.enhance(pred, levels_input[-1][-2], training=training)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
out_levels.append([pred, mean, logvar])
return out_levels
# # 32
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
# pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
# pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# return [ [pred_0, mean_0, logvar_0] ]
# # 32 64
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
# pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
# pred_0 = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
# pred = pred_0 * levels_input[1][-2][:,:,:,tf.newaxis]
# pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
# pred_1 = pred + pred_1
# pred_1 = self.enhance(pred_1, levels_input[-1][-2], training=training)
# pred_1 = tf.image.resize(pred_1, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred_1 = pred_1 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# return [ [pred_0, mean_0, logvar_0], [pred_1, mean_1, logvar_1] ]
# class Pyramid(tf.keras.Model):
# def __init__(self, config):
# super(Pyramid, self).__init__()
# 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['pyramid']['static_model'] != '':
# load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
# 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['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['pyramid']['freeze_encoders']:
# self.static_enc.trainable = False
# self.body_enc.trainable = False
# 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'),
# tf.keras.layers.Dense(128, activation='relu'),
# tf.keras.layers.Dense(config['pyramid']['latent_dim']),])
# self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
# self.upscale_enhancer = None
# if config['pyramid']['final_upscale_enhancer'] != '':
# self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
# load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
# self.upscale_enhancer.trainable = False
# def enhance(self, low, mask, training):
# x = self.upscale_enhancer([low, mask], training=training)
# return x
# def encode_conds(self, uv_template, uv_body_posed, uv_curv):
# z_static = self.static_enc(uv_template)
# z_body = self.body_enc(uv_body_posed)
# z_curv = self.curv_enc(uv_curv)
# z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
# z_cond = self.cond_encoder(z_cond)
# return z_cond
# def call(self, inputs, training):
# levels_input, conds = inputs
# uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
# z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
# out_levels = []
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
# pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
# # pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
# ##################
# # if more than 1 level
# # x = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
# # x = x * levels_input[1][-2][:,:,:,tf.newaxis]
# # if only 1 level
# pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# ##################
# # pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
# # pred_1 = pred_0 + pred_1
# # pred_1 = self.levels[1].enhance(pred, levels_input[1][-2], training=training)
# # for i in range(1, len(self.levels)):
# # pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
# # pred = pred + pred_level
# # if i != len(self.levels) - 1:
# # pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
# # out_levels.append([pred, mean, logvar])
# # pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
# # pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
# # else:
# # if self.upscale_enhancer is not None:
# # pred = self.enhance(pred, levels_input[-1][-2], training=training)
# # pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# # out_levels.append([pred, mean, logvar])
# # pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# out_levels = [
# [pred_0, mean_0, logvar_0]
# ]
# return out_levels
class PyramidNew(tf.keras.Model):
@@ -278,21 +40,6 @@ class PyramidNew(tf.keras.Model):
self.levels = [PyramidLevelNew({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))]
# if len(config['pyramid']['resolutions']) == 2:
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# if len(config['pyramid']['resolutions']) == 3:
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# if len(config['pyramid']['resolutions']) == 4:
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# if len(config['pyramid']['resolutions']) == 5:
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
# self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
self.add_noise = config['pyramid']['add_noise']
self.noise_amount = config['pyramid']['noise_amount']
self.use_offsets = config['use_offsets']
@@ -411,7 +158,7 @@ class PyramidNew(tf.keras.Model):
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
out_levels = []
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
pred, mean, logvar, z_latent = self.levels[0](levels_input[0], z_cond, training=training)
if self.use_offsets:
pred = pred + levels_input[0][2] # uv_body_posed_32
if 1 != len(self.levels):
@@ -429,7 +176,7 @@ class PyramidNew(tf.keras.Model):
for i in range(1, len(self.levels)):
input = levels_input[i]
input[0] = input[0] - pred
pred_level, mean, logvar = self.levels[i](input, z_cond, training=training)
pred_level, mean, logvar, z_latent = self.levels[i](input, z_cond, z_latent, training=training)
pred = pred + pred_level
if i != len(self.levels) - 1:
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)