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

update ae static

Browse Source
This commit is contained in:
luni
2022-10-07 21:56:37 +00:00
parent 3727f341a7
commit 5dd0dd3248
2 changed files with 38 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_static.py
+38 -29
View File
@@ -2,12 +2,24 @@ import os
from configs.config_ae_static import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
# Currently, memory growth needs to be the same across GPUs
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
except RuntimeError as e:
# Memory growth must be set before GPUs have been initialized
print(e)
import time
import argparse
import wandb
import random
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
import plotly.graph_objects as go
import plotly.express as ex
@@ -25,7 +37,7 @@ from models.cond_encoder import AutoEncoder
def init_wandb():
run = wandb.init(project='vae', name=config['name'],
run = wandb.init(project='ae', name=config['name'],
notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
flattened latent dim to 128\
train on 100 garments')
@@ -67,24 +79,19 @@ def train_step(batch):
@tf.function
def valid_step(batch):
pred = model(batch['uv_static'])
uv_mask = batch['uv_mask']
pred = model([batch['uv_static'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 0
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * uv_mask)
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) + 1, max_val=2))
fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_static'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2
}
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_static'] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
loss = 1e3 * (1 - ms_ssim)
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
def plot_everything(batch: dict, type: str):
@@ -100,33 +107,35 @@ def plot_everything(batch: dict, type: str):
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
garment_static, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
body = gt
garment_posed = mesh_pred
garment_posed = garment_static
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((faces[i], faces[i] + np.sum(mask[i])), 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))]
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_static[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **{ type + '/images': imgs } }
def train():
valid_iterator = iter(dataset_valid)
step = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
train_loss = train_step(batch)
# data_valid = valid_iterator.get_next()
# data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
valid_loss = {}#valid_step(data_valid)
data_valid = valid_iterator.get_next()
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
valid_loss = valid_step(data_valid)
if step % 2500 == 0:
if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = {}#plot_everything(data_valid, 'valid')
valid_figs = plot_everything(data_valid, 'valid')
wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
else:
wandb.log({ **train_loss, **valid_loss }, step=step)
@@ -171,10 +180,10 @@ if __name__ == '__main__':
factory = DatasetFactory()
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
model : AutoEncoder = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
configs/config_ae_static.py
+2 -2
View File
@@ -1,8 +1,8 @@
from configs.config_base import config
config['gpu'] = '2'
config['gpu'] = '5'
config['lr'] = 0.0005
config['name'] = 'ae_static_7_2_skirt'
config['name'] = 'static_1'
config['cond_ae']['upscale_output'] = False