forked from cvg/pixloc
-
Notifications
You must be signed in to change notification settings - Fork 0
/
wrappers.py
365 lines (300 loc) · 11.5 KB
/
wrappers.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
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
"""
Convenience classes for an SE3 pose and a pinhole Camera with lens distortion.
Based on PyTorch tensors: differentiable, batched, with GPU support.
"""
import functools
import inspect
import math
from typing import Union, Tuple, List, Dict, NamedTuple
import torch
import numpy as np
from .optimization import skew_symmetric, so3exp_map
from .utils import undistort_points, J_undistort_points
def autocast(func):
"""Cast the inputs of a TensorWrapper method to PyTorch tensors
if they are numpy arrays. Use the device and dtype of the wrapper.
"""
@functools.wraps(func)
def wrap(self, *args):
device = torch.device('cpu')
dtype = None
if isinstance(self, TensorWrapper):
if self._data is not None:
device = self.device
dtype = self.dtype
elif not inspect.isclass(self) or not issubclass(self, TensorWrapper):
raise ValueError(self)
cast_args = []
for arg in args:
if isinstance(arg, np.ndarray):
arg = torch.from_numpy(arg)
arg = arg.to(device=device, dtype=dtype)
cast_args.append(arg)
return func(self, *cast_args)
return wrap
class TensorWrapper:
_data = None
@autocast
def __init__(self, data: torch.Tensor):
self._data = data
@property
def shape(self):
return self._data.shape[:-1]
@property
def device(self):
return self._data.device
@property
def dtype(self):
return self._data.dtype
def __getitem__(self, index):
return self.__class__(self._data[index])
def __setitem__(self, index, item):
self._data[index] = item.data
def to(self, *args, **kwargs):
return self.__class__(self._data.to(*args, **kwargs))
def cpu(self):
return self.__class__(self._data.cpu())
def cuda(self):
return self.__class__(self._data.cuda())
def pin_memory(self):
return self.__class__(self._data.pin_memory())
def float(self):
return self.__class__(self._data.float())
def double(self):
return self.__class__(self._data.double())
def detach(self):
return self.__class__(self._data.detach())
@classmethod
def stack(cls, objects: List, dim=0, *, out=None):
data = torch.stack([obj._data for obj in objects], dim=dim, out=out)
return cls(data)
def __torch_function__(self, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
if func is torch.stack:
return self.stack(*args, **kwargs)
else:
return NotImplemented
class Pose(TensorWrapper):
def __init__(self, data: torch.Tensor):
assert data.shape[-1] == 12
super().__init__(data)
@classmethod
@autocast
def from_Rt(cls, R: torch.Tensor, t: torch.Tensor):
'''Pose from a rotation matrix and translation vector.
Accepts numpy arrays or PyTorch tensors.
Args:
R: rotation matrix with shape (..., 3, 3).
t: translation vector with shape (..., 3).
'''
assert R.shape[-2:] == (3, 3)
assert t.shape[-1] == 3
assert R.shape[:-2] == t.shape[:-1]
data = torch.cat([R.flatten(start_dim=-2), t], -1)
return cls(data)
@classmethod
@autocast
def from_aa(cls, aa: torch.Tensor, t: torch.Tensor):
'''Pose from an axis-angle rotation vector and translation vector.
Accepts numpy arrays or PyTorch tensors.
Args:
aa: axis-angle rotation vector with shape (..., 3).
t: translation vector with shape (..., 3).
'''
assert aa.shape[-1] == 3
assert t.shape[-1] == 3
assert aa.shape[:-1] == t.shape[:-1]
return cls.from_Rt(so3exp_map(aa), t)
@classmethod
def from_4x4mat(cls, T: torch.Tensor):
'''Pose from an SE(3) transformation matrix.
Args:
T: transformation matrix with shape (..., 4, 4).
'''
assert T.shape[-2:] == (4, 4)
R, t = T[..., :3, :3], T[..., :3, 3]
return cls.from_Rt(R, t)
@classmethod
def from_colmap(cls, image: NamedTuple):
'''Pose from a COLMAP Image.'''
return cls.from_Rt(image.qvec2rotmat(), image.tvec)
@property
def R(self) -> torch.Tensor:
'''Underlying rotation matrix with shape (..., 3, 3).'''
rvec = self._data[..., :9]
return rvec.reshape(rvec.shape[:-1]+(3, 3))
@property
def t(self) -> torch.Tensor:
'''Underlying translation vector with shape (..., 3).'''
return self._data[..., -3:]
def inv(self) -> 'Pose':
'''Invert an SE(3) pose.'''
R = self.R.transpose(-1, -2)
t = -(R @ self.t.unsqueeze(-1)).squeeze(-1)
return self.__class__.from_Rt(R, t)
def compose(self, other: 'Pose') -> 'Pose':
'''Chain two SE(3) poses: T_B2C.compose(T_A2B) -> T_A2C.'''
R = self.R @ other.R
t = self.t + (self.R @ other.t.unsqueeze(-1)).squeeze(-1)
return self.__class__.from_Rt(R, t)
@autocast
def transform(self, p3d: torch.Tensor) -> torch.Tensor:
'''Transform a set of 3D points.
Args:
p3d: 3D points, numpy array or PyTorch tensor with shape (..., 3).
'''
assert p3d.shape[-1] == 3
# assert p3d.shape[:-2] == self.shape # allow broadcasting
return p3d @ self.R.transpose(-1, -2) + self.t.unsqueeze(-2)
def __mul__(self, p3D: torch.Tensor) -> torch.Tensor:
'''Transform a set of 3D points: T_A2B * p3D_A -> p3D_B.'''
return self.transform(p3D)
def __matmul__(self, other: 'Pose') -> 'Pose':
'''Chain two SE(3) poses: T_B2C @ T_A2B -> T_A2C.'''
return self.compose(other)
@autocast
def J_transform(self, p3d_out: torch.Tensor):
# [[1,0,0,0,-pz,py],
# [0,1,0,pz,0,-px],
# [0,0,1,-py,px,0]]
J_t = torch.diag_embed(torch.ones_like(p3d_out))
J_rot = -skew_symmetric(p3d_out)
J = torch.cat([J_t, J_rot], dim=-1)
return J # N x 3 x 6
def numpy(self) -> Tuple[np.ndarray]:
return self.R.numpy(), self.t.numpy()
def magnitude(self) -> Tuple[torch.Tensor]:
'''Magnitude of the SE(3) transformation.
Returns:
dr: rotation anngle in degrees.
dt: translation distance in meters.
'''
trace = torch.diagonal(self.R, dim1=-1, dim2=-2).sum(-1)
cos = torch.clamp((trace - 1) / 2, -1, 1)
dr = torch.acos(cos).abs() / math.pi * 180
dt = torch.norm(self.t, dim=-1)
return dr, dt
def __repr__(self):
return f'Pose: {self.shape} {self.dtype} {self.device}'
class Camera(TensorWrapper):
eps = 1e-3
def __init__(self, data: torch.Tensor):
assert data.shape[-1] in {6, 8, 10}
super().__init__(data)
@classmethod
def from_colmap(cls, camera: Union[Dict, NamedTuple]):
'''Camera from a COLMAP Camera tuple or dictionary.
We assume that the origin (0, 0) is the center of the top-left pixel.
This is different from COLMAP.
'''
if isinstance(camera, tuple):
camera = camera._asdict()
model = camera['model']
params = camera['params']
if model in ['OPENCV', 'PINHOLE']:
(fx, fy, cx, cy), params = np.split(params, [4])
elif model in ['SIMPLE_PINHOLE', 'SIMPLE_RADIAL', 'RADIAL']:
(f, cx, cy), params = np.split(params, [3])
fx = fy = f
if model == 'SIMPLE_RADIAL':
params = np.r_[params, 0.]
else:
raise NotImplementedError(model)
data = np.r_[camera['width'], camera['height'],
fx, fy, cx-0.5, cy-0.5, params]
return cls(data)
@property
def size(self) -> torch.Tensor:
'''Size (width height) of the images, with shape (..., 2).'''
return self._data[..., :2]
@property
def f(self) -> torch.Tensor:
'''Focal lengths (fx, fy) with shape (..., 2).'''
return self._data[..., 2:4]
@property
def c(self) -> torch.Tensor:
'''Principal points (cx, cy) with shape (..., 2).'''
return self._data[..., 4:6]
@property
def dist(self) -> torch.Tensor:
'''Distortion parameters, with shape (..., {0, 2, 4}).'''
return self._data[..., 6:]
def scale(self, scales: Union[float, int, Tuple[Union[float, int]]]):
'''Update the camera parameters after resizing an image.'''
if isinstance(scales, (int, float)):
scales = (scales, scales)
s = self._data.new_tensor(scales)
data = torch.cat([
self.size*s,
self.f*s,
(self.c+0.5)*s-0.5,
self.dist], -1)
return self.__class__(data)
def crop(self, left_top: Tuple[float], size: Tuple[int]):
'''Update the camera parameters after cropping an image.'''
left_top = self._data.new_tensor(left_top)
size = self._data.new_tensor(size)
data = torch.cat([
size,
self.f,
self.c - left_top,
self.dist], -1)
return self.__class__(data)
@autocast
def in_image(self, p2d: torch.Tensor):
'''Check if 2D points are within the image boundaries.'''
assert p2d.shape[-1] == 2
# assert p2d.shape[:-2] == self.shape # allow broadcasting
size = self.size.unsqueeze(-2)
valid = torch.all((p2d >= 0) & (p2d <= (size - 1)), -1)
return valid
@autocast
def project(self, p3d: torch.Tensor) -> Tuple[torch.Tensor]:
'''Project 3D points into the camera plane and check for visibility.'''
z = p3d[..., -1]
valid = z > self.eps
z = z.clamp(min=self.eps)
p2d = p3d[..., :-1] / z.unsqueeze(-1)
return p2d, valid
def J_project(self, p3d: torch.Tensor):
x, y, z = p3d[..., 0], p3d[..., 1], p3d[..., 2]
zero = torch.zeros_like(z)
z = z.clamp(min=self.eps)
J = torch.stack([
1/z, zero, -x / z**2,
zero, 1/z, -y / z**2], dim=-1)
J = J.reshape(p3d.shape[:-1]+(2, 3))
return J # N x 2 x 3
@autocast
def undistort(self, pts: torch.Tensor) -> Tuple[torch.Tensor]:
'''Undistort normalized 2D coordinates
and check for validity of the distortion model.
'''
assert pts.shape[-1] == 2
# assert pts.shape[:-2] == self.shape # allow broadcasting
return undistort_points(pts, self.dist)
def J_undistort(self, pts: torch.Tensor):
return J_undistort_points(pts, self.dist) # N x 2 x 2
@autocast
def denormalize(self, p2d: torch.Tensor) -> torch.Tensor:
'''Convert normalized 2D coordinates into pixel coordinates.'''
return p2d * self.f.unsqueeze(-2) + self.c.unsqueeze(-2)
def J_denormalize(self):
return torch.diag_embed(self.f).unsqueeze(-3) # 1 x 2 x 2
@autocast
def world2image(self, p3d: torch.Tensor) -> Tuple[torch.Tensor]:
'''Transform 3D points into 2D pixel coordinates.'''
p2d, visible = self.project(p3d)
p2d, mask = self.undistort(p2d)
p2d = self.denormalize(p2d)
valid = visible & mask & self.in_image(p2d)
return p2d, valid
def J_world2image(self, p3d: torch.Tensor):
p2d_dist, valid = self.project(p3d)
J = (self.J_denormalize()
@ self.J_undistort(p2d_dist)
@ self.J_project(p3d))
return J, valid
def __repr__(self):
return f'Camera {self.shape} {self.dtype} {self.device}'