import os import time import random import torch from torch import nn import torch.nn.functional as F # import numpy as np # import scipy as sp # import plotly.graph_objects as go # import plotly.express as ex from utils.checkpoints import load_model from models.base import AEEncoder from models.pyramid_level import PyramidLevel # class Pyramid(nn.Module): # def __init__(self, config): # super(Pyramid, self).__init__() # self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], kernel_size=config['cond_ae_static']['kernel_size'], activation=config['cond_ae_static']['activation']) # if config['pyramid']['static_model'] != '': # load_model(self.static_enc, config['pyramid']['static_model'] + '_enc') # self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], kernel_size=config['cond_ae_body']['kernel_size'], activation=config['cond_ae_body']['activation']) # if config['pyramid']['body_model'] != '': # load_model(self.body_enc, config['pyramid']['body_model'] + '_enc') # self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], kernel_size=config['cond_ae_curv']['kernel_size'], activation=config['cond_ae_curv']['activation']) # if config['pyramid']['curvature_model'] != '': # load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc') # if config['freeze_encoders']: # for param in self.static_enc.parameters(): param.requires_grad = False # for param in self.body_enc.parameters(): param.requires_grad = False # for param in self.curv_enc.parameters(): param.requires_grad = False # self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae_static': config['cond_ae_static'], 'cond_ae_body': config['cond_ae_body'], 'cond_ae_curv': config['cond_ae_curv'], } }) 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): # x = self.upscale_enhancer([low, mask]) # 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 = torch.concat([z_static, z_body, z_curv], dim=-1) # return z_cond # def call(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) # # 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) # 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]) # levels_input[i][-2]=uv_mask_res # pred = F.interpolate(pred, size=(self.levels[1].size, self.levels[1].size//2), mode='bilinear') # out_levels.append([pred, mean, logvar]) # pred = pred * levels_input[1][-2][:,None,:,:] # else: # pred = self.enhance(pred, levels_input[0][-2]) # pred = F.interpolate(pred, size=(512, 256), mode='bilinear') # pred = pred * levels_input[-1][-1][:,None,:,:] # 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) # pred = pred + pred_level # if i != len(self.levels) - 1: # pred = self.levels[i].enhance(pred, levels_input[i][-2]) # out_levels.append([pred, mean, logvar]) # pred = F.interpolate(pred, size=(self.levels[i+1].size, self.levels[i+1].size//2), mode='bilinear') # pred = pred * levels_input[i+1][-2][:,None,:,:] # else: # pred = self.enhance(pred, levels_input[i][-1]) # pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask # pred = F.interpolate(pred, size=(512, 256), mode='bilinear') # out_levels.append([pred, mean, logvar]) # return out_levels class Pyramid(nn.Module): def __init__(self, config): super(Pyramid, self).__init__() self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], kernel_size=config['cond_ae_static']['kernel_size'], activation=config['cond_ae_static']['activation']) if config['static_model'] != '': load_model(self.static_enc, config['static_model'] + '_enc') self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], kernel_size=config['cond_ae_body']['kernel_size'], activation=config['cond_ae_body']['activation']) if config['body_model'] != '': load_model(self.body_enc, config['body_model'] + '_enc') self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], kernel_size=config['cond_ae_curv']['kernel_size'], activation=config['cond_ae_curv']['activation'], is_label=True, nr_of_labels=7) if config['curvature_model'] != '': load_model(self.curv_enc, config['curvature_model'] + '_enc') if config['freeze_encoders']: for param in self.static_enc.parameters(): param.requires_grad = False for param in self.body_enc.parameters(): param.requires_grad = False for param in self.curv_enc.parameters(): param.requires_grad = False self.levels = [PyramidLevel({ **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.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): z_static = self.static_enc(uv_template) z_body = self.body_enc(uv_body_posed) z_curv = self.curv_enc(uv_curv) z_cond = torch.concat([z_static, z_body, z_curv], dim=1) return z_cond def forward(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) # 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) # 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]) # levels_input[i][-2]=uv_mask_res pred = F.interpolate(pred, size=(self.levels[1].size, self.levels[1].size//2), mode='bilinear') out_levels.append([pred, mean, logvar]) pred = pred * levels_input[1][-2][:,None,:,:] else: # pred = self.enhance(pred, levels_input[0][-2], training=training) pred = self.levels[0].enhance(pred, levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res pred = F.interpolate(pred, size=(512, 256), mode='bilinear') pred = pred * levels_input[-1][-1][:,None,:,:] # 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 # 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) # 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]) out_levels.append([pred, mean, logvar]) pred = F.interpolate(pred, size=(self.levels[i+1].size, self.levels[i+1].size//2), mode='bilinear') pred = pred * levels_input[i+1][-2][:,None,:,:] else: # if self.upscale_enhancer is not None: # pred = self.enhance(pred, levels_input[-1][-2], training=training) pred = self.levels[-1].enhance(pred, levels_input[-1][-2]) pred = F.interpolate(pred, size=(512, 256), mode='bilinear') pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask out_levels.append([pred, mean, logvar]) return out_levels