From 977e49900dd2049bc1d33445230d495addbd0f83 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 2 Sep 2022 14:15:54 +0000 Subject: [PATCH] backup --- .gitignore | 3 +- auto_encoder_body.py | 12 +- auto_encoder_curv.py | 26 +- auto_encoder_static.py | 16 +- custom_utils/custom_utils.py | 8 +- custom_utils/smpl.py | 11 +- cvae.py | 323 +++++++------------ cvae_curv.py | 440 ++++++++++++++++++++++++++ cvae_vae.py | 586 +++++++++++++++++++++++++++++++++++ dataset.py | 28 +- deshape.py | 117 +++++++ models/cvae.py | 191 ++++++++++++ models/cvae_curv.py | 187 +++++++++++ sample.py | 332 ++++++++++++++++++++ verify_poses.py | 160 ++++++++++ verify_scales.py | 73 +++-- 16 files changed, 2241 insertions(+), 272 deletions(-) create mode 100644 cvae_curv.py create mode 100644 cvae_vae.py create mode 100644 deshape.py create mode 100644 models/cvae.py create mode 100644 models/cvae_curv.py create mode 100644 sample.py create mode 100644 verify_poses.py diff --git a/.gitignore b/.gitignore index cfe2520..4f745f0 100644 --- a/.gitignore +++ b/.gitignore @@ -138,9 +138,10 @@ wandb/ output/* *.zip *.out -models/ +models_old/ models_archive/ temp.jpg corrupt/ wandb copy/ checkpoints/ +temp/ \ No newline at end of file diff --git a/auto_encoder_body.py b/auto_encoder_body.py index 232d624..b6f6db0 100644 --- a/auto_encoder_body.py +++ b/auto_encoder_body.py @@ -25,16 +25,18 @@ from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict from custom_utils.smpl import SMPL -name='ae_body_5' +name='ae_body_7' def init_wandb(): - os.environ["CUDA_VISIBLE_DEVICES"] = "3" + os.environ["CUDA_VISIBLE_DEVICES"] = "1" run = wandb.init(project='vae', name=name, notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \ fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \ . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\ two more downscaling than up, so decoder output is resized wto final resolution\ larger batch_size.\ - flattened bottleneck with 128 latent_dim\ remove global orientation') + flattened bottleneck with 128 latent_dim\ remove global orientation\ + fixed gender issue with deoriented bodies, again\ + trained on 100 garments') artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_body.py') artifact.add_file('losses.py') @@ -419,8 +421,8 @@ if __name__ == '__main__': # 'uv_shaped_body': (3, 512, 256,), } - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=32, signature=signature_dynamic) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static) + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static) dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic) static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static) diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index f541c19..f4f3dc7 100644 --- a/auto_encoder_curv.py +++ b/auto_encoder_curv.py @@ -5,6 +5,9 @@ import wandb import random import numpy as np import tensorflow as tf +os.environ["CUDA_VISIBLE_DEVICES"] = "2" + + from custom_utils.custom_utils import UV2Mesh, maskFacesBatch from tensorboardX import SummaryWriter import plotly.graph_objects as go @@ -24,14 +27,15 @@ from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict from custom_utils.smpl import SMPL -name='ae_curvature_6' +name='ae_curvature_7' testing = False def init_wandb(): - os.environ["CUDA_VISIBLE_DEVICES"] = "3" if not testing: - run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128') + run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments') artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_curv.py') + artifact.add_file('dataset.py') + artifact.add_file('losses.py') artifact.add_file('stylegan2_ae/layers/mapping_block.py') artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) @@ -282,19 +286,21 @@ def train(): 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, signature_dynamic)) + valid_loss = {}#valid_step(data_valid) + data_valid = batch to_log = { **train_loss, **valid_loss} - if step % 1000 == 0: + if step % 2500 == 0: uv_curvature = tf.one_hot(data_valid['uv_curvature'], 7) pred = model([uv_curvature, data_valid['uv_mask']]) orig_img = data_valid['uv_curvature'] pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32) img = tf.concat((orig_img, pred_img), axis=2).numpy() - imgs = { 'valid/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] } + # imgs = { 'valid/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] } + imgs = { 'train/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] } to_log = { **to_log, **imgs} wandb.log(to_log, step=step) @@ -351,10 +357,10 @@ if __name__ == '__main__': # 'uv_shaped_body': (3, 512, 256,), } - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static) dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=8, diff --git a/auto_encoder_static.py b/auto_encoder_static.py index e94eb52..48391b4 100644 --- a/auto_encoder_static.py +++ b/auto_encoder_static.py @@ -24,9 +24,9 @@ from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict from custom_utils.smpl import SMPL -name='ae_static_3' +name='ae_static_4' def init_wandb(): - os.environ["CUDA_VISIBLE_DEVICES"] = "2" + os.environ["CUDA_VISIBLE_DEVICES"] = "0" run = wandb.init(project='vae', name=name, notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \ fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \ @@ -389,13 +389,13 @@ def train(): batch = static_dataset_train.expand_batch(batch_to_dict(batch, 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, signature_dynamic)) + valid_loss = {}#valid_step(data_valid) if step % 1000 == 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) @@ -432,8 +432,8 @@ if __name__ == '__main__': # 'uv_shaped_body': (3, 512, 256,), } - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=8, signature=signature_dynamic) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static) + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static) dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=8, signature=signature_dynamic) static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static) diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index fb8feb5..302b5ff 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -196,4 +196,10 @@ def sparse_to_tensor(L): idx = np.mat([L.row, L.col]).transpose().astype(np.float32) data = L.data.astype(dtype=np.float32) shape = L.shape - return tf.SparseTensor(idx, data, shape) \ No newline at end of file + return tf.SparseTensor(idx, data, shape) + + +def load_model(model, checkpoint): + ckpt = tf.train.Checkpoint(model=model) + manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1) + ckpt.restore(manager.latest_checkpoint) diff --git a/custom_utils/smpl.py b/custom_utils/smpl.py index ea87064..68a097e 100644 --- a/custom_utils/smpl.py +++ b/custom_utils/smpl.py @@ -69,6 +69,7 @@ class SMPL: self.nn = np.array(sio.loadmat('custom_utils/NN.mat')) self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1) + self.weights_all = self.weights self.weights = tf.gather(self.weights, self.to_keep, axis=1) def Q2R(self, w, x, y, z): @@ -132,7 +133,7 @@ class SMPL: v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male v_shaped = tf.stack(v_shaped, axis=-1) - def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights, parent=self.parent): + def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights_all, parent=self.parent): R = x[0] v_shaped = x[1] g = x[2] @@ -160,7 +161,7 @@ class SMPL: return v, tf.squeeze(J), v_posed - gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32) + gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32) T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32)) J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3)) @@ -215,7 +216,7 @@ class SMPL: return v, tf.squeeze(J) - gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32) + gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32) T_out, J_out = tf.map_fn(single_run, (R, v_shaped, gender, v_cloth), (tf.float32, tf.float32)) J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3)) @@ -235,7 +236,7 @@ class SMPL: J = tf.matmul(J_regressor[g], v_shaped[:, :, g]) return J - gender = tf.cast(tf.greater(betas[:, 10] , 0), tf.int32) + gender = tf.cast(tf.greater(betas[:, 10] , 0.5), tf.int32) J = tf.map_fn(single_run, (v_shaped, gender), (tf.float32)) J0 = tf.reshape(J[:,0,:], (-1, 1, 3)) @@ -274,7 +275,7 @@ class SMPL: return v, tf.squeeze(J) - gender = tf.cast(tf.greater(betas[:, 10] , 0), tf.int32) + gender = tf.cast(tf.greater(betas[:, 10] , 0.5), tf.int32) T_out, J_out = tf.map_fn(single_run, (R, J, gender, v_cloth), (tf.float32, tf.float32)) J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3)) diff --git a/cvae.py b/cvae.py index 2dce622..f0f125e 100644 --- a/cvae.py +++ b/cvae.py @@ -1,5 +1,5 @@ import os -os.environ["CUDA_VISIBLE_DEVICES"] = "2" +os.environ["CUDA_VISIBLE_DEVICES"] = "3" import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') @@ -24,6 +24,7 @@ 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 utils import str_to_bool from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing @@ -31,26 +32,22 @@ 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 custom_utils.custom_utils import batch_to_dict, load_model 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 models.cvae import CVAE -name='cvae_10_2' +name='cvae_15' 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.\ -new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1000 from 250.\ -offsets over average unposed garment. gamma balancing. remove global orientation (in posed body input and in pose_shape parameters. larger batchsize\ -add edge and collision (fixed) loss. rebalanced l1 loss. increasing edge and coll coeff. retrained body encoder with deoriented bodies \ -don\'t freeze encoders, lower lr') + notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \ +new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \ +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\ +remove shape from unposed garment') artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae.py') artifact.add_file('dataset.py') @@ -61,14 +58,14 @@ don\'t freeze encoders, lower lr') uv2mesh = UV2Mesh(res=256) def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color): - mesh_pred_unposed[:,1] += 1 - collision[:,1] += 2 + mesh_pred_unposed[:,2] += 1 + collision[:,2] += 2 fig = go.Figure(data=[ go.Mesh3d( # 8 vertices of a cube x=mesh_gt[:,0], - y=mesh_gt[:,1], - z=mesh_gt[:,2], + y=mesh_gt[:,2], + z=mesh_gt[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], @@ -77,8 +74,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli go.Mesh3d( # 8 vertices of a cube x=mesh_pred_unposed[:,0], - y=mesh_pred_unposed[:,1], - z=mesh_pred_unposed[:,2], + y=mesh_pred_unposed[:,2], + z=mesh_pred_unposed[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], @@ -87,8 +84,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli go.Mesh3d( # 8 vertices of a cube x=collision[:,0], - y=collision[:,1], - z=collision[:,2], + y=collision[:,2], + z=collision[:,1], # i, j and k give the vertices of triangles i = collision_faces[:,0], j = collision_faces[:,1], @@ -108,151 +105,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli return fig -class Bottleneck(tf.keras.Model): - - def __init__(self, dim): - super(Bottleneck, self).__init__() - self.bottleneck = tf.keras.Sequential( - [Block(dim=dim, drop_path=0, - layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] - ) - - def call(self, inputs, training=None): - x = inputs - - x = self.bottleneck(x, training=training) - - return x - - -class Encoder(tf.keras.Model): - - def __init__(self, latent_dim): - super(Encoder, self).__init__() - - self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False) - self.flatten = tf.keras.layers.Flatten() - self.mean = tf.keras.layers.Dense(latent_dim) - self.log_var = tf.keras.layers.Dense(latent_dim) - - def call(self, inputs, training=None): - x = inputs - - x = self.encoder(x) - x = self.flatten(x) - mean = self.mean(x) - log_var = self.log_var(x) - - return mean, log_var - - -class Decoder(tf.keras.Model): - - def __init__(self): - super(Decoder, self).__init__() - - self.dense = tf.keras.layers.Dense(16 * 8 * 80) - self.reshape = tf.keras.layers.Reshape((16, 8, 80)) - self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16]))) - - - def call(self, inputs, training=None): - x = inputs - - x = self.dense(x) - x = self.reshape(x) - x = self.decoder(x) - - return x - - -class CVAE(tf.keras.Model): - - def __init__(self, latent_dim=128): - super(CVAE, self).__init__() - self.latent_dim = latent_dim - - model : StaticAE = StaticAE() - load_model(model, 'checkpoints/ae_static_3') - self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) - # self.static_enc.trainable = False - - model : BodyAE = BodyAE() - load_model(model, 'checpoints/ae_body_5') - self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) - # self.body_enc.trainable = False - - model : CurvatureAE = CurvatureAE() - load_model(model, 'checpoints/ae_curvature_6') - self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) - # self.curv_enc.trainable = False - - self.encoder = Encoder(latent_dim) - self.decoder = Decoder() - - - # @tf.function - def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, eps=None): - if eps is None: - eps = tf.random.normal(shape=(4, self.latent_dim)) - - uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1]) - uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1]) - uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [4,1,1,1]) - uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [4,1,1]) - - # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [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) - - z_cond = tf.concat([z_static, z_body, z_curv], axis=-1) - z = tf.concat([eps, z_cond], axis=-1) - - x = self.decoder(z) - x = x + avg_jumpsuit[tf.newaxis,:,:,:] - x = x * uv_mask[:,:,:,tf.newaxis] - - return x - - - def reparameterize(self, mean, logvar): - eps = tf.random.normal(shape=mean.shape) - return eps * tf.exp(logvar * .5) + mean - - - def decode(self, z, apply_sigmoid=False): - logits = self.decoder(z) - if apply_sigmoid: - probs = tf.sigmoid(logits) - return probs - return logits - - - def call(self, inputs, training=None): - uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = 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]) - 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) - z = tf.concat([z, z_static, z_body, 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]) - x = x + avg_jumpsuit[tf.newaxis,:,:,:] - x = x * uv_mask[:,:,:,tf.newaxis] - - return x, mean, logvar, z def log_normal_pdf(sample, mean, logvar, raxis=1): @@ -268,7 +121,7 @@ def kl_loss(mean, logvar): return kl_loss -@tf.function +# @tf.function def train_step(batch): with tf.GradientTape() as tape: @@ -280,10 +133,15 @@ def train_step(batch): mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_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] 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.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)) @@ -292,17 +150,18 @@ def train_step(batch): fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed']) tv_loss = tf.image.total_variation(pred) - reconstr_loss = l1_loss / 4 - reconstr_loss = tf.reduce_mean(reconstr_loss) + reg_coeff *(e_loss + c_loss) + 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 / (2 * (gamma ** 2)) + balancing_coeff * kl + loss = reconstr_loss + balancing_coeff * kl # loss = reconstr_loss + balancing_coeff * kl ms_ssim = tf.reduce_mean(ms_ssim) l1_loss = tf.reduce_mean(l1_loss) + l2_loss = tf.reduce_mean(l2_loss) fft_loss = tf.reduce_mean(fft_loss) tv_loss = tf.reduce_mean(tv_loss) mean = tf.reduce_mean(mean) @@ -320,7 +179,7 @@ def train_step(batch): '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/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/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss } @@ -384,7 +243,12 @@ def plot_everything(batch: dict, type: str): body = batch['mesh_body_posed'].numpy() 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])] + + 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])] 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))] @@ -397,25 +261,25 @@ def plot_everything(batch: dict, type: str): return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } -def load_model(model, checkpoint): - ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1) - ckpt.restore(manager.latest_checkpoint) def sample(folder): + # run = wandb.init(project='vae_plots', name=folder.split('/')[-1]) load_model(model,folder) + count = 0 for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) 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]) + 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) pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:] - pose, shape = tf.tile(pose[tf.newaxis,:], [4,1]), tf.tile(shape[tf.newaxis,:], [4,1]) + pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1]) mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy() mesh_orig = batch['mesh_posed'][0].numpy() + mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0] body = batch['mesh_body_posed'][0].numpy() body_color = np.ones_like(body, dtype=np.int32)*100 @@ -424,72 +288,110 @@ def sample(folder): mask = batch['mesh_mask'].numpy().astype(bool)[0] _, faces = _maskFaces(np.ones((14475, 3)), F, mask) - batch = 0 + + 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) + mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2])) + mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy() + + + meshes = [] meshes.append(go.Mesh3d( # 8 vertices of a cube - x=mesh_orig[mask,0], - y=mesh_orig[mask,1], - z=mesh_orig[mask,2], + x=mesh_orig_reposed[mask,0], + y=mesh_orig_reposed[mask,2], + z=mesh_orig_reposed[mask,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] )) - for i in range(4): + for i in range(nr_of_samples): x = mesh_posed[i][mask, :].copy() - x[:, 1] += 1 * (i + 1) + # x[:, 2] += 1 * (i + 1) + x[:, 2] += 1 meshes.append(go.Mesh3d( # 8 vertices of a cube x=x[:,0], - y=x[:,1], - z=x[:,2], + y=x[:,2], + z=x[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] )) - for i in range(5): - x = body.copy() - x[:, 1] += 1 * i + + for i in range(nr_of_samples): + x = mesh_posed_zero[i][mask, :].copy() + # x[:, 2] += 1 * (i + 1) + x[:, 2] += 2 meshes.append(go.Mesh3d( # 8 vertices of a cube x=x[:,0], - y=x[:,1], - z=x[:,2], + y=x[:,2], + z=x[:,1], # i, j and k give the vertices of triangles - i = F[:,0], - j = F[:,1], - k = F[:,2], - vertexcolor=body_color + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] )) + # for i in range(5): + # x = body.copy() + # x[:, 2] += 1 * i + # meshes.append(go.Mesh3d( + # # 8 vertices of a cube + # x=x[:,0], + # y=x[:,2], + # z=x[:,1], + # # i, j and k give the vertices of triangles + # i = F[:,0], + # j = F[:,1], + # k = F[:,2], + # vertexcolor=body_color + # )) + fig = go.Figure(data=meshes) - fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=3, z=1)) + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1)) + # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1)) fig.update_layout( scene = dict( xaxis = dict(nticks=4, range=[-1,1],), - yaxis = dict(nticks=4, range=[-1,5],), + # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],), + yaxis = dict(nticks=4, range=[-1,3],), zaxis = dict(nticks=4, range=[-1,1],),) ) + # wandb.log({'plot_' + str(count): fig}) fig.show() + pred = pred.numpy() + frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])] + frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:], + save_all=True, duration=500, loop=0) + + + if count > 10: + break + + print(count) + count += 1 + print() if __name__ == '__main__': # init_wandb() - # allow_memory_growth() + allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) - batch_size = 12 + batch_size = 4 smpl = SMPL() @@ -511,27 +413,30 @@ if __name__ == '__main__': 'uv_shaped_body': (3, 512, 256,), } - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=batch_size, signature=signature_dynamic) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static) + 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_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() + model = CVAE(avg_garment_path='custom_data/jumpsuit_20_avg_uv.npy', + latent_dim=128, static_model='checkpoints/ae_static_3', + body_model='checpoints/ae_body_6', + curvature_model='checpoints/ae_curvature_6') # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False) optimizer = tf.optimizers.Adam(learning_rate=0.0001) - balancing_coeff = 250 + balancing_coeff = 0.0001 reg_coeff = 1 folder = 'models/ae_logs/' valid_iterator = iter(dataset_valid) gamma = 1 balancing_activated_step = 0 - avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy') + # avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy') - sample('checkpoints/cvae_10') + # sample('checkpoints/cvae_13') ckpt = tf.train.Checkpoint(model=model) @@ -561,9 +466,9 @@ if __name__ == '__main__': if train_loss['train/euclid_error'] < 0.01: if balancing_activated_step == 0: balancing_activated_step = step - balancing_coeff = balancing_coeff + 1000 / 30000 - if train_loss['train/euclid_error'] < 0.015: - reg_coeff = reg_coeff + 100 / 15000 + balancing_coeff = balancing_coeff + 1 / 15000 + # if train_loss['train/euclid_error'] < 0.015: + # reg_coeff = reg_coeff + 100 / 15000 # if step > 15000 and balancing_coeff <= 1e-3: # balancing_coeff = 1e-2 # if step > 20000 and balancing_coeff <= 1e-3: diff --git a/cvae_curv.py b/cvae_curv.py new file mode 100644 index 0000000..76fe407 --- /dev/null +++ b/cvae_curv.py @@ -0,0 +1,440 @@ +import os +os.environ["CUDA_VISIBLE_DEVICES"] = "3" +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 cmapy as cm +import numpy as np +import scipy as sp +from custom_utils.custom_utils import UV2Mesh, UV2MeshLabel, maskFacesBatch, _maskFaces +import plotly.graph_objects as go +import plotly.express as ex +from PIL import Image + +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 dataset import DatasetFactory, StaticDataset +from custom_utils.custom_utils import batch_to_dict +from custom_utils.smpl import SMPL +from models.cvae_curv import CVAE + + +name='cvae_curv_2_2_2' +def init_wandb(): + run = wandb.init(project='vae', name=name, + notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. new implementation for kl loss. \ +no lr schedule. 25 jumpsuits, average reconst_loss before calculation, offsets over average unposed garment. gamma balancing. remove global \ +orientation (in posed body input and in pose_shape parameters.. retrained body encoder with deoriented bodies. lower kl coeff to 0.1. fixed plotting colors\ +fixed deoriented bodies and need retraining (ae_body_6 with 50k iteration), addent MLP for lowering conditional Z diemnsion\ +higher starting kl coeff (0.1)') + artifact = wandb.Artifact('code', type='code') + artifact.add_file('cvae_curv.py') + artifact.add_file('dataset.py') + artifact.add_file('losses.py') + artifact.add_file('convnext/models/convnext_tf.py') + run.log_artifact(artifact) + +uv2mesh = UV2MeshLabel(res=256) + +def plot_3d(curv_gt, curv_pred, faces, mesh): + mesh_2 = mesh.copy() + mesh_2[:,1] += 1 + + colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)]) + v_color_1 = np.zeros((mesh.shape[0], 3)) + v_color_2 = np.zeros((mesh.shape[0], 3)) + for i in range(6+1): + v_color_1[curv_gt == i, :] = colors[i] + v_color_2[curv_pred == i, :] = colors[i] + + fig = go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=mesh[:,0], + y=mesh[:,1], + z=mesh[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + vertexcolor=v_color_1, + text='gt' + ), + go.Mesh3d( + # 8 vertices of a cube + x=mesh_2[:,0], + y=mesh_2[:,1], + z=mesh_2[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + vertexcolor=v_color_2, + text='pred' + ) + ]) + + fig.update_layout(scene_aspectmode='manual', + scene_aspectratio=dict(x=1, y=2, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,3],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + return fig + + +def log_normal_pdf(sample, mean, logvar, raxis=1): + log2pi = tf.math.log(2. * np.pi) + res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis) + return res + + +def kl_loss(mean, logvar): + kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar) + kl_loss = tf.reduce_mean(kl_loss) + kl_loss *= -0.5 + return kl_loss + + + +g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7) +def get_mean_iou(y_true, y_pred): + g_mean_iou.reset_state() + g_mean_iou.update_state(y_true, y_pred) + return g_mean_iou.result() + +# https://github.com/NifTK/NiftyNet/blob/dev/niftynet/layer/loss_segmentation.py +def dice_dense(prediction, ground_truth, weight_map=None): + """ + Computing mean-class Dice similarity. + :param prediction: last dimension should have ``num_classes`` + :param ground_truth: segmentation ground truth (encoded as a binary matrix) + last dimension should be ``num_classes`` + :param weight_map: + :return: ``1.0 - mean(Dice similarity per class)`` + """ + + if weight_map is not None: + raise NotImplementedError + prediction = tf.cast(prediction, dtype=tf.float32) + ground_truth = tf.cast(ground_truth, dtype=tf.float32) + ground_truth = tf.reshape(ground_truth, prediction.shape) + # computing Dice over the spatial dimensions + reduce_axes = list(range(len(prediction.shape) - 1)) + dice_numerator = 2.0 * tf.reduce_sum( + prediction * ground_truth, axis=reduce_axes) + dice_denominator = \ + tf.reduce_sum(tf.square(prediction), axis=reduce_axes) + \ + tf.reduce_sum(tf.square(ground_truth), axis=reduce_axes) + + epsilon = 0.00001 + + dice_score = (dice_numerator + epsilon) / (dice_denominator + epsilon) + return tf.reduce_mean(dice_score) + + +@tf.function +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_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']]) + + ce_uv_loss = scce(batch['uv_curvature'], pred) + ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1)) + + reconstr_loss = ce_uv_loss + reconstr_loss = tf.reduce_mean(reconstr_loss) + kl = kl_loss(mean, logvar) + loss = reconstr_loss + balancing_coeff * kl + ms_ssim = tf.reduce_mean(ms_ssim) + mean = tf.reduce_mean(mean) + logvar = tf.reduce_mean(logvar) + + + # mesh_pred = uv2mesh(tf.tile(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], [1,1,1,3]))[:,:,0] + # mesh_orig = tf.one_hot(batch['mesh_curvature'] + tf.cast(batch['mesh_mask'], tf.int32), 7) + # ce_mesh_loss = scce(mesh_orig, mesh_pred) + + mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1)) + dice = dice_dense(pred, uv_curvature) + + gradients = tape.gradient(loss, model.trainable_variables) + optimizer.apply_gradients(zip(gradients, model.trainable_variables)) + return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/dice': dice, #'ce_mesh_loss': ce_mesh_loss + 'train/reconstr_loss': reconstr_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} + + + + + +@tf.function +def valid_step(batch): + uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + + pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']]) + + ce_uv_loss = scce(batch['uv_curvature'], pred) + ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1)) + loss = ce_uv_loss + + mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1)) + dice = dice_dense(pred, uv_curvature) + + reconstr_loss = tf.reduce_mean(ce_uv_loss) + kl = kl_loss(mean, logvar) + loss = reconstr_loss + balancing_coeff * kl + ms_ssim = tf.reduce_mean(ms_ssim) + mean = tf.reduce_mean(mean) + logvar = tf.reduce_mean(logvar) + + return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou, 'valid/dice': dice, + 'valid/ms_ssim': ms_ssim, 'valid/reconstr_loss': reconstr_loss, + 'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar } + + +def plot_everything(batch: dict, type: str): + uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']]) + + pred = np.argmax(pred.numpy(), axis=-1) + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_curvature'] + mesh_base = batch['mesh_posed'] + + orig_img = batch['uv_curvature'] + pred_img = pred + img = tf.concat((orig_img, pred_img), axis=2) + imgs = [wandb.Image(img[i].numpy().astype(np.float32) / 6.0) for i in range(4)] + + F = np.load("custom_utils/smpl_faces.npy") + mask = batch['mesh_mask'].numpy().astype(bool) + _, faces = maskFacesBatch(np.zeros((4, 14475, 3)), F, mask) + gt_curv = [mesh_orig[i].numpy()[mask[i]] + 1 for i in range(pred.shape[0])] + pred_curv = [mesh_pred[i].numpy()[mask[i]] for i in range(pred.shape[0])] + mesh = [mesh_base[i].numpy()[mask[i]] for i in range(pred.shape[0])] + + figs = { type + "/garment_" + str(i): plot_3d(gt_curv[i], pred_curv[i], faces[i], mesh[i]) 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) + + return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } + + +def load_model(model, checkpoint): + ckpt = tf.train.Checkpoint(model=model) + manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1) + ckpt.restore(manager.latest_checkpoint) + + +def sample(folder): + load_model(model,folder) + for batch in dataset_train: + batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) + uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + nr_of_samples = 16 + + # pred = model.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']) + pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], batch['uv_mask'][0], nr_of_samples=nr_of_samples) + + ce_uv_loss = scce(tf.tile(batch['uv_curvature'][0][tf.newaxis,:,:], [nr_of_samples,1,1]), pred) + + + mesh_pred = uv2mesh(pred).numpy().argmax(axis=-1) + mesh_orig = batch['mesh_curvature'][0].numpy() + + F = np.load("custom_utils/smpl_faces.npy") + mask = batch['mesh_mask'].numpy().astype(bool)[0] + _, faces = _maskFaces(np.ones((14475, 3)), F, mask) + mesh_pred = [mesh_pred[i, mask] for i in range(mesh_pred.shape[0])] + curv_color = [np.zeros((mesh_pred[i].shape[0], 3)) for i in range(len(mesh_pred))] + mesh_orig = mesh_orig[mask] + 1 + + colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)]) + curv_color_gt = np.zeros((mesh_pred[0].shape[0], 3)) + for i in range(6+1): + for j in range(len(mesh_pred)): + curv_color[j][mesh_pred[j] == i, :] = colors[i] + curv_color_gt[mesh_orig == i, :] = colors[i] + + mesh_base = batch['mesh_posed'][0].numpy()[mask, :] + meshes = [] + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=mesh_base[:,0], + y=mesh_base[:,1], + z=mesh_base[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + vertexcolor=curv_color_gt, + text='gt' + )) + + for i in range(nr_of_samples): + x = mesh_base.copy() + x[:, 1] += 1 * (i + 1) + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,1], + z=x[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + vertexcolor=curv_color[i] + )) + + + + + fig = go.Figure(data=meshes) + + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,nr_of_samples + 1],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + # fig.show() + + pred = pred.numpy().argmax(-1) + frames = [Image.fromarray(pred[i]/6*255) for i in range(pred.shape[0])] + frames[0].save("temp/curv2.gif", format="GIF", append_images=frames[1:], + save_all=True, duration=500, loop=0) + + print() + +if __name__ == '__main__': + # init_wandb() + allow_memory_growth() + parser = argparse.ArgumentParser(description='') + parser.add_argument('--name', type=str) + args = vars(parser.parse_args()) + + batch_size = 4 + + smpl = SMPL() + scce = tf.keras.losses.SparseCategoricalCrossentropy() + + factory = DatasetFactory() + signature_dynamic = { 'id': (), + 'pose_shape': (72 + 11,), + # 'uv_posed': (3, 512, 256,), + 'uv_body_posed': (3, 512, 256,), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + # 'mesh_unposed': (14475, 3,), + # 'uv_unposed': (3, 512, 256,), + 'uv_curvature': (512, 256), + 'mesh_curvature': (14475) + } + signature_static = { 'uv_static': (3, 512, 256,), + # 'mesh_static': (14475, 3,), + 'mesh_mask': (14475,), + 'uv_mask': (512, 256,), + 'uv_shaped_body': (3, 512, 256,), + } + + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=batch_size, signature=signature_dynamic) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static) + 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(latent_dim=128, static_model='checkpoints/ae_static_3', + body_model='checpoints/ae_body_6',) + + + # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False) + optimizer = tf.optimizers.Adam(learning_rate=0.0001) + balancing_coeff = 0.1 + reg_coeff = 1 + folder = 'models/ae_logs/' + valid_iterator = iter(dataset_valid) + gamma = 1 + balancing_activated_step = 0 + avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy') + + + # sample('checkpoints/cvae_curv_2') + + + ckpt = tf.train.Checkpoint(model=model) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5) + + step = 0 + for batch in dataset_train: + batch = static_dataset_train.expand_batch(batch_to_dict(batch, 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) + valid_loss = {} + + to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } + if step % 2500 == 0: + train_figs = plot_everything(batch, 'train') + # valid_figs = plot_everything(data_valid, 'valid') + valid_figs = {} + to_log = { **to_log, **train_figs, **valid_figs } + + if step % 25000 == 0: + manager.save(checkpoint_number=step) + + if step > 25000: + # if train_loss['train/euclid_error'] < 0.01: + balancing_coeff = balancing_coeff + 1 / 20000 + # if train_loss['train/euclid_error'] < 0.015: + # reg_coeff = reg_coeff + 100 / 15000 + + + # if step > 15000 and balancing_coeff <= 1e-3: + # balancing_coeff = 1e-2 + # if step > 20000 and balancing_coeff <= 1e-3: + # balancing_coeff = 1e-2 + # if step > 25000 and balancing_coeff <= 1e-2: + # balancing_coeff = 1e-1 + # if step > 30000 and balancing_coeff <= 1e-1: + # balancing_coeff = 1 + + # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss'] + + to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } } + + wandb.log(to_log, step = step) + + print(step) + + step += 1 \ No newline at end of file diff --git a/cvae_vae.py b/cvae_vae.py new file mode 100644 index 0000000..4b97230 --- /dev/null +++ b/cvae_vae.py @@ -0,0 +1,586 @@ +import os +os.environ["CUDA_VISIBLE_DEVICES"] = "2" +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 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 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 CurvatureVAE + + +name='cvae_vae_1' +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.\ +new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1000 from 250.\ +offsets over average unposed garment. gamma balancing. remove global orientation (in posed body input and in pose_shape parameters. \ +rebalanced l1 loss. retrained body encoder with deoriented bodies \ +') + artifact = wandb.Artifact('code', type='code') + artifact.add_file('cvae.py') + artifact.add_file('dataset.py') + artifact.add_file('losses.py') + artifact.add_file('convnext/models/convnext_tf.py') + run.log_artifact(artifact) + +uv2mesh = UV2Mesh(res=256) + +def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color): + mesh_pred_unposed[:,1] += 1 + collision[:,1] += 2 + fig = go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=mesh_gt[:,0], + y=mesh_gt[:,1], + z=mesh_gt[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + ), + go.Mesh3d( + # 8 vertices of a cube + x=mesh_pred_unposed[:,0], + y=mesh_pred_unposed[:,1], + z=mesh_pred_unposed[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + ), + go.Mesh3d( + # 8 vertices of a cube + x=collision[:,0], + y=collision[:,1], + z=collision[:,2], + # i, j and k give the vertices of triangles + i = collision_faces[:,0], + j = collision_faces[:,1], + k = collision_faces[:,2], + vertexcolor=collision_color + ) + ]) + + fig.update_layout(scene_aspectmode='manual', + scene_aspectratio=dict(x=1, y=2, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,3],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + return fig + + +class Bottleneck(tf.keras.Model): + + def __init__(self, dim): + super(Bottleneck, self).__init__() + + self.bottleneck = tf.keras.Sequential( + [Block(dim=dim, drop_path=0, + layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] + ) + + def call(self, inputs, training=None): + x = inputs + + x = self.bottleneck(x, training=training) + + return x + + +class Encoder(tf.keras.Model): + + def __init__(self, latent_dim): + super(Encoder, self).__init__() + + self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False) + self.flatten = tf.keras.layers.Flatten() + self.mean = tf.keras.layers.Dense(latent_dim) + self.log_var = tf.keras.layers.Dense(latent_dim) + + def call(self, inputs, training=None): + x = inputs + + x = self.encoder(x) + x = self.flatten(x) + mean = self.mean(x) + log_var = self.log_var(x) + + return mean, log_var + + +class Decoder(tf.keras.Model): + + def __init__(self): + super(Decoder, self).__init__() + + self.dense = tf.keras.layers.Dense(16 * 8 * 80) + self.reshape = tf.keras.layers.Reshape((16, 8, 80)) + self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16]))) + + + def call(self, inputs, training=None): + x = inputs + + x = self.dense(x) + x = self.reshape(x) + x = self.decoder(x) + + return x + + +class CVAE(tf.keras.Model): + + def __init__(self, latent_dim=128): + super(CVAE, self).__init__() + self.latent_dim = latent_dim + + model : StaticAE = StaticAE() + load_model(model, 'checkpoints/ae_static_3') + self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.static_enc.trainable = False + + model : BodyAE = BodyAE() + load_model(model, 'checpoints/ae_body_5') + self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.body_enc.trainable = False + + model : CurvatureAE = CurvatureVAE() + load_model(model, 'checpoints/cvae_curv_1_3') + self.curv_gen = model + self.curv_gen.trainable = False + + self.encoder = Encoder(latent_dim) + self.decoder = Decoder() + + + # @tf.function + def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, eps=None): + if eps is None: + eps = tf.random.normal(shape=(4, self.latent_dim)) + + uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1]) + uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1]) + uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [4,1,1,1]) + uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [4,1,1]) + + # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1]) + z_static = self.static_enc(uv_template) + z_body = self.body_enc(uv_body_posed) + z_curv = self.curv_gen(uv_curv) + + z_cond = tf.concat([z_static, z_body, z_curv], axis=-1) + z = tf.concat([eps, z_cond], axis=-1) + + x = self.decoder(z) + x = x + avg_jumpsuit[tf.newaxis,:,:,:] + x = x * uv_mask[:,:,:,tf.newaxis] + + return x + + + def reparameterize(self, mean, logvar): + eps = tf.random.normal(shape=mean.shape) + return eps * tf.exp(logvar * .5) + mean + + + def decode(self, z, apply_sigmoid=False): + logits = self.decoder(z) + if apply_sigmoid: + probs = tf.sigmoid(logits) + return probs + return logits + + + def call(self, inputs, training=None): + uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = 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]) + z_static = self.static_enc(uv_template) + z_body = self.body_enc(uv_body_posed) + z_curv = self.curv_gen.sample(uv_template, uv_body_posed, uv_mask) + + 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 = tf.concat([z, z_static, z_body, 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]) + x = x + avg_jumpsuit[tf.newaxis,:,:,:] + x = x * uv_mask[:,:,:,tf.newaxis] + + return x, mean, logvar, z + + +def log_normal_pdf(sample, mean, logvar, raxis=1): + log2pi = tf.math.log(2. * np.pi) + res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis) + return res + + +def kl_loss(mean, logvar): + kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar) + kl_loss = tf.reduce_mean(kl_loss) + kl_loss *= -0.5 + return kl_loss + + +@tf.function +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']]) + + 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_orig = batch['mesh_unposed'] + + + 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]) + # 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.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']) + tv_loss = tf.image.total_variation(pred) + + reconstr_loss = l1_loss / 4 + reconstr_loss = tf.reduce_mean(reconstr_loss) + reg_coeff *(e_loss + c_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 / (2 * (gamma ** 2)) + balancing_coeff * kl + # loss = reconstr_loss + balancing_coeff * kl + ms_ssim = tf.reduce_mean(ms_ssim) + l1_loss = tf.reduce_mean(l1_loss) + fft_loss = tf.reduce_mean(fft_loss) + tv_loss = tf.reduce_mean(tv_loss) + mean = tf.reduce_mean(mean) + logvar = tf.reduce_mean(logvar) + + mesh_orig_2 = uv2mesh(batch['uv_unposed']) + 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']) + euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + gradients = tape.gradient(loss, model.trainable_variables) + # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients] + # 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/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 } + + + +@tf.function +def valid_step(batch): + 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']]) + + mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) + mesh_orig = batch['mesh_unposed'] + e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, 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]) + b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask']) + # 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) + 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]) + # mesh_orig_2 = uv2mesh(batch['uv_unposed']) + 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 = 0 #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']) + + reconstr_loss = l1_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 / (2 * (gamma ** 2)) + kl + loss = reconstr_loss + balancing_coeff * kl + ms_ssim = tf.reduce_mean(ms_ssim) + l1_loss = tf.reduce_mean(l1_loss) + fft_loss = tf.reduce_mean(fft_loss) + b_loss = tf.reduce_mean(b_loss) + reconstr_loss = tf.reduce_mean(reconstr_loss) + mean = tf.reduce_mean(mean) + logvar = tf.reduce_mean(logvar) + + 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, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss, + 'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar } + + +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']]) + + mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) + mesh_orig = batch['mesh_unposed'] + + orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1]) + uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1]) + uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1]) + 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() + 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))] + + 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) } + + 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) + + return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } + + +def load_model(model, checkpoint): + ckpt = tf.train.Checkpoint(model=model) + manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1) + ckpt.restore(manager.latest_checkpoint) + + +def sample(folder): + load_model(model,folder) + for batch in dataset_train: + batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) + uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + + pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0]) + pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:] + pose, shape = tf.tile(pose[tf.newaxis,:], [4,1]), tf.tile(shape[tf.newaxis,:], [4,1]) + + mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) + mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy() + mesh_orig = batch['mesh_posed'][0].numpy() + + body = batch['mesh_body_posed'][0].numpy() + body_color = np.ones_like(body, dtype=np.int32)*100 + + F = np.load("custom_utils/smpl_faces.npy") + mask = batch['mesh_mask'].numpy().astype(bool)[0] + _, faces = _maskFaces(np.ones((14475, 3)), F, mask) + + batch = 0 + + meshes = [] + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=mesh_orig[mask,0], + y=mesh_orig[mask,1], + z=mesh_orig[mask,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + for i in range(4): + x = mesh_posed[i][mask, :].copy() + x[:, 1] += 1 * (i + 1) + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,1], + z=x[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + for i in range(5): + x = body.copy() + x[:, 1] += 1 * i + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,1], + z=x[:,2], + # i, j and k give the vertices of triangles + i = F[:,0], + j = F[:,1], + k = F[:,2], + vertexcolor=body_color + )) + + + fig = go.Figure(data=meshes) + + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=3, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,5],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + fig.show() + + print() + +if __name__ == '__main__': + # init_wandb() + # allow_memory_growth() + parser = argparse.ArgumentParser(description='') + parser.add_argument('--name', type=str) + args = vars(parser.parse_args()) + + batch_size = 12 + + smpl = SMPL() + + factory = DatasetFactory() + signature_dynamic = { 'id': (), + 'pose_shape': (72 + 11,), + # 'uv_posed': (3, 512, 256,), + 'uv_body_posed': (3, 512, 256,), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'mesh_unposed': (14475, 3,), + 'uv_unposed': (3, 512, 256,), + 'uv_curvature': (512, 256), + } + signature_static = { 'uv_static': (3, 512, 256,), + # 'mesh_static': (14475, 3,), + 'mesh_mask': (14475,), + 'uv_mask': (512, 256,), + 'uv_shaped_body': (3, 512, 256,), + } + + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=batch_size, signature=signature_dynamic) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static) + 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() + + + # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False) + optimizer = tf.optimizers.Adam(learning_rate=0.0001) + balancing_coeff = 250 + reg_coeff = 1 + folder = 'models/ae_logs/' + valid_iterator = iter(dataset_valid) + gamma = 1 + balancing_activated_step = 0 + avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy') + + + sample('checkpoints/cvae_10') + + + ckpt = tf.train.Checkpoint(model=model) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) + + step = 0 + for batch in dataset_train: + batch = static_dataset_train.expand_batch(batch_to_dict(batch, 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) + valid_loss = {} + + to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } + if step % 2500 == 0: + train_figs = plot_everything(batch, 'train') + # valid_figs = plot_everything(data_valid, 'valid') + valid_figs = {} + to_log = { **to_log, **train_figs, **valid_figs } + + if step % 10000 == 0: + manager.save(checkpoint_number=step) + + if train_loss['train/euclid_error'] < 0.01: + if balancing_activated_step == 0: + balancing_activated_step = step + balancing_coeff = balancing_coeff + 1000 / 30000 + if train_loss['train/euclid_error'] < 0.015: + reg_coeff = reg_coeff + 100 / 15000 + # if step > 15000 and balancing_coeff <= 1e-3: + # balancing_coeff = 1e-2 + # if step > 20000 and balancing_coeff <= 1e-3: + # balancing_coeff = 1e-2 + # if step > 25000 and balancing_coeff <= 1e-2: + # balancing_coeff = 1e-1 + # if step > 30000 and balancing_coeff <= 1e-1: + # balancing_coeff = 1 + + # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss'] + + to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } } + + wandb.log(to_log, step = step) + + print(step) + + step += 1 \ No newline at end of file diff --git a/dataset.py b/dataset.py index 5b88dc1..4b9107d 100644 --- a/dataset.py +++ b/dataset.py @@ -99,11 +99,12 @@ class StaticDataset(): class DatasetIdx(tf.data.Dataset): - def __new__(cls, folder): + def __new__(cls, folder, shuffle=True): uv_static_folder = os.path.join(folder, 'pose_shape') - frames = sorted(os.listdir(uv_static_folder))#[453:503] + frames = sorted(os.listdir(uv_static_folder)) frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] - random.shuffle(frames) + if shuffle: + random.shuffle(frames) return tf.data.Dataset.from_tensor_slices(frames) @@ -126,9 +127,9 @@ class DatasetFactory(): return tf_signature - def get_dataset(self, folder, batch_size, signature): + def get_dataset(self, folder, batch_size, signature, shuffle=True): # static_data = StaticDataset(folder, signature) - dataset_idx = DatasetIdx(folder) + dataset_idx = DatasetIdx(folder, shuffle) def read_item(frame): frame_split = tf.strings.split(frame, sep='_') @@ -156,6 +157,9 @@ class DatasetFactory(): # if 'mesh_body_posed' in signature.keys(): # mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy')) # item['mesh_body_posed'] = mesh_body_posed + # if 'mesh_posed_body_no_orient' in signature.keys(): + # mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] + # item['mesh_posed_body_no_orient'] = mesh_posed_body_no_orient if 'uv_body_posed' in signature.keys(): uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1]) item['uv_body_posed'] = uv_posed_body_no_orient @@ -165,12 +169,18 @@ class DatasetFactory(): if 'mesh_posed' in signature.keys(): mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'] item['mesh_posed'] = mesh_posed + # if 'mesh_unposed' in signature.keys(): + # mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed'] + # item['mesh_unposed'] = mesh_unposed + # if 'uv_unposed' in signature.keys(): + # uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1]) + # item['uv_unposed'] = uv_unposed if 'mesh_unposed' in signature.keys(): - mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed'] - item['mesh_unposed'] = mesh_unposed + mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + '_' + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] + item['mesh_unposed'] = mesh_unposed_no_shape if 'uv_unposed' in signature.keys(): - uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1]) - item['uv_unposed'] = uv_unposed + uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape', id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1]) + item['uv_unposed'] = uv_unposed_no_shape if 'uv_curvature' in signature.keys(): uv_curvature = np.load(os.path.join(folder, 'uv_curvature', id + '_' + frame_nr + '.npz')); uv_curvature = uv_curvature['curvature_uv'] item['uv_curvature'] = uv_curvature.astype(np.int32) diff --git a/deshape.py b/deshape.py new file mode 100644 index 0000000..8a9ff82 --- /dev/null +++ b/deshape.py @@ -0,0 +1,117 @@ +import os +import sys +import time +import cv2 +import random +import numpy as np +import tensorflow as tf +import cmapy as cm +import plotly.graph_objects as go +import plotly.express as ex + +from dataset import DatasetFactory, StaticDataset +from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces +from custom_utils.smpl import SMPL + +if __name__ == '__main__': + os.environ["CUDA_VISIBLE_DEVICES"] = "5" + + uv2mesh = UV2Mesh() + smpl = SMPL() + + F = np.load("custom_utils/smpl_faces.npy") + + factory = DatasetFactory() + signature_dynamic = { 'id': (), + 'pose_shape': (72 + 11,), + # 'uv_posed': (3, 512, 256,), + # 'uv_body_posed': (3, 512, 256,), + # 'mesh_body_posed': (14475, 3,), + # 'mesh_posed_body_no_orient': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'mesh_unposed': (14475, 3,), + # 'uv_unposed': (3, 512, 256,), + # 'uv_curvature': (512, 256), + # 'mesh_curvature': (14475) + } + signature_static = { #'uv_static': (3, 512, 256,), + # 'mesh_static': (14475, 3,), + 'mesh_mask': (14475,), + # 'uv_mask': (512, 256,), + # 'uv_shaped_body': (3, 512, 256,), + } + + dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', + batch_size=1, + signature=signature_dynamic) + static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', + signature=signature_static) + + + uv_res = 256 + + + + for i, batch in enumerate(dataset): + print(i) + batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) + + mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy()) + mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool) + pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy() + # mesh_posed = np.squeeze(batch['mesh_posed'].numpy()) + gender = int(shape[0, -1] > 0) + print(gender) + + shape_offsets = tf.tensordot(shape[:,:10], smpl.shapedirs[gender], axes=[[1], [2]]).numpy() + + + mesh_unposed[mask, :] = mesh_unposed[mask, :] - np.squeeze(shape_offsets)[smpl.to_keep, :][mask, :] + + np.savez_compressed(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), mesh_posed_body_no_orient = T[i, :, :]) + + + # _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask) + + # x = mesh_unposed[mask, :] + # y = mesh_unposed[mask, :] - np.squeeze(shape_offsets)[smpl.to_keep, :][mask, :] + + # # x = smpl.v_template.numpy()[0, smpl.to_keep, :] + # # R = smpl.run_rodrigues(np.zeros((1, 72), dtype=np.float32)) + # # x = smpl.run(R, shape)[0][0] + # # y = smpl.v_template.numpy()[gender, smpl.to_keep, :] + np.squeeze(shape_offsets)[smpl.to_keep, :] + # # faces = F + + # fig = go.Figure(data=[ + # go.Mesh3d( + # # 8 vertices of a cube + # x=x[:,0], + # y=x[:,1], + # z=x[:,2], + # # i, j and k give the vertices of triangles + # i = faces[:,0], + # j = faces[:,1], + # k = faces[:,2] + # ), + # go.Mesh3d( + # # 8 vertices of a cube + # x=y[:,0], + # y=y[:,1], + # z=y[:,2], + # # i, j and k give the vertices of triangles + # i = faces[:,0], + # j = faces[:,1], + # k = faces[:,2] + # )]) + + # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1)) + # fig.update_layout( + # scene = dict( + # xaxis = dict(nticks=4, range=[-1,1],), + # yaxis = dict(nticks=4, range=[-1,1],), + # zaxis = dict(nticks=4, range=[-1,1],),) + # ) + + # fig.show() + + print() \ No newline at end of file diff --git a/models/cvae.py b/models/cvae.py new file mode 100644 index 0000000..f4d60c7 --- /dev/null +++ b/models/cvae.py @@ -0,0 +1,191 @@ +import os +import time +import random +import tensorflow as tf +import numpy as np +import scipy as sp +# import plotly.graph_objects as go +# import plotly.express as ex + +from custom_utils.custom_utils import load_model +from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block +from auto_encoder_body import AutoEncoder as BodyAE +from auto_encoder_curv import AutoEncoder as CurvatureAE +from auto_encoder_static import AutoEncoder as StaticAE + + +class Bottleneck(tf.keras.Model): + + def __init__(self, dim): + super(Bottleneck, self).__init__() + + self.bottleneck = tf.keras.Sequential( + [Block(dim=dim, drop_path=0, + layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] + ) + + def call(self, inputs, training=None): + x = inputs + + x = self.bottleneck(x, training=training) + + return x + + +class Encoder(tf.keras.Model): + + def __init__(self, latent_dim): + super(Encoder, self).__init__() + + self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False) + self.flatten = tf.keras.layers.Flatten() + self.mean = tf.keras.layers.Dense(latent_dim) + self.log_var = tf.keras.layers.Dense(latent_dim) + + def call(self, inputs, training=None): + x = inputs + + x = self.encoder(x) + x = self.flatten(x) + mean = self.mean(x) + log_var = self.log_var(x) + + return mean, log_var + + +class Decoder(tf.keras.Model): + + def __init__(self): + super(Decoder, self).__init__() + + self.dense = tf.keras.layers.Dense(16 * 8 * 80) + self.reshape = tf.keras.layers.Reshape((16, 8, 80)) + self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16]))) + + + def call(self, inputs, training=None): + x = inputs + + x = self.dense(x) + x = self.reshape(x) + x = self.decoder(x) + + return x + + +class CVAE(tf.keras.Model): + + def __init__(self, avg_garment_path='custom_data/jumpsuit_20_avg_uv.npy', + latent_dim=128, static_model='checkpoints/ae_static_3', + body_model='checpoints/ae_body_6', + curvature_model='checpoints/ae_curvature_6'): + super(CVAE, self).__init__() + self.latent_dim = latent_dim + + model : StaticAE = StaticAE() + load_model(model, static_model) + self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.static_enc.trainable = False + + model : BodyAE = BodyAE() + load_model(model, body_model) + self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.body_enc.trainable = False + + model : CurvatureAE = CurvatureAE() + load_model(model, curvature_model) + self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.curv_enc.trainable = False + + self.encoder = Encoder(latent_dim) + self.decoder = Decoder() + + self.avg_garment = np.load(avg_garment_path) + + self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'), + tf.keras.layers.Dense(256, activation='relu'), + tf.keras.layers.Dense(128, activation='relu'), + tf.keras.layers.Dense(latent_dim),]) + + # 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): + if not fix_noise: + eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim)) + else: + eps = tf.random.normal(shape=(1, self.latent_dim)) + eps = tf.tile(eps, [nr_of_samples,1]) + + uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1]) + uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1]) + if len(uv_curv.shape) == 3: + uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1]) + uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1]) + + # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [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) + + 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) + x = x + self.avg_garment[tf.newaxis,:,:,:] + x = x * uv_mask[:,:,:,tf.newaxis] + + return x + + + def reparameterize(self, mean, logvar): + eps = tf.random.normal(shape=mean.shape) + return eps * tf.exp(logvar * .5) + mean + + + def decode(self, z, apply_sigmoid=False): + logits = self.decoder(z) + if apply_sigmoid: + probs = tf.sigmoid(logits) + return probs + return logits + + + def call(self, inputs, training=None, zero_latent=False): + uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = 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]) + 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) + 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) + + z = tf.concat([z, z_cond], axis=-1) + + # z = self.shuffle(z) + + # 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]) + x = x + self.avg_garment[tf.newaxis,:,:,:] + x = x * uv_mask[:,:,:,tf.newaxis] + + return x, mean, logvar, z \ No newline at end of file diff --git a/models/cvae_curv.py b/models/cvae_curv.py new file mode 100644 index 0000000..0a0e211 --- /dev/null +++ b/models/cvae_curv.py @@ -0,0 +1,187 @@ +import os +import time +import random +import tensorflow as tf +import numpy as np +import scipy as sp +# import plotly.graph_objects as go +# import plotly.express as ex + +from custom_utils.custom_utils import load_model +from convnext.models.convnext_tf import ConvNeXt, Block, ConvNeXtReverseLabels +from auto_encoder_body import AutoEncoder as BodyAE +from auto_encoder_static import AutoEncoder as StaticAE + + + +class Bottleneck(tf.keras.Model): + + def __init__(self, dim): + super(Bottleneck, self).__init__() + + self.bottleneck = tf.keras.Sequential( + [Block(dim=dim, drop_path=0, + layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] + ) + + def call(self, inputs, training=None): + x = inputs + + x = self.bottleneck(x, training=training) + + return x + + +class Encoder(tf.keras.Model): + + def __init__(self, latent_dim): + super(Encoder, self).__init__() + + self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False) + self.flatten = tf.keras.layers.Flatten() + self.mean = tf.keras.layers.Dense(latent_dim) + self.log_var = tf.keras.layers.Dense(latent_dim) + + def call(self, inputs, training=None): + x = inputs + + x = self.encoder(x) + x = self.flatten(x) + mean = self.mean(x) + log_var = self.log_var(x) + + return mean, log_var + + +class Decoder(tf.keras.Model): + + def __init__(self): + super(Decoder, self).__init__() + + self.dense = tf.keras.layers.Dense(16 * 8 * 80) + self.reshape = tf.keras.layers.Reshape((16, 8, 80)) + self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16]))) + + + def call(self, inputs, training=None): + x = inputs + + x = self.dense(x) + x = self.reshape(x) + x = self.decoder(x) + + return x + + +class CVAE(tf.keras.Model): + + def __init__(self, latent_dim=128, static_model='checkpoints/ae_static_3', + body_model='checpoints/ae_body_6',): + super(CVAE, self).__init__() + self.latent_dim = latent_dim + + model : StaticAE = StaticAE() + load_model(model, static_model) + self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.static_enc.trainable = False + + model : BodyAE = BodyAE() + load_model(model, body_model) + self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.body_enc.trainable = False + + self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(256, activation='relu'), + tf.keras.layers.Dense(256, activation='relu'), + tf.keras.layers.Dense(128, activation='relu'), + tf.keras.layers.Dense(latent_dim),]) + + self.encoder = Encoder(latent_dim) + self.decoder = Decoder() + + + # @tf.function + def sample_single(self, uv_template, uv_body_posed, uv_mask, eps=None, nr_of_samples=4): + if eps is None: + eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim)) + + uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1]) + uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1]) + uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1]) + + # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1]) + z_static = self.static_enc(uv_template) + z_body = self.body_enc(uv_body_posed) + + z_cond = tf.concat([z_static, z_body], axis=-1) + z_cond = self.cond_encoder(z_cond) + z = tf.concat([eps, z_cond], axis=-1) + + x = self.decoder(z) + x = x * uv_mask[:,:,:,tf.newaxis] + + return x + + + def sample(self, uv_template, uv_body_posed, uv_mask, eps=None): + eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim)) + + uv_template = tf.transpose(uv_template, [0,2,3,1]) + uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1]) + + z_static = self.static_enc(uv_template) + z_body = self.body_enc(uv_body_posed) + + z_cond = tf.concat([z_static, z_body], axis=-1) + z_cond = self.cond_encoder(z_cond) + z = tf.concat([eps, z_cond], axis=-1) + + x = self.decoder(z) + x = x * uv_mask[:,:,:,tf.newaxis] + + return x + + + # encoders have to be the same for curv_vae and general CVAE + def sample_without_encoding(self, z_static, z_body, uv_mask, eps=None): + eps = tf.random.normal(shape=(z_static.shape[0], self.latent_dim)) + + z_cond = tf.concat([z_static, z_body], axis=-1) + z = tf.concat([eps, z_cond], axis=-1) + + x = self.decoder(z) + x = x * uv_mask[:,:,:,tf.newaxis] + + return x + + + def reparameterize(self, mean, logvar): + eps = tf.random.normal(shape=mean.shape) + return eps * tf.exp(logvar * .5) + mean + + + def decode(self, z, apply_sigmoid=False): + logits = self.decoder(z) + if apply_sigmoid: + probs = tf.sigmoid(logits) + return probs + return logits + + + def call(self, inputs, training=None): + uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs + uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1]) + z_static = self.static_enc(uv_template) + z_body = self.body_enc(uv_body_posed) + + x = tf.concat([uv_template, uv_body_posed, uv_curv], axis=-1) + + mean, logvar = self.encoder(x) + z = self.reparameterize(mean, logvar) + z_cond = tf.concat([z_static, z_body], axis=-1) + z_cond = self.cond_encoder(z_cond) + z = tf.concat([z, z_cond], axis=-1) + x = self.decoder(z) + + x = x * uv_mask[:,:,:,tf.newaxis] + + return x, mean, logvar, z diff --git a/sample.py b/sample.py new file mode 100644 index 0000000..9e38d11 --- /dev/null +++ b/sample.py @@ -0,0 +1,332 @@ +import os +os.environ["CUDA_VISIBLE_DEVICES"] = "0" +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 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 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 import CVAE as CVAEMain + +name='cvae_12_6' +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. \ +new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \ +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.1') + artifact = wandb.Artifact('code', type='code') + artifact.add_file('cvae.py') + artifact.add_file('dataset.py') + artifact.add_file('losses.py') + artifact.add_file('convnext/models/convnext_tf.py') + run.log_artifact(artifact) + +uv2mesh = UV2Mesh(res=256) + +def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color): + mesh_pred_unposed[:,2] += 1 + collision[:,2] += 2 + fig = go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=mesh_gt[:,0], + y=mesh_gt[:,2], + z=mesh_gt[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + ), + go.Mesh3d( + # 8 vertices of a cube + x=mesh_pred_unposed[:,0], + y=mesh_pred_unposed[:,2], + z=mesh_pred_unposed[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + ), + go.Mesh3d( + # 8 vertices of a cube + x=collision[:,0], + y=collision[:,2], + z=collision[:,1], + # i, j and k give the vertices of triangles + i = collision_faces[:,0], + j = collision_faces[:,1], + k = collision_faces[:,2], + vertexcolor=collision_color + ) + ]) + + fig.update_layout(scene_aspectmode='manual', + scene_aspectratio=dict(x=1, y=2, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,3],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + return fig + + +@tf.function +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']]) + + + +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']]) + + mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) + mesh_orig = batch['mesh_unposed'] + + orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1]) + uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1]) + uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1]) + 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() + 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))] + + 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) } + + 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) + + return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } + + +def load_model(model, checkpoint): + ckpt = tf.train.Checkpoint(model=model) + manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1) + ckpt.restore(manager.latest_checkpoint) + + +def sample(folder_main, folder_curv): + # run = wandb.init(project='vae_plots', name=folder.split('/')[-1]) + load_model(model, folder_main) + load_model(curv_model, folder_curv) + nr_of_samples = 2 + count = 0 + for batch in dataset_train: + batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) + + pred = curv_model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], batch['uv_mask'][0], nr_of_samples=nr_of_samples) + uv_curvature = tf.math.argmax(pred, axis=-1) + uv_curvature = tf.one_hot(uv_curvature, 7) + + pred = model.sample(batch['uv_static'][0], + batch['uv_body_posed'][0], + uv_curvature, + batch['uv_mask'][0], + nr_of_samples=nr_of_samples, + fix_noise=True) + 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]) + + mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) + mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy() + mesh_orig = batch['mesh_posed'][0].numpy() + mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0] + + body = batch['mesh_body_posed'][0].numpy() + body_color = np.ones_like(body, dtype=np.int32)*100 + + F = np.load("custom_utils/smpl_faces.npy") + mask = batch['mesh_mask'].numpy().astype(bool)[0] + _, faces = _maskFaces(np.ones((14475, 3)), F, mask) + + + meshes = [] + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=mesh_orig_reposed[mask,0], + y=mesh_orig_reposed[mask,2], + z=mesh_orig_reposed[mask,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + for i in range(nr_of_samples): + x = mesh_posed[i][mask, :].copy() + # x[:, 2] += 1 * (i + 1) + x[:, 2] += 0.5 + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,2], + z=x[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + # for i in range(5): + # x = body.copy() + # x[:, 2] += 1 * i + # meshes.append(go.Mesh3d( + # # 8 vertices of a cube + # x=x[:,0], + # y=x[:,2], + # z=x[:,1], + # # i, j and k give the vertices of triangles + # i = F[:,0], + # j = F[:,1], + # k = F[:,2], + # vertexcolor=body_color + # )) + + + fig = go.Figure(data=meshes) + + # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1)) + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],), + yaxis = dict(nticks=4, range=[-1,1],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + # wandb.log({'plot_' + str(count): fig}) + fig.show() + + uv_curvature_all = np.concatenate(uv_curvature.numpy().argmax(-1), axis=1) + ex.imshow(uv_curvature_all).show() + + # frames = [Image.fromarray(pred[i].numpy()/6*255) for i in range(pred.shape[0])] + # frames[0].save("temp/curv_4.gif", format="GIF", append_images=frames[1:], + # save_all=True, duration=500, loop=0, optimize=False) + + if count > 10: + break + + print(count) + count += 1 + + print() + +if __name__ == '__main__': + # init_wandb() + # allow_memory_growth() + parser = argparse.ArgumentParser(description='') + parser.add_argument('--name', type=str) + args = vars(parser.parse_args()) + + batch_size = 4 + + smpl = SMPL() + + factory = DatasetFactory() + signature_dynamic = { 'id': (), + 'pose_shape': (72 + 11,), + # 'uv_posed': (3, 512, 256,), + 'uv_body_posed': (3, 512, 256,), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'mesh_unposed': (14475, 3,), + 'uv_unposed': (3, 512, 256,), + 'uv_curvature': (512, 256), + } + signature_static = { 'uv_static': (3, 512, 256,), + # 'mesh_static': (14475, 3,), + 'mesh_mask': (14475,), + 'uv_mask': (512, 256,), + 'uv_shaped_body': (3, 512, 256,), + } + + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=batch_size, signature=signature_dynamic) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static) + 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 = CVAEMain(avg_garment_path='custom_data/jumpsuit_20_avg_uv.npy', + latent_dim=128, static_model='checkpoints/ae_static_3', + body_model='checpoints/ae_body_6', + curvature_model='checpoints/ae_curvature_6') + curv_model = CVAECurv() + + + + # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False) + optimizer = tf.optimizers.Adam(learning_rate=0.0001) + balancing_coeff = 0.1 + reg_coeff = 1 + folder = 'models/ae_logs/' + valid_iterator = iter(dataset_valid) + gamma = 1 + balancing_activated_step = 0 + avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy') + + sample('checkpoints/cvae_13', 'checkpoints/cvae_curv_2') + # ckpt = tf.train.Checkpoint(model=model) + # manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) + + # step = 0 + # for batch in dataset_train: + # batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) + diff --git a/verify_poses.py b/verify_poses.py new file mode 100644 index 0000000..865b233 --- /dev/null +++ b/verify_poses.py @@ -0,0 +1,160 @@ +import os +import sys +import time +import cv2 +import random +import numpy as np +import tensorflow as tf +import cmapy as cm +import plotly.graph_objects as go +import plotly.express as ex + +from dataset import DatasetFactory, StaticDataset +from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces +from custom_utils.smpl import SMPL + +if __name__ == '__main__': + os.environ["CUDA_VISIBLE_DEVICES"] = "6" + + uv2mesh = UV2Mesh() + smpl = SMPL() + + F = np.load("custom_utils/smpl_faces.npy") + + factory = DatasetFactory() + signature_dynamic = { 'id': (), + 'pose_shape': (72 + 11,), + # 'uv_posed': (3, 512, 256,), + # 'uv_body_posed': (3, 512, 256,), + 'mesh_body_posed': (14475, 3,), + # 'mesh_posed_body_no_orient': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'mesh_unposed': (14475, 3,), + # 'uv_unposed': (3, 512, 256,), + # 'uv_curvature': (512, 256), + # 'mesh_curvature': (14475) + } + signature_static = { #'uv_static': (3, 512, 256,), + # 'mesh_static': (14475, 3,), + 'mesh_mask': (14475,), + # 'uv_mask': (512, 256,), + # 'uv_shaped_body': (3, 512, 256,), + } + + dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_test/train', + batch_size=1, + signature=signature_dynamic) + static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_test/train', + signature=signature_static) + + + uv_res = 256 + + + + for i, batch in enumerate(dataset): + print(i) + batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) + + mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy()) + mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool) + pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy() + mesh_reposed = smpl.run_garment(pose, shape, batch['mesh_unposed']).numpy()[0] + mesh_posed = np.squeeze(batch['mesh_posed'].numpy()) + mesh_body_posed = np.squeeze(batch['mesh_body_posed'].numpy()) + mesh_posed_body_no_orient = np.squeeze(batch['mesh_posed_body_no_orient'].numpy()) + + # x = mesh_posed[mask, :] + R = smpl.run_rodrigues(pose) + # shape[:, -1] = 1 + x = smpl.run(R, shape)[0][0] + y = mesh_body_posed + z = mesh_posed[mask, :] + _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask) + go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,1], + z=x[:,2], + # i, j and k give the vertices of triangles + i = F[:,0], + j = F[:,1], + k = F[:,2] + ), + go.Mesh3d( + # 8 vertices of a cube + x=y[:,0], + y=y[:,1], + z=y[:,2], + # i, j and k give the vertices of triangles + i = F[:,0], + j = F[:,1], + k = F[:,2] + ), + # go.Mesh3d( + # # 8 vertices of a cube + # x=z[:,0], + # y=z[:,1], + # z=z[:,2], + # # i, j and k give the vertices of triangles + # i = faces[:,0], + # j = faces[:,1], + # k = faces[:,2] + # ), + ]).show() + + pose_2 = pose.copy() + pose_2[:, :3] = 0 + R = smpl.run_rodrigues(pose_2) + x = smpl.run(R, shape)[0][0] + y = mesh_posed_body_no_orient + _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask) + go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,1], + z=x[:,2], + # i, j and k give the vertices of triangles + i = F[:,0], + j = F[:,1], + k = F[:,2] + ), + go.Mesh3d( + # 8 vertices of a cube + x=y[:,0], + y=y[:,1], + z=y[:,2], + # i, j and k give the vertices of triangles + i = F[:,0], + j = F[:,1], + k = F[:,2] + )]).show() + + x = mesh_posed[mask, :] + y = mesh_reposed[mask, :] + _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask) + go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,1], + z=x[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + ), + go.Mesh3d( + # 8 vertices of a cube + x=y[:,0], + y=y[:,1], + z=y[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )]).show() + + print() \ No newline at end of file diff --git a/verify_scales.py b/verify_scales.py index 4ae144b..1389f96 100644 --- a/verify_scales.py +++ b/verify_scales.py @@ -8,6 +8,7 @@ import tensorflow as tf from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces +from custom_utils.smpl import SMPL import plotly.graph_objects as go import plotly.express as ex import cmapy as cm @@ -16,66 +17,90 @@ if __name__ == '__main__': os.environ["CUDA_VISIBLE_DEVICES"] = "6" uv2mesh = UV2Mesh() + smpl = SMPL() F = np.load("custom_utils/smpl_faces.npy") factory = DatasetFactory() signature_dynamic = { 'id': (), - # 'pose_shape': (72 + 11,), + 'pose_shape': (72 + 11,), # 'uv_posed': (3, 512, 256,), # 'uv_body_posed': (3, 512, 256,), # 'mesh_body_posed': (14475, 3,), # 'mesh_posed': (14475, 3,), 'mesh_unposed': (14475, 3,), - 'uv_unposed': (3, 512, 256,), + # 'uv_unposed': (3, 512, 256,), # 'uv_curvature': (512, 256), # 'mesh_curvature': (14475) } signature_static = { #'uv_static': (3, 512, 256,), # 'mesh_static': (14475, 3,), 'mesh_mask': (14475,), - 'uv_mask': (512, 256,), + # 'uv_mask': (512, 256,), # 'uv_shaped_body': (3, 512, 256,), } dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=1, - signature=signature_dynamic) + signature=signature_dynamic, + shuffle=False) static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static) - + meshes = [] + current_id = '' + count = 0 for i, batch in enumerate(dataset): + id = batch[0].numpy()[0].decode('utf-8') + if id == current_id: + continue + current_id = id + print(i) batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) - uv_unposed = np.squeeze(batch['uv_unposed'].numpy()) + # uv_unposed = np.squeeze(batch['uv_unposed'].numpy()) mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy()) mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool) - _, faces = _maskFaces(np.ones((14475, 3)), F, mask) + _, f = _maskFaces(np.ones((14475, 3)), F, mask) - mesh_unposed[:,1] += (1 * i) - fig = go.Figure(data=[ - go.Mesh3d( + # deshaped one + ### + pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy() + gender = int(shape[0, -1] > 0) + if gender == 0: + continue + + shape_offsets = tf.tensordot(shape[:,:10], smpl.shapedirs[gender], axes=[[1], [2]]).numpy() + mesh_unposed -= np.squeeze(shape_offsets)[smpl.to_keep, :] + ### + + mesh_unposed[:,2] += (0.5 * count) + meshes.append(go.Mesh3d( # 8 vertices of a cube x=mesh_unposed[mask,0], y=mesh_unposed[mask,2], z=mesh_unposed[mask,1], # i, j and k give the vertices of triangles - i = faces[:,0], - j = faces[:,1], - k = faces[:,2], - ) - ]) + i = f[:,0], + j = f[:,1], + k = f[:,2], + )) - fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=10, z=1)) - fig.update_layout( - scene = dict( - xaxis = dict(nticks=4, range=[-1,1],), - yaxis = dict(nticks=4, range=[-1,19],), - zaxis = dict(nticks=4, range=[-1,1],),) - ) + count += 1 + if count > 10: + break - fig.show() + fig = go.Figure(data=meshes) - print() \ No newline at end of file + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=10, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,19],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + fig.show() + + print() \ No newline at end of file