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

save all

Browse Source
This commit is contained in:
luni
2023-08-24 11:36:22 +00:00
parent 59f67a2495
commit afd482522d
36 changed files with 2194 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/plot_ablation.py
+43 -28
View File
@@ -25,7 +25,7 @@ import importlib
from data.utils import Frame
from data.dataset import read_frames_file
from utils.utils import save_stl
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
@@ -70,25 +70,21 @@ def step_pyramid(model, batch, config):
return mesh_pred_posed
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/ablation')
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 = {
'final': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) },
'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 195000), '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) },
# '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()
@@ -111,8 +107,23 @@ if __name__ == '__main__':
# frames = ['00054_1_Dress', '00005_1_Jumpsuit']
frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt']
frames = [Frame(frame_str=frame) for frame in frames]
# 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:
@@ -130,8 +141,8 @@ if __name__ == '__main__':
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)
# 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():
@@ -139,20 +150,24 @@ if __name__ == '__main__':
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
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)
# 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()
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