From c3a2f558dead774daec2d4a11fa484406118457f Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 14 Sep 2022 16:54:18 +0000 Subject: [PATCH] create first pyramid model --- cvae_pyramid.py | 506 +++++++++++++++++++++++++++++++++++ models/cvae_model.py | 216 +++++++++++++-- models/cvae_pyramid_model.py | 186 ++----------- 3 files changed, 723 insertions(+), 185 deletions(-) create mode 100644 cvae_pyramid.py diff --git a/cvae_pyramid.py b/cvae_pyramid.py new file mode 100644 index 0000000..a52dbe5 --- /dev/null +++ b/cvae_pyramid.py @@ -0,0 +1,506 @@ +import os +os.environ["CUDA_VISIBLE_DEVICES"] = "1" +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 +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 custom_utils.custom_utils import batch_to_dict, load_model +from dataset import DatasetFactory, StaticDataset +from custom_utils.smpl import SMPL +from models.cvae_pyramid_model import Pyramid + + +name='cvae_pyr_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 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\ +training on 100 garments with offsets over the average of these\ +l2 loss on both 32 and full res, bottom predicts base, high res predicts offsets over this') + artifact = wandb.Artifact('code', type='code') + artifact.add_file('cvae_pyramid.py') + artifact.add_file('models/cvae_model.py') + artifact.add_file('models/cvae_pyramid_model.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, mesh_pred_unposed_32, faces, collision, collision_faces, collision_color): + mesh_pred_unposed[:,2] += 1 + mesh_pred_unposed_32[:,2] += 2 + collision[:,2] += 3 + fig = go.Figure(data=[ + go.Mesh3d( + x=mesh_gt[:,0], + y=mesh_gt[:,2], + z=mesh_gt[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + ), + go.Mesh3d( + x=mesh_pred_unposed[:,0], + y=mesh_pred_unposed[:,2], + z=mesh_pred_unposed[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + ), + go.Mesh3d( + x=mesh_pred_unposed_32[:,0], + y=mesh_pred_unposed_32[:,2], + z=mesh_pred_unposed_32[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + ), + go.Mesh3d( + x=collision[:,0], + y=collision[:,2], + z=collision[:,1], + 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=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],),) + ) + + 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 + + +@tf.function +def train_step(batch): + + with tf.GradientTape() as tape: + uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + uv_curvature_32 = tf.one_hot(batch['uv_curvature_32'], 7) + + pred_full, mean_full, logvar_full, pred_32, mean_32, logvar_32, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, + batch['uv_unposed_32'], batch['uv_static_32'], batch['uv_body_posed_32'], uv_curvature_32]) + + mesh_pred_full = uv2mesh(pred_full) + mesh_pred_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full) + mesh_pred_32 = uv2mesh(pred_32) + mesh_pred_posed_32 = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_32) + 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_full = mesh_pred_posed_full + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] + mesh_pred_posed_32 = mesh_pred_posed_32 + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] + + e_loss = tf.reduce_mean(edge_loss(mesh_pred_full, mesh_orig, batch['mesh_mask'])) + c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed_full, 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_full)), 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_full) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2]) + ms_ssim = tf.image.ssim_multiscale(pred_full + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3) + kl = kl_loss(mean_full, logvar_full) + + e_loss_32 = tf.reduce_mean(edge_loss(mesh_pred_32, mesh_orig, batch['mesh_mask'])) + c_loss_32 = tf.reduce_mean(collision_penalty(mesh_pred_posed_32, batch['mesh_body_posed'], batch['mesh_mask'])) + l1_loss_32 = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed_32'], [0,2,3,1]) - pred_32)), axis=-1), axis=[1, 2]) + l2_loss_32 = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed_32'], [0,2,3,1]) - pred_32) ** 2), axis=-1), axis=[1, 2]) + ms_ssim_32 = tf.image.ssim(pred_32 + 1, tf.transpose(batch['uv_unposed_32'], [0,2,3,1]) + 1, max_val=3) + kl_32 = kl_loss(mean_full, logvar_full) + + reconstr_loss = l2_loss + l2_loss_32 + reconstr_loss = tf.reduce_mean(reconstr_loss) + loss = reconstr_loss + balancing_coeff * (kl + kl_32) + + ms_ssim = tf.reduce_mean(ms_ssim) + l1_loss = tf.reduce_mean(l1_loss) + l2_loss = tf.reduce_mean(l2_loss) + mean = tf.reduce_mean(mean_full) + logvar = tf.reduce_mean(logvar_full) + + ms_ssim_32 = tf.reduce_mean(ms_ssim_32) + l1_loss_32 = tf.reduce_mean(l1_loss_32) + l2_loss_32 = tf.reduce_mean(l2_loss_32) + mean_32 = tf.reduce_mean(mean_32) + logvar_32 = tf.reduce_mean(logvar_32) + + 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_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + metric_32 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_32) ** 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_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + gradients = tape.gradient(loss, model.trainable_variables) + 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/euclid_error_reconst': metric_2, 'train/reconstr_loss': reconstr_loss, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, + 'train/var': np.e ** logvar, 'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss, + 'train/edge_loss_32': e_loss_32, 'train/collision_loss_32': c_loss_32, 'train/l1_loss_32': l1_loss_32, 'train/euclid_error_32': metric_32, 'train/ms_ssim_32': ms_ssim_32, + 'train/mean_32': mean_32, 'train/logvar_32': logvar_32, 'train/kl_loss_32': kl_32, + 'train/var_32': np.e ** logvar_32, 'train/l2_loss_32': l2_loss_32 } + + + +@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) + uv_curvature_32 = tf.one_hot(batch['uv_curvature_32'], 7) + pred_full, mean_full, logvar_full, pred_32, mean_32, logvar_32, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, + batch['uv_unposed_32'], batch['uv_static_32'], batch['uv_body_posed_32'], uv_curvature_32]) + + mesh_pred_full = uv2mesh(pred_full) + mesh_pred_32 = uv2mesh(pred_32) + 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_full + pred_32_up = tf.image.resize(pred_32, (512, 256), tf.image.ResizeMethod.LANCZOS5) + img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_32_up), 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_full.numpy(), F, mask) + garment_unposed_32, faces = maskFacesBatch(mesh_pred_32.numpy(), F, mask) + gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred_full.shape[0])] + + body = batch['mesh_body_posed'].numpy() + mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).numpy() + + shape = batch['pose_shape'][:, 72:] + gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) + shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) + + garment_posed = [mesh_posed_full[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))] + + figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], garment_unposed_32[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) } + + mean_hist = wandb.Histogram(sequence=mean_full.numpy().mean(axis=0)) + logvar_hist = wandb.Histogram(sequence=logvar_full.numpy().mean(axis=0)) + + return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } + + + + +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], nr_of_samples=nr_of_samples) + pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:] + pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1]) + + 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) + + + 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_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] += 1 + 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(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[:,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=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,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() + 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,), + 'uv_body_posed_32': (3, 32, 16,), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'mesh_unposed': (14475, 3,), + 'uv_unposed': (3, 512, 256,), + 'uv_unposed_32': (3, 32, 16,), + 'uv_curvature': (512, 256), + 'uv_curvature_32': (32, 16), + } + signature_static = { 'uv_static': (3, 512, 256,), + 'uv_static_32': (3, 32, 16,), + # '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/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=True, resolutions=[512, 32]) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=[512, 32]) + 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 = Pyramid(latent_dim=128) + + + # 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.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') + + + # sample('checkpoints/cvae_13') + + + 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 + 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: + # 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/models/cvae_model.py b/models/cvae_model.py index f4d60c7..1246e51 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -9,10 +9,12 @@ import scipy as sp 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 - +from models.ae_body_model import AutoEncoder as BodyAE +from models.ae_curv_model import AutoEncoder as CurvatureAE +from models.ae_static_model import AutoEncoder as StaticAE +from models.ae_body_model import AutoEncoder32 as BodyAE32 +from models.ae_curv_model import AutoEncoder32 as CurvatureAE32 +from models.ae_static_model import AutoEncoder32 as StaticAE32 class Bottleneck(tf.keras.Model): @@ -75,32 +77,42 @@ class Decoder(tf.keras.Model): 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'): + def __init__(self, latent_dim=128, + avg_garment_path='', + static_model='', + body_model='', + curvature_model='', + freeze_encoders=False): super(CVAE, self).__init__() self.latent_dim = latent_dim model : StaticAE = StaticAE() - load_model(model, static_model) + if static_model != '': + 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) + if body_model != '': + 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) + if curvature_model != '': + load_model(model, curvature_model) self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) - self.curv_enc.trainable = False + + if freeze_encoders: + self.static_enc.trainable = False + self.body_enc.trainable = False + self.curv_enc.trainable = False self.encoder = Encoder(latent_dim) self.decoder = Decoder() - self.avg_garment = np.load(avg_garment_path) + if avg_garment_path != '': + self.avg_garment = np.load(avg_garment_path) + else: + self.avg_garment = None self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'), tf.keras.layers.Dense(256, activation='relu'), @@ -157,7 +169,7 @@ class CVAE(tf.keras.Model): 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, uv_curv = inputs uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1]) z_static = self.static_enc(uv_template) z_body = self.body_enc(uv_body_posed) @@ -185,7 +197,179 @@ class CVAE(tf.keras.Model): # x = x * uv_mask[:,tf.newaxis,:,:] # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5) # x = x + tf.transpose(uv_shaped_body, [0,2,3,1]) + if self.avg_garment is not None: + x = x + self.avg_garment[tf.newaxis,:,:,:] + # x = x * uv_mask[:,:,:,tf.newaxis] + + return x, mean, logvar, z + + + + +class Encoder32(tf.keras.Model): + + def __init__(self, latent_dim): + super(Encoder32, self).__init__() + + self.encoder = ConvNeXt(depths=[27, 27], dims=[256, 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 Decoder32(tf.keras.Model): + + def __init__(self): + super(Decoder32, self).__init__() + + self.dense = tf.keras.layers.Dense(4 * 2 * 256) + self.reshape = tf.keras.layers.Reshape((4, 2, 256)) + self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([256, 256]))) + + + def call(self, inputs, training=None): + x = inputs + + x = self.dense(x) + x = self.reshape(x) + x = self.decoder(x) + + return x + + +class CVAE32(tf.keras.Model): + + def __init__(self, latent_dim=128, + avg_garment_path='', + static_model='', + body_model='', + curvature_model='', + freeze_encoders=False): + super(CVAE32, self).__init__() + self.latent_dim = latent_dim + + model : StaticAE32 = StaticAE32() + if static_model != '': + load_model(model, static_model) + self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + + model : BodyAE32 = BodyAE32() + if body_model != '': + load_model(model, body_model) + self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + + model : CurvatureAE32 = CurvatureAE32() + if curvature_model != '': + load_model(model, curvature_model) + self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + + if freeze_encoders: + self.static_enc.trainable = False + self.body_enc.trainable = False + self.curv_enc.trainable = False + + self.encoder = Encoder32(latent_dim) + self.decoder = Decoder32() + + if avg_garment_path != '': + self.avg_garment = np.load(avg_garment_path) + else: + self.avg_garment = None + + 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),]) + + + # @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 = 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]) + if self.avg_garment is not None: + 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_pyramid_model.py b/models/cvae_pyramid_model.py index f4d60c7..50f4beb 100644 --- a/models/cvae_pyramid_model.py +++ b/models/cvae_pyramid_model.py @@ -3,189 +3,37 @@ 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 +from models.cvae_model import CVAE, CVAE32 -class Bottleneck(tf.keras.Model): +class Pyramid(tf.keras.Model): - def __init__(self, dim): - super(Bottleneck, self).__init__() + def __init__(self, latent_dim=128, avg_garment_path=''): + super(Pyramid, 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) + self.model_32 = CVAE32(latent_dim=latent_dim) + self.model_full = CVAE(latent_dim=latent_dim) - 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)) + if avg_garment_path != '': + self.avg_garment = np.load(avg_garment_path) else: - eps = tf.random.normal(shape=(1, self.latent_dim)) - eps = tf.tile(eps, [nr_of_samples,1]) + self.avg_garment = None - 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]) + def call(self, inputs, training=None): + uv_unposed, uv_static, uv_body_posed, uv_curv, uv_unposed_32, uv_static_32, uv_body_posed_32, uv_curv_32 = 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) - z_curv = self.curv_enc(uv_curv) + pred_32, mean_32, logvar_32, _ = self.model_32([uv_unposed_32, uv_static_32, uv_body_posed_32, uv_curv_32], training=training) - z_cond = tf.concat([z_static, z_body, z_curv], axis=-1) - z_cond = self.cond_encoder(z_cond) + x = tf.image.resize(pred_32, (512, 256), tf.image.ResizeMethod.LANCZOS5) - if zero_latent: - eps = tf.zeros_like(eps) + pred_full, mean_full, logvar_full, _ = self.model_full([uv_unposed, uv_static, uv_body_posed, uv_curv], training=training) - 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 + x = x + pred_full + + return x, mean_full, logvar_full, pred_32, mean_32, logvar_32, - 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