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 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 tf_utils import check_tf_version, allow_memory_growth # 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_model import CVAE name='cvae_curv_3' 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.. r lower kl coeff to 0.1. fixed plotting colors\ addent MLP for lowering conditional Z diemnsion\ higher starting kl coeff (0.1)\ train on 100 garments') artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae_curv.py') artifact.add_file('models/cvae_curv_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 = 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']]) 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']]) 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']]) 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/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512]) static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512]) 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_5', body_model='checpoints/ae_body_8',) # 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 # sample('checkpoints/cvae_curv_2') ckpt = tf.train.Checkpoint(model=model) manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=20) 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