178 lines
7.8 KiB
Python
178 lines
7.8 KiB
Python
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"] = '3'
|
|
|
|
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
|
|
import importlib
|
|
|
|
from data.utils import Frame
|
|
from data.dataset import read_frames_file
|
|
from utils.utils import save_stl, create_temp_frames
|
|
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_abl import config as config_1
|
|
|
|
|
|
uv2mesh = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config_1['is_cape'])
|
|
uv2mesh_skirt = UV2Mesh(res=256, is_skirt=True, is_smpl_orig=config_1['is_cape'])
|
|
|
|
|
|
def step_pyramid(model, batch, config):
|
|
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 not config['is_skirt'] else uv2mesh_skirt(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]
|
|
if not config['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
|
|
|
|
return mesh_pred_posed
|
|
|
|
|
|
if __name__ == '__main__':
|
|
results_folder = Path('/data-local/127_2/results')
|
|
out_folder = Path('temp/supplementary')
|
|
os.makedirs(out_folder, exist_ok=True)
|
|
|
|
smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
|
|
dataset_resolutions = []
|
|
dataset_resolutions = [64, 128, 256, 512]
|
|
|
|
experiments = {
|
|
# 'inc': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) },
|
|
'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 250000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 200000) },
|
|
# 'skin': { 'nonskirt': ('pyr_64_128_256_512_skin', 225000), 'skirt': ('pyr_64_128_256_512_skin_skirt', 200000) },
|
|
# 'norm': { 'nonskirt': ('pyr_64_128_256_512_norm', 200000), 'skirt': ('pyr_64_128_256_512_norm_skirt', 200000) },
|
|
# 'reg': { 'nonskirt': ('pyr_64_128_256_512_reg', 200000), 'skirt': ('pyr_64_128_256_512_reg_skirt', 200000) },
|
|
}
|
|
|
|
configs = dict()
|
|
models = dict()
|
|
|
|
for experiment in experiments.keys():
|
|
configs[experiment] = dict()
|
|
configs[experiment]['nonskirt'] = importlib.import_module('configs.iccvw.' + experiments[experiment]['nonskirt'][0]).config
|
|
configs[experiment]['skirt'] = importlib.import_module('configs.iccvw.' + experiments[experiment]['skirt'][0]).config
|
|
|
|
models[experiment] = dict()
|
|
models[experiment]['nonskirt'] = Pyramid(configs[experiment]['nonskirt'])
|
|
models[experiment]['skirt'] = Pyramid(configs[experiment]['skirt'])
|
|
|
|
load_model(models[experiment]['nonskirt'], 'checkpoints/iccvw/' + configs[experiment]['nonskirt']['name'], iteration=experiments[experiment]['nonskirt'][1])
|
|
load_model(models[experiment]['skirt'], 'checkpoints/iccvw/' + configs[experiment]['skirt']['name'], iteration=experiments[experiment]['skirt'][1])
|
|
|
|
models[experiment]['nonskirt'].eval()
|
|
models[experiment]['skirt'].eval()
|
|
|
|
|
|
# frames = ['00054_1_Dress', '00005_1_Jumpsuit']
|
|
# frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt']
|
|
# frames = ['01852_70_Tshirt'] # , '03849_27_Dress', '00054_1_Dress'
|
|
# frames = [Frame(frame_str=frame) for frame in frames]
|
|
|
|
frames_file = configs['e2e']['nonskirt']['train_frames']
|
|
with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f:
|
|
frames = f.readlines()
|
|
frames_1 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
|
|
|
|
frames_file = configs['e2e']['skirt']['train_frames']
|
|
with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f:
|
|
frames = f.readlines()
|
|
frames_2 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
|
|
|
|
frames = frames_1 + frames_2
|
|
random.shuffle(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=True, resolutions=dataset_resolutions, is_cape=False)
|
|
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
|
|
|
|
model_picker = 'skirt' if frame.garment in ['Dress', 'Skirt'] else 'nonskirt'
|
|
|
|
fix_data = np.load(Path(dataset.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
|
|
mask = fix_data['mask']
|
|
faces = fix_data['faces']
|
|
|
|
batch = next(iter(dataloader))
|
|
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 experiment in experiments.keys():
|
|
|
|
pred = step_pyramid(models[experiment][model_picker], batch, configs[experiment][model_picker])
|
|
|
|
mesh_posed_fixed = fix_collision(pred, batch['mesh_body_posed'], batch['mesh_mask']).cpu().numpy()
|
|
mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
|
|
try:
|
|
m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
|
|
tm.smoothing.filter_humphrey(m, iterations=2)
|
|
verts = m.vertices
|
|
except:
|
|
verts = mesh_posed_fixed
|
|
print('failed to apply smoothing')
|
|
|
|
# save_stl(out_folder / f'{frame}_{experiment}_body.stl', body, body_faces)
|
|
# save_stl(out_folder / f'{frame}_{experiment}_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.update_layout(title_text=f'{frame} {experiment}')
|
|
fig.show()
|
|
pass
|
|
|
|
pass
|
|
|
|
|
|
|
|
|