import os import sys sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = '3' import time import argparse import wandb from tqdm 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 torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode from pathlib import Path import importlib from scipy.spatial.transform import Rotation as R from data.utils import Frame from utils.utils import save_stl from utils.checkpoints import CheckpointSaver, load_model from utils.postprocess import fix_collision from losses.losses import collision_loss, normal_loss, kl_loss 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 data.dataset import Dataset from utils.checkpoints import CheckpointSaver, load_model from torch.utils.data import DataLoader import trimesh as tm from zipfile import ZipFile from custom_utils.custom_utils import _maskFaces from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_f from configs.iccvw.pyr_64_128_256_512_cape_m_e2e_2 import config as config_m uv2mesh = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config_f['is_cape']) def step_pyramid(model, batch, config): with torch.inference_mode(): levels_input, conds = group_inputs_for_levels(batch, config) out_levels = model([levels_input, conds]) pred, _, _ = out_levels[-1] mesh_pred = uv2mesh(pred) if config['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['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 return mesh_pred_posed def read_frames_file(file): with open(file, 'r') as f: lines = f.readlines() frames = ['_'.join(line.strip().split('.')[:-1]) for line in lines] return frames def create_temp_frames(folder: str, frame: Frame): with open(Path(folder) / 'frames_temp.txt', 'w') as f: f.writelines([f'{frame.id}.{frame.frame_nr}\n' for _ in range(4)]) def add_global_orientation(mesh, pose, inverse=False): rodr = smpl.rodrigues(torch.tensor(pose.T, device=smpl.device).reshape(-1, 1, 3)).reshape(4, -1, 3, 3).cpu().numpy() r = R.from_matrix(rodr[0, 0, :, :]) out = r.apply(mesh, inverse=inverse) return out def read_obj(path): with open(path, 'r') as f: obj = tm.exchange.obj.load_obj(f, skip_materials=True, maintain_order=True) verts = obj['vertices'] return verts if __name__ == '__main__': results_folder = Path('/data-local/127_2/results') out_folder = Path('temp/cape') folder = dict() folder['cape_f'] = '/data-net/code/StyleGAN2_tf2_conditioned/CAPE/resultsOld/CAPE-affineconv_nz18_pose24_clotype8_female/test_reconstruction_objs_dataset_female_4clotypes' folder['cape_m'] = '/data-net/code/StyleGAN2_tf2_conditioned/CAPE/resultsOld/CAPE-affineconv_nz18_pose24_clotype8_male/test_reconstruction_objs_dataset_male_4clotypes' frames_file = dict() frames_file['cape_f'] = '/data-net/dataset/cape_processed/datasets/dataset_female_4clotypes_frames.txt' frames_file['cape_m'] = '/data-net/dataset/cape_processed/datasets/dataset_male_4clotypes_frames.txt' frames = dict() frames_to_i = dict() frames['cape_f'] = read_frames_file(frames_file['cape_f']) frames_to_i['cape_f'] = { frame: i for i, frame in enumerate(frames['cape_f']) } frames['cape_m'] = read_frames_file(frames_file['cape_m']) frames_to_i['cape_m'] = { frame: i for i, frame in enumerate(frames['cape_m']) } random.shuffle(frames['cape_f']) random.shuffle(frames['cape_m']) smpl = SMPL(orig_smpl=True, path='smpl', device=config_f['device']) J0 = dict() J0['cape_f'] = np.load('smpl/root_joint_f.npy') #smpl(torch.zeros((4, 72), device=config_f['device']), torch.zeros((4, 11), device=config_f['device']))[1][0].cpu().numpy() J0['cape_m'] = np.load('smpl/root_joint_m.npy') 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 cape_mask = cape_mask[smpl.to_keep] _, faces = _maskFaces(np.zeros((3638, 3)), np.load('smpl/smpl_faces_orig.npy'), cape_mask) faces = faces.astype(np.int32) cape_template_body = read_obj('/data-net/code/StyleGAN2_tf2_conditioned/CAPE/data/template_mesh.obj') dataset_resolutions = [] dataset_resolutions = [64, 128, 256, 512] # cape_exp = { 'cape_f': ('pyr_64_128_256_512_cape', 100000), 'cape_m': ('pyr_64_128_256_512_skirt', 200000) } ours_exp = { 'cape_f': ('pyr_iccvw_64_128_256_512_cape_f_e2e_fine_3_lvl2', 170000), 'cape_m': ('pyr_iccvw_64_128_256_512_cape_m_e2e_lvl2', 230000) } picker = 'cape_m' config = dict() config['cape_f'] = config_f config['cape_m'] = config_m frames_override = ['00122_shortshort_ballerina_spin_000153', '00127_shortlong_ballerina_spin_000141', '00127_shortlong_soccer_000306', '00134_shortlong_ski_trial2_000230',] #'00134_longlong_climb_trial2_000040', '00134_shortlong_soccer_trial2_000059' _frames = frames_override if frames_override is not None else frames[picker] for _frame in tqdm(_frames): if _frame in frames['cape_f']: picker = 'cape_f' else: picker = 'cape_m' frame = Frame('_'.join(_frame.split('_')[:-1]), _frame.split('_')[-1], 'cape') create_temp_frames(config[picker]['test_dir'], frame) dataset = Dataset(config[picker]['test_dir'], config[picker]['signature'], device=config[picker]['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=True) dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, ) batch = next(iter(dataloader)) gt = batch['mesh_posed'][0].cpu().numpy()[cape_mask] pose = batch['pose_shape'][0, :72].cpu().numpy() pose[:3] = batch['global_orient'][0].cpu().numpy() save_stl(out_folder / f'{frame}_gt.stl', gt, faces) pred = np.load(results_folder / f'{ours_exp[picker][0]}_{ours_exp[picker][1]}' / f'{frame}.npz')['arr_0'] pred = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.from_numpy(pred[None,:,:]).repeat(4,1,1).to('cuda'))[0].cpu().numpy() pred = pred[0, cape_mask, :] cape = read_obj(os.path.join(folder[picker], '{:0>4d}.obj'.format(frames_to_i[picker][f'{frame.id}_{frame.frame_nr}'])))[clothing_verts_idx, :] - J0[picker][np.newaxis, :] cape = add_global_orientation(cape, pose[np.newaxis, :], inverse=True) 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')) cape = cape - cape_template_body[clothing_verts_idx, :] + min_body[clothing_verts_idx, :] save_stl(out_folder / f'{frame}_ours.stl', pred, faces) save_stl(out_folder / f'{frame}_cape.stl', cape, faces) # pred[:,0] = pred[:,0] + 1.5 # cape[:,0] = cape[:,0] + 0.5 # gt[:,0] = gt[:,0] - 0.5 # fig = go.Figure(data=[ # go.Mesh3d(x=gt[:,0], y=gt[:,2], z=gt[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightgreen', opacity=1), # go.Mesh3d(x=pred[:,0], y=pred[:,2], z=pred[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), # go.Mesh3d(x=cape[:,0], y=cape[:,2], z=cape[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=1) # ]) # fig.update_layout(scene=dict(aspectmode='cube')) # fig.update_layout(scene=dict(xaxis=dict(range=[-2, 2], ), yaxis=dict(range=[-2, 2], ), zaxis=dict(range=[-2, 2], ), ), ) # fig.update_layout(title_text=f'{frame} ') # fig.show() pass pass