pyramid pipeline ready
This commit is contained in:
+116
-124
@@ -1,176 +1,168 @@
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import tensorflow as tf
|
||||
import torch
|
||||
from torch import nn
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder, VariationalEncoder, Decoder, VariationalEncoderNew, DecoderNew
|
||||
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
|
||||
from utils.checkpoints import load_model
|
||||
from models.base import AEEncoder, VariationalEncoder, Decoder
|
||||
from models.upscaler import Upscaler
|
||||
|
||||
|
||||
class PyramidLevel(tf.keras.Model):
|
||||
|
||||
# class PyramidLevel(nn.Module):
|
||||
|
||||
# def __init__(self, config):
|
||||
# super(PyramidLevel, self).__init__()
|
||||
# self.latent_dim = config['latent_dim']
|
||||
# self.size = config['resolution']
|
||||
|
||||
# self.encoder: VariationalEncoder = VariationalEncoder(config['encoder_block_settings'], config['input_features'], kernel_size=config['kernel_size'])
|
||||
# self.decoder: Decoder = Decoder(config['decoder_block_settings'], config['kernel_size'])
|
||||
|
||||
# self.upscale_enhancer = None
|
||||
# if config['upscale_enhancer'] != '':
|
||||
# k_size = config['upscale_enhancer_kernel_size']
|
||||
# self.upscale_enhancer = Upscaler(config['resolution'], kernel_size=k_size)
|
||||
# load_model(self.upscale_enhancer, config['upscale_enhancer'])
|
||||
# for param in self.upscale_enhancer.parameters(): param.requires_grad = False
|
||||
|
||||
# if config['avg_garment_path'] != '':
|
||||
# self.avg_garment = np.load(config['avg_garment_path'])
|
||||
# else:
|
||||
# self.avg_garment = None
|
||||
|
||||
|
||||
# def enhance(self, low, mask, training=None):
|
||||
# x = self.upscale_enhancer([low, mask], training=training)
|
||||
# return x
|
||||
|
||||
|
||||
# def sample(self, z_cond, uv_mask, noise=None):
|
||||
# if noise is None:
|
||||
# eps = torch.normal(shape=(uv_mask.shape[0], self.latent_dim))
|
||||
# else:
|
||||
# eps = noise
|
||||
|
||||
# z = torch.concat([eps, z_cond], dim=-1)
|
||||
|
||||
# x = self.decoder(z)
|
||||
|
||||
# if self.avg_garment is not None:
|
||||
# x = x + self.avg_garment[None,:,:,:]
|
||||
|
||||
# x = x * uv_mask[:,None,:,:]
|
||||
|
||||
# return x
|
||||
|
||||
|
||||
# def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
# if training == True:
|
||||
# if not fix_noise:
|
||||
# eps = torch.normal(mean=0.0, std=1.0, size=mean.shape)
|
||||
# else:
|
||||
# eps = torch.normal(mean=0.0, std=1.0, size=mean.shape)
|
||||
# eps = torch.tile(eps, [mean.shape[0],1])
|
||||
# out = eps * torch.exp(logvar * .5) + mean
|
||||
# else:
|
||||
# out = torch.exp(logvar * .5) + mean
|
||||
# return out
|
||||
|
||||
|
||||
# def call(self, inputs, z_cond, training=None):
|
||||
# uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
|
||||
|
||||
# x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
|
||||
# mean, logvar = self.encoder(x, training=training)
|
||||
# z = self.reparameterize(mean, logvar, training)
|
||||
|
||||
# z = torch.concat([z, z_cond], dim=-1)
|
||||
# x = self.decoder(z, training=training)
|
||||
|
||||
# if self.avg_garment is not None:
|
||||
# x = x + self.avg_garment[None,:,:,:]
|
||||
|
||||
# x = x * uv_mask[:,None,:,:]
|
||||
|
||||
# # if self.upscale_enhancer is not None:
|
||||
# # pred_low = self.enhance(x, uv_mask, training=training)
|
||||
# # x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
# # x = x * uv_mask_high[:,:,:,tf.newaxis]
|
||||
|
||||
# return x, mean, logvar
|
||||
|
||||
|
||||
|
||||
class PyramidLevel(nn.Module):
|
||||
|
||||
def __init__(self, config):
|
||||
super(PyramidLevel, self).__init__()
|
||||
self.latent_dim = config['latent_dim']
|
||||
self.size = config['resolution']
|
||||
|
||||
self.encoder = VariationalEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
|
||||
self.decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
|
||||
self.encoder: VariationalEncoder = VariationalEncoder(config['encoder_block_settings'], input_features=config['input_features'], kernel_size=config['kernel_size'])
|
||||
self.decoder: Decoder = Decoder(config['decoder_block_settings'], config['kernel_size'])
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['upscale_enhancer'] != '':
|
||||
k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
|
||||
k_size = config['kernel_size']
|
||||
self.upscale_enhancer = Upscaler(config['resolution'], kernel_size=k_size).to('cuda')
|
||||
load_model(self.upscale_enhancer, config['upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
for param in self.upscale_enhancer.parameters(): param.requires_grad = False
|
||||
|
||||
|
||||
def enhance(self, low, mask, training=None):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
def enhance(self, low, mask):
|
||||
x = self.upscale_enhancer([low, mask])
|
||||
return x
|
||||
|
||||
|
||||
def sample(self, z_cond, uv_mask, noise=None):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
|
||||
eps = torch.normal(shape=(uv_mask.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
z = torch.concat([eps, z_cond], dim=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
x = x + self.avg_garment[None,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
x = x * uv_mask[:,None:,:]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
def reparameterize(self, mean, logvar, fix_noise=False):
|
||||
if self.training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
eps = torch.normal(mean=0.0, std=1.0, size=mean.shape).to(mean.device)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
eps = torch.normal(mean=0.0, std=1.0, size=mean.shape).to(mean.device)
|
||||
eps = torch.tile(eps, [mean.shape[0],1])
|
||||
out = eps * torch.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
out = torch.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def call(self, inputs, z_cond, training=None):
|
||||
def forward(self, inputs, z_cond):
|
||||
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 = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
# if self.upscale_enhancer is not None:
|
||||
# pred_low = self.enhance(x, uv_mask, training=training)
|
||||
# x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
# x = x * uv_mask_high[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
|
||||
|
||||
|
||||
class PyramidLevelNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(PyramidLevelNew, self).__init__()
|
||||
self.size = config['resolution']
|
||||
|
||||
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
|
||||
self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['upscale_enhancer'] != '':
|
||||
k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
|
||||
load_model(self.upscale_enhancer, config['upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
|
||||
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
|
||||
|
||||
|
||||
def enhance(self, low, mask, training=None):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
return x
|
||||
|
||||
|
||||
def sample(self, z_cond, uv_mask, noise=None):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
|
||||
mean, logvar = self.encoder(x)
|
||||
z = self.reparameterize(mean, logvar)
|
||||
z = torch.reshape(z, (mean.shape[0], -1, 8, 4))
|
||||
|
||||
z = torch.concat([z, z_cond], dim=1)
|
||||
x = self.decoder(z)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def call(self, inputs, z_cond, 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 = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
x = x * uv_mask[:,None,:,:]
|
||||
|
||||
# if self.upscale_enhancer is not None:
|
||||
# pred_low = self.enhance(x, uv_mask, training=training)
|
||||
|
||||
@@ -0,0 +1,246 @@
|
||||
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
|
||||
@@ -1,365 +0,0 @@
|
||||
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')
|
||||
|
||||
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
|
||||
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.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)
|
||||
|
||||
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.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
|
||||
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, training=training)
|
||||
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], 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)
|
||||
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
|
||||
Reference in New Issue
Block a user