luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

commit all

Browse Source
This commit is contained in:
luni
2023-01-25 12:25:59 +00:00
parent 1bf224841f
commit 7470ca9fb6
139 changed files with 82069 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_model_optimized.py
+150 -29
View File
@@ -7,9 +7,8 @@ import numpy as np
# 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.base import AEEncoder, AEEncoderNew
from models.pyramid_level import PyramidLevel, PyramidLevelNew
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
@@ -42,57 +41,65 @@ class Pyramid(tf.keras.Model):
self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
# 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)])
# 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 / 10)])
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 / 10)])
if len(config['pyramid']['resolutions']) == 4:
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)])
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
if len(config['pyramid']['resolutions']) == 5:
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)])
self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
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'] != '':
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1])
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=None):
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)
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)
z_cond = self.cond_encoder(z_cond, training=training)
return z_cond
def call(self, inputs, training=None):
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)
z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
if self.add_noise:
z_cond = z_cond + tf.random.uniform(z_cond.shape, maxval=10)
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)
@@ -184,7 +191,7 @@ class Pyramid(tf.keras.Model):
# load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
# self.upscale_enhancer.trainable = False
# def enhance(self, low, mask, training=None):
# def enhance(self, low, mask, training):
# x = self.upscale_enhancer([low, mask], training=training)
# return x
@@ -198,7 +205,7 @@ class Pyramid(tf.keras.Model):
# return z_cond
# def call(self, inputs, training=None):
# 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)
@@ -241,4 +248,118 @@ class Pyramid(tf.keras.Model):
# [pred_0, mean_0, logvar_0]
# ]
# return out_levels
# 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