From 8ef3b1045d9eb018dbc5728c9e71f2d5d38a08d3 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 21 Oct 2022 22:36:38 +0000 Subject: [PATCH] pyramid model without for loops --- models/pyramid_model_optimized.py | 200 ++++++++++++++++++++++++++++++ 1 file changed, 200 insertions(+) create mode 100644 models/pyramid_model_optimized.py diff --git a/models/pyramid_model_optimized.py b/models/pyramid_model_optimized.py new file mode 100644 index 0000000..92c1d1c --- /dev/null +++ b/models/pyramid_model_optimized.py @@ -0,0 +1,200 @@ +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 +from models.cvae_model import CVAE +from models.pyramid_level import PyramidLevel +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']),]) + + 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 = UpscalerBasicConv(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=None): + 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=None): + 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, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training) + # 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]) + + 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] ] + + +# 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=None): +# 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=None): +# 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 \ No newline at end of file