134 lines
5.1 KiB
Python
134 lines
5.1 KiB
Python
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 wandb
|
|
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 tqdm import tqdm
|
|
from data.dataset import Dataset
|
|
from data.utils import Frame
|
|
from data.cloth3d_io import readPC2Frame, readPC2
|
|
from data.preprocessing import deorient
|
|
from utils.checkpoints import CheckpointSaver, load_model
|
|
from torch.utils.data import DataLoader
|
|
from utils.utils import text_to_html, pickle_load
|
|
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 plotly.subplots import make_subplots
|
|
import pytorch3d.loss
|
|
from scipy.spatial.distance import cdist
|
|
from scipy.spatial import cKDTree
|
|
from smpl.smpl_torch import SMPL
|
|
from scipy.spatial.transform import Rotation as R
|
|
|
|
|
|
def chamfer(points1, points2, points2_tree: cKDTree):
|
|
tree1 = cKDTree(points1)
|
|
dist1, _ = tree1.query(points2, workers=40)
|
|
dist2, _ = points2_tree.query(points1, workers=40)
|
|
chamfer_dist = np.mean(dist1) + np.mean(dist2)
|
|
return chamfer_dist
|
|
|
|
|
|
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 calculate():
|
|
errors_hood = []
|
|
errors_ours = []
|
|
garments = sorted(os.listdir(hood_folder))
|
|
garments = [g[:-4] for g in garments]
|
|
|
|
for garment in tqdm(garments):
|
|
picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
|
|
id, garment_type = garment.split('_')
|
|
|
|
pred_hood = pickle_load(f'{hood_folder}/{garment}.pkl')['verts']
|
|
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
|
|
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')[2:]
|
|
gt = deorient(gt, pose, True)
|
|
|
|
|
|
for i in range(pred_hood.shape[0]):
|
|
gt_tree = cKDTree(gt[i])
|
|
|
|
try:
|
|
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
|
|
|
|
errors_hood.append((f'{id}_{i}_{garment_type}', chamfer(pred_hood[i], gt[i], gt_tree)))
|
|
errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask,:], gt[i], gt_tree)))
|
|
|
|
# fig = go.Figure()
|
|
# fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
|
|
# fig.add_trace(go.Scatter3d(x=pred_hood[i][:, 0], y=pred_hood[i][:, 2], z=pred_hood[i][:, 1], mode='markers', marker=dict(size=1)))
|
|
# fig.update_layout(scene_aspectmode='cube')
|
|
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1]), yaxis=dict(range=[-1, 1]), zaxis=dict(range=[-1, 1])))
|
|
# fig.show()
|
|
|
|
# fig = go.Figure()
|
|
# fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
|
|
# fig.add_trace(go.Scatter3d(x=pred_ours[i,mask,:][:, 0], y=pred_ours[i,mask,:][:, 2], z=pred_ours[i,mask,:][:, 1], mode='markers', marker=dict(size=1)))
|
|
# fig.update_layout(scene_aspectmode='cube')
|
|
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1]), yaxis=dict(range=[-1, 1]), zaxis=dict(range=[-1, 1])))
|
|
# fig.show()
|
|
|
|
|
|
# save errors to a csv file
|
|
with open(f'{results_folder}/hood/{name}_hood.csv', 'w') as f:
|
|
for frame, error in errors_hood:
|
|
frame_str = ','.join(str(frame).split('_'))
|
|
f.write(f'{frame_str},{error}\n')
|
|
|
|
with open(f'{results_folder}/hood/{name}_ours.csv', 'w') as f:
|
|
for frame, error in errors_ours:
|
|
frame_str = ','.join(str(frame).split('_'))
|
|
f.write(f'{frame_str},{error}\n')
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
name = 'hood_2'
|
|
|
|
hood_folder = '/data-local/127_2/results/hood/all_2'
|
|
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'
|
|
out_folder = f'{results_folder}/errors/{name}'
|
|
|
|
|
|
calculate() |