-
Notifications
You must be signed in to change notification settings - Fork 3
/
visualize.py
138 lines (119 loc) · 5.7 KB
/
visualize.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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import argparse
from omegaconf import OmegaConf
from sklearn.metrics import pairwise_distances_argmin
from tensorpack.predict import OfflinePredictor, PredictConfig
from tensorpack.tfutils import SaverRestore
from model import Model
from utils import SMPLModel, sample_vertex_from_mesh
import os
from tqdm import tqdm
import numpy as np
from eval_iou import KeypointDataset, name2id, nms_usip
import open3d as o3d
if os.name == 'nt':
from sdv_src.build.Release import sdv
else:
from sdv_src.build import sdv
def gen_mesh_with_color(pc, kp_idxs, color, kp_scale=2):
prim_verts = None
mesh_spheres = []
res = 4
colors = []
for j, pt in enumerate(pc):
if j in kp_idxs:
continue
mesh_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.01, resolution=res)
mesh_sphere.translate(pt)
mesh_sphere.paint_uniform_color(color[j])
mesh_spheres.append(mesh_sphere)
if prim_verts is None:
prim_verts = np.asarray(mesh_sphere.vertices).shape[0]
colors.append(color[j])
for idx in kp_idxs:
pt = pc[idx]
mesh_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.01 * kp_scale, resolution=res)
mesh_sphere.translate(pt)
mesh_sphere.paint_uniform_color(color[idx])
mesh_spheres.append(mesh_sphere)
colors.append(color[idx])
# save to file
vertices = []
triangles = []
total_vertices = 0
for mesh in mesh_spheres:
triangles.append(np.asarray(mesh.triangles) + total_vertices)
vertices.append(np.asarray(mesh.vertices))
total_vertices += np.asarray(mesh.vertices).shape[0]
vertices = np.concatenate(vertices)
triangles = np.concatenate(triangles)
merged_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(vertices), o3d.utility.Vector3iVector(triangles))
mesh_color = np.zeros((vertices.shape[0], 3))
colors = np.stack(colors, 0)
for j in range(pc.shape[0]):
mesh_color[j * prim_verts:(j + 1) * prim_verts] = colors[j]
merged_mesh.vertex_colors = o3d.utility.Vector3dVector(mesh_color.reshape(-1, 3))
return merged_mesh
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--type', default='shapenet', choices=['shapenet', 'smpl', '3dmatch'], help='Which type of evaluation (ShapeNet/SMPL/Real Scene)?')
parser.add_argument('--kpnet_root', default='/kpnet/root', help='KeypointNet data root')
parser.add_argument('--cat_name', default='chair', help='ShapeNet category name')
parser.add_argument('--nms', action='store_true', help='Whether to use NMS')
parser.add_argument('--nms_radius', type=float, default=0.2, help='NMS radius')
parser.add_argument('--kp_num', type=int, default=-1, help='Whether to force a fixed number of keypoints (defaults: -1, no limitation)')
args = parser.parse_args()
if args.type == 'shapenet':
model_path = 'outputs/{}/tflogs/checkpoint'.format(args.cat_name)
elif args.type == 'smpl':
model_path = 'outputs/{}/tflogs/checkpoint'.format('smpl')
elif args.type == '3dmatch':
model_path = 'outputs/{}/tflogs/checkpoint'.format('universal')
else:
raise Exception('unknown type')
cfg = OmegaConf.load('config/config.yaml')
predictor = OfflinePredictor(config=PredictConfig(model=Model(cfg),
input_names=['pc', 'pc_feature'],
output_names=['encoder/z'],
session_init=SaverRestore(model_path)))
if args.type == 'shapenet':
test_dataset = KeypointDataset(args.kpnet_root, name2id[args.cat_name], 'splits/test.txt')
mesh_name, pc, feature, label = test_dataset[np.random.randint(len(test_dataset))]
prob = predictor(pc[None], feature[None])[0][0]
elif args.type == 'smpl':
smpl = SMPLModel('data/model.pkl')
pose = (np.random.rand(*smpl.pose_shape) - 0.5) * 0.4
beta = (np.random.rand(*smpl.beta_shape) - 0.5) * 0.06
trans = np.zeros(smpl.trans_shape)
smpl.set_params(beta=beta, pose=pose, trans=trans)
pc, _, idx, u, v = sample_vertex_from_mesh(smpl.verts, smpl.faces, num_samples=cfg.num_points)
pc_feature = np.array(sdv.compute(pc)).astype(np.float32)
prob = predictor(pc[None], pc_feature[None])[0][0]
elif args.type == '3dmatch':
pcd = o3d.io.read_point_cloud('data/3DMatch/kitchen/cloud_bin_0.ply')
pc_origin = np.array(pcd.points)
_, trace = pcd.voxel_down_sample_and_trace(
3e-2, pcd.get_min_bound(), pcd.get_max_bound(), False
)
sp_idx = trace[:, 0]
sp_idx = sp_idx[sp_idx >= 0]
pc = pc_origin[sp_idx]
pc_feature = np.array(sdv.compute(pc_origin, interest_point_idxs=sp_idx)).astype(np.float32)
prob = []
block_size = 2048
for i in tqdm(range((pc.shape[0] - 1) // block_size + 1)):
prob.append(predictor(pc[i*block_size:(i+1)*block_size][None], pc_feature[i*block_size:(i+1)*block_size][None])[0][0])
prob = np.concatenate(prob)
if args.nms:
kp, prob = nms_usip(pc, prob, args.nms_radius)
kp = kp[prob > 0.5]
prob = prob[prob > 0.5]
else:
kp = pc[prob > 0.5]
prob = prob[prob > 0.5]
if args.kp_num > 0:
kp = kp[np.argsort(prob)[-args.kp_num:]]
kp_idx = pairwise_distances_argmin(kp, pc)
color = np.repeat(np.array([[202 / 255., 202 / 255., 202 / 255.]]), pc.shape[0], 0)
color[kp_idx] = np.array([[1., 0, 0]])
mesh = gen_mesh_with_color(pc, kp_idx, color)
o3d.visualization.draw_geometries([mesh])