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

save all

Browse Source
This commit is contained in:
luni
2023-08-24 11:36:22 +00:00
parent 59f67a2495
commit afd482522d
36 changed files with 2194 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/calculate_chamfer.py
+168 -39
View File
@@ -4,47 +4,40 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
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 data.cloth3d_io import readPC2Frame, readPC2, readOBJ
from utils.utils import pickle_load
from scipy.spatial import cKDTree
from smpl.smpl_torch import SMPL
from scipy.spatial.transform import Rotation as R
import trimesh as tm
from losses.losses import collision_loss, normal_loss, kl_loss, collision_error
parser = argparse.ArgumentParser()
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
parser.add_argument("-e", "--end", help = "end of range", default=550, type=int)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument('--debug', action='store_true')
args = parser.parse_args()
def chamfer(points1, points2, points2_tree: cKDTree):
tree1 = cKDTree(points1)
dist1, _ = tree1.query(points2, workers=40)
dist2, _ = points2_tree.query(points1, workers=40)
dist1, _ = tree1.query(points2, workers=10)
dist2, _ = points2_tree.query(points1, workers=10)
chamfer_dist = np.mean(dist1) + np.mean(dist2)
return chamfer_dist
@@ -60,34 +53,116 @@ def deorient(verts, pose, inverse=True):
return rotated
def get_collision_depth(m1, m2):
cm = tm.collision.CollisionManager()
cm.add_object('m1', m1)
cm.add_object('m2', m2)
data = cm.in_collision_internal(return_data=True)[1]
# c = np.asarray([d.index('m2') for d in data])
# count = np.unique(c).shape[0]
depth_sum = np.sum([d.depth for d in data])
return depth_sum
def get_area_diff(m1, m2):
return np.abs(m1.area - m2.area)
device = 'cuda'
def get_collision_loss(pred, body, mask):
c_loss, c_ratio = collision_error(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device))
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
def get_collision_loss_2(pred, body, mask):
_pred = np.zeros((14475, 3), dtype=np.float32)
_pred[mask] = pred
c_loss, c_ratio = collision_loss(torch.from_numpy(_pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005)
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
def calculate():
errors_hood = []
errors_ours = []
garments = sorted(os.listdir(hood_folder))
collisions_hood = []
areas_hood = []
collision_ratios_hood = []
if calculate_ours:
errors_ours = []
collisions_ours = []
areas_ours = []
collision_ratios_ours = []
garments = sorted(os.listdir(hood_folder))[args.start:args.end]
garments = [g[:-4] for g in garments]
body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
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']
pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl')
pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces']
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
template, template_faces, _, _ = readOBJ(os.path.join(cloth3d_folder, id, garment_type + '.obj'))
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)
body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:]
fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
faces = fix_data['faces']
m_template = tm.Trimesh(vertices=template, faces=template_faces, process=False, validate=False)
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']
if calculate_ours:
# pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0']
except:
print('missed one')
continue
# m_body = tm.Trimesh(vertices=body[i], faces=body_faces, process=False, validate=False)
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)))
m_hood = tm.Trimesh(vertices=pred_hood[i], faces=pred_hood_faces, process=False, validate=False)
# collision_hood = get_collision_depth(m_body, m_hood)
# collisions_hood.append((f'{id}_{i}_{garment_type}', collision_hood))
area_hood = get_area_diff(m_template, m_hood)
areas_hood.append((f'{id}_{i}_{garment_type}', area_hood))
collision_ratio_hood = get_collision_loss(pred_hood[i], body[i], mask)[1]
collision_ratios_hood.append((f'{id}_{i}_{garment_type}', collision_ratio_hood))
if calculate_ours:
errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask], gt[i], gt_tree)))
m_ours = tm.Trimesh(vertices=pred_ours[mask], faces=faces, process=False, validate=False)
# collision_ours = get_collision_depth(m_body, m_ours)
# collisions_ours.append((f'{id}_{i}_{garment_type}', collision_ours))
area_ours = get_area_diff(m_template, m_ours)
areas_ours.append((f'{id}_{i}_{garment_type}', area_ours))
collision_ratio_ours = get_collision_loss(pred_ours[mask], body[i], mask)[1]
# collision_ratio_ours_2 = get_collision_loss_2(pred_ours[mask], body[i], mask)[1]
collision_ratios_ours.append((f'{id}_{i}_{garment_type}', collision_ratio_ours))
# fig = go.Figure(data=[go.Mesh3d(x=pred_ours[mask][:,0], y=pred_ours[mask][:,2], z=pred_ours[mask][:,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=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()
# fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[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
# 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)))
@@ -95,6 +170,7 @@ def calculate():
# 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()
# pass
# 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)))
@@ -105,30 +181,83 @@ def calculate():
# 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')
os.makedirs(f'{results_folder}/hood/{name}_hood', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/chamfer', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/collision', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/area', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/collision_ratio', exist_ok=True)
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')
for frame, error in errors_hood:
np.save(f'{results_folder}/hood/{name}_hood/chamfer/{frame}.npy', error)
for frame, error in collisions_hood:
np.save(f'{results_folder}/hood/{name}_hood/collision/{frame}.npy', error)
for frame, error in areas_hood:
np.save(f'{results_folder}/hood/{name}_hood/area/{frame}.npy', error)
for frame, error in collision_ratios_hood:
np.save(f'{results_folder}/hood/{name}_hood/collision_ratio/{frame}.npy', error)
os.makedirs(f'{results_folder}/hood/{name}_ours', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/chamfer', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/collision', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/area', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/collision_ratio', exist_ok=True)
for frame, error in errors_ours:
np.save(f'{results_folder}/hood/{name}_ours/chamfer/{frame}.npy', error)
for frame, error in collisions_ours:
np.save(f'{results_folder}/hood/{name}_ours/collision/{frame}.npy', error)
for frame, error in areas_ours:
np.save(f'{results_folder}/hood/{name}_ours/area/{frame}.npy', error)
for frame, error in collision_ratios_ours:
np.save(f'{results_folder}/hood/{name}_ours/collision_ratio/{frame}.npy', error)
# with open(f'{results_folder}/hood/{name}/chamfer_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}/collision_hood.csv', 'w') as f:
# for frame, error in collisions_hood:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/area_hood.csv', 'w') as f:
# for frame, error in areas_hood:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# if calculate_ours:
# with open(f'{results_folder}/hood/{name}_chamfer_ours.csv', 'w') as f:
# for frame, error in errors_ours:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/collision_ours.csv', 'w') as f:
# for frame, error in collisions_ours:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/area_ours.csv', 'w') as f:
# for frame, error in areas_ours:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
if __name__ == '__main__':
name = 'hood_2'
name = 'hood_3_var'
hood_folder = '/data-local/127_2/results/hood/all_2'
hood_folder = '/data-local/127_2/results/hood/all_3'
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'
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_var_3_450000'
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_var_2_600000'
cloth3d_folder = '/data-net/dataset/cloth3d/test'
results_folder = '/data-local/127_2/results'
out_folder = f'{results_folder}/errors/{name}'
calculate_ours = True
calculate()