Skip to content

Commit

Permalink
feat: update transformation operation (#17)
Browse files Browse the repository at this point in the history
Signed-off-by: ktro2828 <[email protected]>
ktro2828 authored Oct 21, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent 63e006e commit 8aacfe1
Showing 2 changed files with 297 additions and 55 deletions.
267 changes: 223 additions & 44 deletions t4_devkit/dataclass/transform.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
from __future__ import annotations

from copy import deepcopy
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, overload
from typing import TYPE_CHECKING, Any, overload

import numpy as np
from pyquaternion import Quaternion
@@ -12,7 +13,13 @@
if TYPE_CHECKING:
from t4_devkit.typing import ArrayLike

__all__ = ["TransformBuffer", "HomogeneousMatrix", "TransformLike"]
__all__ = [
"TransformBuffer",
"HomogeneousMatrix",
"TranslateItemLike",
"RotateItemLike",
"TransformItemLike",
]


@dataclass
@@ -36,11 +43,20 @@ def set_transform(self, matrix: HomogeneousMatrix) -> None:

def lookup_transform(self, src: str, dst: str) -> HomogeneousMatrix | None:
if src == dst:
return HomogeneousMatrix(np.zeros(3), Quaternion(), src=src, dst=dst)
return HomogeneousMatrix.as_identity(src)
return self.buffer[(src, dst)] if (src, dst) in self.buffer else None

def do_transform(self, src: str, dst: str, *args: TransformLike) -> TransformLike | None:
return self.buffer[(src, dst)].transform(args) if (src, dst) in self.buffer else None
def do_translate(self, src: str, dst: str, *args, **kwargs) -> TranslateItemLike | None:
tf_matrix = self.lookup_transform(src, dst)
return tf_matrix.translate(*args, **kwargs) if tf_matrix is not None else None

def do_rotate(self, src: str, dst: str, *args, **kwargs) -> RotateItemLike | None:
tf_matrix = self.lookup_transform(src, dst)
return tf_matrix.rotate(*args, **kwargs) if tf_matrix is not None else None

def do_transform(self, src: str, dst: str, *args, **kwargs) -> TransformItemLike | None:
tf_matrix = self.lookup_transform(src, dst)
return tf_matrix.transform(*args, **kwargs) if tf_matrix is not None else None


@dataclass
@@ -73,25 +89,44 @@ def __init__(

self.matrix = _generate_homogeneous_matrix(position, rotation)

@classmethod
def as_identity(cls, frame_id: str) -> Self:
"""Construct a new object with identity.
Args:
frame_id (str): Frame ID.
Returns:
Constructed self instance.
"""
return cls(np.zeros(3), Quaternion(), frame_id, frame_id)

@classmethod
def from_matrix(
cls,
matrix: NDArray | HomogeneousMatrix,
src: str,
dst: str,
src: str | None = None,
dst: str | None = None,
) -> Self:
"""Construct a new object from a homogeneous matrix.
Args:
matrix (NDArray | HomogeneousMatrix): 4x4 homogeneous matrix.
src (str): Source frame ID.
dst (str): Destination frame ID.
src (str | None, optional): Source frame ID.
This must be specified only if the input matrix is `NDArray`.
dst (str | None, optional): Destination frame ID.
This must be specified only if the input matrix is `NDArray`.
Returns:
Self: Constructed instance.
Constructed self instance.
"""
position, rotation = _extract_position_and_rotation_from_matrix(matrix)
return cls(position, rotation, src, dst)
if isinstance(matrix, np.ndarray):
assert matrix.shape == (4, 4)
assert src is not None and dst is not None
return cls(position, rotation, src, dst)
else:
return cls(position, rotation, matrix.src, matrix.dst)

@property
def shape(self) -> tuple[int, ...]:
@@ -154,6 +189,93 @@ def inv(self) -> HomogeneousMatrix:
position, rotation = _extract_position_and_rotation_from_matrix(ret_mat)
return HomogeneousMatrix(position, rotation, src=self.src, dst=self.dst)

@overload
def translate(self, position: ArrayLike) -> NDArray:
"""Translate a position by myself.
Args:
position (ArrayLike): 3D position.
Returns:
Translated position.
"""
pass

@overload
def translate(self, matrix: HomogeneousMatrix) -> HomogeneousMatrix:
"""Translate a homogeneous matrix by myself.
Args:
matrix (HomogeneousMatrix):
Returns:
Translated `HomogeneousMatrix` object.
"""
pass

def translate(self, *args, **kwargs) -> TranslateItemLike:
inputs = _format_transform_args(*args, **kwargs)
if {"position"} == set(inputs.keys()):
return self.position + inputs["position"]
elif {"matrix"} == set(inputs.keys()):
matrix: HomogeneousMatrix = deepcopy(inputs["matrix"])
matrix.position = self.position + matrix.position
return matrix
else:
raise ValueError(f"Unexpected arguments: {list(inputs.keys())}")

@overload
def rotate(self, position: ArrayLike) -> NDArray:
"""Rotate a position by myself.
Args:
position (ArrayLike): 3D position.
Returns:
Rotated position.
"""
pass

@overload
def rotate(self, rotation: RotationType) -> RotationType:
"""Rotate a 3x3 rotation matrix or quaternion by myself.
Args:
rotation (RotationType): 3x3 rotation matrix or quaternion.
Returns:
Rotated quaternion.
"""
pass

@overload
def rotate(self, matrix: HomogeneousMatrix) -> HomogeneousMatrix:
"""Rotate a homogeneous matrix by myself.
Args:
matrix (HomogeneousMatrix): `HomogeneousMatrix` object.
Returns:
Rotated `HomogeneousMatrix` object.
"""
pass

def rotate(self, *args, **kwargs) -> RotateItemLike:
inputs = _format_transform_args(*args, **kwargs)
if {"position"} == set(inputs.keys()):
return np.matmul(self.rotation_matrix, inputs["position"])
elif {"rotation"} == set(inputs.keys()):
rotation_matrix: NDArray = inputs["rotation"].rotation_matrix
return np.matmul(self.rotation_matrix, rotation_matrix)
elif {"matrix"} == set(inputs.keys()):
matrix: HomogeneousMatrix = deepcopy(inputs["matrix"])
matrix.rotation = Quaternion(
matrix=np.matmul(self.rotation_matrix, matrix.rotation_matrix)
)
return matrix
else:
raise ValueError(f"Unexpected arguments: {list(inputs.keys())}")

@overload
def transform(self, position: ArrayLike) -> NDArray:
"""Transform a position by myself.
@@ -166,6 +288,18 @@ def transform(self, position: ArrayLike) -> NDArray:
"""
pass

@overload
def transform(self, rotation: RotationType) -> RotationType:
"""Transform a rotation by myself.
Args:
rotation (RotationType): 3x3 rotation matrix or quaternion.
Returns:
Transformed quaternion.
"""
pass

@overload
def transform(
self,
@@ -195,38 +329,18 @@ def transform(self, matrix: HomogeneousMatrix) -> HomogeneousMatrix:
"""
pass

def transform(self, *args, **kwargs):
# TODO(ktro2828): Refactoring this operations.
s = len(args)
if s == 0:
if not kwargs:
raise ValueError("At least 1 argument specified")

if "position" in kwargs:
position = kwargs["position"]
if "matrix" in kwargs:
raise ValueError("Cannot specify `position` and `matrix` at the same time.")
elif "rotation" in kwargs:
rotation = kwargs["rotation"]
return self.__transform_position_and_rotation(position, rotation)
else:
return self.__transform_position(position)
elif "matrix" in kwargs:
matrix = kwargs["matrix"]
return self.__transform_matrix(matrix)
else:
raise KeyError(f"Unexpected keys are detected: {list(kwargs.keys())}")
elif s == 1:
arg = args[0]
if isinstance(arg, HomogeneousMatrix):
return self.__transform_matrix(matrix=arg)
else:
return self.__transform_position(position=arg)
elif s == 2:
position, rotation = args
return self.__transform_position_and_rotation(position, rotation)
def transform(self, *args, **kwargs) -> TransformItemLike:
inputs = _format_transform_args(*args, **kwargs)
if {"position", "rotation"} == set(inputs.keys()):
return self.__transform_position_and_rotation(**inputs)
elif {"position"} == set(inputs.keys()):
return self.__transform_position(**inputs)
elif {"rotation"} == set(inputs.keys()):
return self.__transform_rotation(**inputs)
elif {"matrix"} == set(inputs.keys()):
return self.__transform_matrix(**inputs)
else:
raise ValueError(f"Unexpected number of arguments {s}")
raise ValueError(f"Unexpected inputs: {list(inputs.keys())}")

def __transform_position(self, position: ArrayLike) -> NDArray:
rotation = Quaternion()
@@ -235,6 +349,13 @@ def __transform_position(self, position: ArrayLike) -> NDArray:
ret_pos, _ = _extract_position_and_rotation_from_matrix(ret_mat)
return ret_pos

def __transform_rotation(self, rotation: RotationType) -> RotationType:
position = np.zeros(3)
matrix = _generate_homogeneous_matrix(position, rotation)
ret_mat = self.matrix.dot(matrix)
_, ret_rot = _extract_position_and_rotation_from_matrix(ret_mat)
return ret_rot

def __transform_position_and_rotation(
self,
position: ArrayLike,
@@ -248,7 +369,66 @@ def __transform_matrix(self, matrix: HomogeneousMatrix) -> HomogeneousMatrix:
return matrix.dot(self)


TransformLike = NDArray | tuple[NDArray, RotationType] | HomogeneousMatrix
TranslateItemLike = NDArray | HomogeneousMatrix
RotateItemLike = NDArray | RotationType | HomogeneousMatrix
TransformItemLike = NDArray | RotationType | tuple[NDArray, RotationType] | HomogeneousMatrix


def _format_transform_args(*args, **kwargs) -> dict[str, Any]:
num_args = len(args)
num_kwargs = len(kwargs)
if num_args == 0 and num_kwargs == 0:
raise ValueError("At least 1 argument specified.")

# >>> (position), (rotation), (position, rotation), (matrix)
if num_args == 0:
if "position" in kwargs:
if "matrix" in kwargs:
raise KeyError("Cannot specify `position` and `matrix` at the same time.")
elif "rotation" in kwargs:
return {"position": kwargs["position"], "rotation": kwargs["rotation"]}
else:
return {"position": kwargs["position"]}
elif "rotation" in kwargs:
if "matrix" in kwargs:
raise KeyError("Cannot specify `rotation` and `matrix` at the same time.")
return {"rotation": kwargs["rotation"]}
elif "matrix" in kwargs:
return {"matrix": kwargs["matrix"]}
else:
raise KeyError(f"Unexpected keys are detected: {list(kwargs.keys())}.")
# >>> (position), (rotation), (position, rotation), (matrix)
elif num_args == 1:
arg0 = args[0]
if num_kwargs == 0:
if isinstance(arg0, HomogeneousMatrix):
return {"matrix": arg0}
elif isinstance(arg0, Quaternion):
return {"rotation": arg0}
else:
arg0 = np.asarray(arg0)
if arg0.ndim == 1:
if len(arg0) == 3:
return {"position": arg0}
elif len(arg0) == 4:
return {"rotation": arg0}
else:
raise ValueError(f"Unexpected argument shape: {arg0.shape}.")
else:
if not arg0.shape != (3, 3):
raise ValueError(f"Unexpected argument shape: {arg0.shape}.")
return {"rotation": arg0}
elif num_kwargs == 1:
if "rotation" not in kwargs:
raise KeyError("Expected two arguments: position and rotation.")
return {"position": arg0, "rotation": kwargs["rotation"]}
else:
raise ValueError(f"Too much arguments {num_args + num_kwargs}.")
# >>> (position, rotation)
elif num_args == 2:
return {"position": args[0], "rotation": args[1]}
else:
raise ValueError(f"Too much arguments {num_args + num_kwargs}.")


def _extract_position_and_rotation_from_matrix(
@@ -302,4 +482,3 @@ def _generate_homogeneous_matrix(
matrix[:3, 3] = position
matrix[:3, :3] = rotation.rotation_matrix
return matrix
return matrix
85 changes: 74 additions & 11 deletions tests/dataclass/test_transform.py
Original file line number Diff line number Diff line change
@@ -1,32 +1,95 @@
from __future__ import annotations

import numpy as np
from pyquaternion import Quaternion

from t4_devkit.dataclass.transform import HomogeneousMatrix


def test_homogeneous_matrix_transform():
ego2map = HomogeneousMatrix((1, 0, 0), (1, 0, 0, 0), src="base_link", dst="map")
pos1 = ego2map.transform((1, 0, 0))
assert np.allclose(pos1, np.array((2, 0, 0)))

pos2 = ego2map.transform(position=(1, 0, 0))
assert np.allclose(pos2, np.array((2, 0, 0)))
def test_tf_buffer(dummy_tf_buffer) -> None:
# (position)
pos1 = dummy_tf_buffer.do_transform(src="base_link", dst="base_link", position=(1, 0, 0))
assert isinstance(pos1, np.ndarray)
assert np.allclose(pos1, (1, 0, 0))

pos1, rot1 = ego2map.transform((1, 0, 0), (1, 0, 0, 0))
assert np.allclose(pos1, np.array((2, 0, 0)))
# (rotation)
rot1 = dummy_tf_buffer.do_transform(
src="base_link",
dst="base_link",
rotation=(1, 0, 0, 0),
)
assert isinstance(rot1, Quaternion)
assert np.allclose(rot1.rotation_matrix, np.eye(3))

pos2, rot2 = ego2map.transform(position=(1, 0, 0), rotation=(1, 0, 0, 0))
assert np.allclose(pos2, np.array((2, 0, 0)))
# (position, rotation)
pos2, rot2 = dummy_tf_buffer.do_transform(
src="base_link",
dst="base_link",
position=(1, 0, 0),
rotation=(1, 0, 0, 0),
)
assert isinstance(pos2, np.ndarray)
assert np.allclose(pos2, (1, 0, 0))
assert isinstance(rot2, Quaternion)
assert np.allclose(rot2.rotation_matrix, np.eye(3))

# matrix
cam2ego = HomogeneousMatrix((2, 2, 2), (1, 0, 0, 0), src="camera", dst="base_link")
mat = dummy_tf_buffer.do_transform(src="base_link", dst="camera", matrix=cam2ego)
assert np.allclose(
mat.matrix,
np.array(
[
[-1, 0, 0, 3],
[0, -1, 0, 3],
[0, 0, 1, 3],
[0, 0, 0, 1],
],
),
)


def test_homogeneous_matrix_transform() -> None:
ego2map = HomogeneousMatrix((1, 0, 0), (1, 0, 0, 0), src="base_link", dst="map")
# transform(position)
pos1_1 = ego2map.transform((1, 0, 0))
assert np.allclose(pos1_1, np.array((2, 0, 0)))

# transform(position=)
pos1_2 = ego2map.transform(position=(1, 0, 0))
assert np.allclose(pos1_2, np.array((2, 0, 0)))

# transform(rotation)
rot1_1 = ego2map.transform((1, 0, 0, 0))
assert np.allclose(rot1_1.rotation_matrix, np.eye(3))

# transform(rotation)
rot1_2 = ego2map.transform(rotation=(1, 0, 0, 0))
assert np.allclose(rot1_2.rotation_matrix, np.eye(3))

# transform(position, rotation)
pos2_1, rot2_1 = ego2map.transform((1, 0, 0), (1, 0, 0, 0))
assert np.allclose(pos2_1, np.array((2, 0, 0)))
assert np.allclose(rot2_1.rotation_matrix, np.eye(3))

# transform(position=, rotation=)
pos2_2, rot2_2 = ego2map.transform(position=(1, 0, 0), rotation=(1, 0, 0, 0))
assert np.allclose(pos2_2, np.array((2, 0, 0)))
assert np.allclose(rot2_2.rotation_matrix, np.eye(3))

# transform(position, rotation=)
pos2_3, rot2_3 = ego2map.transform((1, 0, 0), rotation=(1, 0, 0, 0))
assert np.allclose(pos2_3, np.array((2, 0, 0)))
assert np.allclose(rot2_3.rotation_matrix, np.eye(3))

# transform(matrix)
map2ego = HomogeneousMatrix((-1, 0, 0), (1, 0, 0, 0), src="map", dst="base_link")
mat1 = ego2map.transform(map2ego)
assert np.allclose(mat1.matrix, np.eye(4))
assert np.allclose(mat1.position, np.zeros(3))
assert np.allclose(mat1.rotation_matrix, np.eye(3))

# transform(matrix=)
mat2 = ego2map.transform(matrix=map2ego)
assert np.allclose(mat2.matrix, np.eye(4))
assert np.allclose(mat2.position, np.zeros(3))

0 comments on commit 8aacfe1

Please sign in to comment.