luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

save

Browse Source
This commit is contained in:
luni
2023-07-14 10:12:00 +02:00
parent bb16aaf868
commit 88531da784
16 changed files with 411 additions and 299 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/plot_ablation_pyramid.py
+30 -21
View File
@@ -8,7 +8,7 @@ os.environ["CUDA_VISIBLE_DEVICES"] = '1'
import time
import argparse
import wandb
import tqdm
from tqdm import tqdm
import random
import numpy as np
import plotly.graph_objects as go
@@ -39,30 +39,32 @@ import trimesh as tm
from zipfile import ZipFile
from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_1
from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_2
uv2mesh = dict()
uv2mesh['nonskirt'] = 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'])
uv2mesh = UV2Mesh(res=256, is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
def step_pyramid(model, batch):
def step_pyramid(model, batch, config):
with torch.inference_mode():
levels_input, conds = group_inputs_for_levels(batch, config_1)
levels_input, conds = group_inputs_for_levels(batch, config)
out = []
out_levels = model([levels_input, conds])
for i in range(len(out_levels)):
pred, _, _ = out_levels[i]
mesh_pred = uv2mesh(pred)
mesh_pred = uv2mesh[picker](pred)
if config_1['use_skinning']:
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_1['is_cape']:
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
@@ -82,31 +84,39 @@ def create_temp_frames(folder: str, frame: Frame):
if __name__ == '__main__':
results_folder = Path('/data-local/127_2/results')
out_folder = Path('temp/pyramid_ablation')
iteration = '100000'
iteration = {'nonskirt': 190000, 'skirt': 200000 }
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()
model = { 'nonskirt': Pyramid(config_1), 'skirt': Pyramid(config_2) }
load_model(model['nonskirt'], 'checkpoints/iccvw/' + config_1['name'], iteration=iteration['nonskirt'])
load_model(model['skirt'], 'checkpoints/iccvw/' + config_2['name'], iteration=iteration['skirt'])
model['nonskirt'].eval()
model['skirt'].eval()
frames = ['00054_1_Dress']
frames = ['00054_1_Dress', '01070_13_Jumpsuit', '03849_27_Dress', '06057_152_Dress',
'03179_274_Dress', '04247_102_Tshirt', '01852_70_Tshirt']
frames_corrupt = ['00300_69_Jumpsuit', '07608_77_Jumpsuit', '06955_117_Dress']
frames = [Frame(frame_str=frame) for frame in frames]
body_faces = smpl.f
for frame in frames:
config = { 'nonskirt': config_1, 'skirt': config_2 }
for frame in tqdm(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'])
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=config_1['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
picker = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'nonskirt'
batch = next(iter(dataloader))
pred = step_pyramid(model, batch)
pred = step_pyramid(model[picker], batch, config[picker])
fix_data = np.load(Path(dataset.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
mask = fix_data['mask']
@@ -124,19 +134,18 @@ if __name__ == '__main__':
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
pass