import os import time import random import tensorflow as tf import numpy as np # import plotly.graph_objects as go # import plotly.express as ex from custom_utils.custom_utils import load_model from models.base import AEEncoder, AEEncoderNew 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): def __init__(self, config): super(PyramidNew, self).__init__() self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation']) if config['pyramid']['static_model'] != '': load_model(self.static_enc, config['pyramid']['static_model'] + '_enc') self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size']) if config['pyramid']['body_model'] != '': load_model(self.body_enc, config['pyramid']['body_model'] + '_enc') if config['use_curvature']: self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['encoder_depths'], config['cond_ae_curv']['encoder_dims'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation']) 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 if config['use_curvature']: self.curv_enc.trainable = False 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'] 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) 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) if self.use_curvature: z_curv = self.curv_enc(uv_curv, 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 def sample(self, inputs): 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=False) pred = [None, None, None, None, None] for i in range(len(self.levels)): pred[i] = self.levels[i].sample(z_cond, levels_input[i][-2]) pred_1 = pred[0] pred_2 = pred[1] pred_3 = pred[2] pred_4 = pred[3] pred_5 = pred[4] pred_1 = self.levels[0].enhance(pred_1, levels_input[0][-2], training=False) pred_1 = tf.image.resize(pred_1, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR) pred_1 = pred_1 * levels_input[1][-2][:,:,:,tf.newaxis] pred_2 = pred_1 + pred_2 pred_2 = self.levels[1].enhance(pred_2, levels_input[1][-2], training=False) pred_2 = tf.image.resize(pred_2, (self.levels[2].size, self.levels[2].size//2), tf.image.ResizeMethod.BILINEAR) pred_2 = pred_2 * levels_input[2][-2][:,:,:,tf.newaxis] pred_3 = pred_2 + pred_3 pred_3 = self.levels[2].enhance(pred_3, levels_input[2][-2], training=False) pred_3 = tf.image.resize(pred_3, (self.levels[3].size, self.levels[3].size//2), tf.image.ResizeMethod.BILINEAR) pred_3 = pred_3 * levels_input[3][-2][:,:,:,tf.newaxis] pred_4 = pred_3 + pred_4 pred_4 = self.levels[3].enhance(pred_4, levels_input[3][-2], training=False) pred_4 = tf.image.resize(pred_4, (self.levels[4].size, self.levels[4].size//2), tf.image.ResizeMethod.BILINEAR) pred_4 = pred_4 * levels_input[4][-2][:,:,:,tf.newaxis] pred_5 = pred_4 + pred_5 pred_5 = pred_5 * levels_input[4][-2][:,:,:,tf.newaxis] return pred_5 def sample_seq(self, inputs): 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=False) # level 1 pred_1 = self.levels[0].sample(z_cond, levels_input[0][-2]) pred_1 = self.levels[0].enhance(pred_1, levels_input[0][-2], training=False) pred_1 = tf.image.resize(pred_1, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR) pred_1 = pred_1 * levels_input[1][-2][:,:,:,tf.newaxis] # level 2 pred_2 = self.levels[1].sample(z_cond, levels_input[1][-2]) pred_2 = pred_1 + pred_2 pred_2 = self.levels[1].enhance(pred_2, levels_input[1][-2], training=False) pred_2 = tf.image.resize(pred_2, (self.levels[2].size, self.levels[2].size//2), tf.image.ResizeMethod.BILINEAR) pred_2 = pred_2 * levels_input[2][-2][:,:,:,tf.newaxis] # level 3 pred_3 = self.levels[2].sample(z_cond, levels_input[2][-2]) pred_3 = pred_2 + pred_3 pred_3 = self.levels[2].enhance(pred_3, levels_input[2][-2], training=False) pred_3 = tf.image.resize(pred_3, (self.levels[3].size, self.levels[3].size//2), tf.image.ResizeMethod.BILINEAR) pred_3 = pred_3 * levels_input[3][-2][:,:,:,tf.newaxis] # level 4 pred_4 = self.levels[3].sample(z_cond, levels_input[3][-2]) pred_4 = pred_3 + pred_4 pred_4 = self.levels[3].enhance(pred_4, levels_input[3][-2], training=False) pred_4 = tf.image.resize(pred_4, (self.levels[4].size, self.levels[4].size//2), tf.image.ResizeMethod.BILINEAR) pred_4 = pred_4 * levels_input[4][-2][:,:,:,tf.newaxis] # level 5 pred_5 = self.levels[4].sample(z_cond, levels_input[4][-2]) pred_5 = pred_5 * levels_input[4][-2][:,:,:,tf.newaxis] pred_5 = pred_4 + pred_5 return pred_5 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.BILINEAR) 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 = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR) 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)): input = levels_input[i] input[0] = input[0] - pred pred_level, mean, logvar = self.levels[i](input, 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.BILINEAR) 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) if len(self.levels) != 5: pred = self.levels[-1].enhance(pred, levels_input[-1][-2], training=training) pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR) pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask out_levels.append([pred, mean, logvar]) return out_levels