import os import sys from configs.iccv.config_cvae_curv import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] 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, load_model, get_summary_html, get_summary_str import plotly.graph_objects as go import plotly.express as ex from plotly.subplots import make_subplots from PIL import Image # 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 from losses import kl_loss from custom_utils.custom_utils import Accumulator def init_wandb(): run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae_curv.py') artifact.add_file('models/cvae_curv_model.py') artifact.add_file('models/cvae_model.py') artifact.add_file('dataset.py') artifact.add_file('losses.py') artifact.add_file('models/base.py') artifact.add_file('models/cond_encoder.py') artifact.add_file('dataset.py') artifact.add_file('convnext/models/convnext_tf.py') config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')] for config_file_name in config_file_names: file_name = '/'.join(config_file_name.split('.')) + '.py' if os.path.exists(file_name): artifact.add_file(file_name) run.log_artifact(artifact) uv2mesh = UV2MeshLabel(res=256) def plot_3d(curv_gt, curv_pred, faces, mesh): 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 = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ]) fig.add_trace(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' ), row=1, col=1) fig.add_trace(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_2, text='pred' ), row=1, col=2) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), ) helper = dict(xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,1],), zaxis = dict(nticks=4, range=[-1,1],),) fig.update_layout(scene = helper, scene2 = helper, ) return fig 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: pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True) ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred) dice = dice_dense(pred, batch['uv_curvature']) mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis] mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask'] ce_mesh_loss = scce(mesh_orig, mesh_pred) reconstr_loss = ce_uv_loss + ce_mesh_loss # reconstr_loss = 1 - dice reconstr_loss = tf.reduce_mean(reconstr_loss) kl = kl_loss(mean, logvar) loss = reconstr_loss + g_balancing_coeff * kl mean = tf.reduce_mean(mean) logvar = tf.reduce_mean(logvar) mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1)) 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/mean_iou': mean_iou, 'train/dice': dice, 'train/reconstr_loss': reconstr_loss, 'train/ce_mesh_loss': ce_mesh_loss, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl,'train/var': np.e ** logvar, } @tf.function def valid_step(batch): pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred) loss = ce_uv_loss mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1)) dice = dice_dense(pred, batch['uv_curvature']) reconstr_loss = ce_uv_loss # reconstr_loss = 1 - dice reconstr_loss = tf.reduce_mean(reconstr_loss) kl = kl_loss(mean, logvar) loss = reconstr_loss + g_balancing_coeff * kl mean = tf.reduce_mean(mean) logvar = tf.reduce_mean(logvar) mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis] mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask'] ce_mesh_loss = scce(mesh_orig, mesh_pred) return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/mean_iou': mean_iou, 'valid/dice': dice, 'valid/reconstr_loss': reconstr_loss, 'valid/ce_mesh_loss': ce_mesh_loss, 'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar } def plot_everything(batch: dict, type: str): pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) pred = np.argmax(pred.numpy(), axis=-1) mesh_pred = uv2mesh(pred) mesh_orig = uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1)) # batch['mesh_curvature'] mesh_base = batch['mesh_posed'] orig_img = batch['uv_curvature'] pred_img = pred img = tf.concat((tf.math.argmax(orig_img, axis=-1), pred_img), axis=2) colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)]) color_img = np.zeros((img.shape[0], img.shape[1], img.shape[2], 3), np.uint8) for i in range(7): color_img[img == i, :] = colors[i] imgs = [wandb.Image(img[i].numpy().astype(np.float32) / 6.0) for i in range(4)] color_img = [wandb.Image(color_img[i]) 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]] 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 + '/images_color': color_img, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } def sample(): iteration = load_model(model, f'checkpoints/iccv/cvae_curv_3', return_iteration=True) for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) nr_of_samples = 4 for i in range(4): # 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) pred = model.sample(tf.tile(batch['uv_static'][i][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][i][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][i][None,:,:], [nr_of_samples,1,1])) # pred,_,_ = model([tf.tile(batch['uv_static'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][0][None,:,:], [nr_of_samples,1,1])], training=False) # pred = model.call_sample([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']]) orig_img = np.argmax(batch['uv_curvature'][i].numpy(), axis=-1) pred_img = np.argmax(pred.numpy(), axis=-1) img = np.concatenate((orig_img, np.concatenate(pred_img, axis=1)), axis=1) colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)]) color_img = np.zeros((img.shape[0], img.shape[1], 3), np.uint8) for i in range(7): color_img[img == i, :] = colors[i] ex.imshow(color_img).show() print() print() def sample_2(batch): nr_of_samples = 4 # 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) # pred = model.sample(tf.tile(batch['uv_static'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][0][None,:,:], [nr_of_samples,1,1])) pred,_,_ = model([tf.tile(batch['uv_static'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][0][None,:,:,:], [nr_of_samples,1,1,1]),tf.tile(batch['uv_curvature'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][0][None,:,:], [nr_of_samples,1,1])], training=False) # pred = model.call_sample([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']]) orig_img = np.argmax(batch['uv_curvature'][0].numpy(), axis=-1) pred_img = np.argmax(pred.numpy(), axis=-1) img = np.concatenate((orig_img, np.concatenate(pred_img, axis=1)), axis=1) colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)]) color_img = np.zeros((img.shape[0], img.shape[1], 3), np.uint8) for i in range(7): color_img[img == i, :] = colors[i] ex.imshow(color_img).show() print() if __name__ == '__main__': # init_wandb() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) smpl = SMPL() skirt_modifier = '_skirt' if config['is_skirt'] else '' scce = tf.keras.losses.SparseCategoricalCrossentropy() factory = DatasetFactory() dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], use_one_hot=False, use_no_orient=True, frames_file=config['frames_file'] if 'frames_file' in config.keys() else None) static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], resolutions=[config['resolution']], in_memory=True) dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], use_one_hot=False, use_no_orient=True) static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], resolutions=[config['resolution']], in_memory=True) model = CVAE(config) sample() # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False) optimizer = tf.optimizers.Adam(learning_rate=config['lr']) #0.0001 g_balancing_coeff = config['starting_balancing_coeff'] folder = 'models/ae_logs/' valid_iterator = iter(dataset_valid) ckpt = tf.train.Checkpoint(model=model) manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + config['name'], max_to_keep=20) train_accumulator = Accumulator() valid_accumulator = Accumulator() if config['pretrained_model'] != '': load_model(model, config['pretrained_model'], iteration=config['pretrained_model_iteration']) step = 0 for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) train_loss = train_step(batch) train_accumulator.update(train_loss) if step % 25 == 0: data_valid = valid_iterator.get_next() data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic'])) valid_loss = valid_step(data_valid) train_accumulator.update(train_loss) if step % 1000 == 0 or (step < 1000 and step % 100 == 0): to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } }) if step % 5000 == 0 or (step < 5000 and step % 1000 == 0): train_figs = plot_everything(batch, 'train') valid_figs = plot_everything(data_valid, 'valid') to_log = { **to_log, **train_figs, **valid_figs } if step == 0: to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } } wandb.log(to_log, step=step) train_accumulator.reset() valid_accumulator.reset() if step % 10000 == 0: manager.save(checkpoint_number=step) if step > 40000: g_balancing_coeff = g_balancing_coeff + 1 / 20000 print(step) step += 1