Quaternion conversions

kornia/geometry/quaternion.py

@@ -12,6 +12,7 @@
  euler_from_quaternion,
  normalize_quaternion,
  quaternion_from_euler,
+ quaternion_to_axis_angle,
  quaternion_to_rotation_matrix,
  rotation_matrix_to_quaternion,
 )
@@ -262,7 +263,7 @@ def from_matrix(cls, matrix: Tensor) -> "Quaternion":
 
  @classmethod
  def from_euler(cls, roll: Tensor, pitch: Tensor, yaw: Tensor) -> "Quaternion":
- """Create a quaternion from euler angles.
+ """Create a quaternion from Euler angles.
 
  Args:
  roll: the roll euler angle.
@@ -281,10 +282,7 @@ def from_euler(cls, roll: Tensor, pitch: Tensor, yaw: Tensor) -> "Quaternion":
  return cls(q)
 
  def to_euler(self) -> Tuple[Tensor, Tensor, Tensor]:
- """Create a quaternion from euler angles.
-
- Args:
- matrix: the rotation matrix to convert of shape :math:`(B, 3, 3)`.
+ """Convert the quaternion to a triple of Euler angles (roll, pitch, yaw).
 
  Example:
  >>> q = Quaternion(tensor([2., 0., 1., 1.]))
@@ -314,6 +312,17 @@ def from_axis_angle(cls, axis_angle: Tensor) -> "Quaternion":
  """
  return cls(axis_angle_to_quaternion(axis_angle))
 
+ def to_axis_angle(self) -> Tensor:
+ """Converts the quaternion to an axis-angle representation.
+
+ Example:
+ >>> q = Quaternion.identity()
+ >>> axis_angle = q.to_axis_angle()
+ >>> axis_angle
+ tensor([0., 0., 0.], grad_fn=<AsStridedBackward0>)
+ """
+ return quaternion_to_axis_angle(self.data)
+
  @classmethod
  def identity(
  cls, batch_size: Optional[int] = None, device: Optional[Device] = None, dtype: Dtype = None

tests/geometry/test_quaternion.py

@@ -208,6 +208,39 @@ def test_axis_angle(self, device, dtype, batch_size):
  q2 = q2.to(device, dtype)
  self.assert_close(q1, q2)
 
+ @pytest.mark.parametrize("batch_size", (None, 1, 2, 5))
+ def test_to_axis_angle(self, device, dtype, batch_size):
+ # batch_s = 5
+ # random_coefs = Quaternion.random(batch_s).data
+ random_coefs = torch.tensor(
+ [
+ [2.5398e-04, -2.2677e-01, -8.3897e-01, 4.9467e-01],
+ [-1.7005e-01, -1.0974e-01, 3.7635e-01, -9.0410e-01],
+ [9.1273e-01, 4.8935e-02, -6.2994e-03, 4.0558e-01],
+ [-9.8316e-01, 5.4078e-03, 1.4471e-01, 1.1145e-01],
+ [4.5794e-02, -7.0831e-01, 6.7577e-01, 1.9883e-01],
+ ],
+ device=device,
+ dtype=dtype,
+ )
+
+ q = Quaternion(random_coefs)
+ axis_angle_actual = q.to_axis_angle()
+
+ axis_angle_expected = torch.tensor(
+ [
+ [-0.7123, -2.6353, 1.5538],
+ [0.3118, -1.0693, 2.5687],
+ [0.1008, -0.0130, 0.8356],
+ [-0.0109, -0.2911, -0.2242],
+ [-2.1626, 2.0632, 0.6071],
+ ],
+ device=device,
+ dtype=dtype,
+ )
+
+ self.assert_close(axis_angle_expected, axis_angle_actual, 1e-4, 1e-4)
+
  @pytest.mark.parametrize("batch_size", (None, 1, 2, 5))
  def test_slerp(self, device, dtype, batch_size):
  for axis in torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]):