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

progress save

Browse Source
This commit is contained in:
luni
2023-02-22 21:06:41 +00:00
parent c5d576f11f
commit c1113379ec
4 changed files with 36 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/iccv/config_abl_all.py
+2 -2
View File
@@ -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'] == '')
cvae.py
-6
View File
@@ -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
models/cvae_model.py
+8 -9
View File
@@ -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
models/gnn.py
+26 -1
View File
@@ -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