new scripts

scripts/calculate_error.py

@@ -0,0 +1,176 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('--iteration', type=str, default='')
+parser.add_argument('-g', '--gpu', type=str)
+args = vars(parser.parse_args())
+
+if args['debug']:
+    from configs.iccvw.pyr_64_128_cape_f_2 import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+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 utils.checkpoints import CheckpointSaver, load_model
+from torch.utils.data import DataLoader
+from utils.utils import text_to_html
+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
+
+
+
+def calculate():
+    errors = []
+
+    for frame in tqdm(frames):
+
+        # import plotly.graph_objects as go
+        # mesh_pred = mesh_pred.cpu().numpy()
+        # mask = batch['mesh_mask'].cpu().numpy()
+        # i = 3
+        # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,mask[i],0], y=mesh_pred[i,mask[i],2], z=mesh_pred[i,mask[i],1], mode='markers', marker=dict(size=1))])
+        # fig.update_layout(scene_aspectmode='cube')
+        # fig.show()
+
+        pred = np.load(f'{out_folder}/{frame}.npz')['arr_0']
+        gt = readPC2Frame(f'{config["test_dir"]}/mesh_posed_no_orient/{frame.id}_{frame.garment}.pc16', frame.frame_nr, True)
+        # mask = np.load(f'{config["test_dir"]}/mesh_mask/{frame.id}_{frame.garment}.npy')
+
+        fix_data = np.load(f'{config["test_dir"]}/mesh_fixed/{frame.id}_{frame.garment}.npz')
+        mask_convert = fix_data['mas_convert']
+        new_mask = fix_data['mask']
+
+        pred = pred[new_mask]
+        gt = gt[mask_convert]
+
+        error = np.linalg.norm(pred - gt, axis=1).mean()
+        errors.append((str(frame), error))
+
+    # save errors to a csv file
+    with open(f'{results_folder}/errors/{name}.csv', 'w') as f:
+        for frame, error in errors:
+            frame_str = ','.join(str(frame).split('_'))
+            f.write(f'{frame_str},{error}\n')
+
+
+def calculate_cape():
+    from smpl.smpl_torch import SMPL
+    errors = []
+
+    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
+
+    smpl = SMPL(path='smpl', orig_smpl=True, device='cpu')
+
+    for frame in tqdm(frames):
+
+        # import plotly.graph_objects as go
+        # mesh_pred = mesh_pred.cpu().numpy()
+        # mask = batch['mesh_mask'].cpu().numpy()
+        # i = 3
+        # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,mask[i],0], y=mesh_pred[i,mask[i],2], z=mesh_pred[i,mask[i],1], mode='markers', marker=dict(size=1))])
+        # fig.update_layout(scene_aspectmode='cube')
+        # fig.show()
+
+        if not os.path.exists(f'{out_folder}/{frame}.npz'):
+            print('missed one')
+            continue
+
+
+
+        pred = np.load(f'{out_folder}/{frame}.npz')['arr_0']
+        gt = np.load(f'{config["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
+        pose_shape = np.load(f'{config["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32)
+        pose, shape = pose_shape[:72], pose_shape[72:]
+        gender = int(shape[10] >= 0.5)
+        shape_offsets = torch.tensordot(torch.from_numpy(shape[:10][None, :]), smpl.shapedirs[gender], dims=[[1], [2]])[0,smpl.to_keep,:].numpy()
+
+        gt_shaped = gt + shape_offsets 
+        
+        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'))
+        offsets_pred = pred[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:]
+        offsets_orig =  gt_shaped[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:]
+
+        error = np.linalg.norm(offsets_pred - offsets_orig, axis=1).mean()
+        errors.append((str(frame), error))
+
+        # import plotly.graph_objects as go
+        # fig = go.Figure(data=[go.Scatter3d(x=pred[:,0], y=pred[:,2], z=pred[:,1], mode='markers', marker=dict(size=1))])
+        # fig.add_trace(go.Scatter3d(x=gt[:,0], y=gt[:,2], z=gt[:,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}/errors/{name}.csv', 'w') as f:
+        for frame, error in errors:
+            split = str(frame).split('_')
+            if not config['is_cape']:
+                frame_str = ','.join(split)
+            else:
+                id = '_'.join(split[:-2])
+                frame_str = f'{id},{split[-2]},{split[-1]}'
+            f.write(f'{frame_str},{error}\n')
+
+
+if __name__ == '__main__':
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    name = config['name']
+
+    results_folder = '/data-local/127_2/results'
+    if args['iteration'] == '':
+        folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config['name'])]
+        assert len(folders) == 1
+        name = folders[0]
+    else:
+        name = f'{name}_{args["iteration"]}'
+    out_folder = f'{results_folder}/{name}'
+
+    os.makedirs(f'{results_folder}/errors', exist_ok=True)
+
+    if config['is_cape']:
+        from data.dataset_cape import read_frames
+        frames = read_frames(config['test_dir'], shuffle=False, is_cape=True)
+        calculate_cape()
+    else:
+        frames_file = config['test_frames']
+        with open(os.path.join(config['test_dir'], frames_file), 'r') as f:
+            frames = f.readlines()
+        frames = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
+
+        calculate()

scripts/get_error.py

@@ -0,0 +1,82 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('--iteration', type=str, default='')
+parser.add_argument('-g', '--gpu', type=str)
+args = vars(parser.parse_args())
+
+if args['debug']:
+    from configs.iccvw.pyr_64_128_cape_f_2 import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
+
+import time
+import argparse
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from data.utils import Frame
+
+
+
+
+def calculate():
+    errors = dict()
+    errors_mean = dict()
+
+    with open(f'{results_folder}/errors/{name}.csv', 'r') as f:
+        lines = f.readlines()
+
+    for line in lines:
+        line = line.split(',')
+        category = line[2]
+        if category not in errors:
+            errors[category] = []
+        errors[category].append(float(line[3]))
+
+    for category in errors:
+        errors_mean[category] = np.mean(errors[category])
+
+    print(errors_mean) 
+
+    mean = np.mean(list(errors_mean.values()))
+
+    print(f'Final error: {mean}')
+
+
+
+if __name__ == '__main__':
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    name = config['name']
+    if not config['is_cape']:
+        frames_file = config['test_frames']
+        with open(os.path.join(config['test_dir'], frames_file), 'r') as f:
+            frames = f.readlines()
+        frames = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
+    else:
+        from data.dataset_cape import read_frames
+        frames = read_frames(config['test_dir'], shuffle=False, is_cape=True)
+
+    results_folder = '/data-local/127_2/results'
+    if args['iteration'] == '':
+        folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config['name'])]
+        assert len(folders) == 1
+        name = folders[0]
+    else:
+        name = f'{name}_{args["iteration"]}'
+    out_folder = f'{results_folder}/{name}'
+
+
+    calculate()

utils/checkpoints.py

@@ -33,12 +33,16 @@ class CheckpointSaver:
 
 def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=None, scaler: torch.nn.Module=None, iteration=''):
     # list of checkpoint files in the folder checkpoints
-    checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
+    if iteration == '':
+        checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith('.pt')]
+    else:
+        checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'iteration_{iteration}.pt')]
+        assert len(checkpoints) == 1
     iterations = np.asarray([int(f.split('_')[-1].split('.')[0]) for f in checkpoints])
     last_iteration = np.argmax(iterations)
     checkpoint = torch.load(checkpoints[last_iteration])
     iteration = iterations[last_iteration]
-    model.load_state_dict(checkpoint['model'])
+    model.load_state_dict(checkpoint['model'], strict=False)
     if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict())
     if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler'])
     print(f'Loaded model from {checkpoints[last_iteration]} at iteration {iteration}')