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"] = '0' import time import argparse import random import tqdm import numpy as np import plotly.graph_objects as go import plotly.express as ex from utils.postprocess import fix_collision import torch from tqdm import tqdm from data.utils import Frame from data.cloth3d_io import readPC2Frame, readPC2, readOBJ from utils.utils import pickle_load, save_stl from scipy.spatial import cKDTree from smpl.smpl_torch import SMPL from scipy.spatial.transform import Rotation as R import trimesh as tm from plotly.subplots import make_subplots smpl = SMPL(path='smpl', device='cpu') def deorient(verts, pose, inverse=True): assert pose.shape[1] == 72 rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy() rotated = np.zeros_like(verts) for frame in range(verts.shape[0]): r = R.from_matrix(rodr[frame, 0, :, :]) rotated[frame] = r.apply(verts[frame], inverse=inverse) return rotated def readBIN(fname): if '.' not in fname: fname += '.bin' elif '.' in fname and not fname.endswith('.bin'): print("File name extension should be '.bin'") return with open(fname, 'rb') as f: F = np.frombuffer(f.read(), dtype=np.uint16).astype(np.int32) return F.reshape((-1,3)) def calculate(): garments = sorted(os.listdir('/data-local/127_2/cloth3d_processed_test/mesh_mask')) garments = [g[:-4] for g in garments] # random.shuffle(garments) # garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)] body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32) for garment in tqdm(garments): # garment = garment_data[0] picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt' id, garment_type = garment.split('_') if not os.path.exists(f'temp/cloth3d/{id}/outfit.pc16'): continue if not os.path.exists(os.path.join(cloth3d_folder, id, garment_type + '.pc16')): continue pred_cloth3d = readPC2(f'temp/cloth3d/{id}/outfit.pc16', True)['V'] pred_cloth3d_faces = readBIN(f'temp/cloth3d/{id}/outfit.bin') gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'] mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask'] pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy') gt = deorient(gt, pose, True) # gt = deorient(gt, pose, True) body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'] fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz') faces = fix_data['faces'] i = np.random.randint(0, len(pred_cloth3d)) # i = garment_data[1] try: if calculate_ours: pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3'] except: print('missed one') continue mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy() mesh_posed_fixed = mesh_posed_fixed[0, mask, :] try: m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False) tm.smoothing.filter_humphrey(m, iterations=2) pred_ours = m.vertices except: pred_ours = mesh_posed_fixed print('failed to apply smoothing') # fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')]) # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), ) # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1)) # fig.show() fig = go.Figure(data=[go.Mesh3d(x=pred_cloth3d[i][:,0], y=pred_cloth3d[i][:,1], z=pred_cloth3d[i][:,2], i=pred_cloth3d_faces[:,0], j=pred_cloth3d_faces[:,1], k=pred_cloth3d_faces[:,2], color='blue')]) fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), ) fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1)) fig.show() # fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')]) # fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray')) # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), ) # fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1)) # fig.show() pass # save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces) # save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces) # save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces) # save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces) # save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces) # save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces) if __name__ == '__main__': hood_folder = '/data-local/127_2/results/hood/all_2_105k' ours_folder = dict() ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000' ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000' cloth3d_folder = '/data-net/dataset/cloth3d/test' results_folder = '/data-local/127_2/results' calculate_ours = True calculate()