diff --git a/aaa_helpers.txt b/aaa_helpers.txt index f5d2261..32abb84 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -1,4 +1,6 @@ nohup python cvae.py --name="lhunor_weights_test_6" > nohup/iccv/weights_test_6.out 2>&1 & +nohup python evaluate_ae.py --name iccv_weights_test_6 > nohup/iccv/evaluate_iccv_weights_test_6.out 2>&1 & +nohup python calculate_error.py --name iccv_weights_test_6 > nohup/iccv/error_iccv_weights_test_6.out 2>&1 & one final convolutional layer for smoothing, on the final coordinates diff --git a/compare_models_saved.py b/compare_models_saved.py new file mode 100644 index 0000000..b95268f --- /dev/null +++ b/compare_models_saved.py @@ -0,0 +1,188 @@ +import os +import sys +from configs.config_abl_all_post import config as config_a + +os.environ["CUDA_VISIBLE_DEVICES"] = '5' +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 glob +import time +import argparse +import wandb +import tqdm +import random +import numpy as np +import plotly.graph_objects as go +import plotly.express as ex +from PIL import Image +from postprocessing import fix_collision + +if config_a['is_skirt']: + from losses import edge_loss_skirt as edge_loss +else: + from losses import edge_loss + + + +from losses import collision_ratio, kl_loss #, concept_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, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision +from dataset import DatasetFactory, StaticDataset +from custom_utils.smpl import SMPL +from models.cvae_model import CVAE +from custom_utils.supersmpl_io import writeOBJ +from plotly.subplots import make_subplots + + +# uv2mesh = UV2Mesh(res=256, is_skirt=config_a['is_skirt'], is_adaptive=False) + + + +def plot(mesh_gt, pred_a, pred_b, body, mask): + fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ]) + + _, faces_garment = _maskFaces(np.zeros((14475, 3), np.float32), F, mask) + faces = np.concatenate((F, faces_garment + 14475), axis=0) + mesh_gt = np.concatenate((body, mesh_gt[mask, :]), axis=0) + pred_a = np.concatenate((body, pred_a), axis=0) + pred_b = np.concatenate((body, pred_b), axis=0) + colors = np.concatenate((np.ones_like(body, dtype=np.int32)*100, np.ones_like(mesh_gt, dtype=np.int32)*230), axis=0) + # colors = np.ones_like(mesh_gt, dtype=np.int32)*230 + + fig.add_trace(go.Mesh3d( + x=mesh_gt[:,0], + y=mesh_gt[:,2], + z=mesh_gt[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + text="Ground Truth", + vertexcolor=colors, + ), + row=1, col=1) + + fig.add_trace(go.Mesh3d( + x=pred_a[:,0], + y=pred_a[:,2], + z=pred_a[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + text="Pred A", + vertexcolor=colors, + ), + row=1, col=2) + + fig.add_trace(go.Mesh3d( + x=pred_b[:,0], + y=pred_b[:,2], + z=pred_b[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + text="Pred B", + vertexcolor=colors, + ), + row=1, col=3) + + + + 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), + scene3_aspectmode='manual', scene3_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, + scene3 = helper, + ) + + return fig + + + +if __name__ == '__main__': + run = wandb.init(project='plots', name=config_a['name'], notes=config_a['notes'], dir='/wandb') + + parser = argparse.ArgumentParser(description='') + parser.add_argument('--name', type=str) + args = vars(parser.parse_args()) + + skirt_modifier = '_skirt' if config_a['is_skirt'] else '' + + folder = config_a['test_dir'] + # smpl = SMPL(use_one_hot=False) + + F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32) + + iteration = '' + + # load_model(model, 'checkpoints/' + 'abl_all_120k') + + + factory = DatasetFactory() + dataset_a = factory.get_dataset(folder=config_a['test_dir'], + batch_size=config_a['batch_size'], + signature=config_a['signature_dynamic'], + resolutions=[config_a['resolution']], + in_memory=False, + do_repeat=False, + use_one_hot=config_a['use_one_hot'], + use_no_orient=config_a['remove_global_orientation'], + shuffle=False, + testing=True, + # frames_file='frames_abl.txt', + ) + static_dataset_a = StaticDataset(folder=config_a['test_dir'], + signature=config_a['signature_static'], + resolutions=[config_a['resolution']], + in_memory=True) + + iterator_a = iter(dataset_a) + + + faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config_a['is_skirt'] else '') + ids = os.listdir(faces_folder) + ids = [id.split('.')[0] for id in ids] + faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids } + + results_folder = '/home/workspace/data/results/post' + + count = 0 + while True: + batch = iterator_a.get_next() + + batch = static_dataset_a.expand_batch(batch_to_dict(batch, config_a['signature_dynamic'])) + + mask = batch['mesh_mask'].numpy().astype(bool) + gt = batch['mesh_posed'].numpy() + body = batch['mesh_body_posed'].numpy() + + for i in range(mask.shape[0]): + id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8") + pred_a = np.load(results_folder + '/al_all_post-/' + id + '_' + frame_nr + '.npy') + pred_b = np.load(results_folder + '/iccv_weights_test_6-/' + id + '_' + frame_nr + '.npy') + # pred_c = np.load(results_folder + '/iccv_weights_test_6-/' + id + '_' + frame_nr + '.npy') + + fig = plot(gt[i], pred_a, pred_b, body[i], mask[i]) + + wandb.log({f'{id}_{frame_nr}': fig}, step = count) + count += 1 + + print(count) + diff --git a/configs/config_abl_all_post.py b/configs/config_abl_all_post.py index cc5d1a8..2f3373a 100644 --- a/configs/config_abl_all_post.py +++ b/configs/config_abl_all_post.py @@ -21,7 +21,7 @@ assert ablation_count == 1 or ablation_count == 0 config['gpu'] = '5' config['name'] = 'iccv_weights_test_8' -config['notes'] = 'smaller weight loss coeff, 1e-6, convolutional weight learning, continued from _5' +config['notes'] = 'smaller weight loss coeff, 1e-6, convolutional weight learning, changed convnext reverse' config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv if config['use_skinning']: @@ -53,7 +53,7 @@ if config['is_cape']: config['smpl_res'] = 3638 -config['pretrained_model'] = ''#'checkpoints/post/iccv_weights_test_5' +config['pretrained_model'] = '' ############### config['lr'] = 0.0001# 0.0001 @@ -91,9 +91,9 @@ config['cvae']['noise_amount'] = 0.00001 config['cvae']['kernel_size'] = 7 config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3] config['cvae']['encoder_dims'] = [32, 64, 128, 256, 64] -config['cvae']['decoder_depths'] = config['cvae']['encoder_depths'] -config['cvae']['decoder_dims'] = config['cvae']['encoder_dims'] -config['cvae']['decoder_weights_depths'] = [3, 3, 3, 5, 3] +config['cvae']['decoder_depths'] = [3, 9, 3, 3, 3, 3] +config['cvae']['decoder_dims'] = [3, 32, 128, 128, 256, 512] +config['cvae']['decoder_weights_depths'] = [3, 9, 3, 3, 3] config['cvae']['decoder_weights_dims'] = [24, 32, 64, 128, 256] config['cond_ae_static'] = dict() diff --git a/convnext/models/convnext_tf.py b/convnext/models/convnext_tf.py index e0fc3a4..f9380bf 100644 --- a/convnext/models/convnext_tf.py +++ b/convnext/models/convnext_tf.py @@ -225,26 +225,28 @@ class ConvNeXtReverse(tf.keras.Model): drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., learn_weights=False, ): super().__init__() - self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers + self.upsample_layers = [] + stem = tf.keras.Sequential([ - Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'), - # UpSampling2D((4, 4), interpolation='nearest'), - # Conv2D(dims[0], kernel_size=4, strides=1, padding='same'), - # Dense(dims[0]), - LayerNormalization(epsilon=1e-6)] - ) + LayerNormalization(epsilon=1e-6), + Conv2DTranspose(dims[0], kernel_size=2, strides=2, padding='same'), + ]) self.upsample_layers.append(stem) - for i in range(len(depths) - 1): + + for i in range(len(depths) - 2): downsample_layer = tf.keras.Sequential([ LayerNormalization(epsilon=1e-6), - Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same') - # UpSampling2D((2,2), interpolation='nearest'), - # Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'), - # Dense(dims[i+1]), + Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same'), ]) self.upsample_layers.append(downsample_layer) - self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks + downsample_layer = tf.keras.Sequential([ + Conv2DTranspose(dims[i+2], kernel_size=2, strides=2, padding='same'), + LayerNormalization(epsilon=1e-6), + ]) + self.upsample_layers.append(downsample_layer) + + self.stages = [] dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))] cur = 0 for i in range(len(depths)): @@ -257,6 +259,7 @@ class ConvNeXtReverse(tf.keras.Model): if not learn_weights: self.last = Dense(3, activation='tanh') + # self.last = Activation('tanh') else: self.last = Activation('softmax') # self.last = Dense(24, activation='softmax') diff --git a/evaluate_ae.py b/evaluate_ae.py index 28bb889..c160969 100644 --- a/evaluate_ae.py +++ b/evaluate_ae.py @@ -39,7 +39,7 @@ else: from losses import collision_ratio, kl_loss #, concept_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, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision +from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, UV2MeshLabel, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision from dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL from models.cvae_model import CVAENew as CVAE @@ -48,6 +48,7 @@ from custom_utils.supersmpl_io import writeOBJ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape']) +uv2mesh_label = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) @tf.function def step(batch): @@ -56,9 +57,10 @@ def step(batch): else: uv_orig = batch['uv_posed'] - pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) + pred, _, _, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) mesh_pred = uv2mesh(pred) + mesh_weigths = uv2mesh_label(weights) if config['use_skinning']: shape = batch['pose_shape'][:, 72:] @@ -68,7 +70,13 @@ def step(batch): mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped) else: - mesh_pred_posed = mesh_pred + 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_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] + mesh_pred_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped, mesh_weigths) + # mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])