forked from shanice-l/gdrnpp_bop2022
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathlmo_full.py
128 lines (111 loc) · 3.38 KB
/
lmo_full.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
# encoding: utf-8
"""This file includes necessary params, info."""
import os
import mmcv
import os.path as osp
import numpy as np
# ---------------------------------------------------------------- #
# ROOT PATH INFO
# ---------------------------------------------------------------- #
cur_dir = osp.abspath(osp.dirname(__file__))
root_dir = osp.normpath(osp.join(cur_dir, ".."))
# directory storing experiment data (result, model checkpoints, etc).
output_dir = osp.join(root_dir, "output")
data_root = osp.join(root_dir, "datasets")
bop_root = osp.join(data_root, "BOP_DATASETS/")
# ---------------------------------------------------------------- #
# LINEMOD OCCLUSION DATASET
# ---------------------------------------------------------------- #
dataset_root = osp.join(bop_root, "lmo")
train_dir = osp.join(dataset_root, "train")
test_dir = osp.join(dataset_root, "test")
model_dir = osp.join(dataset_root, "models")
model_eval_dir = osp.join(dataset_root, "models_eval")
model_scaled_simple_dir = osp.join(dataset_root, "models_scaled_f5k")
vertex_scale = 0.001
# for surfemb
surface_samples_dir = osp.join(dataset_root, "surface_samples")
surface_samples_normals_dir = osp.join(dataset_root, "surface_samples_normals")
test_scenes = [2]
# object info
objects = [
"ape",
"can",
"cat",
"driller",
"duck",
"eggbox",
"glue",
"holepuncher",
]
id2obj = {
1: "ape",
# 2: 'benchvise',
# 3: 'bowl',
# 4: 'camera',
5: "can",
6: "cat",
# 7: 'cup',
8: "driller",
9: "duck",
10: "eggbox",
11: "glue",
12: "holepuncher",
# 13: 'iron',
# 14: 'lamp',
# 15: 'phone'
}
obj_num = len(id2obj)
obj2id = {_name: _id for _id, _name in id2obj.items()}
model_paths = [osp.join(model_dir, "obj_{:06d}.ply").format(_id) for _id in id2obj]
texture_paths = None
model_colors = [((i + 1) * 10, (i + 1) * 10, (i + 1) * 10) for i in range(obj_num)] # for renderer
diameters = (
np.array(
[
102.099, # 1
# 247.506,
# 167.355,
# 172.492,
201.404, # 5
154.546, # 6
# 124.264,
261.472, # 8
108.999, # 9
164.628, # 10
175.889, # 11
145.543, # 12
# 278.078,
# 282.601,
# 212.358,
]
)
/ 1000.0
)
# Camera info
width = 640
height = 480
zNear = 0.25
zFar = 6.0
center = (height / 2, width / 2)
camera_matrix = np.array([[572.4114, 0, 325.2611], [0, 573.57043, 242.04899], [0, 0, 1]])
def get_models_info():
"""key is str(obj_id)"""
models_info_path = osp.join(model_dir, "models_info.json")
assert osp.exists(models_info_path), models_info_path
models_info = mmcv.load(models_info_path) # key is str(obj_id)
return models_info
def get_fps_points():
"""key is str(obj_id) generated by
core/gdrn_modeling/tools/lmo/lmo_1_compute_fps.py."""
fps_points_path = osp.join(model_dir, "fps_points.pkl")
assert osp.exists(fps_points_path), fps_points_path
fps_dict = mmcv.load(fps_points_path)
return fps_dict
def get_keypoints_3d():
"""key is str(obj_id) generated by
core/roi_pvnet/tools/lmo/lmo_1_compute_keypoints_3d.py."""
keypoints_3d_path = osp.join(model_dir, "keypoints_3d.pkl")
assert osp.exists(keypoints_3d_path), keypoints_3d_path
kpts_dict = mmcv.load(keypoints_3d_path)
return kpts_dict