import os import sys sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') import importlib import argparse parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) parser.add_argument('--config', type=str) parser.add_argument('--iteration', type=str, default='') parser.add_argument('--debug', action='store_true') parser.add_argument('-g', '--gpu', type=str, default='1') args = vars(parser.parse_args()) if args['debug']: from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_f else: config_module = importlib.import_module(args['config']) config_f = config_module.config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args['gpu'] if args['debug']: os.environ["CUDA_LAUNCH_BLOCKING"] = '1' 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 import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from tqdm import tqdm from data.dataset import Dataset from utils.checkpoints import CheckpointSaver, load_model from torch.utils.data import DataLoader from utils.utils import text_to_html from utils.postprocess import fix_collision from losses.losses import collision_loss, normal_loss, kl_loss from pytorch_msssim import MS_SSIM from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels from smpl.smpl_torch import SMPL from plotly.subplots import make_subplots from models.pyramid_model import Pyramid from models.cvae_model import CVAE from data.utils import Frame uv2mesh = UV2Mesh(res=256, is_skirt=config_f['is_skirt'], is_smpl_orig=config_f['is_cape']) def step_pyramid(batch): model.eval() with torch.inference_mode(): levels_input, conds = group_inputs_for_levels(batch, config_f) out = [] out_levels = model([levels_input, conds]) for i in range(len(out_levels)): pred, _, _ = out_levels[i] mesh_pred = uv2mesh(pred) if config_f['use_skinning']: shape = batch['pose_shape'][:, 72:] gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]) mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] mesh_pred_posed = mesh_pred_shaped else: mesh_pred_posed = mesh_pred out.append(mesh_pred_posed) return out from smpl.smpl_torch import SMPL smpl = SMPL(path='smpl', orig_smpl=True, device='cpu') clothing_verts_idx = np.load(os.path.join('custom_data', 'clothing_verts_idx.npy')) cape_mask = np.zeros((6890), bool) cape_mask[clothing_verts_idx] = True def get_error(pred, frame: Frame): gt = np.load(f'{config_f["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape'] pose_shape = np.load(f'{config_f["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32) pose, shape = pose_shape[:72], pose_shape[72:] gender = int(shape[10] >= 0.5) shape_offsets = torch.tensordot(torch.from_numpy(shape[:10][None, :]), smpl.shapedirs[gender].to('cpu'), dims=[[1], [2]])[0,smpl.to_keep,:].numpy() gt_shaped = gt + shape_offsets min_body = np.load(os.path.join('/data-local/127_2/cape_processed', 'minimal_body_shape', frame.id.split('_')[0], frame.id.split('_')[0] + '_minimal.npy')) offsets_pred = pred[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:] offsets_orig = gt_shaped[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:] error = np.linalg.norm(offsets_pred - offsets_orig, axis=1).mean() return error def generate(): errors = {0: [], 1: [], 2: [], 3: []} for batch in tqdm(test_dataloader): with torch.inference_mode(): mesh_pred = step_pyramid(batch) ids = batch['id'] frame_nrs = batch['frame_nr'] garments = batch['garment'] for i in range(len(mesh_pred)): frame = Frame(ids[i], frame_nrs[i], garments[i]) for j in range(4): error = get_error(mesh_pred[j][i].cpu().numpy(), frame) errors[j].append(error) # print(errors) for j in range(4): print(f'Level {j}: {np.mean(errors[j])}') # import plotly.graph_objects as go # mesh_pred = mesh_pred.cpu().numpy() # mask = batch['mesh_mask'].cpu().numpy() # i = 3 # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,mask[i],0], y=mesh_pred[i,mask[i],2], z=mesh_pred[i,mask[i],1], mode='markers', marker=dict(size=1))]) # fig.update_layout(scene_aspectmode='cube') # fig.show() # import plotly.graph_objects as go # mesh_pred = mesh_pred.cpu().numpy() # i = 3 # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,:,0], y=mesh_pred[i,:,2], z=mesh_pred[i,:,1], mode='markers', marker=dict(size=1))]) # fig.update_layout(scene_aspectmode='cube') # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) # fig.show() if __name__ == '__main__': smpl = SMPL(orig_smpl=config_f['is_cape'], path='smpl', device=config_f['device']) skirt_modifier = '_skirt' if config_f['is_skirt'] else '' dataset_resolutions = [] if 'pyramid' in config_f: dataset_resolutions = config_f['pyramid']['resolutions'].copy() if 512 not in dataset_resolutions: dataset_resolutions.append(512) # remove all values from config['signature] that start with mesh_ config_f['signature'] = set([x for x in config_f['signature'] if not x.startswith('mesh_')]) dataset_test = Dataset(config_f['test_dir'], config_f['signature'], device=config_f['device'], frames_file=config_f['test_frames'], use_cache=config_f['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_f['is_cape']) # dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape']) test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, ) is_pyramid = config_f['name'].startswith('pyr') model: Pyramid = Pyramid(config_f) it = load_model(model, 'checkpoints/iccvw/' + config_f['name'], iteration=args['iteration']) model.eval() # it = load_model(model, 'checkpoints/iccvw/' + config['name']) name = config_f['name'] + '_' + str(it) folder = '/data-local/127_2/results' out_folder = f'{folder}/{name}' os.makedirs(out_folder, exist_ok=True) generate()