new scripts

scripts/calculate_chamfer.py

@@ -0,0 +1,134 @@
+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()