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')
models/cvae_model.py
+47 -15
View File
@@ -126,7 +126,7 @@ class CVAE(tf.keras.Model):
# self.shuffle = tf.keras.layers.Lambda(lambda x: tf.stack([tf.gather(x[i], tf.random.shuffle(tf.range(x[i].shape[0]))) for i in range(x.shape[0])]))
# @tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False):
if not fix_noise:
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
else:
@@ -147,9 +147,6 @@ class CVAE(tf.keras.Model):
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
eps = tf.zeros_like(eps)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
@@ -159,6 +156,26 @@ class CVAE(tf.keras.Model):
return x
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
if noise is None:
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
else:
eps = noise
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if uv_mask is not None:
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean, logvar, fix_noise=False):
if not fix_noise:
eps = tf.random.normal(shape=mean.shape)
@@ -176,22 +193,37 @@ class CVAE(tf.keras.Model):
return logits
def call(self, inputs, training=None, zero_latent=False, fix_noise=False):
uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
def encode_conds(self, uv_template, uv_body_posed, uv_curv):
if uv_template.shape[1] == 3:
uv_template = tf.transpose(uv_template, [0,2,3,1])
if uv_body_posed.shape[1] == 3:
uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
if uv_curv.shape[1] == 3:
uv_curv = tf.transpose(uv_curv, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar, fix_noise)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
z = tf.zeros_like(z)
return z_cond
def call(self, inputs, training=None, use_mask=False):
if use_mask:
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
else:
uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = tf.concat([z, z_cond], axis=-1)
@@ -200,14 +232,14 @@ class CVAE(tf.keras.Model):
# z = tf.concat([z, z_curv], axis=-1)
x = self.decoder(z)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
# x = x * uv_mask[:,:,:,tf.newaxis]
if use_mask:
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar, z
sample.py
+1 -1
View File
@@ -174,7 +174,7 @@ def sample(folder_main, folder_curv):
uv_curvature = tf.math.argmax(pred, axis=-1)
uv_curvature = tf.one_hot(uv_curvature, 7)
pred = model.sample(batch['uv_static'][0],
pred = model.sample_single(batch['uv_static'][0],
batch['uv_body_posed'][0],
uv_curvature,
batch['uv_mask'][0],
sample_pro.py
+218 -67
View File
@@ -1,4 +1,3 @@
from lib2to3.refactor import get_fixers_from_package
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow as tf
@@ -22,27 +21,12 @@ 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
from plotly.subplots import make_subplots
from utils import str_to_bool
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
from load_models import load_generator, load_discriminator
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
from auto_encoder_body import AutoEncoder as BodyAE
from auto_encoder_curv import AutoEncoder as CurvatureAE
from auto_encoder_static import AutoEncoder as StaticAE
from cvae_curv import CVAE as CVAECurv
from models.cvae_model import CVAE as CVAEMain
@@ -163,84 +147,251 @@ def get_frame(folder, nr):
frames = sorted(os.listdir(pose_folder))
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
random.shuffle(frames)
return frames[0], frames[1:nr]
out = [frames[0]]
out.extend(frames[1:nr])
return out
def get_bodies(folder, frame_main, frame_rand, nr=4, random=True):
def get_bodies(folder, frames, random=True):
bodies = []
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', frame_main + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; bodies.append(uv_posed_body_no_orient)
pass
def get_curvs(folder, body, static, mask, frame=None, nr=4, random=True, model=None):
assert (random and model is not None) or (not random and model is None)
if random:
curv = model([body, static, mask])
uv_curvature = tf.math.argmax(pred, axis=-1)
uv_curvature = tf.one_hot(uv_curvature, 7)
pass
poses_shapes = []
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', frames[0] + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; bodies.append(uv_posed_body_no_orient)
pose_shape = np.load(os.path.join(folder, 'pose_shape', frames[0] + '.npy')); poses_shapes.append(pose_shape)
if not random:
for _ in range(len(frames)-1):
bodies.append(uv_posed_body_no_orient)
poses_shapes.append(pose_shape)
else:
curv = np.load(os.path.join(folder, 'uv_curvature')) #...
pass
for i in range(1, len(frames)):
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', frames[i] + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; bodies.append(uv_posed_body_no_orient)
pose_shape = np.load(os.path.join(folder, 'pose_shape', frames[i] + '.npy')); poses_shapes.append(pose_shape)
bodies = np.stack(bodies)
poses_shapes = np.stack(poses_shapes).astype(np.float32)
return bodies, poses_shapes
def get_statics(folder, nr=4, random=True):
pass
def get_unposed(folder, nr=4, random=True):
pass
def get_curvs(folder, body, static, mask, frames, noise, random=True, model: CVAECurv=None, use_original=True):
assert (random and model is not None) or (not random)
uv_curvature_orig = np.load(os.path.join(folder, 'uv_curvature', frames[0] + '.npz')); uv_curvature_orig = uv_curvature_orig['curvature_uv']
def plot_inputs(unposed, static, body, curv):
pass
noise = get_noise(not random, body.shape[0], latent_dim=128)
curv = model.sample(static, body, mask, noise)
def plot_variable_condition(condition, is_curv=False):
pass
uv_curvature = tf.math.argmax(curv, axis=-1).numpy()
if use_original:
uv_curvature[0,:,:] = uv_curvature_orig
out = tf.one_hot(uv_curvature, 7).numpy()
def plot_results(unposed, posed, body, mesh_mask):
body_color = np.ones_like(body, dtype=np.int32)*100
return out
def get_statics(folder, frames, random=True):
statics = []
uv_masks = []
mesh_masks = []
uv_static = np.load(os.path.join(folder, 'uv_static', frames[0].split('_')[0] + '.npz')); uv_static = uv_static['uv_static']; statics.append(uv_static)
uv_mask = np.load(os.path.join(folder, 'uv_mask', frames[0].split('_')[0] + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); uv_masks.append(uv_mask)
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', frames[0].split('_')[0] + '.npy')); mesh_masks.append(mesh_mask)
if not random:
for _ in range(len(frames)-1):
statics.append(uv_static)
uv_masks.append(uv_mask)
mesh_masks.append(mesh_mask)
else:
for i in range(1, len(frames)):
uv_static = np.load(os.path.join(folder, 'uv_static', frames[i].split('_')[0] + '.npz')); uv_static = uv_static['uv_static']; statics.append(uv_static)
uv_mask = np.load(os.path.join(folder, 'uv_mask', frames[i].split('_')[0] + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); uv_masks.append(uv_mask)
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', frames[i].split('_')[0] + '.npy')); mesh_masks.append(mesh_mask)
statics = np.stack(statics)
uv_masks = np.stack(uv_masks)
mesh_masks = np.stack(mesh_masks)
return statics, uv_masks, mesh_masks
def get_unposed(folder, frames, random=True):
uv_unposeds = []
mesh_unposeds = []
uv_unposed = np.load(os.path.join(folder, 'uv_posed', frames[0] + '.npz')); uv_unposed = uv_unposed['uv_gt']; uv_unposeds.append(uv_unposed)
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed_no_shape', frames[0] + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed_no_shape']; mesh_unposeds.append(mesh_unposed)
if not random:
for _ in range(len(frames)-1):
uv_unposeds.append(uv_unposed)
mesh_unposeds.append(mesh_unposed)
else:
for i in range(1, len(frames)):
uv_static = np.load(os.path.join(folder, 'uv_static', frames[i].split('_')[0] + '.npz')); uv_static = uv_static['uv_static']; uv_unposeds.append(uv_static)
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed_no_shape', frames[i] + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed_no_shape']; mesh_unposeds.append(mesh_unposed)
uv_unposeds = np.stack(uv_unposeds)
mesh_unposeds = np.stack(mesh_unposeds)
return uv_unposeds, mesh_unposeds
def make_square_2(figure):
figure.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1))
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,1],),
zaxis = dict(nticks=4, range=[-1,1],),)
figure.update_layout(scene = helper,
scene2 = helper)
return figure
def make_square_4(figure):
figure.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1))
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,1],),
zaxis = dict(nticks=4, range=[-1,1],),)
figure.update_layout(scene = helper,
scene2 = helper,
scene3 = helper,
scene4 = helper,)
return figure
def plot_inputs(unposed, static, body, curv, mesh_mask):
curv = np.tile(curv.argmax(-1)[:,:,np.newaxis], [1,1,3])/6
img = np.concatenate([unposed, static, body, curv], axis=-2)
ex.imshow((img + 1) / 2).show()
mesh_static = uv2mesh(static[np.newaxis,:,:,:]).numpy()[0]
mesh_body = uv2mesh(body[np.newaxis,:,:,:]).numpy()[0]
fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ])
F = np.load("custom_utils/smpl_faces.npy")
_, faces = _maskFaces(np.ones((14475, 3)), F, mesh_mask)
pass
x = mesh_static[mesh_mask, :]
y = mesh_body
fig.add_trace(go.Mesh3d(
x=x[:,0],
y=x[:,1],
z=x[:,2],
i=faces[:,0],
j=faces[:,1],
k=faces[:,2]
),
row=1, col=1)
fig.add_trace(go.Mesh3d(
x=y[:,0],
y=y[:,2],
z=y[:,1],
i=F[:,0],
j=F[:,1],
k=F[:,2]
),
row=1, col=2)
make_square_2(fig)
fig.show()
def plot_variable_condition(condition, mesh_mask, is_curv=False, is_body=False):
if is_curv:
img = np.concatenate(condition.argmax(-1), axis=-1)
ex.imshow(img / img.max()).show()
else:
mesh = uv2mesh(condition).numpy()
fig = make_subplots(rows=1, cols=condition.shape[0], start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(condition.shape[0])] ])
F = np.load("custom_utils/smpl_faces.npy")
if is_body:
faces = [F for _ in range(condition.shape[0])]
else:
_, faces = maskFacesBatch(np.ones((condition.shape[0], 14475, 3)), F, mesh_mask)
for i, m in enumerate(mesh):
if not is_body:
m = m[mesh_mask[i], :]
if is_body:
x = m.copy()
m[:,1] = x[:,2]
m[:,2] = x[:,1]
fig.add_trace(go.Mesh3d(
# 8 vertices of a cube
x=m[:,0],
y=m[:,1],
z=m[:,2],
# i, j and k give the vertices of triangles
i = faces[i][:,0],
j = faces[i][:,1],
k = faces[i][:,2]
),
row=1, col=i+1)
make_square_4(fig)
fig.show()
def plot_results(unposed, posed, body, mesh_mask):
fig = make_subplots(rows=1, cols=unposed.shape[0], start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(unposed.shape[0])] ])
F = np.load("custom_utils/smpl_faces.npy")
_, faces = maskFacesBatch(np.ones((unposed.shape[0], 14475, 3)), F, mesh_mask)
body_mesh = uv2mesh(body).numpy()
collision = [np.concatenate((body_mesh[i], posed[i][mesh_mask[i], :]), axis=0) for i in range(unposed.shape[0])]
collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(unposed.shape[0])]
collision_color = [np.concatenate((np.ones_like(body_mesh[i], dtype=np.int32)*100, np.ones_like(unposed[i], dtype=np.int32)*230), axis=0) for i in range(unposed.shape[0])]
for i, c in enumerate(collision):
fig.add_trace(go.Mesh3d(
x=c[:,0],
y=c[:,1],
z=c[:,2],
i=collision_faces[i][:,0],
j=collision_faces[i][:,1],
k=collision_faces[i][:,2],
vertexcolor=collision_color[i]
),
row=1, col=i+1)
make_square_4(fig)
fig.show()
def get_noise(fix_noise, nr, latent_dim=128):
eps = tf.random.normal(shape=(1 if fix_noise else nr, latent_dim))
if fix_noise:
eps = tf.tile(eps, [nr,1])
return eps
def sample(folder, variable_condition, fix_noise, nr=4):
# run = wandb.init(project='vae_plots', name=)
random_static, random_body, random_curv, random_unposed = variable_condition == 'static', variable_condition == 'body', variable_condition == 'curv', variable_condition == 'unposed'
frame_main, frame_rand = get_frame(folder, nr)
unposed, uv_mask, mesh_mask = get_unposed(folder, frame, nr, random_unposed)
static = get_statics(folder, nr, random_static)
body, pose, shape = get_bodies(folder, nr, random_body)
curv = get_curvs(folder, body, static, uv_mask, frame, nr, random_curv, curv_model)
frames = get_frame(folder, nr)
noise = get_noise(fix_noise, nr)
noise_curv = get_noise(fix_noise, nr)
body, pose_shape = get_bodies(folder, frames, random_body)
static, uv_mask, mesh_mask = get_statics(folder, frames, random_static)
unposed, mesh_unposed = get_unposed(folder, frames, random_unposed)
curv = get_curvs(folder, body, static, uv_mask, frames, noise_curv, random_curv, curv_model, use_original=True)
pred = model([unposed, static, body, curv, uv_mask], fix_noise=fix_noise)
pred = model.sample(static, body, curv, uv_mask, noise=noise)
mesh_pred = uv2mesh(pred)
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])])
gender = tf.cast(tf.greater_equal(pose_shape[:, -1] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(pose_shape[i,72:-1][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(len(frames))])
mesh_pred_shaped = (mesh_pred + shape_offsets).numpy()
mesh_posed = smpl.run_garment(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped).numpy()
mesh_pred_shaped = mesh_pred + shape_offsets
plot_inputs(unposed[0], static[0], body[0], curv[0], mesh_mask[0])
mesh_posed = smpl.run_garment(pose, shape, mesh_pred_shaped).numpy()
plot_inputs(unposed[0], static[0], body[0], curv[0])
variable_cond = None
variable_cond_data = None
if random_static:
variable_cond = static
variable_cond_data = static
elif random_body:
variable_cond = body
variable_cond_data = body
elif random_curv:
variable_cond = curv
variable_cond_data = curv
elif random_unposed:
variable_cond = unposed
variable_cond_data = unposed
else:
assert False
plot_variable_condition(variable_cond, is_curv=random_curv)
plot_variable_condition(variable_cond_data, mesh_mask, is_curv=random_curv, is_body=random_body)
plot_results(mesh_pred, mesh_posed, body, mesh_mask)
plot_results(mesh_pred_shaped, mesh_posed, body, mesh_mask)
print()
@@ -256,8 +407,8 @@ if __name__ == '__main__':
folder = '/home/workspace/local-data/train_100'
folder_main = 'checkpoints/cvae_13'
folder_curv = 'checkpoints/cvae_curv_2'
variable_condition = 'curv'
folder_curv = 'checkpoints/cvae_curv_3'
variable_condition = 'static'
batch_size = 4
smpl = SMPL()