plot cape

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
+
+
+
+

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