From a72e71c8ff4a7a7aaf08681394f1699ecf5ff871 Mon Sep 17 00:00:00 2001
From: ktro2828 <kotaro.uetake@tier4.jp>
Date: Thu, 21 Nov 2024 19:26:32 +0900
Subject: [PATCH] feat: add support of rendering geocoordinate

Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>
---
 t4_devkit/schema/tables/ego_pose.py |  6 ++--
 t4_devkit/viewer/viewer.py          | 43 +++++++++++++++++++++++++----
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/t4_devkit/schema/tables/ego_pose.py b/t4_devkit/schema/tables/ego_pose.py
index e5b0d31..9963d20 100644
--- a/t4_devkit/schema/tables/ego_pose.py
+++ b/t4_devkit/schema/tables/ego_pose.py
@@ -45,12 +45,12 @@ class EgoPose(SchemaBase):
     translation: TranslationType = field(converter=np.asarray)
     rotation: RotationType = field(converter=as_quaternion)
     timestamp: int
-    twist: TwistType = field(
+    twist: TwistType | None = field(
         default=None, converter=lambda x: np.asarray(x) if x is not None else x
     )
-    acceleration: AccelerationType = field(
+    acceleration: AccelerationType | None = field(
         default=None, converter=lambda x: np.asarray(x) if x is not None else x
     )
-    geocoordinate: GeoCoordinateType = field(
+    geocoordinate: GeoCoordinateType | None = field(
         default=None, converter=lambda x: np.asarray(x) if x is not None else x
     )
diff --git a/t4_devkit/viewer/viewer.py b/t4_devkit/viewer/viewer.py
index db4d0c1..29643fb 100644
--- a/t4_devkit/viewer/viewer.py
+++ b/t4_devkit/viewer/viewer.py
@@ -12,7 +12,13 @@
 
 from t4_devkit.common.timestamp import us2sec
 from t4_devkit.schema import CalibratedSensor, EgoPose, Sensor, SensorModality
-from t4_devkit.typing import CamIntrinsicType, NDArrayU8, RotationType, TranslationType
+from t4_devkit.typing import (
+    CamIntrinsicType,
+    GeoCoordinateType,
+    NDArrayU8,
+    RotationType,
+    TranslationType,
+)
 
 from .color import distance_color
 from .rendering_data import BoxData2D, BoxData3D, SegmentationData2D
@@ -59,6 +65,7 @@ class Tier4Viewer:
     # entity paths
     map_entity = "map"
     ego_entity = "map/base_link"
+    geocoordinate_entity = "geocoordinate"
     timeline = "timestamp"
 
     def __init__(
@@ -95,11 +102,14 @@ def __init__(
 
         view_container = []
         if not without_3d:
-            view_container.append(
-                rrb.Horizontal(
-                    rrb.Spatial3DView(name="3D", origin=self.map_entity),
-                    column_shares=[3, 1],
-                )
+            view_container.extend(
+                [
+                    rrb.Horizontal(
+                        rrb.Spatial3DView(name="3D", origin=self.map_entity),
+                        rrb.Horizontal(rrb.MapView(name="Map", origin=self.geocoordinate_entity)),
+                        column_shares=[3, 1],
+                    ),
+                ]
             )
 
         if self.cameras is not None:
@@ -308,12 +318,20 @@ def _render_ego_with_schema(self, ego_pose: EgoPose) -> None:
             ),
         )
 
+        if ego_pose.geocoordinate is not None:
+            latitude, longitude, _ = ego_pose.geocoordinate
+            rr.log(
+                self.geocoordinate_entity,
+                rr.GeoPoints(lat_lon=(latitude, longitude), radii=rr.Radius.ui_points(10.0)),
+            )
+
     @render_ego.register
     def _render_ego_without_schema(
         self,
         seconds: float,
         translation: TranslationType,
         rotation: RotationType,
+        geocoordinate: GeoCoordinateType | None = None,
     ) -> None:
         rr.set_time_seconds(self.timeline, seconds)
 
@@ -327,6 +345,13 @@ def _render_ego_without_schema(
             ),
         )
 
+        if geocoordinate is not None:
+            latitude, longitude, _ = geocoordinate
+            rr.log(
+                self.geocoordinate_entity,
+                rr.GeoPoints(lat_lon=(latitude, longitude)),
+            )
+
     @singledispatchmethod
     def render_calibration(self, *args, **kwargs) -> None:
         raise TypeError("Unexpected parameter types.")
@@ -394,3 +419,9 @@ def _render_calibration_without_schema(
                 rr.Pinhole(image_from_camera=camera_intrinsic),
                 static=True,
             )
+        if modality == SensorModality.CAMERA:
+            rr.log(
+                format_entity(self.ego_entity, channel),
+                rr.Pinhole(image_from_camera=camera_intrinsic),
+                static=True,
+            )