From fb639e6c1d166dd8b5fe31609752e8de7962b039 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 29 Oct 2022 20:10:11 +0000 Subject: [PATCH] cape evaluation script --- evaluate_cape.py | 413 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 413 insertions(+) create mode 100644 evaluate_cape.py diff --git a/evaluate_cape.py b/evaluate_cape.py new file mode 100644 index 0000000..0336497 --- /dev/null +++ b/evaluate_cape.py @@ -0,0 +1,413 @@ +import os +from configs.config_pyramid 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 tqdm +import argparse +import wandb +import random +import numpy as np +import scipy as sp +from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, get_summary_str +import plotly.graph_objects as go +import plotly.express as ex +from plotly.subplots import make_subplots +from PIL import Image + +if config['is_skirt']: + from losses import edge_loss_skirt as edge_loss +else: + from losses import edge_loss + +from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing +from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio +# 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.pyramid_model_optimized import Pyramid + + +def init_wandb(): + run = wandb.init(project='cvae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp') + + artifact = wandb.Artifact('code', type='code') + artifact.add_file('cvae_pyramid.py') + artifact.add_file('models/cvae_model.py') + artifact.add_file('models/pyramid_model.py') + artifact.add_file('models/pyramid_level.py') + artifact.add_file('models/upscaler.py') + artifact.add_file('losses.py') + artifact.add_file('configs/config_base.py') + artifact.add_file('configs/config_pyramid.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') + run.log_artifact(artifact) + +uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + +def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, collision_faces, collision_color): + fig = make_subplots(rows=1, cols=1+len(mesh_pred_unposed_levels)+1+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(1+len(mesh_pred_unposed_levels)+1+1)] ]) + + 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" + ), + row=1, col=1) + + for i, mesh in enumerate(mesh_pred_unposed_levels): + fig.add_trace(go.Mesh3d( + x=mesh[:,0], + y=mesh[:,2], + z=mesh[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + text='Unposed level_' + str(i) + ), + row=1, col=i+2) + + fig.add_trace(go.Mesh3d( + x=mesh_posed[:,0], + y=mesh_posed[:,1], + z=mesh_posed[:,2], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + text='Posed GT' + ), + row=1, col=1+len(mesh_pred_unposed_levels)+1) + + fig.add_trace(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, + text='Posed output on body' + ), + row=1, col=1+len(mesh_pred_unposed_levels)+1+1) + + if len(mesh_pred_unposed_levels) == 1: + 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), + scene4_aspectmode='manual', scene4_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, + scene4 = helper, + ) + if len(mesh_pred_unposed_levels) == 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), + scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1), + scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), + scene5_aspectmode='manual', scene5_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, + scene4 = helper, + scene5 = helper, + ) + if len(mesh_pred_unposed_levels) == 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), + scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), + scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1), + scene6_aspectmode='manual', scene6_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, + scene4 = helper, + scene5 = helper, + scene6 = helper, + ) + if len(mesh_pred_unposed_levels) == 4: + 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), + scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), + scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1), + scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,), + scene7_aspectmode='manual', scene7_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, + scene4 = helper, + scene5 = helper, + scene6 = helper, + scene7 = helper, + ) + if len(mesh_pred_unposed_levels) == 5: + 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), + scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), + scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1), + scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,), + scene7_aspectmode='manual', scene7_aspectratio=dict(x=1, y=1, z=1,), + scene8_aspectmode='manual', scene8_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, + scene4 = helper, + scene5 = helper, + scene6 = helper, + scene7 = helper, + scene8 = helper, + ) + return fig + + +def group_inputs_for_levels(batch): + levels_input = [] + for res in config['pyramid']['resolutions']: + res_modifier = '_' + str(res) if res != 512 else '' + levels_input.append([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']]) + conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']] + return levels_input, conds + + + +@tf.function +def test_step(batch): + levels_input, conds = group_inputs_for_levels(batch) + out_levels = model([levels_input, conds], training=False) + + 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_orig = batch['mesh_unposed'] + + output = dict() + reconstr_loss_total = 0 + kl_total = 0 + + #for i, res in enumerate(config['pyramid']['resolutions']): + i = -1 + res = config['pyramid']['resolutions'][-1] + res_modifier = '_' + str(res) if res != 512 else '' + pred, mean, logvar = out_levels[i] + + mesh_pred = uv2mesh(pred) + 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) + + # is the output already masked? + + 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'])) + c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])) + l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2])) + # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3)) + + reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']) + + # 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']) + + output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss, + 'test/ms_ssim' + res_modifier: ms_ssim, 'test/l2_loss' + res_modifier: l2_loss, + 'test/collision_ratio' + res_modifier: c_ratio } } + + loss = reconstr_loss_total + balancing_coeff * kl_total + + return { **{ 'test/loss': loss, 'test/reconstr_loss_total': reconstr_loss_total }, **output }, out_levels + + + +def plot_everything(batch: dict, type: str): + levels_input = group_inputs_for_levels(batch) + output_levels = model(levels_input, training=False) + + mesh_pred_levels = [uv2mesh(output_levels[i][0]).numpy() for i in range(len(output_levels))] + mesh_orig = batch['mesh_unposed'].numpy() + mesh_orig_posed = batch['mesh_posed'].numpy() + body = batch['mesh_body_posed'].numpy() + + orig_img = (batch['uv_unposed'].numpy() + 1) / 2 + uv_static_img = (batch['uv_static'].numpy() + 1) / 2 + uv_body_posed_img = (batch['uv_body_posed'].numpy() + 1) / 2 + uv_curvature = (np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 6) + pred_img = np.concatenate([tf.image.resize(output_levels[i][0], (512, 256), tf.image.ResizeMethod.LANCZOS5) for i in range(len(output_levels))], axis=2) + pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) + img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2) + imgs = [wandb.Image(img[i]) for i in range(4)] + + 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])]).numpy() + + F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") + mask = batch['mesh_mask'].numpy().astype(bool) + + _, faces = maskFacesBatch(np.ones_like(mesh_orig), F, mask) + gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])] + gt_posed = [mesh_orig_posed[i, mask[i], :] for i in range(mask.shape[0])] + + mesh_unposed_shaped_levels = [[] for _ in range(mask.shape[0])] + for i in range(len(output_levels)): + mesh_unposed_shaped = [mesh_pred_levels[i][j, mask[j], :] + shape_offsets[j, mask[j]] for j in range(mask.shape[0])] + for k in range(mask.shape[0]): + mesh_unposed_shaped_levels[k].append(mesh_unposed_shaped[k]) + + # full res version + mesh_unposed_shaped_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32) + for i in range(mask.shape[0]): + mesh_unposed_shaped_full[i, mask[i], :] = mesh_unposed_shaped[i] + mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_unposed_shaped_full).numpy() + mesh_posed = [mesh_posed[j, mask[j], :] for j in range(mask.shape[0])] + + collision = [np.concatenate((body[i], mesh_posed[i]), axis=0) for i in range(len(mesh_posed))] + collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(mesh_posed))] + collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(mesh_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(mesh_posed))] + + # body_2 = [b.copy() for b in body] + # for i in range(len(gt_posed)): + # body_2[i][:, 1] = body[i][:,2].copy() + # body_2[i][:, 2] = body[i][:,1].copy() + # gt_posed_with_body = [np.concatenate((body_2[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))] + + figs = { type + "/garment_" + str(i): plot_3d(gt[i], mesh_unposed_shaped_levels[i], faces[i], gt_posed[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) } + + return { **figs, **{ type + '/images': imgs, } } + + + + +def test(): + str_split = config['name'].split('_') + test_idx = str_split.index('test') + checkpoint_name = '_'.join(str_split[:test_idx]) + load_model(model, 'checkpoints/' + checkpoint_name) + + dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False, is_cape=config['is_cape']) + static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, is_cape=config['is_cape']) + + step = 0 + losses = [] + + + for batch in tqdm.tqdm(dataset_test): + batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) + + test_loss, out_levels = test_step(batch) + all_errors = [] + + errors = [] + for i, id in enumerate(batch['id']): + id_str = id.numpy().decode("utf-8") + pred = uv2mesh(out_levels[-1][0]) + min_body = np.load(os.path.join('/home/workspace/data/dataset/cape/cape_release', 'minimal_body_shape', id_str.split('_')[0], id_str.split('_')[0] + '_minimal.npy'))[np.newaxis, :, :].astype(np.float32) + offsets_pred = pred[i] - min_body[0,smpl.to_keep,:] + offsets_orig = batch['mesh_unposed'][i] - min_body[0,smpl.to_keep,:] + error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((offsets_orig - offsets_pred) ** 2, axis=-1))) + errors.append(error) + + test_loss = { **{ 'test/error': tf.reduce_mean(errors) }, **test_loss } + losses.append(test_loss) + all_errors.extend(errors) + to_log = { **{ 'current_step': step, 'test/error': np.asarray(errors).mean() }, **test_loss } + if step % 1000 == 0: + figs = plot_everything(batch, 'test') + to_log = { **to_log, **figs } + + wandb.log(to_log, step = step) + + # print(step) + step += 1 + + # averages = dict() + path_to_plotly_html = "temp/plotly_figure.html" + + errors = np.asarray(all_errors) + print('mean: ' + str(errors.mean())) + print('min: ' + str(errors.min())) + print('max: ' + str(errors.max())) + print('median: ' + str(np.median(errors))) + np.save('temp/cape_f_pyr_64_128_256_error.npy', errors) + + histograms = dict() + for key in test_loss.keys(): + accumulated = [losses[i][key].numpy() for i in range(len(losses))] + fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean())) + fig.write_html(path_to_plotly_html, auto_play = False) + + histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html) + # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean() + + wandb.log({ **histograms }, step = step) + wandb.finish() + pass + + + +if __name__ == '__main__': + init_wandb() + allow_memory_growth() + parser = argparse.ArgumentParser(description='') + parser.add_argument('--name', type=str) + args = vars(parser.parse_args()) + + smpl = SMPL(orig_smpl=config['is_cape']) + skirt_modifier = '_skirt' if config['is_skirt'] else '' + dataset_resolutions = config['pyramid']['resolutions'].copy() + if 512 not in dataset_resolutions: + dataset_resolutions.append(512) + + factory = DatasetFactory() + + model: Pyramid = Pyramid(config) + + optimizer = tf.optimizers.Adam(learning_rate=config['lr']) + balancing_coeff = config['starting_balancing_coeff'] + balancing_activated_step = 0 + + test() +