diff --git a/scripts/plot_cape.py b/scripts/plot_cape.py new file mode 100644 index 0000000..fad3046 --- /dev/null +++ b/scripts/plot_cape.py @@ -0,0 +1,196 @@ +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', 170000), 'cape_m': ('pyr_iccvw_64_128_256_512_cape_m_e2e', 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_longlong_climb_trial2_000040', + '00134_shortlong_ski_trial2_000230', '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 + + + + diff --git a/smpl/root_joint_f.npy b/smpl/root_joint_f.npy new file mode 100644 index 0000000..0152fdf Binary files /dev/null and b/smpl/root_joint_f.npy differ diff --git a/smpl/root_joint_m.npy b/smpl/root_joint_m.npy new file mode 100644 index 0000000..d5d0b9d Binary files /dev/null and b/smpl/root_joint_m.npy differ diff --git a/smpl/smpl_torch.py b/smpl/smpl_torch.py index b18dfb7..c505d43 100644 --- a/smpl/smpl_torch.py +++ b/smpl/smpl_torch.py @@ -174,7 +174,7 @@ class SMPL(Module): # J = torch.matmul(self.J_regressor, v_shaped) if verts is None: - J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])]) + J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])]) else: J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])]) J0 = J[:,0,:].view((-1, 1, 3)) @@ -222,6 +222,8 @@ class SMPL(Module): weights = self.weights if verts is not None else self.weights_all rest_shape_h_ones_local = self.rest_shape_h_ones if verts is not None else self.rest_shape_h_ones_full + if weights.shape[1] == 6890: + rest_shape_h_ones_local = torch.ones((batch_num, 6890, 1), dtype=torch.float32).to(self.device) T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) rest_shape_h = torch.cat( (v_posed, rest_shape_h_ones_local), dim=2