forked from Daisy-1227/SSCF
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathlsnrr_opt_mugs.py
107 lines (84 loc) · 3.62 KB
/
lsnrr_opt_mugs.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
import os
import sys
sys.path.append('/mnt/d/projects/sscf')
import torch
from src.datasets.mugs import MugsDataset
from src.datasets.shapenet import ShapeNetTestDataset
from src.optimizer.mug.lsnrr_opt import LSNRROptimizer
from src.human.hand import load_mano, param2mesh
import numpy as np
import pyvista as pv
import trimesh
from tqdm import tqdm
from src.utils.rot_utils import rotate_np
H_model = load_mano(mano_assets_root='data/human_model/mano')
dataset_src = MugsDataset(f'data/hoi/mugs')
dataset_tgt = ShapeNetTestDataset('data/shapenet','03797390')
optimizer = LSNRROptimizer('data/ckpts/mug/cpdpca.pt', vnndif_path='data/ckpts/mug/checkpoint.ckpt')
color_anc = 'w'
color_src = '#15c4ed'
color_tgt = '#fe7a89'
visualize = True
save_results = True
save_folder = f'results/lsnrr/mug'
if not os.path.exists(save_folder):
os.makedirs(save_folder, exist_ok=True)
src_list = [i for i in range(len(dataset_src))]
tgt_list = [i for i in range(len(dataset_tgt))]
for i_src in src_list:
i_tgt = tgt_list[np.random.randint(len(tgt_list))]
print(f'i_src: {i_src}, i_tgt: {i_tgt}')
# data_src: src_interaction
data_src = dataset_src[i_src]
# O_tgt: tgt_object
O_tgt = dataset_tgt[i_tgt][:2048]
# O_src: src_object
O_src = data_src['O_src'][:2048]
# H_src_param: src_human
H_src_param = data_src['H_src_param']
# transfer
out_dict = optimizer.transfer(H_src_param, O_src, O_tgt)
# visualize results
if visualize:
O_src_mesh = dataset_src.get_mesh(i_src)
O_src_loc = out_dict['O_src_loc']
O_src_scale = out_dict['O_src_scale']
if 'O_src_rot' not in out_dict:
O_src_rot = np.eye(3)
else:
O_src_rot = out_dict['O_src_rot']
O_src_mesh.vertices = rotate_np((O_src_mesh.vertices - O_src_loc) / O_src_scale, O_src_rot)
H_src_norm = out_dict['H_src_norm']
H_src = trimesh.Trimesh(H_src_norm, H_model.get_mano_closed_faces().numpy(), process=False)
pl = pv.Plotter(shape=(1,2))
pl.subplot(0,0)
pl.add_mesh(O_src_mesh, color_src)
pl.add_mesh(H_src, color_anc)
pl.subplot(0,1)
O_tgt_mesh = dataset_tgt.get_mesh_wt(i_tgt)
O_tgt_mesh.vertices = (O_tgt_mesh.vertices - out_dict['O_tgt_loc']) / out_dict['O_tgt_scale']
pl.add_mesh(O_tgt_mesh, color_tgt)
pl.add_mesh(trimesh.Trimesh(out_dict['H_tgt_norm'], H_model.get_mano_closed_faces().numpy(), process=False), color_anc)
pl.link_views()
pl.show()
# save results
if save_results:
sub_dir = f"{save_folder}/{i_src}_{i_tgt}"
os.makedirs(sub_dir, exist_ok=True)
O_tgt_mesh = dataset_tgt.get_mesh_wt(i_tgt)
O_tgt_mesh.vertices = (O_tgt_mesh.vertices - out_dict['O_tgt_loc']) / out_dict['O_tgt_scale']
O_tgt_mesh.export(f'{sub_dir}/{i_src}_{i_tgt}_tgt_obj.obj')
H_tgt = trimesh.Trimesh(out_dict['H_tgt_norm'], H_model.get_mano_closed_faces().numpy(), process=False)
H_tgt.export(f'{sub_dir}/{i_src}_{i_tgt}_tgt_H.obj')
O_src_mesh = dataset_src.get_mesh(i_src)
O_src_loc = out_dict['O_src_loc']
O_src_scale = out_dict['O_src_scale']
if 'O_src_rot' not in out_dict:
O_src_rot = np.eye(3)
else:
O_src_rot = out_dict['O_src_rot']
O_src_mesh.vertices = rotate_np((O_src_mesh.vertices - O_src_loc) / O_src_scale, O_src_rot)
O_src_mesh.export(f'{sub_dir}/{i_src}_{i_tgt}_src_obj.obj')
H_src.export(f'{sub_dir}/{i_src}_{i_tgt}_src_H.obj')
torch.save(out_dict,f'{sub_dir}/{i_src}_{i_tgt}.pt')
print(f'Saved to {sub_dir}/{i_src}_{i_tgt}.pt')