luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

plot cape

Browse Source
This commit is contained in:
luni
2023-07-13 17:11:57 +02:00
parent 2f77e2af18
commit 955e87a0c6
4 changed files with 199 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/plot_cape.py
+196
View File
@@ -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
smpl/root_joint_f.npy
BIN
View File
Binary file not shown.
smpl/root_joint_m.npy
BIN
View File
Binary file not shown.
smpl/smpl_torch.py
+3 -1
View File
@@ -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