plot ablation pyramid

scripts/plot_ablation_pyramid.py

@@ -0,0 +1,145 @@
+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"] = '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 torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
+from pathlib import Path
+
+from data.utils import Frame
+from data.dataset import read_frames_file
+from utils.utils import save_stl
+from utils.checkpoints import CheckpointSaver, load_model
+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 smpl.smpl_torch import SMPL
+from plotly.subplots import make_subplots
+from models.pyramid_model import Pyramid
+from data.dataset import Dataset
+from utils.checkpoints import CheckpointSaver, load_model
+from torch.utils.data import DataLoader
+import trimesh as tm
+from zipfile import ZipFile
+
+
+from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_1
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
+
+
+def step_pyramid(model, batch):
+    with torch.inference_mode():
+        levels_input, conds = group_inputs_for_levels(batch, config_1)
+
+        out = []
+        out_levels = model([levels_input, conds])
+        for i in range(len(out_levels)):
+            pred, _, _ = out_levels[i]
+
+            mesh_pred = uv2mesh(pred)
+
+            if config_1['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]
+                if not config_1['is_cape']:
+                    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+                else:
+                    mesh_pred_posed = mesh_pred_shaped
+            else:
+                mesh_pred_posed = mesh_pred
+
+            out.append(mesh_pred_posed)
+
+    return out
+
+
+def create_temp_frames(folder: str, frame: Frame):
+    with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+        f.writelines([f'{frame}\n' for _ in range(4)])
+
+
+if __name__ == '__main__':
+    results_folder = Path('/data-local/127_2/results')
+    out_folder = Path('temp/pyramid_ablation')
+    iteration = '100000'
+
+    smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
+    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
+    dataset_resolutions = []
+    if 'pyramid' in config_1:
+        dataset_resolutions = config_1['pyramid']['resolutions'].copy()
+    if 512 not in dataset_resolutions:
+        dataset_resolutions.append(512)
+
+    model: Pyramid = Pyramid(config_1)
+    it = load_model(model, 'checkpoints/iccvw/' + config_1['name'], iteration=iteration)
+    model.eval()
+
+    frames = ['00054_1_Dress']
+    frames = [Frame(frame_str=frame) for frame in frames]
+    body_faces = smpl.f
+
+    for frame in frames:
+        create_temp_frames(config_1['train_dir'], frame)
+        dataset = Dataset(config_1['train_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+
+        batch = next(iter(dataloader))
+        pred = step_pyramid(model, batch)
+        
+        fix_data = np.load(Path(dataset.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+        mask = fix_data['mask']
+        faces = fix_data['faces']
+
+        body = batch['mesh_body_posed'][0].cpu().numpy()
+        gt = batch['mesh_posed'][0].cpu().numpy()
+
+        save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces)
+        save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces)
+
+        for i, _pred in enumerate(pred):
+            mesh_posed_fixed = fix_collision(_pred, batch['mesh_body_posed'], batch['mesh_mask']).cpu().numpy()
+
+            mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
+            m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
+            tm.smoothing.filter_humphrey(m, iterations=2)
+
+            verts = m.vertices
+
+            save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
+            save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)
+
+
+            # fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
+            #                       go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
+            #                       ])
+            # fig.update_layout(scene=dict(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
+
+
+
+