188 lines
7.1 KiB
Python
188 lines
7.1 KiB
Python
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_1
|
|
else:
|
|
config_module = importlib.import_module(args['config'])
|
|
config_1 = 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_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
|
|
|
|
|
|
def step_pyramid(batch):
|
|
model.eval()
|
|
with torch.inference_mode():
|
|
levels_input, conds = group_inputs_for_levels(batch, config_1)
|
|
|
|
out = []
|
|
|
|
out_levels = model([levels_input, conds])
|
|
for i in range(len(out_levels)):
|
|
pred, _, _ = out_levels[i]
|
|
|
|
mesh_pred = uv2mesh(pred)
|
|
|
|
if config_1['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]
|
|
if not config_1['is_cape']:
|
|
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
|
|
else:
|
|
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_1["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
|
|
pose_shape = np.load(f'{config_1["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_1['is_cape'], path='smpl', device=config_1['device'])
|
|
skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
|
|
dataset_resolutions = []
|
|
if 'pyramid' in config_1:
|
|
dataset_resolutions = config_1['pyramid']['resolutions'].copy()
|
|
if 512 not in dataset_resolutions:
|
|
dataset_resolutions.append(512)
|
|
|
|
# remove all values from config['signature] that start with mesh_
|
|
config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
|
|
|
|
dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['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_1['name'].startswith('pyr')
|
|
|
|
model: Pyramid = Pyramid(config_1)
|
|
|
|
it = load_model(model, 'checkpoints/iccvw/' + config_1['name'], iteration=args['iteration'])
|
|
model.eval()
|
|
# it = load_model(model, 'checkpoints/iccvw/' + config['name'])
|
|
|
|
name = config_1['name'] + '_' + str(it)
|
|
|
|
folder = '/data-local/127_2/results'
|
|
out_folder = f'{folder}/{name}'
|
|
|
|
os.makedirs(out_folder, exist_ok=True)
|
|
|
|
generate() |