diff --git a/convnext/models/convnext_tf.py b/convnext/models/convnext_tf.py index 0387db0..6a36bb1 100644 --- a/convnext/models/convnext_tf.py +++ b/convnext/models/convnext_tf.py @@ -45,10 +45,10 @@ class Block(Layer): layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. """ - def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6, prefix=''): + def __init__(self, dim, kernel_size=7, drop_path=0., layer_scale_init_value=1e-6, prefix=''): super().__init__() self.dwconv = DepthwiseConv2D( - kernel_size=7, padding='same') # depthwise conv + kernel_size=kernel_size, padding='same') # depthwise conv self.norm = LayerNormalization(epsilon=1e-6) # pointwise/1x1 convs, implemented with linear layers self.pwconv1 = Dense(4 * dim) @@ -97,7 +97,7 @@ class ConvNeXt(tf.keras.Model): head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. """ - def __init__(self, num_classes=1000, + def __init__(self, num_classes=1000, kernel_size=7, depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True, drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., ): @@ -121,7 +121,7 @@ class ConvNeXt(tf.keras.Model): cur = 0 for i in range(len(depths)): stage = tf.keras.Sequential( - [Block(dim=dims[i], drop_path=dp_rates[cur + j], + [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size, layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] ) self.stages.append(stage) @@ -167,7 +167,7 @@ class ConvNeXtWithSkip(tf.keras.Model): """ def __init__(self, - depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], + depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7, drop_path_rate=0., layer_scale_init_value=1e-6, ): super().__init__() @@ -189,7 +189,7 @@ class ConvNeXtWithSkip(tf.keras.Model): cur = 0 for i in range(len(depths)): stage = tf.keras.Sequential( - [Block(dim=dims[i], drop_path=dp_rates[cur + j], + [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size, layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] ) self.stages.append(stage) @@ -221,7 +221,7 @@ class ConvNeXtReverse(tf.keras.Model): """ def __init__(self, - depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], + depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7, drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., ): super().__init__() @@ -249,7 +249,7 @@ class ConvNeXtReverse(tf.keras.Model): cur = 0 for i in range(len(depths)): stage = tf.keras.Sequential( - [Block(dim=dims[i], drop_path=dp_rates[cur + j], + [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size, layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] ) self.stages.append(stage) @@ -281,7 +281,7 @@ class ConvNeXtReverseWithSkip(tf.keras.Model): """ def __init__(self, - depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], + depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7, drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., ): super().__init__() @@ -304,7 +304,7 @@ class ConvNeXtReverseWithSkip(tf.keras.Model): cur = 0 for i in range(len(depths)): stage = tf.keras.Sequential( - [Block(dim=dims[i], drop_path=dp_rates[cur + j], + [Block(dim=dims[i], kernel_size=kernel_size, drop_path=dp_rates[cur + j], layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] ) self.stages.append(stage) diff --git a/cvae_low.py b/cvae_low.py new file mode 100644 index 0000000..b27067e --- /dev/null +++ b/cvae_low.py @@ -0,0 +1,498 @@ +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 +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 custom_utils.custom_utils import batch_to_dict, load_model +from dataset import DatasetFactory, StaticDataset +from custom_utils.smpl import SMPL +from models.cvae_model import CVAE, CVAE128, CVAE32 + + +name='cvae_low_1_d128' +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 , absolute, not offsets\ +using filled versions') + artifact = wandb.Artifact('code', type='code') + artifact.add_file('cvae.py') + artifact.add_file('models/cvae_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, 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 + + + + +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_low = tf.one_hot(batch['uv_curvature' + res_modifier], 7) + pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low, batch['uv_static'], batch['uv_body_posed'], uv_curvature]) + # pred = pred * batch['uv_mask' + res_modifier][:,:,:,tf.newaxis] + + 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'] + + 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' + res_modifier], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask' + res_modifier], axis=[1, 2]) + l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred)), axis=-1), axis=[1, 2]) + # l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred) ** 2), axis=-1) * batch['uv_mask' + res_modifier], axis=[1, 2]) + l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred) ** 2), axis=-1), 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' + res_modifier], [0,2,3,1]) + 1, max_val=3, power_factors=(0.6, 0.4)) #(0.0448, 0.2856, 0.3001, 0.2363, 0.1333) + ms_ssim = tf.image.ssim(pred + 1, tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) + 1, max_val=3) #(0.0448, 0.2856, 0.3001, 0.2363, 0.1333) + # 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' + res_modifier], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed' + res_modifier], [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 = l2_loss + reconstr_loss = tf.reduce_mean(reconstr_loss) + # logpx_z = -reconstr_loss + # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.) + # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar) + # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x)) + kl = kl_loss(mean, logvar) + loss = reconstr_loss + balancing_coeff * kl + # loss = reconstr_loss + balancing_coeff * kl + 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) + logvar = tf.reduce_mean(logvar) + + mesh_orig_2 = uv2mesh(batch['uv_unposed' + res_modifier]) + 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, 'train/l2_loss': l2_loss } + + + +@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_low = tf.one_hot(batch['uv_curvature' + res_modifier], 7) + pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low, batch['uv_static'], batch['uv_body_posed'], uv_curvature]) + # pred = pred * batch['uv_mask' + res_modifier][:,:,:,tf.newaxis] + + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_unposed'] + + orig_img = tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) + uv_static_img = tf.transpose(batch['uv_static' + res_modifier], [0,2,3,1]) + uv_body_posed_img = tf.transpose(batch['uv_body_posed' + res_modifier], [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() + + 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))] + + 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 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' + res_modifier][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' + res_modifier][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() + + resolution = 128 + res_modifier = '' + if resolution != 512: + res_modifier = '_' + str(resolution) + + factory = DatasetFactory() + signature_dynamic = { 'id': (), + 'pose_shape': (72 + 11,), + # 'uv_posed': (3, 512, 256,), + 'uv_body_posed': (3, 512, 256,), + 'uv_body_posed' + res_modifier: (3, resolution, resolution//2,), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'mesh_unposed': (14475, 3,), + 'uv_unposed': (3, 512, 256,), + 'uv_unposed' + res_modifier: (3, resolution, resolution//2,), + 'uv_curvature': (512, 256,), + 'uv_curvature' + res_modifier: (resolution, resolution//2), + } + signature_static = { 'uv_static': (3, 512, 256,), + 'uv_static' + res_modifier: (3, resolution, resolution//2,), + # 'mesh_static': (14475, 3,), + 'mesh_mask': (14475,), + 'uv_mask' + res_modifier: (resolution, resolution//2,), + # '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, resolution]) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=[512, resolution]) + 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 = CVAE128(#avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy', + latent_dim=128, static_model='checkpoints/ae_static_4', + body_model='checpoints/ae_body_8', + curvature_model='checpoints/ae_curvature_7', + freeze_encoders=True, + use_full_res_encoders=True) + + # 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 + + # 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_pyramid.py b/cvae_pyramid.py index fc9dada..0833510 100644 --- a/cvae_pyramid.py +++ b/cvae_pyramid.py @@ -1,5 +1,5 @@ import os -os.environ["CUDA_VISIBLE_DEVICES"] = "0" +os.environ["CUDA_VISIBLE_DEVICES"] = "2" import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') @@ -37,7 +37,7 @@ from custom_utils.smpl import SMPL from models.cvae_pyramid_model import Pyramid -name='cvae_pyr_2_d128' +name='cvae_pyr_3_d32' 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. \ @@ -49,11 +49,13 @@ 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 128 and full res, bottom predicts base, high res predicts offsets over this\ -low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input') +low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input\ +not using filled in versions, but using custom MLP upscaling') 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('models/upscaler.py') artifact.add_file('dataset.py') artifact.add_file('losses.py') artifact.add_file('convnext/models/convnext_tf.py') @@ -135,8 +137,8 @@ def train_step(batch): uv_curvature = tf.one_hot(batch['uv_curvature'], 7) uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7) - pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, - batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low]) + pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], + batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier],]) mesh_pred_full = uv2mesh(pred_full) mesh_pred_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full) @@ -152,11 +154,12 @@ def train_step(batch): 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]) + l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred_full)), axis=-1), 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), 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) + pred_low = pred_low * batch['uv_mask' + resolution_modifier][:,:,:,tf.newaxis] e_loss_low = tf.reduce_mean(edge_loss(mesh_pred_low, mesh_orig, batch['mesh_mask'])) c_loss_low = tf.reduce_mean(collision_penalty(mesh_pred_posed_low, batch['mesh_body_posed'], batch['mesh_mask'])) l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) - pred_low)), axis=-1), axis=[1, 2]) @@ -236,8 +239,8 @@ def valid_step(batch): def plot_everything(batch: dict, type: str): uv_curvature = tf.one_hot(batch['uv_curvature'], 7) uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7) - pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, - batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low]) + pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], + batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier]]) mesh_pred_full = uv2mesh(pred_full) mesh_pred_low = uv2mesh(pred_low) @@ -409,7 +412,7 @@ if __name__ == '__main__': args = vars(parser.parse_args()) batch_size = 4 - low_resolution = 128 + low_resolution = 32 resolution_modifier = '_' + str(low_resolution) smpl = SMPL() @@ -433,11 +436,12 @@ if __name__ == '__main__': # 'mesh_static': (14475, 3,), 'mesh_mask': (14475,), 'uv_mask': (512, 256,), + 'uv_mask' + resolution_modifier: (low_resolution, low_resolution//2,), # '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, low_resolution]) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=[512, low_resolution]) + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512, low_resolution]) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_resolution]) 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) @@ -447,7 +451,8 @@ if __name__ == '__main__': freeze_encoders=True, static_model='checkpoints/ae_static_4', body_model='checpoints/ae_body_7', - curvature_model='checpoints/ae_curvature_7') + curvature_model='checpoints/ae_curvature_7', + upscaler_model='checkpoints/up_2_3') # 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) diff --git a/models/cvae_model.py b/models/cvae_model.py index 51458e3..efb9a3e 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -161,7 +161,7 @@ class CVAE(tf.keras.Model): def reparameterize(self, mean, logvar, fix_noise=False): if not fix_noise: - eps = tf.random.normal(shape=mean.shape) + eps = tf.random.normal(shape=mean.shape) else: eps = tf.random.normal(shape=(1, mean.shape[1])) eps = tf.tile(eps, [mean.shape[0],1]) @@ -358,11 +358,11 @@ class CVAE32(tf.keras.Model): 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, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = 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) + z_static = self.static_enc(uv_template_high) + z_body = self.body_enc(uv_body_posed_high) + z_curv = self.curv_enc(uv_curv_high) x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1) diff --git a/models/cvae_pyramid_model.py b/models/cvae_pyramid_model.py index ff12319..922b9fc 100644 --- a/models/cvae_pyramid_model.py +++ b/models/cvae_pyramid_model.py @@ -8,6 +8,7 @@ import numpy as np from custom_utils.custom_utils import load_model from models.cvae_model import CVAE, CVAE128, CVAE32 +from models.upscaler import UpscalerBasic as Upscaler class Pyramid(tf.keras.Model): @@ -17,7 +18,8 @@ class Pyramid(tf.keras.Model): avg_garment_path='', static_model='', body_model='', - curvature_model='', + curvature_model='', + upscaler_model='', freeze_encoders=False, use_full_res_encoders=True): super(Pyramid, self).__init__() @@ -42,6 +44,18 @@ class Pyramid(tf.keras.Model): body_model=body_model, curvature_model=curvature_model, freeze_encoders=freeze_encoders) + + self.upscale_enhancer = None + if upscaler_model != '': + self.upscale_enhancer = Upscaler(512, low_res) + load_model(self.upscale_enhancer, upscaler_model) + self.upscale_enhancer = tf.keras.Sequential([ + self.upscale_enhancer.flatten, + self.upscale_enhancer.seq, + self.upscale_enhancer.reshape + ]) + + self.upscale_enhancer.trainable = False self.use_full_res_encoders = use_full_res_encoders @@ -51,18 +65,27 @@ class Pyramid(tf.keras.Model): self.avg_garment = None def call(self, inputs, training=None): - uv_unposed, uv_static, uv_body_posed, uv_curv, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low = inputs + uv_unposed, uv_static, uv_body_posed, uv_curv, uv_mask, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_mask_low = inputs if self.use_full_res_encoders: pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_static, uv_body_posed, uv_curv], training=training) else: pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training) - x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5) + pred_low = pred_low * uv_mask_low[:,:,:,tf.newaxis] + + if self.upscale_enhancer is not None: + x = self.upscale_enhancer(pred_low, training=training) + x = x * (1 - uv_mask_low[:,:,:,tf.newaxis]) + pred_low = x + pred_low + + x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5) + x = x * uv_mask[:,:,:,tf.newaxis] pred_full, mean_full, logvar_full, _ = self.model_full([uv_unposed, uv_static, uv_body_posed, uv_curv], training=training) x = x + pred_full + x = x * uv_mask[:,:,:,tf.newaxis] return x, mean_full, logvar_full, pred_low, mean_low, logvar_low,