feat: add translate and rotate methods

Signed-off-by: ktro2828 <[email protected]>

t4_devkit/dataclass/box.py

@@ -170,6 +170,32 @@ def area(self) -> float:
     def volume(self) -> float:
         return self.area * self.size[2]
 
+    def translate(self, x: TranslationType) -> None:
+        """Apply a translation.
+
+        Args:
+            x (TranslationType): 3D translation vector in the order of (x, y, z).
+        """
+        self.position += x
+
+        if self.future is not None:
+            self.future = [t.translate(x) for t in self.future]
+
+    def rotate(self, q: RotationType) -> None:
+        """Apply a rotation.
+
+        Args:
+            q (RotationType): Rotation quaternion.
+        """
+        self.position = np.dot(q.rotation_matrix, self.position)
+        self.rotation = q * self.rotation
+
+        if self.velocity is not None:
+            self.velocity = np.dot(q.rotation_matrix, self.velocity)
+
+        if self.future is not None:
+            self.future = [t.rotate(q) for t in self.future]
+
     def corners(self, box_scale: float = 1.0) -> NDArrayF64:
         """Return the bounding box corners.
 

t4_devkit/dataclass/trajectory.py

@@ -6,7 +6,7 @@
 from attrs import define, field
 
 if TYPE_CHECKING:
-    from t4_devkit.typing import TrajectoryType, TranslationType
+    from t4_devkit.typing import RotationType, TrajectoryType, TranslationType
 
 __all__ = ["Trajectory", "to_trajectories"]
 
@@ -66,6 +66,23 @@ def shape(self) -> tuple[int, ...]:
         """
         return self.waypoints.shape
 
+    def translate(self, x: TranslationType) -> None:
+        """Apply a translation.
+
+        Args:
+            x (TranslationType): 3D translation vector.
+        """
+        self.waypoints += x
+
+    def rotate(self, q: RotationType) -> None:
+        """Apply a rotation.
+
+        Args:
+            q (RotationType): Rotation quaternion.
+        """
+        # NOTE: R * X = X * R^T
+        self.waypoints = np.dot(self.waypoints, q.rotation_matrix.T)
+
 
 def to_trajectories(
     waypoints: list[TrajectoryType],

t4_devkit/tier4.py

@@ -293,18 +293,21 @@ def get_sample_data_path(self, sample_data_token: str) -> str:
     def get_sample_data(
         self,
         sample_data_token: str,
-        selected_ann_tokens: list[str] | None = None,
         *,
+        selected_ann_tokens: list[str] | None = None,
         as_3d: bool = True,
+        as_sensor_coord: bool = True,
         visibility: VisibilityLevel = VisibilityLevel.NONE,
     ) -> tuple[str, list[BoxType], CamIntrinsicType | None]:
         """Return the data path as well as all annotations related to that `sample_data`.
+        Note that output boxes is w.r.t base link or sensor coordinate system.
 
         Args:
             sample_data_token (str): Token of `sample_data`.
             selected_ann_tokens (list[str] | None, optional):
                 Specify if you want to extract only particular annotations.
             as_3d (bool, optional): Whether to return 3D or 2D boxes.
+            as_sensor_coord (bool, optional): Whether to transform boxes as sensor origin coordinate system.
             visibility (VisibilityLevel, optional): If `sample_data` is an image,
                 this sets required visibility for only 3D boxes.
 
@@ -350,10 +353,13 @@ def get_sample_data(
             # Move box to ego vehicle coord system.
             box.translate(-pose_record.translation)
             box.rotate(pose_record.rotation.inverse)
+            box.frame_id = "base_link"
 
-            #  Move box to sensor coord system.
-            box.translate(-cs_record.translation)
-            box.rotate(cs_record.rotation.inverse)
+            if as_sensor_coord:
+                #  Move box to sensor coord system.
+                box.translate(-cs_record.translation)
+                box.rotate(cs_record.rotation.inverse)
+                box.frame_id = sensor_record.channel
 
             if sensor_record.modality == SensorModality.CAMERA and not is_box_in_image(
                 box,

tests/conftest.py

@@ -8,6 +8,7 @@
     SemanticLabel,
     Shape,
     ShapeType,
+    Trajectory,
     TransformBuffer,
 )
 
@@ -17,7 +18,7 @@ def label2id() -> dict[str, int]:
     return {"car": 0, "bicycle": 1, "pedestrian": 2}
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="function")
 def dummy_box3d() -> Box3D:
     """Return a dummy 3D box.
 
@@ -29,15 +30,15 @@ def dummy_box3d() -> Box3D:
         frame_id="base_link",
         semantic_label=SemanticLabel("car"),
         position=(1.0, 1.0, 1.0),
-        rotation=Quaternion([0.0, 0.0, 0.0, 1.0]),
+        rotation=Quaternion([1.0, 0.0, 0.0, 0.0]),
         shape=Shape(shape_type=ShapeType.BOUNDING_BOX, size=(1.0, 1.0, 1.0)),
         velocity=(1.0, 1.0, 1.0),
         confidence=1.0,
         uuid="car3d_0",
     )
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="function")
 def dummy_box3ds() -> list[Box3D]:
     """Return a list of dummy 3D boxes.
 
@@ -50,7 +51,7 @@ def dummy_box3ds() -> list[Box3D]:
             frame_id="base_link",
             semantic_label=SemanticLabel("car"),
             position=(1.0, 1.0, 1.0),
-            rotation=Quaternion([0.0, 0.0, 0.0, 1.0]),
+            rotation=Quaternion([1.0, 0.0, 0.0, 0.0]),
             shape=Shape(shape_type=ShapeType.BOUNDING_BOX, size=(1.0, 1.0, 1.0)),
             velocity=(1.0, 1.0, 1.0),
             confidence=1.0,
@@ -61,7 +62,7 @@ def dummy_box3ds() -> list[Box3D]:
             frame_id="base_link",
             semantic_label=SemanticLabel("bicycle"),
             position=(-1.0, -1.0, 1.0),
-            rotation=Quaternion([0.0, 0.0, 0.0, 1.0]),
+            rotation=Quaternion([1.0, 0.0, 0.0, 0.0]),
             shape=Shape(shape_type=ShapeType.BOUNDING_BOX, size=(1.0, 1.0, 1.0)),
             velocity=(1.0, 1.0, 1.0),
             confidence=1.0,
@@ -72,7 +73,7 @@ def dummy_box3ds() -> list[Box3D]:
             frame_id="base_link",
             semantic_label=SemanticLabel("pedestrian"),
             position=(-1.0, 1.0, 1.0),
-            rotation=Quaternion([0.0, 0.0, 0.0, 1.0]),
+            rotation=Quaternion([1.0, 0.0, 0.0, 0.0]),
             shape=Shape(shape_type=ShapeType.BOUNDING_BOX, size=(1.0, 1.0, 1.0)),
             velocity=(1.0, 1.0, 1.0),
             confidence=1.0,
@@ -81,7 +82,7 @@ def dummy_box3ds() -> list[Box3D]:
     ]
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="function")
 def dummy_box2d() -> Box2D:
     """Return a dummy 2D box.
 
@@ -98,7 +99,7 @@ def dummy_box2d() -> Box2D:
     )
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="function")
 def dummy_box2ds() -> list[Box2D]:
     """Return a list of dummy 2D boxes.
 
@@ -133,6 +134,20 @@ def dummy_box2ds() -> list[Box2D]:
     ]
 
 
+@pytest.fixture(scope="function")
+def dummy_trajectory() -> Trajectory:
+    """Return a dummy trajectory.
+
+    Returns:
+        A trajectory.
+    """
+    # list item is converted to NDArray internally
+    return Trajectory(
+        waypoints=[[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]],
+        confidence=1.0,
+    )
+
+
 @pytest.fixture(scope="module")
 def dummy_tf_buffer() -> TransformBuffer:
     """Return a dummy transformation buffer.
@@ -145,7 +160,7 @@ def dummy_tf_buffer() -> TransformBuffer:
     tf_buffer.set_transform(
         HomogeneousMatrix(
             [1.0, 1.0, 1.0],
-            Quaternion([0.0, 0.0, 0.0, 1.0]),
+            Quaternion([1.0, 0.0, 0.0, 0.0]),
             src="base_link",
             dst="map",
         )
@@ -154,11 +169,10 @@ def dummy_tf_buffer() -> TransformBuffer:
     tf_buffer.set_transform(
         HomogeneousMatrix(
             [1.0, 1.0, 1.0],
-            Quaternion([0.0, 0.0, 0.0, 1.0]),
+            Quaternion([1.0, 0.0, 0.0, 0.0]),
             src="base_link",
             dst="camera",
         )
     )
 
     return tf_buffer
-    return tf_buffer

tests/dataclass/test_box.py

@@ -1,6 +1,7 @@
 import math
 
 import numpy as np
+from pyquaternion import Quaternion
 
 from t4_devkit.dataclass import distance_box
 
@@ -20,14 +21,37 @@ def test_box3d(dummy_box3d) -> None:
         dummy_box3d.corners(box_scale=1.0),
         np.array(
             [
-                [0.5, 0.5, 1.5],
-                [0.5, 1.5, 1.5],
-                [0.5, 1.5, 0.5],
-                [0.5, 0.5, 0.5],
-                [1.5, 0.5, 1.5],
                 [1.5, 1.5, 1.5],
-                [1.5, 1.5, 0.5],
+                [1.5, 0.5, 1.5],
                 [1.5, 0.5, 0.5],
+                [1.5, 1.5, 0.5],
+                [0.5, 1.5, 1.5],
+                [0.5, 0.5, 1.5],
+                [0.5, 0.5, 0.5],
+                [0.5, 1.5, 0.5],
+            ],
+        ),
+    )
+
+
+def test_box3d_translate(dummy_box3d) -> None:
+    dummy_box3d.translate(x=(1.0, 2.0, 3.0))
+
+    assert np.allclose(dummy_box3d.position, (2.0, 3.0, 4.0))
+
+
+def test_box3d_rotate(dummy_box3d) -> None:
+    # +90 [deg]
+    dummy_box3d.rotate(q=Quaternion([0.7071067811865475, 0.0, 0.0, -0.7071067811865475]))
+
+    assert np.allclose(dummy_box3d.position, (1.0, -1.0, 1.0))
+    assert np.allclose(
+        dummy_box3d.rotation.rotation_matrix,
+        np.array(
+            [
+                [0, 1, 0],
+                [-1, 0, 0],
+                [0, 0, 1],
             ]
         ),
     )

tests/dataclass/test_trajectory.py

@@ -1,33 +1,46 @@
 import numpy as np
 import pytest
+from pyquaternion import Quaternion
 
-from t4_devkit.dataclass.trajectory import Trajectory, to_trajectories
+from t4_devkit.dataclass.trajectory import to_trajectories
 
 
-def test_trajectory() -> None:
+def test_trajectory(dummy_trajectory) -> None:
     """Test `Trajectory` class including its initialization and methods."""
-    # list item is converted to NDArray internally
-    trajectory = Trajectory(
-        waypoints=[[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]],
-        confidence=1.0,
-    )
 
-    assert trajectory.confidence == 1.0
+    assert dummy_trajectory.confidence == 1.0
 
     # test __len__()
-    assert len(trajectory) == 2
+    assert len(dummy_trajectory) == 2
 
     # test __getitem__()
-    assert np.allclose(trajectory[0], [1.0, 1.0, 1.0])
-    assert np.allclose(trajectory[1], [2.0, 2.0, 2.0])
+    assert np.allclose(dummy_trajectory[0], (1.0, 1.0, 1.0))
+    assert np.allclose(dummy_trajectory[1], (2.0, 2.0, 2.0))
 
     # test __iter__()
-    for point in trajectory:
+    for point in dummy_trajectory:
         assert isinstance(point, np.ndarray)
         assert point.shape == (3,)
 
     # test shape property
-    assert trajectory.shape == (2, 3)
+    assert dummy_trajectory.shape == (2, 3)
+
+
+def test_trajectory_translate(dummy_trajectory) -> None:
+    """Test `translate` methods of `Trajectory` class."""
+    dummy_trajectory.translate(x=(1.0, 2.0, 3.0))
+
+    assert np.allclose(dummy_trajectory[0], (2.0, 3.0, 4.0))
+    assert np.allclose(dummy_trajectory[1], (3.0, 4.0, 5.0))
+
+
+def test_trajectory_rotate(dummy_trajectory) -> None:
+    """Test `rotate` methods of `Trajectory` class."""
+    # +90 [deg]
+    dummy_trajectory.rotate(q=Quaternion([0.7071067811865475, 0.0, 0.0, -0.7071067811865475]))
+
+    assert np.allclose(dummy_trajectory[0], (1.0, -1.0, 1.0))
+    assert np.allclose(dummy_trajectory[1], (2.0, -2.0, 2.0))
 
 
 def test_to_trajectories() -> None:

tests/dataclass/test_transform.py

@@ -40,8 +40,8 @@ def test_tf_buffer(dummy_tf_buffer) -> None:
         mat.matrix,
         np.array(
             [
-                [-1, 0, 0, 3],
-                [0, -1, 0, 3],
+                [1, 0, 0, 3],
+                [0, 1, 0, 3],
                 [0, 0, 1, 3],
                 [0, 0, 0, 1],
             ],