rename

scenarionet/converter/__init__.py

@@ -1,3 +1,3 @@
-from scenarionet.converter.nuscenes.utils import convert_one_nuscenes_scenario
-from scenarionet.converter.nuplan.utils import convert_one_nuplan_scenario
+from scenarionet.converter.nuscenes.utils import convert_nuscenes_scenario
+from scenarionet.converter.nuplan.utils import convert_nuplan_scenario
 from scenarionet.converter.utils import write_to_directory

scenarionet/converter/nuplan/utils.py

@@ -388,7 +388,7 @@ def extract_traffic(scenario: NuPlanScenario, center):
     return tracks
 
 
-def convert_one_nuplan_scenario(scenario: NuPlanScenario):
+def convert_nuplan_scenario(scenario: NuPlanScenario):
     """
     Data will be interpolated to 0.1s time interval, while the time interval of original key frames are 0.5s.
     """

scenarionet/converter/nuscenes/utils.py

@@ -345,7 +345,7 @@ def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=250, poi
     return ret
 
 
-def convert_one_nuscenes_scenario(scene, nuscenes: NuScenes):
+def convert_nuscenes_scenario(scene, nuscenes: NuScenes):
     """
     Data will be interpolated to 0.1s time interval, while the time interval of original key frames are 0.5s.
     """
@@ -386,4 +386,4 @@ def convert_one_nuscenes_scenario(scene, nuscenes: NuScenes):
     map_center = result[SD.TRACKS]["ego"]["state"]["position"][0]
     result[SD.MAP_FEATURES] = get_map_features(scene_info, nuscenes, map_center, 250)
 
-    return result
+    return result, scene_token

scenarionet/converter/scripts/convert_nuplan.py

@@ -10,7 +10,7 @@ import shutil
 import tqdm
 from metadrive.scenario.scenario_description import ScenarioDescription
 
-from scenarionet.converter.nuplan.utils import get_nuplan_scenarios, convert_one_nuplan_scenario
+from scenarionet.converter.nuplan.utils import get_nuplan_scenarios, convert_nuplan_scenario
 from scenarionet.converter.utils import dict_recursive_remove_array
 
 
@@ -41,7 +41,7 @@ def convert_nuplan(dataset_params, output_path, worker_index=None, force_overwri
     # Init.
     scenarios = get_nuplan_scenarios(dataset_params)
     for scenario in tqdm.tqdm(scenarios):
-        sd_scenario = convert_one_nuplan_scenario(scenario)
+        sd_scenario = convert_nuplan_scenario(scenario)
         sd_scenario = sd_scenario.to_dict()
         ScenarioDescription.sanity_check(sd_scenario, check_self_type=True)
         export_file_name = "sd_{}_{}.pkl".format("nuplan", scenario.scenario_name)

scenarionet/converter/scripts/convert_nuscenes.py

@@ -2,28 +2,26 @@
 This script aims to convert nuscenes scenarios to ScenarioDescription, so that we can load any nuscenes scenarios into
 MetaDrive.
 """
+import os.path
 
+from nuscenes import NuScenes
 from scenarionet import SCENARIONET_DATASET_PATH
-from scenarionet.converter.nuscenes.utils import convert_one_nuscenes_scenario
+from scenarionet.converter.nuscenes.utils import convert_nuscenes_scenario
 from scenarionet.converter.utils import write_to_directory
 
-try:
-    from nuscenes import NuScenes
-except ImportError:
-    print("Can not find nuscenes-devkit")
-
 #
 if __name__ == "__main__":
-    output_path = SCENARIONET_DATASET_PATH
+    output_path = os.path.join(SCENARIONET_DATASET_PATH, "nuscenes")
     version = 'v1.0-mini'
     dataroot = '/home/shady/data/nuscenes'
     force_overwrite = True
     nusc = NuScenes(version=version, dataroot=dataroot)
     scenarios = nusc.scene
-    write_to_directory(convert_func=convert_one_nuscenes_scenario,
+
+    write_to_directory(convert_func=convert_nuscenes_scenario,
                        scenarios=scenarios,
                        output_path=output_path,
-                       version=version,
+                       dataset_version=version,
                        dataset_name="nuscenes",
-                       force_overwrite=True,
+                       force_overwrite=force_overwrite,
                        nuscenes=nusc)

scenarionet/converter/utils.py

@@ -1,7 +1,7 @@
-import copy
-import math
 import ast
+import copy
 import inspect
+import math
 import os
 import pickle
 import shutil
@@ -112,7 +112,16 @@ def contains_explicit_return(f):
     return any(isinstance(node, ast.Return) for node in ast.walk(ast.parse(inspect.getsource(f))))
 
 
-def write_to_directory(convert_func, scenarios, output_path, version, dataset_name, force_overwrite=False, **kwargs):
+def write_to_directory(convert_func,
+                       scenarios,
+                       output_path,
+                       dataset_version,
+                       dataset_name,
+                       force_overwrite=False,
+                       **kwargs):
+    """
+    Convert a batch of scenarios.
+    """
     if not contains_explicit_return(convert_func):
         raise RuntimeError("The convert function should return a metadata dict")
 
@@ -135,10 +144,29 @@ def write_to_directory(convert_func, scenarios, output_path, version, dataset_na
 
     metadata_recorder = {}
     for scenario in tqdm.tqdm(scenarios):
-        sd_scenario = convert_func(scenario, **kwargs)
+        # convert scenario
+        sd_scenario, scenario_id = convert_func(scenario, **kwargs)
+        export_file_name = "sd_{}_{}.pkl".format(dataset_name + "_" + dataset_version, scenario_id)
+
+        # add agents summary
+        summary_dict = {}
+        ego_car_id = sd_scenario[SD.METADATA][SD.SDC_ID]
+        summary_dict[ego_car_id] = get_agent_summary(
+            state_dict=sd_scenario.get_sdc_track()["state"], id=ego_car_id, type=sd_scenario.get_sdc_track()["type"]
+        )
+        for track_id, track in sd_scenario[SD.TRACKS].items():
+            summary_dict[track_id] = get_agent_summary(state_dict=track["state"], id=track_id, type=track["type"])
+        sd_scenario[SD.METADATA]["object_summary"] = summary_dict
+
+        # count some objects occurrence
+        sd_scenario[SD.METADATA]["number_summary"] = get_number_summary(sd_scenario)
+        metadata_recorder[export_file_name] = copy.deepcopy(sd_scenario[SD.METADATA])
+
+        # sanity check
         sd_scenario = sd_scenario.to_dict()
         ScenarioDescription.sanity_check(sd_scenario, check_self_type=True)
-        export_file_name = "sd_{}_{}.pkl".format(dataset_name+"_" + version, scenario["token"])
+
+        # dump
         p = os.path.join(output_path, export_file_name)
         with open(p, "wb") as f:
             pickle.dump(sd_scenario, f)