google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
import numpy as np
17

18

19
def generate_camera_noise(num_poses=100, noise_std=0.15):
20
  """Generate camera noise."""
21
  se3_noise = np.random.randn(num_poses, 6) * noise_std
22
  return se3_noise
23

24

25
def add_camera_noise(poses, se3_noise):
26
  """Add camera noise."""
27
  num_poses = len(poses)
28
  se3_noise = se3_exp(se3_noise[:num_poses])
29
  noisy_poses = compose([se3_noise, poses[:, :3]])
30
  new_poses = np.copy(poses)
31
  new_poses[:, :3] = noisy_poses
32
  return new_poses
33

34

35
def define_pose(rot=None, trans=None):
36
  """Construct pose matrix."""
37
  assert(rot is not None or trans is not None)
38
  if rot is None:
39
    rot = np.repeat(np.eye(3)[None, Ellipsis], repeats=trans.shape[0], axis=0)
40
  if trans is None:
41
    trans = np.zeros(rot.shape[:-1])  # Bx3
42
  pose = np.concatenate([rot, trans[Ellipsis, None]], axis=-1)  # [..., 3, 4]
43
  assert pose.shape[-2:] == (3, 4)
44
  return pose
45

46

47
def invert_pose(pose):
48
  """Invert pose."""
49
  # invert a camera pose
50
  rot, trans = pose[Ellipsis, :3], pose[Ellipsis, 3:]
51
  r_inv = np.swapaxes(rot, -1, -2)
52
  t_inv = (-r_inv@trans)[Ellipsis, 0]
53
  pose_inv = define_pose(rot=r_inv, trans=t_inv)
54
  return pose_inv
55

56

57
def compose(pose_list):
58
  """Compose sequence of poses."""
59
  # compose a sequence of poses together
60
  # pose_new(x) = poseN o ... o pose2 o pose1(x)
61
  pose_new = pose_list[0]
62
  for pose in pose_list[1:]:
63
    pose_new = compose_pair(pose_new, pose)
64
  return pose_new
65

66

67
def compose_pair(pose_a, pose_b):
68
  """Compose poses."""
69
  # pose_new(x) = pose_b o pose_a(x)
70
  r_a, t_a = pose_a[Ellipsis, :3], pose_a[Ellipsis, 3:]
71
  r_b, t_b = pose_b[Ellipsis, :3], pose_b[Ellipsis, 3:]
72
  r_new = r_b@r_a
73
  t_new = (r_b@t_a+t_b)[Ellipsis, 0]
74
  pose_new = define_pose(rot=r_new, trans=t_new)
75
  return pose_new
76

77

78
def se3_exp(wu):  # [..., 3]
79
  """Exponential map."""
80
  w_vec, u_vec = np.split(wu, 2, axis=1)
81
  wx = skew_symmetric(w_vec)
82
  theta = np.linalg.norm(w_vec, axis=-1)[Ellipsis, None, None]
83
  i_mat = np.eye(3)
84
  a_mat = taylor_a(theta)
85
  b_mat = taylor_b(theta)
86
  c_mat = taylor_c(theta)
87
  r_mat = i_mat+a_mat*wx+b_mat*wx@wx
88
  v_mat = i_mat+b_mat*wx+c_mat*wx@wx
89
  rt = np.concatenate([r_mat, (v_mat@u_vec[Ellipsis, None])], axis=-1)  # Bx3x4
90
  return rt
91

92

93
def skew_symmetric(w_vec):
94
  """Skey symmetric matrix from vectors."""
95
  w0, w1, w2 = np.split(w_vec, 3, axis=-1)
96
  w0 = w0.squeeze()
97
  w1 = w1.squeeze()
98
  w2 = w2.squeeze()
99
  zeros = np.zeros(w0.shape, np.float32)
100
  wx = np.stack([np.stack([zeros, -w2, w1], axis=-1),
101
                 np.stack([w2, zeros, -w0], axis=-1),
102
                 np.stack([-w1, w0, zeros], axis=-1)], axis=-2)
103
  return wx
104

105

106
def taylor_a(x_val, nth=10):
107
  """Taylor expansion."""
108
  # Taylor expansion of sin(x)/x
109
  ans = np.zeros_like(x_val)
110
  denom = 1.
111
  for i in range(nth+1):
112
    if i > 0:
113
      denom *= (2*i)*(2*i+1)
114
    ans = ans+(-1)**i*x_val**(2*i)/denom
115
  return ans
116

117

118
def taylor_b(x_val, nth=10):
119
  """Taylor expansion."""
120
  # Taylor expansion of (1-cos(x))/x**2
121
  ans = np.zeros_like(x_val)
122
  denom = 1.
123
  for i in range(nth+1):
124
    denom *= (2*i+1)*(2*i+2)
125
    ans = ans+(-1)**i*x_val**(2*i)/denom
126
  return ans
127

128

129
def taylor_c(x_val, nth=10):
130
  """Taylor expansion."""
131
  # Taylor expansion of (x-sin(x))/x**3
132
  ans = np.zeros_like(x_val)
133
  denom = 1.
134
  for i in range(nth+1):
135
    denom *= (2*i+2)*(2*i+3)
136
    ans = ans+(-1)**i*x_val**(2*i)/denom
137
  return ans
138

139

140
def to_hom(x_mat):
141
  """Homogenous coordinates."""
142
  # get homogeneous coordinates of the input
143
  x_hom = np.concatenate([x_mat, np.ones_like(x_mat[Ellipsis, :1])], axis=-1)
144
  return x_hom
145

146

147
def world2cam(x_mat, pose):  # [B, N, 3]
148
  """Coordinate transformations."""
149
  x_hom = to_hom(x_mat)
150
  return x_hom@np.swapaxes(pose, -1, -2)
151

152

153
def cam2img(x_mat, cam_intr):
154
  """Coordinate transformations."""
155
  return x_mat@np.swapaxes(cam_intr, -1, -2)
156

157

158
def img2cam(x_mat, cam_intr):
159
  """Coordinate transformations."""
160
  return x_mat@np.swapaxes(np.linalg.inv(cam_intr), -1, -2)
161

162

163
def cam2world(x_mat, pose):
164
  """Coordinate transformations."""
165
  x_hom = to_hom(x_mat)
166
  pose_inv = invert_pose(pose)
167
  return x_hom@np.swapaxes(pose_inv, -1, -2)
168

169

170
def angle_to_rotation_matrix(angle, axis):
171
  """Angle to rotation matrix."""
172
  # get the rotation matrix from Euler angle around specific axis
173
  roll = dict(X=1, Y=2, Z=0)[axis]
174
  zeros = np.zeros_like(angle)
175
  eye = np.ones_like(angle)
176
  mat = np.stack([np.stack([angle.cos(), -angle.sin(), zeros], axis=-1),
177
                  np.stack([angle.sin(), angle.cos(), zeros], axis=-1),
178
                  np.stack([zeros, zeros, eye], axis=-1)], axis=-2)
179
  mat = mat.roll((roll, roll), axis=(-2, -1))
180
  return mat
181

182

183
def rotation_distance(r1, r2, eps=1e-7):
184
  """Rotation distance."""
185
  # http://www.boris-belousov.net/2016/12/01/quat-dist/
186
  r_diff = r1@np.swapaxes(r2, -2, -1)
187
  trace = r_diff[Ellipsis, 0, 0]+r_diff[Ellipsis, 1, 1]+r_diff[Ellipsis, 2, 2]
188
  angle = np.arccos(((trace-1)/2).clip(-1+eps, 1-eps))
189
  return angle
190

191

192
def procrustes_analysis(x0, x1):  # [N, 3]
193
  """Procrustes."""
194
  # translation
195
  t0 = x0.mean(axis=0, keepdims=True)
196
  t1 = x1.mean(axis=0, keepdims=True)
197
  x0c = x0-t0
198
  x1c = x1-t1
199
  # scale
200
  s0 = np.sqrt((x0c**2).sum(axis=-1).mean())
201
  s1 = np.sqrt((x1c**2).sum(axis=-1).mean())
202
  x0cs = x0c/s0
203
  x1cs = x1c/s1
204
  # rotation (use double for SVD, float loses precision)
205
  # TODO do we need float64?
206
  u_mat, _, v_mat = np.linalg.svd((x0cs.transpose()@x1cs).astype(np.float64),
207
                                  full_matrices=False)
208
  rot = (u_mat@v_mat).astype(np.float32)
209
  if np.linalg.det(rot) < 0:
210
    rot[2] *= -1
211
  # align X1 to X0: X1to0 = (X1-t1)/s1@R.t()*s0+t0
212
  sim3 = {"t0": t0[0], "t1": t1[0], "s0": s0, "s1": s1, "R": rot}
213
  return sim3
214

215

216
def prealign_cameras(poses_pred, poses_gt):
217
  """Prealign cameras."""
218
  center = np.zeros((1, 1, 3))
219
  centers_pred = cam2world(center, poses_pred)[:, 0]
220
  centers_gt = cam2world(center, poses_gt)[:, 0]
221
  sim3 = procrustes_analysis(centers_gt, centers_pred)
222
  centers_aligned = (centers_pred-sim3["t1"])/(
223
      sim3["s1"]@sim3["R"].transpose()*sim3["s0"]+sim3["t0"])
224
  r_aligned = poses_pred[Ellipsis, :3]@sim3["R"].transpose()
225
  t_aligned = (-r_aligned@centers_aligned[Ellipsis, None])[Ellipsis, 0]
226
  poses_aligned = define_pose(rot=r_aligned, trans=t_aligned)
227
  return poses_aligned, sim3
228

229

230
def refine_test_cameras(poses_init, sim3):
231
  """Refine test cameras."""
232
  center = np.zeros((1, 1, 3))
233
  centers = cam2world(center, poses_init)[:, 0]
234
  centers_aligned = (
235
      centers-sim3["t0"])/sim3["s0"]@sim3["R"]*sim3["s1"]+sim3["t1"]
236
  r_aligned = poses_init[Ellipsis, :3]@sim3["R"]
237
  t_aligned = (-r_aligned@centers_aligned[Ellipsis, None])[Ellipsis, 0]
238
  poses_aligned = define_pose(rot=r_aligned, trans=t_aligned)
239
  return poses_aligned
240

241

242
def evaluate_camera(pose, pose_gt):
243
  """Evaluate cameras."""
244
  # evaluate the distance between pose and pose_ref; pose_ref is fixed
245
  pose_aligned, _ = prealign_cameras(pose, pose_gt)
246
  r_aligned, t_aligned = np.split(pose_aligned, [3,], axis=-1)
247
  r_gt, t_gt = np.split(pose_gt, [3,], axis=-1)
248
  r_error = rotation_distance(r_aligned, r_gt)
249
  t_error = np.linalg.norm((t_aligned-t_gt)[Ellipsis, 0], axis=-1)
250
  return r_error, t_error
251

252

253
def evaluate_aligned_camera(pose_aligned, pose_gt):
254
  """Evaluate aligned cameras."""
255
  # evaluate the distance between pose and pose_ref; pose_ref is fixed
256
  r_aligned, t_aligned = np.split(pose_aligned, [3,], axis=-1)
257
  r_gt, t_gt = np.split(pose_gt, [3,], axis=-1)
258
  r_error = rotation_distance(r_aligned, r_gt)
259
  t_error = np.linalg.norm((t_aligned-t_gt)[Ellipsis, 0], axis=-1)
260
  return r_error, t_error
261