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

commit all

Browse Source
This commit is contained in:
luni
2023-01-25 12:25:59 +00:00
parent 1bf224841f
commit 7470ca9fb6
139 changed files with 82069 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae.py
+39 -25
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.config_abl_norm import config
from configs.config_abl_all_post import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
@@ -41,11 +41,11 @@ from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAE
from models.cvae_model import CVAENew as CVAE
def init_wandb():
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae.py')
artifact.add_file('models/cvae_model.py')
@@ -61,7 +61,7 @@ def init_wandb():
artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
def kl_loss(mean, logvar):
@@ -71,7 +71,9 @@ def kl_loss(mean, logvar):
return kl_loss
@tf.function
# 3d loss on posed !!!!!!!!!!!!
# @tf.function
def train_step(batch):
with tf.GradientTape() as tape:
@@ -100,14 +102,17 @@ def train_step(batch):
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred, mesh_orig, batch['mesh_mask'], is_skirt=config['is_skirt'])
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
# conc_loss = concept_loss(pred, batch['uv_unposed'])
kl = kl_loss(mean, logvar)
reconstr_loss = l2_loss
if config['use_extra_losses']:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff']
else:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
if config['use_norm_loss'] and step > 2000:
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
@@ -155,8 +160,8 @@ def valid_step(batch):
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred, mesh_orig, batch['mesh_mask'], is_skirt=config['is_skirt'])
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
# conc_loss = concept_loss(pred, batch['uv_unposed'])
kl = kl_loss(mean, logvar)
@@ -256,7 +261,7 @@ def plot_everything(batch: dict, type: str):
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)]
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
mask = batch['mesh_mask'].numpy().astype(bool)
if config['use_skinning']:
@@ -284,7 +289,7 @@ def plot_everything(batch: dict, type: str):
garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(garment_posed))]
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
@@ -294,8 +299,8 @@ def plot_everything(batch: dict, type: str):
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0))
logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0))
return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
@@ -323,7 +328,7 @@ def sample(folder):
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
_, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
mesh_pred_zero = uv2mesh(pred_zero)
@@ -471,12 +476,12 @@ def test():
if __name__ == '__main__':
init_wandb()
# init_wandb()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL(use_one_hot=config['use_one_hot'])
smpl = SMPL(orig_smpl=config['is_cape'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
factory = DatasetFactory()
@@ -487,11 +492,14 @@ if __name__ == '__main__':
resolutions=[config['resolution']],
in_memory=config['dataset_in_memory'],
use_one_hot=config['use_one_hot'],
use_no_orient=config['remove_global_orientation'])
use_no_orient=config['remove_global_orientation'],
frames_file=config['frames_file'] if 'frames_file' in config.keys() else None,
is_cape=config['is_cape'])
static_dataset_train = StaticDataset(folder=config['train_dir'],
signature=config['signature_static'],
resolutions=[config['resolution']],
in_memory=True)
in_memory=True,
is_cape=config['is_cape'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'],
batch_size=config['batch_size'],
@@ -499,11 +507,14 @@ if __name__ == '__main__':
resolutions=[config['resolution']],
in_memory=config['dataset_in_memory'],
use_one_hot=config['use_one_hot'],
use_no_orient=config['remove_global_orientation'])
use_no_orient=config['remove_global_orientation'],
frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None,
is_cape=config['is_cape'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'],
signature=config['signature_static'],
resolutions=[config['resolution']],
in_memory=True)
in_memory=True,
is_cape=config['is_cape'])
model = CVAE(config)
@@ -515,11 +526,14 @@ if __name__ == '__main__':
# test()
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
if config['pretrained_model'] != '':
load_model(model, config['pretrained_model'])
# load_model(model, 'checkpoints/' + config['name'])
valid_iterator = iter(dataset_valid)
step = 0
for batch in dataset_train:
@@ -537,15 +551,15 @@ if __name__ == '__main__':
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step % 10000 == 0:
if step % 5000 == 0:
manager.save(checkpoint_number=step)
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if train_loss['train/euclid_error'] < 0.017:
if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
if balancing_activated_step == 0:
balancing_activated_step = step
g_balancing_coeff = g_balancing_coeff + 1 / 20000
g_balancing_coeff = g_balancing_coeff + 1 / 10000
to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }