luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

weight test with gnn

Browse Source
This commit is contained in:
luni
2023-02-24 22:14:16 +00:00
parent 9fbd09386c
commit 5062ed95d8
3 changed files with 51 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/config_abl_all_post.py
+2 -2
View File
@@ -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']:
cvae.py
+15 -5
View File
@@ -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
models/gnn.py
+34
View File
@@ -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