feat: add support of rendering future trajectories

t4_devkit/tier4.py

@@ -217,10 +217,15 @@ def __make_reverse_index__(self, verbose: bool) -> None:
                 sample_record: Sample = self.get("sample", record.sample_token)
                 sample_record.data[record.channel] = record.token
 
+        self._sample_and_instance_to_ann3d: dict[tuple[str, str], str] = {}
         for ann_record in self.sample_annotation:
             sample_record: Sample = self.get("sample", ann_record.sample_token)
             sample_record.ann_3ds.append(ann_record.token)
 
+            self._sample_and_instance_to_ann3d[(sample_record.token, ann_record.instance_token)] = (
+                ann_record.token
+            )
+
         for ann_record in self.object_ann:
             sd_record: SampleData = self.get("sample_data", ann_record.sample_data_token)
             sample_record: Sample = self.get("sample", sd_record.sample_token)
@@ -301,6 +306,7 @@ def get_sample_data(
         selected_ann_tokens: list[str] | None = None,
         as_3d: bool = True,
         as_sensor_coord: bool = True,
+        future_seconds: float = 0.0,
         visibility: VisibilityLevel = VisibilityLevel.NONE,
     ) -> tuple[str, list[BoxType], CamIntrinsicType | None]:
         """Return the data path as well as all annotations related to that `sample_data`.
@@ -339,13 +345,18 @@ def get_sample_data(
         boxes: list[BoxType]
         if selected_ann_tokens is not None:
             boxes = (
-                list(map(self.get_box3d, selected_ann_tokens))
+                [
+                    self.get_box3d(token, future_seconds=future_seconds)
+                    for token in selected_ann_tokens
+                ]
                 if as_3d
                 else list(map(self.get_box2d, selected_ann_tokens))
             )
         else:
             boxes = (
-                self.get_box3ds(sample_data_token) if as_3d else self.get_box2ds(sample_data_token)
+                self.get_box3ds(sample_data_token, future_seconds=future_seconds)
+                if as_3d
+                else self.get_box2ds(sample_data_token)
             )
 
         if not as_3d:
@@ -399,11 +410,12 @@ def get_semantic_label(
 
         return SemanticLabel(category.name, attributes)
 
-    def get_box3d(self, sample_annotation_token: str) -> Box3D:
+    def get_box3d(self, sample_annotation_token: str, *, future_seconds: float = 0.0) -> Box3D:
         """Return a Box3D class from a `sample_annotation` record.
 
         Args:
             sample_annotation_token (str): Token of `sample_annotation`.
+            future_seconds (float, optional): Future time in [s].
 
         Returns:
             Instantiated Box3D.
@@ -423,7 +435,7 @@ def get_box3d(self, sample_annotation_token: str) -> Box3D:
         # velocity
         velocity = self.box_velocity(sample_annotation_token=sample_annotation_token)
 
-        return Box3D(
+        box = Box3D(
             unix_time=sample.timestamp,
             frame_id="map",
             semantic_label=semantic_label,
@@ -435,6 +447,51 @@ def get_box3d(self, sample_annotation_token: str) -> Box3D:
             uuid=instance.token,  # TODO(ktro2828): extract uuid from `instance_name`.
         )
 
+        if future_seconds > 0.0:
+            # NOTE: Future trajectory is map coordinate frame
+            anns: list[SampleAnnotation] = self.get_sample_annotations_until(
+                ann.instance_token, ann.sample_token, future_seconds
+            )
+            if len(anns) == 0:
+                return box
+            waypoints = [ann.translation for ann in anns]
+            return box.with_future(waypoints=[waypoints], confidences=[1.0])
+        else:
+            return box
+
+    def get_sample_annotations_until(
+        self,
+        instance_token: str,
+        sample_token: str,
+        seconds: float,
+    ) -> list[SampleAnnotation]:
+        """Return a list of sample annotations until the specified seconds.
+
+        Args:
+            instance_token (str): Instance token.
+            sample_token (str): Start sample token.
+            seconds (float): Time seconds until.
+
+        Returns:
+            list[SampleAnnotation]: List of sample annotation records.
+        """
+        outputs = []
+        start_sample: Sample = self.get("sample", sample_token)
+
+        current_sample = start_sample
+        while current_sample.next != "":
+            next_sample: Sample = self.get("sample", current_sample.next)
+            if us2sec(next_sample.timestamp - start_sample.timestamp) > seconds:
+                break
+
+            ann_token = self._sample_and_instance_to_ann3d.get((next_sample.token, instance_token))
+            if ann_token is not None:
+                outputs.append(self.get("sample_annotation", ann_token))
+
+            current_sample = next_sample
+
+        return outputs
+
     def get_box2d(self, object_ann_token: str) -> Box2D:
         """Return a Box2D class from a `object_ann` record.
 
@@ -462,12 +519,13 @@ def get_box2d(self, object_ann_token: str) -> Box2D:
             uuid=instance.token,  # TODO(ktro2828): extract uuid from `instance_name`.
         )
 
-    def get_box3ds(self, sample_data_token: str) -> list[Box3D]:
+    def get_box3ds(self, sample_data_token: str, *, future_seconds: float = 0.0) -> list[Box3D]:
         """Rerun a list of Box3D classes for all annotations of a particular `sample_data` record.
         It the `sample_data` is a keyframe, this returns annotations for the corresponding `sample`.
 
         Args:
             sample_data_token (str): Token of `sample_data`.
+            future_seconds (float, optional): Future time in [s].
 
         Returns:
             List of instantiated Box3D classes.
@@ -478,7 +536,10 @@ def get_box3ds(self, sample_data_token: str) -> list[Box3D]:
 
         if curr_sample_record.prev == "" or sd_record.is_key_frame:
             # If no previous annotations available, or if sample_data is keyframe just return the current ones.
-            boxes = list(map(self.get_box3d, curr_sample_record.ann_3ds))
+            boxes = [
+                self.get_box3d(token, future_seconds=future_seconds)
+                for token in curr_sample_record.ann_3ds
+            ]
 
         else:
             prev_sample_record: Sample = self.get("sample", curr_sample_record.prev)
@@ -542,7 +603,7 @@ def get_box3ds(self, sample_data_token: str) -> list[Box3D]:
                     )
                 else:
                     # If not, simply grab the current annotation.
-                    box = self.get_box3d(curr_ann.token)
+                    box = self.get_box3d(curr_ann.token, future_seconds=future_seconds)
                 boxes.append(box)
 
         return boxes
@@ -708,6 +769,7 @@ def render_scene(
         scene_token: str,
         *,
         max_time_seconds: float = np.inf,
+        future_seconds: float = 0.0,
         save_dir: str | None = None,
         show: bool = True,
     ) -> None:
@@ -716,6 +778,7 @@ def render_scene(
         Args:
             scene_token (str): Unique identifier of scene.
             max_time_seconds (float, optional): Max time length to be rendered [s].
+            future_seconds (float, optional): Future time in [s].
             save_dir (str | None, optional): Directory path to save the recording.
             show (bool, optional): Whether to spawn rendering viewer.
         """
@@ -756,7 +819,9 @@ def render_scene(
         self._render_cameras(viewer, first_camera_tokens, max_timestamp_us)
 
         # render annotation
-        self._render_annotation_3ds(viewer, scene.first_sample_token, max_timestamp_us)
+        self._render_annotation_3ds(
+            viewer, scene.first_sample_token, max_timestamp_us, future_seconds
+        )
         self._render_annotation_2ds(viewer, scene.first_sample_token, max_timestamp_us)
 
         if save_dir is not None:
@@ -766,13 +831,15 @@ def render_instance(
         self,
         instance_token: str,
         *,
+        future_seconds: float = 0.0,
         save_dir: str | None = None,
         show: bool = True,
     ) -> None:
         """Render particular instance.
 
         Args:
             instance_token (str): Instance token.
+            future_seconds (float, optional): Future time in [s].
             save_dir (str | None, optional): Directory path to save the recording.
             show (bool, optional): Whether to spawn rendering viewer.
         """
@@ -822,6 +889,7 @@ def render_instance(
             viewer,
             first_sample.token,
             max_timestamp_us,
+            future_seconds=future_seconds,
             instance_token=instance_token,
         )
         self._render_annotation_2ds(
@@ -1094,6 +1162,7 @@ def _render_annotation_3ds(
         viewer: RerunViewer,
         first_sample_token: str,
         max_timestamp_us: float,
+        future_seconds: float = 0.0,
         instance_token: str | None = None,
     ) -> None:
         """Render annotated 3D boxes.
@@ -1116,10 +1185,12 @@ def _render_annotation_3ds(
                 for ann_token in sample.ann_3ds:
                     ann: SampleAnnotation = self.get("sample_annotation", ann_token)
                     if ann.instance_token == instance_token:
-                        boxes.append(self.get_box3d(ann_token))
+                        boxes.append(self.get_box3d(ann_token, future_seconds=future_seconds))
                         break
             else:
-                boxes = list(map(self.get_box3d, sample.ann_3ds))
+                boxes = [
+                    self.get_box3d(token, future_seconds=future_seconds) for token in sample.ann_3ds
+                ]
             viewer.render_box3ds(us2sec(sample.timestamp), boxes)
 
             current_sample_token = sample.next

t4_devkit/viewer/rendering_data/box.py

@@ -7,8 +7,14 @@
 from attrs import define, field
 
 if TYPE_CHECKING:
-    from t4_devkit.dataclass import Box2D, Box3D
-    from t4_devkit.typing import RoiType, RotationType, SizeType, TranslationType, VelocityType
+    from t4_devkit.dataclass import Box2D, Box3D, Trajectory
+    from t4_devkit.typing import (
+        RoiType,
+        RotationType,
+        SizeType,
+        TranslationType,
+        VelocityType,
+    )
 
 __all__ = ["BoxData3D", "BoxData2D"]
 
@@ -35,6 +41,7 @@ class BoxData3D:
     class_ids: list[int] = field(init=False, factory=list)
     uuids: list[str] = field(init=False, factory=list)
     velocities: list[VelocityType] = field(init=False, factory=list)
+    future: list[list[Trajectory]] = field(init=False, factory=list)
 
     @overload
     def append(self, box: Box3D) -> None:
@@ -54,6 +61,7 @@ def append(
         class_id: int,
         uuid: str | None = None,
         velocity: VelocityType | None = None,
+        future: list[Trajectory] | None = None,
     ) -> None:
         """Append a 3D box data with its elements.
 
@@ -63,7 +71,8 @@ def append(
             size (SizeType): Box size in the order of (width, height, length).
             class_id (int): Class ID.
             uuid (str | None, optional): Unique identifier.
-            velocity (VelocityType | None, optional): Box velocity. Defaults to None.
+            velocity (VelocityType | None, optional): Box velocity.
+            future (list[Trajectory] | None, optional): Future trajectory.
         """
         pass
 
@@ -93,6 +102,9 @@ def _append_with_box(self, box: Box3D) -> None:
         if box.uuid is not None:
             self.uuids.append(box.uuid[:6])
 
+        if box.future is not None:
+            self.future.append(box.future)
+
     def _append_with_elements(
         self,
         center: TranslationType,
@@ -101,6 +113,7 @@ def _append_with_elements(
         class_id: int,
         velocity: VelocityType | None = None,
         uuid: str | None = None,
+        future: list[Trajectory] | None = None,
     ) -> None:
         self.centers.append(center)
 
@@ -118,6 +131,9 @@ def _append_with_elements(
         if uuid is not None:
             self.uuids.append(uuid)
 
+        if future is not None:
+            self.future.append(future)
+
     def as_boxes3d(self) -> rr.Boxes3D:
         """Return 3D boxes data as a `rr.Boxes3D`.
 
@@ -145,6 +161,19 @@ def as_arrows3d(self) -> rr.Arrows3D:
             class_ids=self.class_ids,
         )
 
+    def as_linestrips3d(self) -> rr.LineStrips3D:
+        """Return future trajectories data as a list of `rr.LineStrips3D`.
+
+        Returns:
+            `rr.LineStrips3D` object for each box.
+        """
+        stripes = []
+        class_ids = []
+        for class_id, modes in zip(self.class_ids, self.future):
+            class_ids += [class_id] * len(modes)
+            stripes += [x.waypoints for x in modes]
+        return rr.LineStrips3D(strips=stripes, class_ids=class_ids)
+
 
 @define
 class BoxData2D:

t4_devkit/viewer/viewer.py

@@ -17,7 +17,7 @@
 from .rendering_data import BoxData2D, BoxData3D, SegmentationData2D
 
 if TYPE_CHECKING:
-    from t4_devkit.dataclass import Box2D, Box3D, PointCloudLike
+    from t4_devkit.dataclass import Box2D, Box3D, PointCloudLike, Trajectory
     from t4_devkit.schema import CalibratedSensor, EgoPose, Sensor
     from t4_devkit.typing import (
         CamIntrinsicType,
@@ -259,6 +259,11 @@ def _render_box3ds_with_boxes(self, seconds: float, boxes: Sequence[Box3D]) -> N
                 format_entity(self.map_entity, frame_id, "velocity"),
                 data.as_arrows3d(),
             )
+            # record futures
+            rr.log(
+                format_entity(self.map_entity, frame_id, "future"),
+                data.as_linestrips3d(),
+            )
 
     def _render_box3ds_with_elements(
         self,
@@ -269,6 +274,7 @@ def _render_box3ds_with_elements(
         class_ids: Sequence[int],
         velocities: Sequence[VelocityType] | None = None,
         uuids: Sequence[str] | None | None = None,
+        future: Sequence[Sequence[Trajectory]] | None = None,
     ) -> None:
         if uuids is None:
             uuids = [None] * len(centers)
@@ -279,9 +285,22 @@ def _render_box3ds_with_elements(
         else:
             show_arrows = True
 
+        if future is None:
+            futures = [None] * len(centers)
+            show_futures = False
+        else:
+            show_futures = True
+
         box_data = BoxData3D(label2id=self.label2id)
-        for center, rotation, size, class_id, velocity, uuid in zip(
-            centers, rotations, sizes, class_ids, velocities, uuids, strict=True
+        for center, rotation, size, class_id, velocity, uuid, future in zip(
+            centers,
+            rotations,
+            sizes,
+            class_ids,
+            velocities,
+            uuids,
+            futures,
+            strict=True,
         ):
             box_data.append(
                 center=center,
@@ -290,6 +309,7 @@ def _render_box3ds_with_elements(
                 class_id=class_id,
                 velocity=velocity,
                 uuid=uuid,
+                future=future,
             )
 
         rr.set_time_seconds(self.timeline, seconds)
@@ -299,6 +319,9 @@ def _render_box3ds_with_elements(
         if show_arrows:
             rr.log(format_entity(self.ego_entity, "velocity"), box_data.as_arrows3d())
 
+        if show_futures:
+            rr.log(format_entity(self.ego_entity, "future"), box_data.as_linestrips3d())
+
     @overload
     def render_box2ds(self, seconds: float, boxes: Sequence[Box2D]) -> None:
         """Render 2D boxes. Note that if the viewer initialized without `cameras=None`,