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

pro sampling avaiable, updated cvae

Browse Source
This commit is contained in:
luni
2022-09-22 16:32:37 +00:00
parent c5ae37aacc
commit a124f95f1f
4 changed files with 297 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae.py
+31 -35
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -21,7 +21,6 @@ import random
import numpy as np
import scipy as sp
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
@@ -38,7 +37,7 @@ from custom_utils.smpl import SMPL
from models.cvae_model import CVAE
name='cvae_16'
name='cvae_17'
def init_wandb():
run = wandb.init(project='vae', name=name,
notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
@@ -46,7 +45,7 @@ new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_lo
offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
retrained body encoder with deoriented bodies \
use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
fixed gender issue and retrained body encoder\
fixed gender issue and retrained body encoder, FINAL\
remove shape from unposed garment\
training on 100 garments with offsets over the average of these')
artifact = wandb.Artifact('code', type='code')
@@ -129,35 +128,29 @@ def train_step(batch):
with tf.GradientTape() as tape:
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], training=True, use_mask=True)
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_pred_posed = mesh_pred_posed + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
# ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
fft_loss = tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
tv_loss = tf.image.total_variation(pred)
reconstr_loss = l2_loss
reconstr_loss = tf.reduce_mean(reconstr_loss)
# logpx_z = -reconstr_loss
# logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
# logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
# loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
kl = kl_loss(mean, logvar)
loss = reconstr_loss + balancing_coeff * kl
# loss = reconstr_loss + balancing_coeff * kl
@@ -178,8 +171,8 @@ def train_step(batch):
# gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2,
'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
@@ -226,9 +219,9 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], training=False, use_mask=True)
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_pred = uv2mesh(pred).numpy()
mesh_orig = batch['mesh_unposed']
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
@@ -237,20 +230,23 @@ 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)]
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
body = batch['mesh_body_posed'].numpy()
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
garment_posed = [mesh_posed[i, mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(body.shape[0])]
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(pred.shape[0])]
body = batch['mesh_body_posed'].numpy()
for i in range(body.shape[0]):
mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
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_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))]
@@ -274,7 +270,7 @@ def sample(folder):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
nr_of_samples=2
pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pred = 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)
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
@@ -291,7 +287,7 @@ def sample(folder):
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
pred_zero = model.sample(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)
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(tf.transpose(pred_zero, [0,3,1,2]))
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
@@ -412,18 +408,18 @@ if __name__ == '__main__':
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
'uv_shaped_body': (3, 512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static)
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
model = CVAE(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
latent_dim=128, static_model='checkpoints/ae_static_4',
body_model='checpoints/ae_body_7',
body_model='checpoints/ae_body_8',
curvature_model='checpoints/ae_curvature_7')