diff --git a/scripts/plot_ablation_pyramid.py b/scripts/plot_ablation_pyramid.py new file mode 100644 index 0000000..e8d4d75 --- /dev/null +++ b/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 + + + +