From c5d576f11f77582716531544a4bd8844f5310f28 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 22 Feb 2023 20:39:24 +0000 Subject: [PATCH] try smoothing with gnn --- aaa_helpers.txt | 1 + configs/iccv/config_abl_all.py | 84 +++++++++++++++++++--------------- cvae.py | 14 ++++-- models/gnn.py | 41 +++++++++++++++++ 4 files changed, 99 insertions(+), 41 deletions(-) create mode 100644 models/gnn.py diff --git a/aaa_helpers.txt b/aaa_helpers.txt index 63aa28e..f86226b 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -8,6 +8,7 @@ nohup python auto_encoder_body.py --name="lhunor_body_iccv_1_skirt" > nohup/iccv nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1 " > nohup/iccv/encoders/curv_iccv_1.out 2>&1 & nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_skirt" > nohup/iccv/encoders/curv_iccv_1_skirt.out 2>&1 & nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_cape_f" > nohup/iccv/encoders/curv_iccv_1_cape_f.out 2>&1 & +nohup python cvae.py --name="lhunor_cvae_iccv_gnn_1" > nohup/iccv/cvae_iccv_gnn_1.out 2>&1 & skirts!!!! pyramid next week diff --git a/configs/iccv/config_abl_all.py b/configs/iccv/config_abl_all.py index 1ec8467..ef7cdc6 100644 --- a/configs/iccv/config_abl_all.py +++ b/configs/iccv/config_abl_all.py @@ -14,13 +14,15 @@ config['use_extra_losses'] = True # edge, collision, 3d mesh config['use_offsets'] = True config['use_norm_loss'] = True +config['only_curvatures'] = False ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False) assert ablation_count == 1 or ablation_count == 0 -config['gpu'] = '0' -config['name'] = 'abl_all' +config['gpu'] = '3' +config['name'] = 'iccv_gnn_2' config['notes'] = '' +config['tags'] = ['iccv', 'gnn'] config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv if config['use_skinning']: @@ -38,39 +40,42 @@ config['train_dir'] = '/home/workspace/local-data/train' config['valid_dir'] = '/home/workspace/local-data/valid' config['test_dir'] = '/home/workspace/data/dataset/test_ready' -config['static_model'] = 'checkpoints/static_1' -config['body_model'] = 'checkpoints/body_1' -config['curvature_model'] = 'checkpoints/curv_1' +# config['frames_file'] = 'frames_dress_30k.txt' +# config['frames_file_valid'] = 'frames_dress_30k.txt' +# config['frames_file_test'] = 'frames_dress_30k.txt' + +config['static_model'] = 'checkpoints/post/static_post_2' +config['body_model'] = 'checkpoints/post/body_post_2' +config['curvature_model'] = 'checkpoints/post/curv_post_1' config['smpl_res'] = 14475 if config['is_cape']: config['smpl_res'] = 3638 -config['latent_dim'] = 128 config['pretrained_model'] = '' ############### config['lr'] = 0.0001# 0.0001 ############### -config['starting_balancing_coeff'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 config['signature_dynamic'] = { 'id': (), 'frame_nr': (), 'pose_shape': (72 + 11,), 'uv_body_posed': (512, 256, 3), - 'mesh_body_posed': (14475, 3,), - 'mesh_posed': (14475, 3,), + 'mesh_body_posed': (config['smpl_res'], 3,), + 'mesh_posed': (config['smpl_res'], 3,), 'uv_curvature': (512, 256, 7), } config['signature_static'] = { 'uv_static': (512, 256, 3), - 'mesh_mask': (14475,), + 'mesh_mask': (config['smpl_res'],), 'uv_mask': (512, 256,), } if config['use_skinning']: config['signature_dynamic']['uv_unposed'] = (512, 256, 3) - config['signature_dynamic']['mesh_unposed'] = (14475, 3,) + config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,) else: config['signature_dynamic']['uv_posed'] = (512, 256, 3) @@ -81,37 +86,42 @@ config['use_one_hot'] = False config['cvae'] = dict() config['cvae']['is_label'] = False -config['cvae']['latent_dim'] = config['latent_dim'] config['cvae']['add_noise'] = True -config['cvae']['noise_amount'] = 0.025 +config['cvae']['noise_amount'] = 0.001 config['cvae']['kernel_size'] = 7 -config['cvae']['encoder_depths'] = [3, 3, 9, 3] -config['cvae']['encoder_dims'] = [32, 64, 128, 256] -config['cvae']['decoder_input'] = (16, 8, 80) -config['cvae']['decoder_depths'] = [3, 3, 9, 3] -config['cvae']['decoder_dims'] = [256, 128, 64, 16] +config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3] +config['cvae']['encoder_dims'] = [32, 64, 128, 256, 64] +config['cvae']['decoder_depths'] = config['cvae']['encoder_depths'] +config['cvae']['decoder_dims'] = config['cvae']['encoder_dims'] -config['cond_ae'] = dict() -config['cond_ae']['is_label'] = False -config['cond_ae']['upscale_output'] = False -config['cond_ae']['latent_dim'] = config['latent_dim'] -config['cond_ae']['kernel_size'] = 7 -config['cond_ae']['bottleneck_size'] = 64 -config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3] -config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']] -config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size']) -config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3] -config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16] +config['cond_ae_static'] = dict() +config['cond_ae_static']['is_label'] = False +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 +config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3] +config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']] +config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths'] +config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims'] +config['cond_ae_static']['activation'] = 'tanh' + +config['cond_ae_curv'] = dict() +config['cond_ae_curv']['is_label'] = True +config['cond_ae_curv']['kernel_size'] = 7 +config['cond_ae_curv']['bottleneck_size'] = 40 +config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3] +config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']] +config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths'] +config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims'] +config['cond_ae_curv']['activation'] = 'tanh' config['cond_ae_body'] = dict() -config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['is_label'] = False config['cond_ae_body']['kernel_size'] = 7 -config['cond_ae_body']['upscale_output'] = True -config['cond_ae_body']['bottleneck_size'] = 512 -config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3] -config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']] -config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size']) -config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3] -config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16] +config['cond_ae_body']['bottleneck_size'] = 16 +config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3] +config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']] +config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths'] +config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims'] +config['cond_ae_body']['activation'] = 'tanh' \ No newline at end of file diff --git a/cvae.py b/cvae.py index f79cdd1..f35f68d 100644 --- a/cvae.py +++ b/cvae.py @@ -1,6 +1,6 @@ import os import sys -from configs.config_abl_all_post import config +from configs.iccv.config_abl_all import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf @@ -42,6 +42,7 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mes from dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL from models.cvae_model import CVAENew as CVAE +from models.gnn import Smoother def init_wandb(): @@ -53,6 +54,7 @@ def init_wandb(): # artifact.add_file('configs/config_base.py') # artifact.add_file('configs/config_cvae.py') artifact.add_file('models/base.py') + artifact.add_file('models/gnn.py') artifact.add_file('models/cond_encoder.py') artifact.add_file('dataset.py') artifact.add_file('convnext/models/convnext_tf.py') @@ -75,7 +77,7 @@ def kl_loss(mean, logvar): # 3d loss on posed !!!!!!!!!!!! -# @tf.function +@tf.function def train_step(batch): with tf.GradientTape() as tape: @@ -87,6 +89,7 @@ def train_step(batch): pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True) mesh_pred = uv2mesh(pred) + mesh_pred = smoother(mesh_pred, training=True) if config['use_skinning']: shape = batch['pose_shape'][:, 72:] @@ -145,6 +148,7 @@ def valid_step(batch): pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) mesh_pred = uv2mesh(pred) + mesh_pred = smoother(mesh_pred, training=False) if config['use_skinning']: shape = batch['pose_shape'][:, 72:] @@ -254,7 +258,8 @@ def plot_everything(batch: dict, type: str): pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) - mesh_pred = uv2mesh(pred).numpy() + mesh_pred = uv2mesh(pred) + mesh_pred = smoother(mesh_pred, training=False).numpy() orig_img = uv_orig uv_static_img = batch['uv_static'] @@ -478,7 +483,7 @@ def test(): if __name__ == '__main__': - # init_wandb() + init_wandb() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) @@ -519,6 +524,7 @@ if __name__ == '__main__': is_cape=config['is_cape']) model = CVAE(config) + smoother = Smoother() optimizer = tf.optimizers.Adam(learning_rate=config['lr']) balancing_activated_step = 0 diff --git a/models/gnn.py b/models/gnn.py new file mode 100644 index 0000000..ad9d495 --- /dev/null +++ b/models/gnn.py @@ -0,0 +1,41 @@ +import os +import sys + +import scipy.io as sio +import tensorflow as tf +from spektral.layers import GCNConv, GeneralConv +from spektral.utils import sp_matrix_to_sp_tensor +from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D, + GlobalAveragePooling2D, Layer, + LayerNormalization, Flatten, Reshape, UpSampling2D) + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +class Smoother(tf.keras.Model): + def __init__(self, ): + super().__init__() + self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()) + self.conv1 = GeneralConv(channels=32, batch_norm=True, dropout=0.1, aggregate='mean', activation='prelu', use_bias=True) + self.conv2 = GeneralConv(channels=128, batch_norm=True, dropout=0.1, aggregate='mean', activation='prelu', use_bias=True) + self.conv3 = GeneralConv(channels=64, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True) + self.conv4 = GeneralConv(channels=16, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True) + + self.conv5 = GCNConv(3, activation='tanh', use_bias=False) + self.conv6 = GCNConv(3, activation='tanh', use_bias=False) + + def call(self, inputs, training=None): + x = inputs + x = self.conv1([x, self.a], training=training) + x = self.conv2([x, self.a], training=training) + x = self.conv3([x, self.a], training=training) + x = self.conv4([x, self.a], training=training) + x = self.conv5([x, self.a], training=training) + x = self.conv6([x, self.a], training=training) + return x + + + + + +