From 1f61a4b9d14db1bc85fa2827fafa46edd217a0ee Mon Sep 17 00:00:00 2001
From: Alan <36124025+Alan-LanFeng@users.noreply.github.com>
Date: Thu, 18 Jan 2024 22:47:15 +0100
Subject: [PATCH] 1. add back map center to restore original coordinate in
 nuScnes (#51)

* 1. add back map center to restore the original coordinate in nuScenes
---
 scenarionet/converter/nuscenes/utils.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/scenarionet/converter/nuscenes/utils.py b/scenarionet/converter/nuscenes/utils.py
index 84392f2..9ef9d93 100644
--- a/scenarionet/converter/nuscenes/utils.py
+++ b/scenarionet/converter/nuscenes/utils.py
@@ -485,6 +485,16 @@ def convert_nuscenes_scenario(
     }
     # map
     result[SD.MAP_FEATURES] = get_map_features(scene_info, nuscenes, map_center, map_radius, only_lane=only_lane)
+    # add back map center
+    map_center = map_center[np.newaxis]
+    for k, v in result[SD.TRACKS].items():
+        v['state']['position'] += map_center
+
+    for k, v in result[SD.MAP_FEATURES].items():
+        if 'polygon' in v:
+            v['polygon'] += map_center[:, :v['polygon'].shape[-1]]
+        else:
+            v['polyline'] += map_center[:, :v['polyline'].shape[-1]]
     del frames_scene_info
     del frames
     del scene_info