-
Notifications
You must be signed in to change notification settings - Fork 46
/
Copy pathbase_utils.py
335 lines (276 loc) · 12.3 KB
/
base_utils.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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
import cv2
import lmdb
import numpy as np
import os
from PIL import Image
from plyfile import PlyData
from config import cfg
from transforms3d.euler import mat2euler
import pickle
def read_pickle(pkl_path):
with open(pkl_path, 'rb') as f:
return pickle.load(f)
def save_pickle(data, pkl_path):
os.system('mkdir -p {}'.format(os.path.dirname(pkl_path)))
with open(pkl_path, 'wb') as f:
pickle.dump(data, f)
def read_pose(rot_path, tra_path):
rot = np.loadtxt(rot_path, skiprows=1)
tra = np.loadtxt(tra_path, skiprows=1) / 100.
return np.concatenate([rot, np.reshape(tra, newshape=[3, 1])], axis=-1)
class ModelAligner(object):
rotation_transform = np.array([[1., 0., 0.],
[0., -1., 0.],
[0., 0., -1.]])
translation_transforms = {
# 'cat': np.array([-0.00577495, -0.01259045, -0.04062323])
}
intrinsic_matrix = {
'linemod': np.array([[572.4114, 0., 325.2611],
[0., 573.57043, 242.04899],
[0., 0., 1.]]),
# 'blender': np.array([[280.0, 0.0, 128.0],
# [0.0, 280.0, 128.0],
# [0.0, 0.0, 1.0]]),
'blender': np.array([[700., 0., 320.],
[0., 700., 240.],
[0., 0., 1.]])
}
def __init__(self, class_type='cat'):
self.class_type = class_type
self.blender_model_path = os.path.join(cfg.LINEMOD,'{}/{}.ply'.format(class_type, class_type))
self.orig_model_path = os.path.join(cfg.LINEMOD_ORIG,'{}/mesh.ply'.format(class_type))
self.orig_old_model_path = os.path.join(cfg.LINEMOD_ORIG,'{}/OLDmesh.ply'.format(class_type))
self.transform_dat_path = os.path.join(cfg.LINEMOD_ORIG,'{}/transform.dat'.format(class_type))
self.R_p2w,self.t_p2w,self.s_p2w=self.setup_p2w_transform()
@staticmethod
def setup_p2w_transform():
transform1 = np.array([[0.161513626575, -0.827108919621, 0.538334608078, -0.245206743479],
[-0.986692547798, -0.124983474612, 0.104004733264, -0.050683632493],
[-0.018740313128, -0.547968924046, -0.836288750172, 0.387638419867]])
transform2 = np.array([[0.976471602917, 0.201606079936, -0.076541729271, -0.000718327821],
[-0.196746662259, 0.978194475174, 0.066531419754, 0.000077120210],
[0.088285841048, -0.049906700850, 0.994844079018, -0.001409600372]])
R1 = transform1[:, :3]
t1 = transform1[:, 3]
R2 = transform2[:, :3]
t2 = transform2[:, 3]
# printer system to world system
t_p2w = np.dot(R2, t1) + t2
R_p2w = np.dot(R2, R1)
s_p2w = 0.85
return R_p2w,t_p2w,s_p2w
def pose_p2w(self,RT):
t,R=RT[:,3],RT[:,:3]
R_w2c=np.dot(R, self.R_p2w.T)
t_w2c=-np.dot(R_w2c,self.t_p2w)+self.s_p2w*t
return np.concatenate([R_w2c,t_w2c[:,None]],1)
@staticmethod
def load_ply_model(model_path):
ply = PlyData.read(model_path)
data = ply.elements[0].data
x = data['x']
y = data['y']
z = data['z']
return np.stack([x, y, z], axis=-1)
def read_transform_dat(self):
transform_dat = np.loadtxt(self.transform_dat_path, skiprows=1)[:, 1]
transform_dat = np.reshape(transform_dat, newshape=[3, 4])
return transform_dat
def load_orig_model(self):
if os.path.exists(self.orig_model_path):
return self.load_ply_model(self.orig_model_path) / 1000.
else:
transform = self.read_transform_dat()
old_model = self.load_ply_model(self.orig_old_model_path) / 1000.
old_model = np.dot(old_model, transform[:, :3].T) + transform[:, 3]
return old_model
def get_translation_transform(self):
if self.class_type in self.translation_transforms:
return self.translation_transforms[self.class_type]
blender_model = self.load_ply_model(self.blender_model_path)
orig_model = self.load_orig_model()
blender_model = np.dot(blender_model, self.rotation_transform.T)
translation_transform = np.mean(orig_model, axis=0) - np.mean(blender_model, axis=0)
self.translation_transforms[self.class_type] = translation_transform
return translation_transform
def align_model(self):
blender_model = self.load_ply_model(self.blender_model_path)
orig_model = self.load_orig_model()
blender_model = np.dot(blender_model, self.rotation_transform.T)
blender_model += (np.mean(orig_model, axis=0) - np.mean(blender_model, axis=0))
np.savetxt(os.path.join(cfg.DATA_DIR, 'blender_model.txt'), blender_model)
np.savetxt(os.path.join(cfg.DATA_DIR, 'orig_model.txt'), orig_model)
def project_model(self, model, pose, camera_type):
camera_points_3d = np.dot(model, pose[:, :3].T) + pose[:, 3]
camera_points_3d = np.dot(camera_points_3d, self.intrinsic_matrix[camera_type].T)
return camera_points_3d[:, :2] / camera_points_3d[:, 2:]
def validate(self, idx):
model = self.load_ply_model(self.blender_model_path)
pose = read_pickle('/home/pengsida/Datasets/LINEMOD/renders/{}/{}_RT.pkl'.format(self.class_type, idx))['RT']
model_2d = self.project_model(model, pose, 'blender')
img = np.array(Image.open('/home/pengsida/Datasets/LINEMOD/renders/{}/{}.jpg'.format(self.class_type, idx)))
import matplotlib.pyplot as plt
plt.imshow(img)
plt.plot(model_2d[:, 0], model_2d[:, 1], 'r.')
plt.show()
class PoseTransformer(object):
rotation_transform = np.array([[1., 0., 0.],
[0., -1., 0.],
[0., 0., -1.]])
translation_transforms = {}
class_type_to_number = {
'ape': '001',
'can': '004',
'cat': '005',
'driller': '006',
'duck': '007',
'eggbox': '008',
'glue': '009',
'holepuncher': '010'
}
blender_models={}
def __init__(self, class_type):
self.class_type = class_type
self.blender_model_path = os.path.join(cfg.LINEMOD,'{}/{}.ply'.format(class_type, class_type))
self.orig_model_path = os.path.join(cfg.LINEMOD_ORIG,'{}/mesh.ply'.format(class_type))
self.xyz_pattern = os.path.join(cfg.OCCLUSION_LINEMOD,'models/{}/{}.xyz')
self.model_aligner = ModelAligner(class_type)
def orig_pose_to_blender_pose(self, pose):
rot, tra = pose[:, :3], pose[:, 3]
tra = tra + np.dot(rot, self.model_aligner.get_translation_transform())
rot = np.dot(rot, self.rotation_transform)
return np.concatenate([rot, np.reshape(tra, newshape=[3, 1])], axis=-1)
@staticmethod
def blender_pose_to_blender_euler(pose):
euler = [r / np.pi * 180 for r in mat2euler(pose, axes='szxz')]
euler[0] = -(euler[0] + 90) % 360
euler[1] = euler[1] - 90
return np.array(euler)
def orig_pose_to_blender_euler(self, pose):
blender_pose = self.orig_pose_to_blender_pose(pose)
return self.blender_pose_to_blender_euler(blender_pose)
@staticmethod
def load_ply_model(model_path):
ply = PlyData.read(model_path)
data = ply.elements[0].data
x = data['x']
y = data['y']
z = data['z']
return np.stack([x, y, z], axis=-1)
def get_blender_model(self):
if self.class_type in self.blender_models:
return self.blender_models[self.class_type]
blender_model = self.load_ply_model(self.blender_model_path.format(self.class_type, self.class_type))
self.blender_models[self.class_type] = blender_model
return blender_model
def get_translation_transform(self):
if self.class_type in self.translation_transforms:
return self.translation_transforms[self.class_type]
model = self.get_blender_model()
xyz = np.loadtxt(self.xyz_pattern.format(
self.class_type.title(), self.class_type_to_number[self.class_type]))
rotation = np.array([[0., 0., 1.],
[1., 0., 0.],
[0., 1., 0.]])
xyz = np.dot(xyz, rotation.T)
translation_transform = np.mean(xyz, axis=0) - np.mean(model, axis=0)
self.translation_transforms[self.class_type] = translation_transform
return translation_transform
def occlusion_pose_to_blender_pose(self, pose):
rot, tra = pose[:, :3], pose[:, 3]
rotation = np.array([[0., 1., 0.],
[0., 0., 1.],
[1., 0., 0.]])
rot = np.dot(rot, rotation)
tra[1:] *= -1
translation_transform = np.dot(rot, self.get_translation_transform())
rot[1:] *= -1
translation_transform[1:] *= -1
tra += translation_transform
pose = np.concatenate([rot, np.reshape(tra, newshape=[3, 1])], axis=-1)
return pose
class Projector(object):
intrinsic_matrix = {
'linemod': np.array([[572.4114, 0., 325.2611],
[0., 573.57043, 242.04899],
[0., 0., 1.]]),
'blender': np.array([[700., 0., 320.],
[0., 700., 240.],
[0., 0., 1.]]),
'pascal': np.asarray([[-3000.0, 0.0, 0.0],
[0.0, 3000.0, 0.0],
[0.0, 0.0, 1.0]])
}
def project(self,pts_3d,RT,K_type):
pts_2d=np.matmul(pts_3d,RT[:,:3].T)+RT[:,3:].T
pts_2d=np.matmul(pts_2d,self.intrinsic_matrix[K_type].T)
pts_2d=pts_2d[:,:2]/pts_2d[:,2:]
return pts_2d
def project_h(self,pts_3dh,RT,K_type):
'''
:param pts_3dh: [n,4]
:param RT: [3,4]
:param K_type:
:return: [n,3]
'''
K=self.intrinsic_matrix[K_type]
return np.matmul(np.matmul(pts_3dh,RT.transpose()),K.transpose())
def project_pascal(self,pts_3d,RT,principle):
'''
:param pts_3d: [n,3]
:param principle: [2,2]
:return:
'''
K=self.intrinsic_matrix['pascal'].copy()
K[:2,2]=principle
cam_3d=np.matmul(pts_3d,RT[:,:3].T)+RT[:,3:].T
cam_3d[np.abs(cam_3d[:,2])<1e-5,2]=1e-5 # revise depth
pts_2d=np.matmul(cam_3d,K.T)
pts_2d=pts_2d[:,:2]/pts_2d[:,2:]
return pts_2d, cam_3d
def project_pascal_h(self, pts_3dh,RT,principle):
K=self.intrinsic_matrix['pascal'].copy()
K[:2,2]=principle
return np.matmul(np.matmul(pts_3dh,RT.transpose()),K.transpose())
@staticmethod
def project_K(pts_3d,RT,K):
pts_2d=np.matmul(pts_3d,RT[:,:3].T)+RT[:,3:].T
pts_2d=np.matmul(pts_2d,K.T)
pts_2d=pts_2d[:,:2]/pts_2d[:,2:]
return pts_2d
def randomly_read_background():
background_lmdb = lmdb.open(os.path.join(cfg.YCB, 'background'), max_readers=1,
readonly=True, lock=False, readahead=False, meminit=False)
with background_lmdb.begin() as txn:
background_len = txn.stat()['entries']
background_idx=np.random.randint(0,background_len)
with background_lmdb.begin(write=False) as txn:
background=cv2.imdecode(np.fromstring(
txn.get('{:08d}'.format(background_idx).encode('ascii')),np.uint8),cv2.IMREAD_COLOR)
return background
def vertex_layer_reshape(vertex):
b,vn,h,w=vertex.shape
vertex=vertex.permute(0,2,3,1)
vn//=2
vertex=vertex.view(b,h,w,vn,2)
return vertex
def mask_depth_to_point_cloud(mask,depth,K):
ys, xs=np.nonzero(mask)
dpts=depth[ys,xs]
xs,ys=np.asarray(xs,np.float32),np.asarray(ys,np.float32)
xys=np.concatenate([xs[:,None],ys[:,None]],1)
xys*=dpts[:,None]
xyds=np.concatenate([xys,dpts[:,None]],1)
pts=np.matmul(xyds,np.linalg.inv(K).transpose())
return pts
def mask_depth_to_pts(mask,depth,K,output_2d=False):
hs,ws=np.nonzero(mask)
pts_2d=np.asarray([ws,hs],np.float32).transpose()
depth=depth[hs,ws]
pts=np.asarray([ws,hs,depth],np.float32).transpose()
pts[:,:2]*=pts[:,2:]
if output_2d:
return np.dot(pts,np.linalg.inv(K).transpose()), pts_2d
else:
return np.dot(pts,np.linalg.inv(K).transpose())