diff --git a/aaa_helpers.txt b/aaa_helpers.txt index a1ed97d..3f66e3e 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -11,6 +11,7 @@ nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1" > nohup/iccv/encod 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 auto_encoder_curv.py --name="lhunor_curv_iccv_1_cape_m" > nohup/iccv/encoders/curv_iccv_1_cape_m.out 2>&1 & +nohup python cvae.py --name="lhunor_cvae_iccv_gnn_1" > nohup/iccv/cvae_iccv_gnn_1.out 2>&1 & nohup python cvae.py --name="lhunor_abl_all" > nohup/iccv/ablation/abl_all.out 2>&1 & diff --git a/configs/iccv/config_abl_all.py b/configs/iccv/config_abl_all.py index f3d1bfb..5c9a412 100644 --- a/configs/iccv/config_abl_all.py +++ b/configs/iccv/config_abl_all.py @@ -14,6 +14,7 @@ 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 @@ -43,35 +44,38 @@ config['static_model'] = 'checkpoints/iccv/static_iccv_1' config['body_model'] = 'checkpoints/iccv/body_iccv_1' config['curvature_model'] = 'checkpoints/iccv/curv_iccv_1' +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) diff --git a/cvae.py b/cvae.py index 7964455..846e14d 100644 --- a/cvae.py +++ b/cvae.py @@ -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 from custom_utils.custom_utils import Accumulator @@ -54,6 +55,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') @@ -88,6 +90,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:] @@ -146,6 +149,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:] @@ -255,7 +259,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'] @@ -520,6 +525,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 + + + + + +