generating results

scripts/generate_results.py

@@ -0,0 +1,129 @@
+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')
+args = vars(parser.parse_args())
+
+if args['debug']:
+    from configs.iccvw.pyr_64 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 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 pytorch_msssim import MS_SSIM
+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 smpl.smpl_torch import SMPL
+from plotly.subplots import make_subplots
+from models.pyramid_model import Pyramid
+from models.cvae_model import CVAE
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+
+def step_pyramid(batch):
+    model.eval()
+    with torch.inference_mode():
+        levels_input, conds = group_inputs_for_levels(batch, config)
+
+        out_levels = model([levels_input, conds])
+        pred, _, _ = out_levels[-1]
+
+        mesh_pred = uv2mesh(pred)
+
+        if config['use_skinning']:
+            shape = batch['pose_shape'][:, 72:]
+            gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+            shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+            mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+        else:
+            mesh_pred_posed = mesh_pred
+
+    return mesh_pred_posed
+
+
+def step(batch):
+    return None
+
+
+def generate():
+    for batch in tqdm(test_dataloader):
+        if is_pyramid:
+            mesh_pred = step_pyramid(batch)
+        else: 
+            mesh_pred = step(batch)
+
+        ids = batch['id']
+        frame_nrs = batch['frame_nr']
+        garments = batch['garment']
+
+        for i in range(len(ids)):
+            np.savez_compressed(f'{out_folder}/{ids[i]}_{frame_nrs[i]}_{garments[i]}', mesh_pred[i].cpu().numpy())
+
+
+
+
+if __name__ == '__main__':
+
+    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+    dataset_resolutions = config['pyramid']['resolutions'].copy()
+    if 512 not in dataset_resolutions:
+        dataset_resolutions.append(512)
+
+    dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file=config['test_frames'], use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
+
+    is_pyramid = config['name'].startswith('pyr')
+
+    if is_pyramid:
+        model: Pyramid = Pyramid(config)
+    else:
+        model: CVAE = CVAE(config)
+
+    it = load_model(model, 'checkpoints/iccv/torch/' + config['name'])
+    # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
+
+    name = config['name'] + '_' + str(it)
+
+    folder = '/data-local/127_2/results'
+    out_folder = f'{folder}/{name}'
+
+    os.makedirs(out_folder, exist_ok=True)
+
+    generate()