Change pose sampling from Euler to quaternion

lifters/robust_pose_lifter.py

@@ -7,6 +7,7 @@
 from poly_matrix.poly_matrix import PolyMatrix
 from utils.geometry import (
     get_C_r_from_theta,
+    get_C_r_from_unknown,
     get_C_r_from_xtheta,
     get_xtheta_from_C_r,
     get_xtheta_from_theta,
@@ -285,10 +286,8 @@ def get_cost(self, theta, y):
         else:
             x = theta
 
-        try:
-            R, t = get_C_r_from_theta(x, self.d)
-        except Exception:
-            R, t = get_C_r_from_xtheta(x, self.d)
+        # TODO(FD) below is a bit hacky, should be changed eventually.
+        R, t = get_C_r_from_unknown(x, self.d)
 
         cost = 0
         for i in range(self.n_landmarks):
@@ -366,10 +365,8 @@ def euclidean_gradient_unused(R, t):
         optimizer = Optimizer(**solver_kwargs)
 
         if self.robust:
-            try:
-                R_0, t_0 = get_C_r_from_theta(t0[: -self.n_landmarks], self.d)
-            except ValueError:
-                R_0, t_0 = get_C_r_from_xtheta(t0[: -self.n_landmarks], self.d)
+            x = t0[: -self.n_landmarks]
+            R_0, t_0 = get_C_r_from_unknown(x, self.d)
         else:
             R_0, t_0 = get_C_r_from_xtheta(t0[: self.d + self.d**2], self.d)
         res = optimizer.run(problem, initial_point=(R_0, t_0))

lifters/stereo3d_lifter.py

@@ -1,11 +1,7 @@
 import autograd.numpy as np
 
-from lifters.stereo_lifter import StereoLifter, NORMALIZE
-from utils.geometry import (
-    get_T,
-    get_xtheta_from_T,
-    get_xtheta_from_theta,
-)
+from lifters.stereo_lifter import NORMALIZE, StereoLifter
+from utils.geometry import get_T, get_xtheta_from_T, get_xtheta_from_theta
 
 
 def change_dimensions(a, y):
@@ -57,7 +53,7 @@ def get_cost(self, t, y, W=None):
 
         p_w, y = change_dimensions(a, y)
 
-        if len(t) == 6:
+        if len(t) == 7:
             xtheta = get_xtheta_from_theta(t, 3)
         else:
             xtheta = t

lifters/stereo_lifter.py

@@ -173,15 +173,18 @@ def get_x(self, theta=None, parameters=None, var_subset=None):
         if self.d == 2:
             if len(theta) == 3:
                 C, r = get_C_r_from_theta(theta, self.d)
-            elif len(theta) == 6:
+            elif len(theta) == 7:
                 C, r = get_C_r_from_xtheta(theta, self.d)
+            else:
+                raise ValueError(len(theta))
         elif self.d == 3:
-            if len(theta) == 6:
-                # theta is (x, y, C.flatten()), technically this should be called xtheta!
+            if len(theta) == 7:
                 C, r = get_C_r_from_theta(theta, self.d)
             elif len(theta) == 12:
                 # theta is (x, y, z, C.flatten()), technically this should be called xtheta!
                 C, r = get_C_r_from_xtheta(theta, self.d)
+            else:
+                raise ValueError(len(theta))
 
         if (self.param_level != "no") and (len(parameters) > 1):
             landmarks = np.array(parameters[1:]).reshape((self.n_landmarks, self.d))
@@ -255,25 +258,9 @@ def get_A_known(self, var_dict=None, output_poly=False):
             assert abs(u_z1 * (cy1 * p1[0] + cy2 * p1[1]) + u_z1 * ty - 1) < 1e-10
         elif self.d == 3:
             x = self.get_x()
-            (
-                _,
-                tx,
-                ty,
-                tz,
-                cx1,
-                cx2,
-                cx3,
-                cy1,
-                cy2,
-                cy3,
-                cz1,
-                cz2,
-                cz3,
-                u_x1,
-                u_y1,
-                u_z1,
-                *_,
-            ) = x
+            # fmt: off
+            (_, tx, ty, tz, cx1, cx2, cx3, cy1, cy2, cy3, cz1, cz2, cz3, u_x1, u_y1, u_z1, *_) = x
+            # fmt: on
             p1 = self.landmarks[0]
             assert (
                 abs(u_z1 * (cx1 * p1[0] + cx2 * p1[1] + cx3 * p1[2]) + u_z1 * tx - u_x1)

utils/geometry.py

@@ -1,22 +1,29 @@
+""" Tools related to geometry.
+
+Naming conventions:
+- C and r are rotation matrix and translation. T is the transformation matrix.
+- theta contains the translation and orientation (in quaternion form)
+- xtheta contains the translation and vectorzied rotation (i.e., r and vec(C))
+"""
+
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 
 
 def generate_random_pose(d=2, size=1):
     if d == 2:
-        n_angles = 1
+        rot = np.random.rand() * 2 * np.pi
     elif d == 3:
-        n_angles = 3
+        rot = R.random().as_quat()
     else:
         raise ValueError("d has to be 2 or 3.")
     trans = np.random.rand(d) * size - size / 2
-    rot = np.random.rand(n_angles) * 2 * np.pi
     return np.r_[trans, rot]
 
 
-def get_euler(C):
+def get_quat(C):
     r = R.from_matrix(C)
-    return r.as_euler("xyz")
+    return r.as_quat()
 
 
 def get_rot_matrix(rot):
@@ -26,8 +33,8 @@ def get_rot_matrix(rot):
             rot = float(rot[0])
         r = R.from_euler("z", rot)
         return r.as_matrix()[:2, :2]
-    elif len(rot) == 3:
-        r = R.from_euler("xyz", rot)
+    elif len(rot) == 4:
+        r = R.from_quat(rot)
         return r.as_matrix()
     else:
         raise ValueError(rot)
@@ -43,9 +50,25 @@ def get_C_r_from_theta(theta, d):
 def get_C_r_from_xtheta(xtheta, d):
     r = xtheta[:d]
     C = xtheta[d:].reshape((d, d))
+    np.testing.assert_allclose(C.T @ C, np.eye(d), atol=1e-10)
     return C, r
 
 
+def get_C_r_from_unknown(x, d):
+    """Chooses the correct function depending of dimension of x."""
+    if (d == 2) and (len(x) == 3):  # r, alpha
+        R, t = get_C_r_from_theta(x, d)
+    elif (d == 2) and (len(x) == 6):  # r, vec(C)
+        R, t = get_C_r_from_xtheta(x, d)
+    elif (d == 3) and (len(x) == 7):  # r, alpha
+        R, t = get_C_r_from_theta(x, d)
+    elif (d == 3) and (len(x) == 12):  # r, vec(C)
+        R, t = get_C_r_from_xtheta(x, d)
+    else:
+        raise ValueError(f"Inconsistent dimension: {len(x)}, {d}")
+    return R, t
+
+
 def get_xtheta_from_C_r(C, r):
     # row-wise flatten
     return np.r_[r, C.flatten("C")]
@@ -64,7 +87,7 @@ def get_T(xtheta=None, d=None, theta=None):
 
 
 def get_xtheta_from_theta(theta, d):
-    n_rot = d * (d - 1) // 2
+    n_rot = 1 if d == 2 else 4
     pos = theta[:d]
     alpha = theta[d : d + n_rot]
     C = get_rot_matrix(alpha)
@@ -74,13 +97,11 @@ def get_xtheta_from_theta(theta, d):
 
 
 def get_theta_from_xtheta(xtheta, d):
-    from utils.geometry import get_euler
-
     pos = xtheta[1 : d + 1]
     c = xtheta[d + 1 : d + 1 + d**2]
     C = c.reshape(d, d)
 
-    alpha = get_euler(C)
+    alpha = get_quat(C)
     theta = np.r_[pos, alpha]
     return theta