Spaces:
Running on Zero
Running on Zero
multimodalart HF Staff
Learn2Splat ZeroGPU demo: app, vendored optgs, prebuilt sm_120 wheels
9196862 verified | # Author: True Price <jtprice at cs.unc.edu> | |
| import numpy as np | |
| #------------------------------------------------------------------------------- | |
| # | |
| # Axis-Angle Functions | |
| # | |
| #------------------------------------------------------------------------------- | |
| # returns the cross product matrix representation of a 3-vector v | |
| def cross_prod_matrix(v): | |
| return np.array(((0., -v[2], v[1]), (v[2], 0., -v[0]), (-v[1], v[0], 0.))) | |
| #------------------------------------------------------------------------------- | |
| # www.euclideanspace.com/maths/geometry/rotations/conversions/angleToMatrix/ | |
| # if angle is None, assume ||axis|| == angle, in radians | |
| # if angle is not None, assume that axis is a unit vector | |
| def axis_angle_to_rotation_matrix(axis, angle=None): | |
| if angle is None: | |
| angle = np.linalg.norm(axis) | |
| if np.abs(angle) > np.finfo('float').eps: | |
| axis = axis / angle | |
| cp_axis = cross_prod_matrix(axis) | |
| return np.eye(3) + ( | |
| np.sin(angle) * cp_axis + (1. - np.cos(angle)) * cp_axis.dot(cp_axis)) | |
| #------------------------------------------------------------------------------- | |
| # after some deliberation, I've decided the easiest way to do this is to use | |
| # quaternions as an intermediary | |
| def rotation_matrix_to_axis_angle(R): | |
| return Quaternion.FromR(R).ToAxisAngle() | |
| #------------------------------------------------------------------------------- | |
| # | |
| # Quaternion | |
| # | |
| #------------------------------------------------------------------------------- | |
| class Quaternion: | |
| # create a quaternion from an existing rotation matrix | |
| # euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/ | |
| def FromR(R): | |
| trace = np.trace(R) | |
| if trace > 0: | |
| qw = 0.5 * np.sqrt(1. + trace) | |
| qx = (R[2,1] - R[1,2]) * 0.25 / qw | |
| qy = (R[0,2] - R[2,0]) * 0.25 / qw | |
| qz = (R[1,0] - R[0,1]) * 0.25 / qw | |
| elif R[0,0] > R[1,1] and R[0,0] > R[2,2]: | |
| s = 2. * np.sqrt(1. + R[0,0] - R[1,1] - R[2,2]) | |
| qw = (R[2,1] - R[1,2]) / s | |
| qx = 0.25 * s | |
| qy = (R[0,1] + R[1,0]) / s | |
| qz = (R[0,2] + R[2,0]) / s | |
| elif R[1,1] > R[2,2]: | |
| s = 2. * np.sqrt(1. + R[1,1] - R[0,0] - R[2,2]) | |
| qw = (R[0,2] - R[2,0]) / s | |
| qx = (R[0,1] + R[1,0]) / s | |
| qy = 0.25 * s | |
| qz = (R[1,2] + R[2,1]) / s | |
| else: | |
| s = 2. * np.sqrt(1. + R[2,2] - R[0,0] - R[1,1]) | |
| qw = (R[1,0] - R[0,1]) / s | |
| qx = (R[0,2] + R[2,0]) / s | |
| qy = (R[1,2] + R[2,1]) / s | |
| qz = 0.25 * s | |
| return Quaternion(np.array((qw, qx, qy, qz))) | |
| # if angle is None, assume ||axis|| == angle, in radians | |
| # if angle is not None, assume that axis is a unit vector | |
| def FromAxisAngle(axis, angle=None): | |
| if angle is None: | |
| angle = np.linalg.norm(axis) | |
| if np.abs(angle) > np.finfo('float').eps: | |
| axis = axis / angle | |
| qw = np.cos(0.5 * angle) | |
| axis = axis * np.sin(0.5 * angle) | |
| return Quaternion(np.array((qw, axis[0], axis[1], axis[2]))) | |
| #--------------------------------------------------------------------------- | |
| def __init__(self, q=np.array((1., 0., 0., 0.))): | |
| if isinstance(q, Quaternion): | |
| self.q = q.q.copy() | |
| else: | |
| q = np.asarray(q) | |
| if q.size == 4: | |
| self.q = q.copy() | |
| elif q.size == 3: # convert from a 3-vector to a quaternion | |
| self.q = np.empty(4) | |
| self.q[0], self.q[1:] = 0., q.ravel() | |
| else: | |
| raise Exception('Input quaternion should be a 3- or 4-vector') | |
| def __add__(self, other): | |
| return Quaternion(self.q + other.q) | |
| def __iadd__(self, other): | |
| self.q += other.q | |
| return self | |
| # conjugation via the ~ operator | |
| def __invert__(self): | |
| return Quaternion( | |
| np.array((self.q[0], -self.q[1], -self.q[2], -self.q[3]))) | |
| # returns: self.q * other.q if other is a Quaternion; otherwise performs | |
| # scalar multiplication | |
| def __mul__(self, other): | |
| if isinstance(other, Quaternion): # quaternion multiplication | |
| return Quaternion(np.array(( | |
| self.q[0] * other.q[0] - self.q[1] * other.q[1] - | |
| self.q[2] * other.q[2] - self.q[3] * other.q[3], | |
| self.q[0] * other.q[1] + self.q[1] * other.q[0] + | |
| self.q[2] * other.q[3] - self.q[3] * other.q[2], | |
| self.q[0] * other.q[2] - self.q[1] * other.q[3] + | |
| self.q[2] * other.q[0] + self.q[3] * other.q[1], | |
| self.q[0] * other.q[3] + self.q[1] * other.q[2] - | |
| self.q[2] * other.q[1] + self.q[3] * other.q[0]))) | |
| else: # scalar multiplication (assumed) | |
| return Quaternion(other * self.q) | |
| def __rmul__(self, other): | |
| return self * other | |
| def __imul__(self, other): | |
| self.q[:] = (self * other).q | |
| return self | |
| def __irmul__(self, other): | |
| self.q[:] = (self * other).q | |
| return self | |
| def __neg__(self): | |
| return Quaternion(-self.q) | |
| def __sub__(self, other): | |
| return Quaternion(self.q - other.q) | |
| def __isub__(self, other): | |
| self.q -= other.q | |
| return self | |
| def __str__(self): | |
| return str(self.q) | |
| def copy(self): | |
| return Quaternion(self) | |
| def dot(self, other): | |
| return self.q.dot(other.q) | |
| # assume the quaternion is nonzero! | |
| def inverse(self): | |
| return Quaternion((~self).q / self.q.dot(self.q)) | |
| def norm(self): | |
| return np.linalg.norm(self.q) | |
| def normalize(self): | |
| self.q /= np.linalg.norm(self.q) | |
| return self | |
| # assume x is a Nx3 numpy array or a numpy 3-vector | |
| def rotate_points(self, x): | |
| x = np.atleast_2d(x) | |
| return x.dot(self.ToR().T) | |
| # convert to a rotation matrix | |
| def ToR(self): | |
| return np.eye(3) + 2 * np.array(( | |
| (-self.q[2] * self.q[2] - self.q[3] * self.q[3], | |
| self.q[1] * self.q[2] - self.q[3] * self.q[0], | |
| self.q[1] * self.q[3] + self.q[2] * self.q[0]), | |
| ( self.q[1] * self.q[2] + self.q[3] * self.q[0], | |
| -self.q[1] * self.q[1] - self.q[3] * self.q[3], | |
| self.q[2] * self.q[3] - self.q[1] * self.q[0]), | |
| ( self.q[1] * self.q[3] - self.q[2] * self.q[0], | |
| self.q[2] * self.q[3] + self.q[1] * self.q[0], | |
| -self.q[1] * self.q[1] - self.q[2] * self.q[2]))) | |
| # convert to axis-angle representation, with angle encoded by the length | |
| def ToAxisAngle(self): | |
| # recall that for axis-angle representation (a, angle), with "a" unit: | |
| # q = (cos(angle/2), a * sin(angle/2)) | |
| # below, for readability, "theta" actually means half of the angle | |
| sin_sq_theta = self.q[1:].dot(self.q[1:]) | |
| # if theta is non-zero, then we can compute a unique rotation | |
| if np.abs(sin_sq_theta) > np.finfo('float').eps: | |
| sin_theta = np.sqrt(sin_sq_theta) | |
| cos_theta = self.q[0] | |
| # atan2 is more stable, so we use it to compute theta | |
| # note that we multiply by 2 to get the actual angle | |
| angle = 2. * ( | |
| np.arctan2(-sin_theta, -cos_theta) if cos_theta < 0. else | |
| np.arctan2(sin_theta, cos_theta)) | |
| return self.q[1:] * (angle / sin_theta) | |
| # otherwise, the result is singular, and we avoid dividing by | |
| # sin(angle/2) = 0 | |
| return np.zeros(3) | |
| # euclideanspace.com/maths/geometry/rotations/conversions/quaternionToEuler | |
| # this assumes the quaternion is non-zero | |
| # returns yaw, pitch, roll, with application in that order | |
| def ToEulerAngles(self): | |
| qsq = self.q**2 | |
| k = 2. * (self.q[0] * self.q[3] + self.q[1] * self.q[2]) / qsq.sum() | |
| if (1. - k) < np.finfo('float').eps: # north pole singularity | |
| return 2. * np.arctan2(self.q[1], self.q[0]), 0.5 * np.pi, 0. | |
| if (1. + k) < np.finfo('float').eps: # south pole singularity | |
| return -2. * np.arctan2(self.q[1], self.q[0]), -0.5 * np.pi, 0. | |
| yaw = np.arctan2(2. * (self.q[0] * self.q[2] - self.q[1] * self.q[3]), | |
| qsq[0] + qsq[1] - qsq[2] - qsq[3]) | |
| pitch = np.arcsin(k) | |
| roll = np.arctan2(2. * (self.q[0] * self.q[1] - self.q[2] * self.q[3]), | |
| qsq[0] - qsq[1] + qsq[2] - qsq[3]) | |
| return yaw, pitch, roll | |
| #------------------------------------------------------------------------------- | |
| # | |
| # DualQuaternion | |
| # | |
| #------------------------------------------------------------------------------- | |
| class DualQuaternion: | |
| # DualQuaternion from an existing rotation + translation | |
| def FromQT(q, t): | |
| return DualQuaternion(qe=(0.5 * np.asarray(t))) * DualQuaternion(q) | |
| def __init__(self, q0=np.array((1., 0., 0., 0.)), qe=np.zeros(4)): | |
| self.q0, self.qe = Quaternion(q0), Quaternion(qe) | |
| def __add__(self, other): | |
| return DualQuaternion(self.q0 + other.q0, self.qe + other.qe) | |
| def __iadd__(self, other): | |
| self.q0 += other.q0 | |
| self.qe += other.qe | |
| return self | |
| # conguation via the ~ operator | |
| def __invert__(self): | |
| return DualQuaternion(~self.q0, ~self.qe) | |
| def __mul__(self, other): | |
| if isinstance(other, DualQuaternion): | |
| return DualQuaternion( | |
| self.q0 * other.q0, | |
| self.q0 * other.qe + self.qe * other.q0) | |
| elif isinstance(other, complex): # multiplication by a dual number | |
| return DualQuaternion( | |
| self.q0 * other.real, | |
| self.q0 * other.imag + self.qe * other.real) | |
| else: # scalar multiplication (assumed) | |
| return DualQuaternion(other * self.q0, other * self.qe) | |
| def __rmul__(self, other): | |
| return self.__mul__(other) | |
| def __imul__(self, other): | |
| tmp = self * other | |
| self.q0, self.qe = tmp.q0, tmp.qe | |
| return self | |
| def __neg__(self): | |
| return DualQuaternion(-self.q0, -self.qe) | |
| def __sub__(self, other): | |
| return DualQuaternion(self.q0 - other.q0, self.qe - other.qe) | |
| def __isub__(self, other): | |
| self.q0 -= other.q0 | |
| self.qe -= other.qe | |
| return self | |
| # q^-1 = q* / ||q||^2 | |
| # assume that q0 is nonzero! | |
| def inverse(self): | |
| normsq = complex(q0.dot(q0), 2. * self.q0.q.dot(self.qe.q)) | |
| inv_len_real = 1. / normsq.real | |
| return ~self * complex( | |
| inv_len_real, -normsq.imag * inv_len_real * inv_len_real) | |
| # returns a complex representation of the real and imaginary parts of the norm | |
| # assume that q0 is nonzero! | |
| def norm(self): | |
| q0_norm = self.q0.norm() | |
| return complex(q0_norm, self.q0.dot(self.qe) / q0_norm) | |
| # assume that q0 is nonzero! | |
| def normalize(self): | |
| # current length is ||q0|| + eps * (<q0, qe> / ||q0||) | |
| # writing this as a + eps * b, the inverse is | |
| # 1/||q|| = 1/a - eps * b / a^2 | |
| norm = self.norm() | |
| inv_len_real = 1. / norm.real | |
| self *= complex(inv_len_real, -norm.imag * inv_len_real * inv_len_real) | |
| return self | |
| # return the translation vector for this dual quaternion | |
| def getT(self): | |
| return 2 * (self.qe * ~self.q0).q[1:] | |
| def ToQT(self): | |
| return self.q0, self.getT() | |