diff --git a/configs/config_ae_body.py b/configs/config_ae_body.py index e2f9285..9930124 100644 --- a/configs/config_ae_body.py +++ b/configs/config_ae_body.py @@ -1,8 +1,21 @@ from configs.config_base import config -config['gpu'] = '0' +config['gpu'] = '1' config['lr'] = 0.0001 -config['name'] = 'ae_body_9' +config['name'] = 'ae_body_10_3' + +config['notes'] = 'convnext + convnext, add dense with tanh at the end of reverse convnext,\ + flattened bottleneck with 128 latent_dim\ remove global orientation\ + fixed gender issue with deoriented bodies, again, aaand again\ + trained on 100 garments\ + remove multiplications with garment uv mask\ + larger bottleneck, low res output upscaled\ + larger batch size, 32 instead of 4\ + ' + +config['dataset_in_memory'] = False +config['batch_size'] = 32 + config['signature_dynamic'] = { 'id': (), 'pose_shape': (72 + 11,), @@ -11,4 +24,12 @@ config['signature_dynamic'] = { 'id': (), } config['signature_static'] = { 'mesh_mask': (14475,), - } \ No newline at end of file + } + +config['cond_ae']['upscale_output'] = True +config['cond_ae']['bottleneck_size'] = 512 +config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3, 3] +config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae']['bottleneck_size']] +config['cond_ae']['decoder_input'] = (4, 2, config['cond_ae']['bottleneck_size']) +config['cond_ae']['decoder_depths'] = [3, 3, 9, 3] +config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 128, 64, 16] \ No newline at end of file diff --git a/configs/config_ae_curv.py b/configs/config_ae_curv.py index 461b1ba..9b34116 100644 --- a/configs/config_ae_curv.py +++ b/configs/config_ae_curv.py @@ -2,12 +2,14 @@ from configs.config_base import config config['gpu'] = '2' config['lr'] = 0.001 -config['name'] = 'ae_curv_9' +config['name'] = 'x' #'ae_curv_9' config['cond_ae']['is_label'] = True +config['cond_ae']['upscale_output'] = False + config['signature_dynamic'] = { 'id': (), - 'uv_curvature': (512, 256), + 'uv_curvature': (512, 256, 7), 'mesh_curvature': (14475) } config['signature_static'] = { 'mesh_mask': (14475,), diff --git a/configs/config_ae_static.py b/configs/config_ae_static.py index 00c5d82..85adca4 100644 --- a/configs/config_ae_static.py +++ b/configs/config_ae_static.py @@ -4,6 +4,9 @@ config['gpu'] = '1' config['lr'] = 0.0005 config['name'] = 'ae_static_7_2' +config['cond_ae']['upscale_output'] = False + + config['signature_dynamic'] = { 'id': (), } config['signature_static'] = { 'uv_static': (512, 256, 3), 'mesh_static': (14475, 3,), diff --git a/configs/config_cvae.py b/configs/config_cvae.py index e71f6ed..eccece7 100644 --- a/configs/config_cvae.py +++ b/configs/config_cvae.py @@ -1,7 +1,7 @@ from configs.config_base import config -config['gpu'] = '0' -config['name'] = 'cvae_19_2' +config['gpu'] = '4' +config['name'] = 'cvae_low_4_3_d32' config['notes'] = '\ offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \ retrained body encoder with deoriented bodies \ @@ -11,10 +11,12 @@ remove shape from unposed garment\ training on 100 garments with offsets over the average of these\ fixed autoencoders!!!!!!!!!\ cvae only loads encoders instead AEs\ -with edge and collision loss\ +full res loss\ ' +config['dataset_in_memory'] = False -# config['resolution'] = 32 + +config['resolution'] = 32 if config['resolution'] != 512: @@ -32,7 +34,7 @@ config['signature_dynamic'] = { 'id': (), 'mesh_posed': (14475, 3,), 'mesh_unposed': (14475, 3,), 'uv_unposed': (512, 256, 3), - 'uv_curvature': (512, 256), + 'uv_curvature': (512, 256, 7), } config['signature_static'] = { 'uv_static': (512, 256, 3), 'mesh_mask': (14475,), diff --git a/configs/config_pyramid.py b/configs/config_pyramid.py new file mode 100644 index 0000000..513368d --- /dev/null +++ b/configs/config_pyramid.py @@ -0,0 +1,108 @@ +from configs.config_base import config + +config['gpu'] = '1' +config['name'] = 'cvae_pyr_4_4' +config['notes'] = '\ +offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \ +retrained body encoder with deoriented bodies \ +use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\ +fixed gender issue and retrained body encoder, FINAL\ +remove shape from unposed garment\ +training on 100 garments with offsets over the average of these\ +fixed autoencoders!!!!!!!!!\ +common condition encoders for all levels, only run once\ +full res loss on each level\ +all higher levels are added to the output of first level\ +using custom MLP upscale enhancement\ +layers 32 and 512\ +try loading weights from previous trainings\ +' +config['dataset_in_memory'] = False +config['resolution'] = None + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.0001 + +config['signature_dynamic'] = { 'id': (), + 'pose_shape': (72 + 11,), + 'uv_body_posed': (512, 256, 3), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'mesh_unposed': (14475, 3,), + 'uv_unposed': (512, 256, 3), + 'uv_curvature': (512, 256, 7), + } +config['signature_static'] = { 'uv_static': (512, 256, 3), + 'mesh_mask': (14475,), + 'uv_mask': (512, 256,), + } + +config['pyramid'] = dict() +config['pyramid']['is_label'] = False +config['pyramid']['latent_dim'] = config['latent_dim'] + +config['pyramid']['resolutions'] = [32, 512] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3) + config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3) + config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7) + config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3) + config['signature_static']['uv_mask' + res_modifier] = (res, res//2) + + +config['pyramid']['static_model'] = 'checkpoints/ae_static_7_2_x' +config['pyramid']['body_model'] = 'checkpoints/ae_body_9_x' +config['pyramid']['curvature_model'] = 'checkpoints/ae_curv_9_x' +config['pyramid']['freeze_encoders'] = True + +config['avg_garment_path'] = '' + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['latent_dim'] = config['latent_dim'] + level['kernel_size'] = 7 + level['encoder_depths'] = [3, 3, 9, 3] + level['encoder_dims'] = [32, 64, 128, 256] + level['decoder_input'] = (16, 8, 80) + level['decoder_depths'] = [3, 3, 9, 3] + level['decoder_dims'] = [256, 128, 64, 16] + level['avg_garment_path'] = '' + level['upscale_enhancer'] = '' + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['latent_dim'] = config['latent_dim'] + level['kernel_size'] = 5 + level['encoder_depths'] = [16, 16, 9, 9] + level['encoder_dims'] = [64, 128, 256, 256] + level['decoder_input'] = (4, 2, 128) + level['decoder_depths'] = [16, 16, 9, 9] + level['decoder_dims'] = [256, 256, 128, 64] + level['avg_garment_path'] = '' + level['upscale_enhancer'] = 'checkpoints/up_2_3_d128' + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['latent_dim'] = config['latent_dim'] + level['kernel_size'] = 3 + level['encoder_depths'] = [27, 27] + level['encoder_dims'] = [256, 256] + level['decoder_input'] = (4, 2, 256) + level['decoder_depths'] = [27, 27] + level['decoder_dims'] = [256, 256] + level['avg_garment_path'] = '' + level['upscale_enhancer'] = 'checkpoints/up_2_3' + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + diff --git a/cvae_pyramid.py b/cvae_pyramid.py index bbdbc10..ea32530 100644 --- a/cvae_pyramid.py +++ b/cvae_pyramid.py @@ -1,5 +1,6 @@ import os -os.environ["CUDA_VISIBLE_DEVICES"] = "0" +from configs.config_pyramid import config +os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') @@ -23,178 +24,180 @@ import scipy as sp from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces import plotly.graph_objects as go import plotly.express as ex +from plotly.subplots import make_subplots from PIL import Image -from utils import str_to_bool from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing -from load_models import load_generator, load_discriminator -from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss +from losses import edge_loss, uv_mask_loss, bend_loss, kl_loss # from losses import collision_penalty as collision_penalty from losses import collision_loss as collision_penalty from custom_utils.custom_utils import batch_to_dict, load_model from dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL -from models.cvae_pyramid_model import Pyramid +from models.pyramid_model import Pyramid -name='cvae_pyr_4_d32' def init_wandb(): - run = wandb.init(project='vae', name=name, - notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \ -new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \ -offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \ -retrained body encoder with deoriented bodies \ -use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\ -fixed gender issue and retrained body encoder\ -remove shape from unposed garment\ -training on 100 garments with offsets over the average of these\ -l2 loss on both 128 and full res, bottom predicts base, high res predicts offsets over this\ -low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input\ -not using filled in versions, but using custom MLP upscaling') + run = wandb.init(project='vae', name=config['name'], notes=config['notes']) + artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae_pyramid.py') artifact.add_file('models/cvae_model.py') - artifact.add_file('models/cvae_pyramid_model.py') + artifact.add_file('models/pyramid_model.py') + artifact.add_file('models/pyramid_level.py') artifact.add_file('models/upscaler.py') - artifact.add_file('dataset.py') artifact.add_file('losses.py') + artifact.add_file('configs/config_base.py') + artifact.add_file('configs/config_cvae.py') + artifact.add_file('models/base.py') + artifact.add_file('models/cond_encoder.py') + artifact.add_file('dataset.py') artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) uv2mesh = UV2Mesh(res=256) -def plot_3d(mesh_gt, mesh_pred_unposed, mesh_pred_unposed_low, faces, collision, collision_faces, collision_color): - mesh_pred_unposed[:,2] += 1 - mesh_pred_unposed_low[:,2] += 2 - collision[:,2] += 3 - fig = go.Figure(data=[ - go.Mesh3d( +def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, collision_faces, collision_color): + fig = make_subplots(rows=1, cols=1+len(mesh_pred_unposed_levels)+1+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(1+len(mesh_pred_unposed_levels)+1+1)] ]) + + fig.add_trace(go.Mesh3d( x=mesh_gt[:,0], y=mesh_gt[:,2], z=mesh_gt[:,1], i = faces[:,0], j = faces[:,1], - k = faces[:,2] + k = faces[:,2], + text="Ground Truth" ), - go.Mesh3d( - x=mesh_pred_unposed[:,0], - y=mesh_pred_unposed[:,2], - z=mesh_pred_unposed[:,1], + row=1, col=1) + + for i, mesh in enumerate(mesh_pred_unposed_levels): + fig.add_trace(go.Mesh3d( + x=mesh[:,0], + y=mesh[:,2], + z=mesh[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2], + text='Unposed level_' + str(i) ), - go.Mesh3d( - x=mesh_pred_unposed_low[:,0], - y=mesh_pred_unposed_low[:,2], - z=mesh_pred_unposed_low[:,1], + row=1, col=i+2) + + fig.add_trace(go.Mesh3d( + x=mesh_posed[:,0], + y=mesh_posed[:,1], + z=mesh_posed[:,2], i = faces[:,0], j = faces[:,1], k = faces[:,2], + text='Posed GT' ), - go.Mesh3d( + row=1, col=1+len(mesh_pred_unposed_levels)+1) + + fig.add_trace(go.Mesh3d( x=collision[:,0], y=collision[:,2], z=collision[:,1], i = collision_faces[:,0], j = collision_faces[:,1], k = collision_faces[:,2], - vertexcolor=collision_color - ) - ]) + vertexcolor=collision_color, + text='Posed output on body' + ), + row=1, col=1+len(mesh_pred_unposed_levels)+1+1) - fig.update_layout(scene_aspectmode='manual', - scene_aspectratio=dict(x=1, y=3, z=1)) - fig.update_layout( - scene = dict( - xaxis = dict(nticks=4, range=[-1,1],), - yaxis = dict(nticks=4, range=[-1,5],), - zaxis = dict(nticks=4, range=[-1,1],),) + if len(mesh_pred_unposed_levels) == 3: + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), + scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), + scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1), + scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), + scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1), + scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,) ) - + helper = dict(xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,1],), + zaxis = dict(nticks=4, range=[-1,1],),) + fig.update_layout(scene = helper, + scene2 = helper, + scene3 = helper, + scene4 = helper, + scene5 = helper, + scene6 = helper, + ) + if len(mesh_pred_unposed_levels) == 2: + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), + scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), + scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1), + scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), + scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1), + ) + helper = dict(xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,1],), + zaxis = dict(nticks=4, range=[-1,1],),) + fig.update_layout(scene = helper, + scene2 = helper, + scene3 = helper, + scene4 = helper, + scene5 = helper, + ) return fig +def group_inputs_for_levels(batch): + levels_input = [] + for res in config['pyramid']['resolutions']: + res_modifier = '_' + str(res) if res != 512 else '' + levels_input.append([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']]) + return levels_input -def log_normal_pdf(sample, mean, logvar, raxis=1): - log2pi = tf.math.log(2. * np.pi) - res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis) - return res - - -def kl_loss(mean, logvar): - kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar) - kl_loss = tf.reduce_mean(kl_loss) - kl_loss *= -0.5 - return kl_loss - - -# @tf.function +@tf.function def train_step(batch): with tf.GradientTape() as tape: - uv_curvature = tf.one_hot(batch['uv_curvature'], 7) - uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7) - - pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], - batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier],]) - - mesh_pred_full = uv2mesh(pred_full) - mesh_pred_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full) - mesh_pred_low = uv2mesh(pred_low) - mesh_pred_posed_low = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_low) - mesh_orig = batch['mesh_unposed'] + levels_input = group_inputs_for_levels(batch) + out_levels = model(levels_input, training=True) shape = batch['pose_shape'][:, 72:] gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) - mesh_pred_posed_full = mesh_pred_posed_full + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] - mesh_pred_posed_low = mesh_pred_posed_low + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] + mesh_orig = batch['mesh_unposed'] - e_loss = tf.reduce_mean(edge_loss(mesh_pred_full, mesh_orig, batch['mesh_mask'])) - c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed_full, batch['mesh_body_posed'], batch['mesh_mask'])) - l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred_full)), axis=-1), axis=[1, 2]) - l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full) ** 2), axis=-1), axis=[1, 2]) - ms_ssim = tf.image.ssim_multiscale(pred_full + 1, batch['uv_unposed'] + 1, max_val=3) - kl = kl_loss(mean_full, logvar_full) + output = dict() + reconstr_loss_total = 0 + kl_total = 0 + for i, res in enumerate(config['pyramid']['resolutions']): + res_modifier = '_' + str(res) if res != 512 else '' + pred, mean, logvar = out_levels[i] - pred_low = pred_low * batch['uv_mask' + resolution_modifier][:,:,:,tf.newaxis] - e_loss_low = tf.reduce_mean(edge_loss(mesh_pred_low, mesh_orig, batch['mesh_mask'])) - c_loss_low = tf.reduce_mean(collision_penalty(mesh_pred_posed_low, batch['mesh_body_posed'], batch['mesh_mask'])) - l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + resolution_modifier] - pred_low)), axis=-1), axis=[1, 2]) - l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + resolution_modifier] - pred_low) ** 2), axis=-1), axis=[1, 2]) - ms_ssim_low = tf.image.ssim(pred_low + 1, batch['uv_unposed' + resolution_modifier] + 1, max_val=3) - kl_low = kl_loss(mean_full, logvar_full) + mesh_pred = uv2mesh(pred) + mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] + mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped) - reconstr_loss = l2_loss + l2_loss_low - reconstr_loss = tf.reduce_mean(reconstr_loss) - loss = reconstr_loss + balancing_coeff * (kl + kl_low) + e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])) + l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2])) + ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + kl = kl_loss(mean, logvar) - ms_ssim = tf.reduce_mean(ms_ssim) - l1_loss = tf.reduce_mean(l1_loss) - l2_loss = tf.reduce_mean(l2_loss) - mean = tf.reduce_mean(mean_full) - logvar = tf.reduce_mean(logvar_full) + reconstr_loss_total += l2_loss #+ e_loss + c_loss + kl_total += kl - ms_ssim_low = tf.reduce_mean(ms_ssim_low) - l1_loss_low = tf.reduce_mean(l1_loss_low) - l2_loss_low = tf.reduce_mean(l2_loss_low) - mean_low = tf.reduce_mean(mean_low) - logvar_low = tf.reduce_mean(logvar_low) + mean = tf.reduce_mean(mean) + logvar = tf.reduce_mean(logvar) + + metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + + output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, + 'train/ms_ssim' + res_modifier: ms_ssim, 'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl, + 'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, + 'train/euclid_error': metric } } + + loss = reconstr_loss_total + balancing_coeff * kl_total - mesh_orig_2 = uv2mesh(batch['uv_unposed']) - metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) - metric_low = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_low) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) - metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) - return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, - 'train/euclid_error_reconst': metric_2, 'train/reconstr_loss': reconstr_loss, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, - 'train/var': np.e ** logvar, 'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss, - 'train/edge_loss' + resolution_modifier: e_loss_low, 'train/collision_loss' + resolution_modifier: c_loss_low, 'train/l1_loss' + resolution_modifier: l1_loss_low, 'train/euclid_error' + resolution_modifier: metric_low, 'train/ms_ssim' + resolution_modifier: ms_ssim_low, - 'train/mean' + resolution_modifier: mean_low, 'train/logvar' + resolution_modifier: logvar_low, 'train/kl_loss' + resolution_modifier: kl_low, - 'train/var' + resolution_modifier: np.e ** logvar_low, 'train/l2_loss' + resolution_modifier: l2_loss_low } + return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output } @tf.function @@ -237,51 +240,53 @@ def valid_step(batch): def plot_everything(batch: dict, type: str): - uv_curvature = tf.one_hot(batch['uv_curvature'], 7) - uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7) - pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], - batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier]]) + levels_input = group_inputs_for_levels(batch) + output_levels = model(levels_input, training=False) - mesh_pred_full = uv2mesh(pred_full) - mesh_pred_low = uv2mesh(pred_low) - mesh_orig = batch['mesh_unposed'] + mesh_pred_levels = [uv2mesh(output_levels[i][0]).numpy() for i in range(len(output_levels))] + mesh_orig = batch['mesh_unposed'].numpy() + mesh_orig_posed = batch['mesh_posed'].numpy() + body = batch['mesh_body_posed'].numpy() orig_img = batch['uv_unposed'] uv_static_img = batch['uv_static'] uv_body_posed_img = batch['uv_body_posed'] - pred_img = pred_full - pred_low_up = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5) - img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_low_up), axis=2) + uv_curvature = np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 2 - 1 + pred_img = np.concatenate([output_levels[i][0] for i in range(len(output_levels))], axis=2) + img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2) imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] shape = batch['pose_shape'][:, 72:] gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) - shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) + shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy() F = np.load("custom_utils/smpl_faces.npy") mask = batch['mesh_mask'].numpy().astype(bool) - garment_unposed, faces = maskFacesBatch(mesh_pred_full.numpy(), F, mask) - garment_unposed_low, faces = maskFacesBatch(mesh_pred_low.numpy(), F, mask) - garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])] - garment_unposed_low = [garment_unposed_low[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])] - gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])] + + _, faces = maskFacesBatch(np.ones_like(mesh_orig), F, mask) + gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])] + gt_posed = [mesh_orig_posed[i, mask[i], :] for i in range(mask.shape[0])] + mesh_unposed_shaped_levels = [[] for _ in range(mask.shape[0])] + for i in range(len(output_levels)): + mesh_unposed_shaped = [mesh_pred_levels[i][j, mask[j], :] + shape_offsets[j, mask[j]] for j in range(mask.shape[0])] + for k in range(mask.shape[0]): + mesh_unposed_shaped_levels[k].append(mesh_unposed_shaped[k]) + + # full res version + mesh_unposed_shaped_full = np.zeros((mask.shape[0], 14475, 3), dtype=np.float32) for i in range(mask.shape[0]): - mesh_pred_full[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :] - body = batch['mesh_body_posed'].numpy() - mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).numpy() + mesh_unposed_shaped_full[i, mask[i], :] = mesh_unposed_shaped[i] + mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_unposed_shaped_full).numpy() + mesh_posed = [mesh_posed[j, mask[j], :] for j in range(mask.shape[0])] - garment_posed = [mesh_posed_full[i, mask[i], :] for i in range(body.shape[0])] - collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))] - collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))] - collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))] + collision = [np.concatenate((body[i], mesh_posed[i]), axis=0) for i in range(len(mesh_posed))] + collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(mesh_posed))] + collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(mesh_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(mesh_posed))] - figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], garment_unposed_low[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) } - - mean_hist = wandb.Histogram(sequence=mean_full.numpy().mean(axis=0)) - logvar_hist = wandb.Histogram(sequence=logvar_full.numpy().mean(axis=0)) + figs = { type + "/garment_" + str(i): plot_3d(gt[i], mesh_unposed_shaped_levels[i], faces[i], gt_posed[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) } - return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } + return { **figs, **{ type + '/images': imgs, } } def sample(folder): @@ -289,7 +294,7 @@ def sample(folder): load_model(model,folder) count = 0 for batch in dataset_train: - batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) + batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) uv_curvature = tf.one_hot(batch['uv_curvature'], 7) nr_of_samples=2 @@ -407,75 +412,41 @@ def sample(folder): if __name__ == '__main__': - # init_wandb() + init_wandb() allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) - batch_size = 4 - low_resolution = 32 - resolution_modifier = '_' + str(low_resolution) - smpl = SMPL() factory = DatasetFactory() - signature_dynamic = { 'id': (), - 'pose_shape': (72 + 11,), - # 'uv_posed': (512, 256, 3), - 'uv_body_posed': (512, 256, 3), - 'uv_body_posed' + resolution_modifier: (low_resolution, low_resolution//2, 3), - 'mesh_body_posed': (14475, 3,), - 'mesh_posed': (14475, 3,), - 'mesh_unposed': (14475, 3,), - 'uv_unposed': (512, 256, 3), - 'uv_unposed' + resolution_modifier: (low_resolution, low_resolution//2, 3), - 'uv_curvature': (512, 256), - 'uv_curvature' + resolution_modifier: (low_resolution, low_resolution//2), - } - signature_static = { 'uv_static': (512, 256, 3), - 'uv_static' + resolution_modifier: (low_resolution, low_resolution//2, 3), - # 'mesh_static': (14475, 3,), - 'mesh_mask': (14475,), - 'uv_mask': (512, 256,), - 'uv_mask' + resolution_modifier: (low_resolution, low_resolution//2,), - # 'uv_shaped_body': (512, 256, 3), - } + dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=config['pyramid']['resolutions']) + static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=config['pyramid']['resolutions']) + dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=config['pyramid']['resolutions']) + static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static']) - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512, low_resolution]) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_resolution]) - dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic) - static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static) + model: Pyramid = Pyramid(config) - - model = Pyramid(latent_dim=128, - low_res=low_resolution, - freeze_encoders=True, - static_model='checkpoints/ae_static_6', - body_model='checkpoints/ae_body_8', - curvature_model='checkpoints/ae_curvature_7', - upscaler_model='checkpoints/up_2_3') - - # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False) - optimizer = tf.optimizers.Adam(learning_rate=0.0001) - balancing_coeff = 0.0001 + optimizer = tf.optimizers.Adam(learning_rate=config['lr']) + balancing_coeff = config['starting_balancing_coeff'] reg_coeff = 1 - folder = 'models/ae_logs/' valid_iterator = iter(dataset_valid) gamma = 1 balancing_activated_step = 0 - # avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy') - # sample('checkpoints/cvae_13') - ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3) + + load_model(model.cond_encoder, 'checkpoints/cvae_19') + load_model(model.levels[0], 'checkpoints/cvae_low_4_3_d32') + load_model(model.levels[-1], 'checkpoints/cvae_19') step = 0 for batch in dataset_train: - batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) + batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) train_loss = train_step(batch) @@ -498,18 +469,7 @@ if __name__ == '__main__': if balancing_activated_step == 0: balancing_activated_step = step balancing_coeff = balancing_coeff + 1 / 15000 - # if train_loss['train/euclid_error'] < 0.015: - # reg_coeff = reg_coeff + 100 / 15000 - # if step > 15000 and balancing_coeff <= 1e-3: - # balancing_coeff = 1e-2 - # if step > 20000 and balancing_coeff <= 1e-3: - # balancing_coeff = 1e-2 - # if step > 25000 and balancing_coeff <= 1e-2: - # balancing_coeff = 1e-1 - # if step > 30000 and balancing_coeff <= 1e-1: - # balancing_coeff = 1 - # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss'] to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } } diff --git a/cvae_pyramid_old.py b/cvae_pyramid_old.py new file mode 100644 index 0000000..bbdbc10 --- /dev/null +++ b/cvae_pyramid_old.py @@ -0,0 +1,520 @@ +import os +os.environ["CUDA_VISIBLE_DEVICES"] = "0" +import tensorflow as tf + +gpus = tf.config.experimental.list_physical_devices('GPU') +if gpus: + try: + # Currently, memory growth needs to be the same across GPUs + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + logical_gpus = tf.config.experimental.list_logical_devices('GPU') + print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs") + except RuntimeError as e: + # Memory growth must be set before GPUs have been initialized + print(e) + +import time +import argparse +import wandb +import random +import numpy as np +import scipy as sp +from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces +import plotly.graph_objects as go +import plotly.express as ex +from PIL import Image + +from utils import str_to_bool +from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing +from load_models import load_generator, load_discriminator +from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss +# from losses import collision_penalty as collision_penalty +from losses import collision_loss as collision_penalty +from custom_utils.custom_utils import batch_to_dict, load_model +from dataset import DatasetFactory, StaticDataset +from custom_utils.smpl import SMPL +from models.cvae_pyramid_model import Pyramid + + +name='cvae_pyr_4_d32' +def init_wandb(): + run = wandb.init(project='vae', name=name, + notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \ +new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \ +offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \ +retrained body encoder with deoriented bodies \ +use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\ +fixed gender issue and retrained body encoder\ +remove shape from unposed garment\ +training on 100 garments with offsets over the average of these\ +l2 loss on both 128 and full res, bottom predicts base, high res predicts offsets over this\ +low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input\ +not using filled in versions, but using custom MLP upscaling') + artifact = wandb.Artifact('code', type='code') + artifact.add_file('cvae_pyramid.py') + artifact.add_file('models/cvae_model.py') + artifact.add_file('models/cvae_pyramid_model.py') + artifact.add_file('models/upscaler.py') + artifact.add_file('dataset.py') + artifact.add_file('losses.py') + artifact.add_file('convnext/models/convnext_tf.py') + run.log_artifact(artifact) + +uv2mesh = UV2Mesh(res=256) + +def plot_3d(mesh_gt, mesh_pred_unposed, mesh_pred_unposed_low, faces, collision, collision_faces, collision_color): + mesh_pred_unposed[:,2] += 1 + mesh_pred_unposed_low[:,2] += 2 + collision[:,2] += 3 + fig = go.Figure(data=[ + go.Mesh3d( + x=mesh_gt[:,0], + y=mesh_gt[:,2], + z=mesh_gt[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + ), + go.Mesh3d( + x=mesh_pred_unposed[:,0], + y=mesh_pred_unposed[:,2], + z=mesh_pred_unposed[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + ), + go.Mesh3d( + x=mesh_pred_unposed_low[:,0], + y=mesh_pred_unposed_low[:,2], + z=mesh_pred_unposed_low[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + ), + go.Mesh3d( + x=collision[:,0], + y=collision[:,2], + z=collision[:,1], + i = collision_faces[:,0], + j = collision_faces[:,1], + k = collision_faces[:,2], + vertexcolor=collision_color + ) + ]) + + fig.update_layout(scene_aspectmode='manual', + scene_aspectratio=dict(x=1, y=3, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,5],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + return fig + + + + +def log_normal_pdf(sample, mean, logvar, raxis=1): + log2pi = tf.math.log(2. * np.pi) + res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis) + return res + + +def kl_loss(mean, logvar): + kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar) + kl_loss = tf.reduce_mean(kl_loss) + kl_loss *= -0.5 + return kl_loss + + +# @tf.function +def train_step(batch): + + with tf.GradientTape() as tape: + uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7) + + pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], + batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier],]) + + mesh_pred_full = uv2mesh(pred_full) + mesh_pred_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full) + mesh_pred_low = uv2mesh(pred_low) + mesh_pred_posed_low = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_low) + mesh_orig = batch['mesh_unposed'] + + shape = batch['pose_shape'][:, 72:] + gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) + shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) + mesh_pred_posed_full = mesh_pred_posed_full + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] + mesh_pred_posed_low = mesh_pred_posed_low + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] + + e_loss = tf.reduce_mean(edge_loss(mesh_pred_full, mesh_orig, batch['mesh_mask'])) + c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed_full, batch['mesh_body_posed'], batch['mesh_mask'])) + l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred_full)), axis=-1), axis=[1, 2]) + l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full) ** 2), axis=-1), axis=[1, 2]) + ms_ssim = tf.image.ssim_multiscale(pred_full + 1, batch['uv_unposed'] + 1, max_val=3) + kl = kl_loss(mean_full, logvar_full) + + pred_low = pred_low * batch['uv_mask' + resolution_modifier][:,:,:,tf.newaxis] + e_loss_low = tf.reduce_mean(edge_loss(mesh_pred_low, mesh_orig, batch['mesh_mask'])) + c_loss_low = tf.reduce_mean(collision_penalty(mesh_pred_posed_low, batch['mesh_body_posed'], batch['mesh_mask'])) + l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + resolution_modifier] - pred_low)), axis=-1), axis=[1, 2]) + l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + resolution_modifier] - pred_low) ** 2), axis=-1), axis=[1, 2]) + ms_ssim_low = tf.image.ssim(pred_low + 1, batch['uv_unposed' + resolution_modifier] + 1, max_val=3) + kl_low = kl_loss(mean_full, logvar_full) + + reconstr_loss = l2_loss + l2_loss_low + reconstr_loss = tf.reduce_mean(reconstr_loss) + loss = reconstr_loss + balancing_coeff * (kl + kl_low) + + ms_ssim = tf.reduce_mean(ms_ssim) + l1_loss = tf.reduce_mean(l1_loss) + l2_loss = tf.reduce_mean(l2_loss) + mean = tf.reduce_mean(mean_full) + logvar = tf.reduce_mean(logvar_full) + + ms_ssim_low = tf.reduce_mean(ms_ssim_low) + l1_loss_low = tf.reduce_mean(l1_loss_low) + l2_loss_low = tf.reduce_mean(l2_loss_low) + mean_low = tf.reduce_mean(mean_low) + logvar_low = tf.reduce_mean(logvar_low) + + mesh_orig_2 = uv2mesh(batch['uv_unposed']) + metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + metric_low = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_low) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + gradients = tape.gradient(loss, model.trainable_variables) + optimizer.apply_gradients(zip(gradients, model.trainable_variables)) + return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, + 'train/euclid_error_reconst': metric_2, 'train/reconstr_loss': reconstr_loss, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, + 'train/var': np.e ** logvar, 'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss, + 'train/edge_loss' + resolution_modifier: e_loss_low, 'train/collision_loss' + resolution_modifier: c_loss_low, 'train/l1_loss' + resolution_modifier: l1_loss_low, 'train/euclid_error' + resolution_modifier: metric_low, 'train/ms_ssim' + resolution_modifier: ms_ssim_low, + 'train/mean' + resolution_modifier: mean_low, 'train/logvar' + resolution_modifier: logvar_low, 'train/kl_loss' + resolution_modifier: kl_low, + 'train/var' + resolution_modifier: np.e ** logvar_low, 'train/l2_loss' + resolution_modifier: l2_loss_low } + + +@tf.function +def valid_step(batch): + uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']]) + + mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) + mesh_orig = batch['mesh_unposed'] + e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])) + l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2]) + b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask']) + # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2) + ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3) + fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3]) + # mesh_orig_2 = uv2mesh(batch['uv_unposed']) + metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + + reconstr_loss = l1_loss + # logpx_z = -reconstr_loss + # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.) + # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar) + # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x)) + kl = kl_loss(mean, logvar) + # loss = reconstr_loss / (2 * (gamma ** 2)) + kl + loss = reconstr_loss + balancing_coeff * kl + ms_ssim = tf.reduce_mean(ms_ssim) + l1_loss = tf.reduce_mean(l1_loss) + fft_loss = tf.reduce_mean(fft_loss) + b_loss = tf.reduce_mean(b_loss) + reconstr_loss = tf.reduce_mean(reconstr_loss) + mean = tf.reduce_mean(mean) + logvar = tf.reduce_mean(logvar) + + return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, + 'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss, + 'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar } + + +def plot_everything(batch: dict, type: str): + uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7) + pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], + batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier]]) + + mesh_pred_full = uv2mesh(pred_full) + mesh_pred_low = uv2mesh(pred_low) + mesh_orig = batch['mesh_unposed'] + + orig_img = batch['uv_unposed'] + uv_static_img = batch['uv_static'] + uv_body_posed_img = batch['uv_body_posed'] + pred_img = pred_full + pred_low_up = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5) + img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_low_up), axis=2) + imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] + + shape = batch['pose_shape'][:, 72:] + gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) + shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) + + F = np.load("custom_utils/smpl_faces.npy") + mask = batch['mesh_mask'].numpy().astype(bool) + garment_unposed, faces = maskFacesBatch(mesh_pred_full.numpy(), F, mask) + garment_unposed_low, faces = maskFacesBatch(mesh_pred_low.numpy(), F, mask) + garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])] + garment_unposed_low = [garment_unposed_low[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])] + gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])] + + for i in range(mask.shape[0]): + mesh_pred_full[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :] + body = batch['mesh_body_posed'].numpy() + mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).numpy() + + garment_posed = [mesh_posed_full[i, mask[i], :] for i in range(body.shape[0])] + collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))] + collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))] + collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))] + + figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], garment_unposed_low[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) } + + mean_hist = wandb.Histogram(sequence=mean_full.numpy().mean(axis=0)) + logvar_hist = wandb.Histogram(sequence=logvar_full.numpy().mean(axis=0)) + + return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } + + +def sample(folder): + # run = wandb.init(project='vae_plots', name=folder.split('/')[-1]) + load_model(model,folder) + count = 0 + for batch in dataset_train: + batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) + uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + nr_of_samples=2 + + pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples) + pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:] + pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1]) + + mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) + mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy() + mesh_orig = batch['mesh_posed'][0].numpy() + mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0] + + body = batch['mesh_body_posed'][0].numpy() + body_color = np.ones_like(body, dtype=np.int32)*100 + + F = np.load("custom_utils/smpl_faces.npy") + mask = batch['mesh_mask'].numpy().astype(bool)[0] + _, faces = _maskFaces(np.ones((14475, 3)), F, mask) + + + pred_zero = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True) + mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2])) + mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy() + + + + + meshes = [] + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=mesh_orig_reposed[mask,0], + y=mesh_orig_reposed[mask,2], + z=mesh_orig_reposed[mask,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + for i in range(nr_of_samples): + x = mesh_posed[i][mask, :].copy() + # x[:, 2] += 1 * (i + 1) + x[:, 2] += 1 + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,2], + z=x[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + + for i in range(nr_of_samples): + x = mesh_posed_zero[i][mask, :].copy() + # x[:, 2] += 1 * (i + 1) + x[:, 2] += 2 + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,2], + z=x[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + # for i in range(5): + # x = body.copy() + # x[:, 2] += 1 * i + # meshes.append(go.Mesh3d( + # # 8 vertices of a cube + # x=x[:,0], + # y=x[:,2], + # z=x[:,1], + # # i, j and k give the vertices of triangles + # i = F[:,0], + # j = F[:,1], + # k = F[:,2], + # vertexcolor=body_color + # )) + + + fig = go.Figure(data=meshes) + + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1)) + # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],), + yaxis = dict(nticks=4, range=[-1,3],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + # wandb.log({'plot_' + str(count): fig}) + fig.show() + + pred = pred.numpy() + frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])] + frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:], + save_all=True, duration=500, loop=0) + + + if count > 10: + break + + print(count) + count += 1 + + print() + + +if __name__ == '__main__': + # init_wandb() + allow_memory_growth() + parser = argparse.ArgumentParser(description='') + parser.add_argument('--name', type=str) + args = vars(parser.parse_args()) + + batch_size = 4 + low_resolution = 32 + resolution_modifier = '_' + str(low_resolution) + + smpl = SMPL() + + factory = DatasetFactory() + signature_dynamic = { 'id': (), + 'pose_shape': (72 + 11,), + # 'uv_posed': (512, 256, 3), + 'uv_body_posed': (512, 256, 3), + 'uv_body_posed' + resolution_modifier: (low_resolution, low_resolution//2, 3), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'mesh_unposed': (14475, 3,), + 'uv_unposed': (512, 256, 3), + 'uv_unposed' + resolution_modifier: (low_resolution, low_resolution//2, 3), + 'uv_curvature': (512, 256), + 'uv_curvature' + resolution_modifier: (low_resolution, low_resolution//2), + } + signature_static = { 'uv_static': (512, 256, 3), + 'uv_static' + resolution_modifier: (low_resolution, low_resolution//2, 3), + # 'mesh_static': (14475, 3,), + 'mesh_mask': (14475,), + 'uv_mask': (512, 256,), + 'uv_mask' + resolution_modifier: (low_resolution, low_resolution//2,), + # 'uv_shaped_body': (512, 256, 3), + } + + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512, low_resolution]) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_resolution]) + dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic) + static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static) + + + model = Pyramid(latent_dim=128, + low_res=low_resolution, + freeze_encoders=True, + static_model='checkpoints/ae_static_6', + body_model='checkpoints/ae_body_8', + curvature_model='checkpoints/ae_curvature_7', + upscaler_model='checkpoints/up_2_3') + + # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False) + optimizer = tf.optimizers.Adam(learning_rate=0.0001) + balancing_coeff = 0.0001 + reg_coeff = 1 + folder = 'models/ae_logs/' + valid_iterator = iter(dataset_valid) + gamma = 1 + balancing_activated_step = 0 + # avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy') + + + # sample('checkpoints/cvae_13') + + + ckpt = tf.train.Checkpoint(model=model) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) + + step = 0 + for batch in dataset_train: + batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) + + train_loss = train_step(batch) + + # data_valid = valid_iterator.get_next() + # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic)) + # valid_loss = valid_step(data_valid) + valid_loss = {} + + to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } + if step % 2500 == 0: + train_figs = plot_everything(batch, 'train') + # valid_figs = plot_everything(data_valid, 'valid') + valid_figs = {} + to_log = { **to_log, **train_figs, **valid_figs } + + if step % 10000 == 0: + manager.save(checkpoint_number=step) + + if train_loss['train/euclid_error'] < 0.01: + if balancing_activated_step == 0: + balancing_activated_step = step + balancing_coeff = balancing_coeff + 1 / 15000 + # if train_loss['train/euclid_error'] < 0.015: + # reg_coeff = reg_coeff + 100 / 15000 + # if step > 15000 and balancing_coeff <= 1e-3: + # balancing_coeff = 1e-2 + # if step > 20000 and balancing_coeff <= 1e-3: + # balancing_coeff = 1e-2 + # if step > 25000 and balancing_coeff <= 1e-2: + # balancing_coeff = 1e-1 + # if step > 30000 and balancing_coeff <= 1e-1: + # balancing_coeff = 1 + + # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss'] + + to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } } + + wandb.log(to_log, step = step) + + print(step) + + step += 1 \ No newline at end of file diff --git a/models/cvae_model.py b/models/cvae_model.py index ddd13bc..6d71307 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -15,7 +15,7 @@ class CVAE(tf.keras.Model): def __init__(self, config): super(CVAE, self).__init__() - self.latent_dim = config['latent_dim'] + self.latent_dim = config['cvae']['latent_dim'] self.size = config['resolution'] 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']) @@ -35,6 +35,11 @@ class CVAE(tf.keras.Model): 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['latent_dim']),]) + self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size']) self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size']) @@ -43,10 +48,6 @@ class CVAE(tf.keras.Model): else: self.avg_garment = None - 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['latent_dim']),]) # @tf.function diff --git a/models/pyramid_level.py b/models/pyramid_level.py new file mode 100644 index 0000000..c73a8a6 --- /dev/null +++ b/models/pyramid_level.py @@ -0,0 +1,97 @@ +import os +import time +import random +import tensorflow as tf +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 +from models.upscaler import UpscalerBasic + + +class PyramidLevel(tf.keras.Model): + + 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.upscale_enhancer = None + if config['upscale_enhancer'] != '': + self.upscale_enhancer = UpscalerBasic(512, config['resolution']) + load_model(self.upscale_enhancer, config['upscale_enhancer']) + self.upscale_enhancer = tf.keras.Sequential([ + self.upscale_enhancer.flatten, + self.upscale_enhancer.seq, + self.upscale_enhancer.reshape + ]) + + 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, training=training) + x = x * (1 - mask[:,:,:,tf.newaxis]) + low = x + low + return low + + + 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 = 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, fix_noise=False): + 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]) + return eps * tf.exp(logvar * .5) + mean + + + 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) + + 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 + diff --git a/models/pyramid_model.py b/models/pyramid_model.py new file mode 100644 index 0000000..dc3d056 --- /dev/null +++ b/models/pyramid_model.py @@ -0,0 +1,70 @@ +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 + + +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']))] + + + 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): + uv_static, uv_body, uv_curv = inputs[-1][1], inputs[-1][2], inputs[-1][3] + z_cond = self.encode_conds(uv_static, uv_body, uv_curv) + + out_levels = [] + pred, mean, logvar = self.levels[0](inputs[0], z_cond, training=training) + out_levels.append([pred, mean, logvar]) + + for i in range(1, len(self.levels)): + pred_level, mean, logvar = self.levels[i](inputs[i], z_cond, training=training) + pred = pred + pred_level + out_levels.append([pred, mean, logvar]) + + return out_levels + +