init commit

workflows/simbox/solver/kpam/SE3_utils.py

@@ -0,0 +1,131 @@
+import numpy as np
+
+
+def transform_point_cloud(T_homo, pc):
+    # type: (np.ndarray, np.ndarray) -> np.ndarray
+    """
+    Transform a point cloud using homogeous transform
+    :param T_homo: 4x4 transformation
+    :param pc: (N, 3) point cloud
+    :return: (N, 3) point cloud
+    """
+    transformed_pc = T_homo[0:3, 0:3].dot(pc.T)
+    transformed_pc = transformed_pc.T
+    transformed_pc[:, 0] += T_homo[0, 3]
+    transformed_pc[:, 1] += T_homo[1, 3]
+    transformed_pc[:, 2] += T_homo[2, 3]
+    return transformed_pc
+
+
+def xyzrpy_to_matrix(xyzrpy):
+    """
+    Create 4x4 homogeneous transform matrix from pos and rpy
+    """
+    xyz = xyzrpy[0:3]
+    rpy = xyzrpy[3:6]
+
+    T = np.zeros([4, 4], dtype=xyzrpy.dtype)
+    T[0:3, 0:3] = rpy_to_rotation_matrix(rpy)
+    T[3, 3] = 1
+    T[0:3, 3] = xyz
+    return T
+
+
+def rpy_to_rotation_matrix(rpy):
+    """
+    Creates 3x3 rotation matrix from rpy
+    See http://danceswithcode.net/engineeringnotes/rotations_in_3d/rotations_in_3d_part1.html
+    """
+    u = rpy[0]
+    v = rpy[1]
+    w = rpy[2]
+
+    R = np.zeros([3, 3], dtype=rpy.dtype)
+
+    # first row
+    R[0, 0] = np.cos(v) * np.cos(w)
+    R[0, 1] = np.sin(u) * np.sin(v) * np.cos(w) - np.cos(u) * np.sin(w)
+    R[0, 2] = np.sin(u) * np.sin(w) + np.cos(u) * np.sin(v) * np.cos(w)
+
+    # second row
+    R[1, 0] = np.cos(v) * np.sin(w)
+    R[1, 1] = np.cos(u) * np.cos(w) + np.sin(u) * np.sin(v) * np.sin(w)
+    R[1, 2] = np.cos(u) * np.sin(v) * np.sin(w) - np.sin(u) * np.cos(w)
+
+    # third row
+    R[2, 0] = -np.sin(v)
+    R[2, 1] = np.sin(u) * np.cos(v)
+    R[2, 2] = np.cos(u) * np.cos(v)
+
+    return R
+
+
+def transform_point(T, p):
+    """
+    Transform a point via a homogeneous transform matrix T
+
+    :param: T 4x4 numpy array
+    """
+
+    p_homog = np.concatenate((p, np.array([1])))
+    q = np.dot(T, p_homog)
+    q = q.squeeze()
+    q = q[0:3]
+    return q
+
+
+def transform_vec(T, v):
+    """
+    Transform a vector via a homogeneous transform matrix T
+
+    :param: T 4x4 numpy array
+    """
+    v = np.array(v)
+    R = T[:3, :3]
+    u = np.dot(R, v)
+    return u
+
+
+def xyzrpy_to_matrix_symbolic(xyzrpy):
+    """
+    Create 4x4 homogeneous transform matrix from pos and rpy
+    """
+    xyz = xyzrpy[0:3]
+    rpy = xyzrpy[3:6]
+
+    T = np.zeros([4, 4], dtype=xyzrpy.dtype)
+    T[0:3, 0:3] = rpy_to_rotation_matrix_symbolic(rpy)
+    T[3, 3] = 1
+    T[0:3, 3] = xyz
+    return T
+
+
+def rpy_to_rotation_matrix_symbolic(rpy):
+    """
+    Creates 3x3 rotation matrix from rpy
+    See http://danceswithcode.net/engineeringnotes/rotations_in_3d/rotations_in_3d_part1.html
+    """
+    from pydrake.math import cos, sin
+
+    u = rpy[0]
+    v = rpy[1]
+    w = rpy[2]
+
+    R = np.zeros([3, 3], dtype=rpy.dtype)
+
+    # first row
+    R[0, 0] = cos(v) * cos(w)
+    R[0, 1] = sin(u) * sin(v) * cos(w) - cos(u) * sin(w)
+    R[0, 2] = sin(u) * sin(w) + cos(u) * sin(v) * cos(w)
+
+    # second row
+    R[1, 0] = cos(v) * sin(w)
+    R[1, 1] = cos(u) * cos(w) + sin(u) * sin(v) * sin(w)
+    R[1, 2] = cos(u) * sin(v) * sin(w) - sin(u) * cos(w)
+
+    # third row
+    R[2, 0] = -sin(v)
+    R[2, 1] = sin(u) * cos(v)
+    R[2, 2] = cos(u) * cos(v)
+
+    return R