From c1113379ecae0d7dd59b72b2af8e81cbd4a00111 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 22 Feb 2023 21:06:41 +0000 Subject: [PATCH] progress save --- configs/iccv/config_abl_all.py | 4 ++-- cvae.py | 6 ------ models/cvae_model.py | 17 ++++++++--------- models/gnn.py | 27 ++++++++++++++++++++++++++- 4 files changed, 36 insertions(+), 18 deletions(-) diff --git a/configs/iccv/config_abl_all.py b/configs/iccv/config_abl_all.py index ef7cdc6..6062ffa 100644 --- a/configs/iccv/config_abl_all.py +++ b/configs/iccv/config_abl_all.py @@ -20,11 +20,11 @@ ablation_count = [config['use_skinning'], config['remove_global_orientation'], c assert ablation_count == 1 or ablation_count == 0 config['gpu'] = '3' -config['name'] = 'iccv_gnn_2' +config['name'] = 'iccv_gnn_3' 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 +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']: assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '') diff --git a/cvae.py b/cvae.py index f35f68d..b02cf9f 100644 --- a/cvae.py +++ b/cvae.py @@ -89,7 +89,6 @@ 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:] @@ -148,7 +147,6 @@ 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:] @@ -259,7 +257,6 @@ 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) - mesh_pred = smoother(mesh_pred, training=False).numpy() orig_img = uv_orig uv_static_img = batch['uv_static'] @@ -432,8 +429,6 @@ def sample(folder): - - def test(): str_split = config['name'].split('_') test_idx = str_split.index('test') @@ -524,7 +519,6 @@ 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/cvae_model.py b/models/cvae_model.py index 3ed4571..4e4e152 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -10,7 +10,7 @@ import scipy as sp from custom_utils.custom_utils import load_model from models.base import AEEncoder, VariationalEncoder, Decoder from models.base import AEEncoderNew, VariationalEncoderNew, DecoderNew - +from models.gnn import GNNDecoder class CVAE(tf.keras.Model): @@ -150,8 +150,6 @@ class CVAE(tf.keras.Model): return x, mean, logvar - - class CVAENew(tf.keras.Model): def __init__(self, config): @@ -180,7 +178,8 @@ class CVAENew(tf.keras.Model): self.curv_enc.trainable = False self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size']) - self.decoder = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size']) + self.decoder = GNNDecoder() + self.reshape = tf.keras.layers.Reshape((8, 4, 64)) @@ -273,13 +272,13 @@ class CVAENew(tf.keras.Model): 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,:,:,:] + # if self.avg_garment is not None: + # x = x + self.avg_garment[tf.newaxis,:,:,:] - if self.use_offsets_over_body: - x = x + uv_body_posed + # if self.use_offsets_over_body: + # x = x + uv_body_posed - x = x * uv_mask[:,:,:,tf.newaxis] + # x = x * uv_mask[:,:,:,tf.newaxis] return x, mean, logvar diff --git a/models/gnn.py b/models/gnn.py index ad9d495..ad94aaa 100644 --- a/models/gnn.py +++ b/models/gnn.py @@ -3,7 +3,7 @@ import sys import scipy.io as sio import tensorflow as tf -from spektral.layers import GCNConv, GeneralConv +from spektral.layers import GCNConv, GeneralConv, GATConv from spektral.utils import sp_matrix_to_sp_tensor from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D, GlobalAveragePooling2D, Layer, @@ -35,7 +35,32 @@ class Smoother(tf.keras.Model): return x +class GNNDecoder(tf.keras.Model): + def __init__(self, ): + super().__init__() + self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()) + self.a_dense = sio.loadmat("custom_utils/ssmpl0.mat")['L'].todense() + self.flatten = Flatten() + self.conv1 = GATConv(channels=1024, attn_heads=4, concat_heads=True, dropout_rate=0.0, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True) + self.conv2 = GATConv(channels=512, attn_heads=4, concat_heads=True, dropout_rate=0.1, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True) + self.conv3 = GATConv(channels=256, attn_heads=4, concat_heads=True, dropout_rate=0.1, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True) + self.conv4 = GeneralConv(channels=128, batch_norm=True, dropout=0.1, aggregate='sum', activation='prelu', use_bias=True) + self.conv5 = GeneralConv(channels=64, batch_norm=True, dropout=0.0, aggregate='sum', activation='prelu', use_bias=True) + self.conv6 = GeneralConv(channels=3, batch_norm=False, dropout=0.0, aggregate='sum', activation='tanh', use_bias=True) + + def call(self, inputs, training=None): + x = inputs + x = self.flatten(x) + x = tf.tile(x[:,None,:], [1,14475,1]) + + x = self.conv1([x, self.a_dense], training=training) + x = self.conv2([x, self.a_dense], training=training) + x = self.conv3([x, self.a_dense], 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