From 5062ed95d8dcf5cfda2198bb1581d184ed70a45a Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 24 Feb 2023 22:14:16 +0000 Subject: [PATCH] weight test with gnn --- configs/config_abl_all_post.py | 4 ++-- cvae.py | 20 +++++++++++++++----- models/gnn.py | 34 ++++++++++++++++++++++++++++++++++ 3 files changed, 51 insertions(+), 7 deletions(-) create mode 100644 models/gnn.py diff --git a/configs/config_abl_all_post.py b/configs/config_abl_all_post.py index 954a55f..5a960fc 100644 --- a/configs/config_abl_all_post.py +++ b/configs/config_abl_all_post.py @@ -20,8 +20,8 @@ ablation_count = [config['use_skinning'], config['remove_global_orientation'], c assert ablation_count == 1 or ablation_count == 0 config['gpu'] = '4' -config['name'] = 'iccv_weights_test_6' -config['notes'] = 'smaller weight loss coeff, 1e-6, convolutional weight learning, continued from _5' +config['name'] = 'iccv_weights_test_10' +config['notes'] = 'smaller weight loss coeff, 1e-6, convolutional weight learning, add 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']: diff --git a/cvae.py b/cvae.py index 8bf7754..3901dfd 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 def init_wandb(): @@ -92,7 +93,9 @@ def train_step(batch): pred, mean, logvar, weights = 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) mesh_weigths = uv2mesh_label(weights) + mesh_weigths = smoother_weights(mesh_weigths, training=True) if config['use_skinning']: shape = batch['pose_shape'][:, 72:] @@ -137,8 +140,8 @@ def train_step(batch): # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_posed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) - gradients = tape.gradient(loss, model.trainable_variables) - optimizer.apply_gradients(zip(gradients, model.trainable_variables)) + gradients = tape.gradient(loss, model.trainable_variables + smoother.trainable_variables + smoother_weights.trainable_variables) + optimizer.apply_gradients(zip(gradients, model.trainable_variables + smoother.trainable_variables + smoother_weights.trainable_variables)) return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric, 'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar, @@ -157,7 +160,9 @@ def valid_step(batch): pred, mean, logvar, weights = 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) mesh_weigths = uv2mesh_label(weights) + mesh_weigths = smoother_weights(mesh_weigths, training=False) if config['use_skinning']: shape = batch['pose_shape'][:, 72:] @@ -273,8 +278,10 @@ def plot_everything(batch: dict, type: str): pred, mean, logvar, weights = 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() mesh_weigths = uv2mesh_label(weights) + mesh_weigths = smoother_weights(mesh_weigths, training=False).numpy() orig_img = uv_orig uv_static_img = batch['uv_static'] @@ -282,7 +289,8 @@ def plot_everything(batch: dict, type: str): pred_img = pred img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2) imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] - + weights_imgs = [wandb.Image(np.concatenate((np.concatenate([weights[b,:,:,i].numpy() for i in range(12)], axis=-1), np.concatenate([weights[b,:,:,i].numpy() for i in range(12,24)], axis=-1)), axis=0)) for b in range(4)] + F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") mask = batch['mesh_mask'].numpy().astype(bool) @@ -332,7 +340,7 @@ def plot_everything(batch: dict, type: str): mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0)) logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0)) - return { **figs, **figs_coll, **figs_weights, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } + return { **figs, **figs_coll, **figs_weights, **{ type + '/images': imgs, type + '/weight_images': weights_imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } def sample(folder): @@ -547,6 +555,8 @@ if __name__ == '__main__': is_cape=config['is_cape']) model = CVAE(config) + smoother = Smoother(features=16, activation_middle='tanh', activation_final='tanh', learn_weights=False) + smoother_weights = Smoother(features=16, activation_middle=None, activation_final='softmax', learn_weights=True) 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..dc203a6 --- /dev/null +++ b/models/gnn.py @@ -0,0 +1,34 @@ +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, features=64, activation_middle='tanh', activation_final='tanh', learn_weights=False): + super().__init__() + self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()) + self.conv1 = GCNConv(features, activation=activation_middle, use_bias=False) + self.conv2 = GCNConv(features, activation=activation_middle, use_bias=False) + self.conv3 = GCNConv(features, activation=activation_middle, use_bias=False) + # self.conv4 = GCNConv(64, activation='tanh', use_bias=False) + # self.conv5 = GCNConv(64, activation='tanh', use_bias=False) + self.conv6 = GCNConv(3 if not learn_weights else 24, activation=activation_final, 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 \ No newline at end of file