pyramid model without for loops
This commit is contained in:
@@ -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
|
||||||
Reference in New Issue
Block a user