From 7501a046184afae69c141e4b2b42adb1b30e0180 Mon Sep 17 00:00:00 2001 From: QuanyiLi Date: Sat, 6 May 2023 21:06:51 +0100 Subject: [PATCH] all convert is ready now --- scenarionet/converter/nuplan/utils.py | 1 + scenarionet/converter/nuscenes/utils.py | 1 + scenarionet/converter/utils.py | 2 +- scenarionet/converter/waymo/utils.py | 5 +++-- scenarionet/examples/convert_nuplan.py | 4 ++-- 5 files changed, 8 insertions(+), 5 deletions(-) diff --git a/scenarionet/converter/nuplan/utils.py b/scenarionet/converter/nuplan/utils.py index 77830cf..2ff411e 100644 --- a/scenarionet/converter/nuplan/utils.py +++ b/scenarionet/converter/nuplan/utils.py @@ -415,6 +415,7 @@ def convert_nuplan_scenario(scenario: NuPlanScenario, version): result[SD.METADATA]["coordinate"] = "right-handed" result[SD.METADATA]["scenario_token"] = scenario.scenario_name result[SD.METADATA]["scenario_id"] = scenario.scenario_name + result[SD.METADATA][SD.ID] = scenario.scenario_name result[SD.METADATA]["scenario_type"] = scenario.scenario_type result[SD.METADATA]["sample_rate"] = scenario_log_interval result[SD.METADATA][SD.TIMESTEP] = np.asarray([i * scenario_log_interval for i in range(result[SD.LENGTH])]) diff --git a/scenarionet/converter/nuscenes/utils.py b/scenarionet/converter/nuscenes/utils.py index 9b4bd7b..c22eef4 100644 --- a/scenarionet/converter/nuscenes/utils.py +++ b/scenarionet/converter/nuscenes/utils.py @@ -370,6 +370,7 @@ def convert_nuscenes_scenario(scene, version, nuscenes: NuScenes): result[SD.METADATA]["date"] = nuscenes.get("log", scene_info["log_token"])["date_captured"] result[SD.METADATA]["coordinate"] = "right-handed" result[SD.METADATA]["scenario_token"] = scene_token + result[SD.METADATA][SD.ID] = scene_info["name"] result[SD.METADATA]["scenario_id"] = scene_info["name"] result[SD.METADATA]["sample_rate"] = scenario_log_interval result[SD.METADATA][SD.TIMESTEP] = np.arange(0., (len(frames) - 1) * 0.5 + 0.1, 0.1) diff --git a/scenarionet/converter/utils.py b/scenarionet/converter/utils.py index df0108c..d25a32e 100644 --- a/scenarionet/converter/utils.py +++ b/scenarionet/converter/utils.py @@ -79,7 +79,7 @@ def write_to_directory(convert_func, if not contains_explicit_return(convert_func): raise RuntimeError("The convert function should return a metadata dict") - if "version" not in kwargs: + if "version" in kwargs: kwargs.pop("version") logger.info("the specified version in kwargs is replaced by argument: 'dataset_version'") diff --git a/scenarionet/converter/waymo/utils.py b/scenarionet/converter/waymo/utils.py index 24546c2..d055cb2 100644 --- a/scenarionet/converter/waymo/utils.py +++ b/scenarionet/converter/waymo/utils.py @@ -382,13 +382,14 @@ def convert_waymo_scenario(scenario, version): compute_width(md_scenario[SD.MAP_FEATURES]) md_scenario[SD.METADATA] = {} + md_scenario[SD.METADATA][SD.ID] = md_scenario[SD.ID] md_scenario[SD.METADATA][SD.COORDINATE] = MetaDriveType.COORDINATE_WAYMO md_scenario[SD.METADATA][SD.TIMESTEP] = np.asarray(list(scenario_pb2.timestamps_seconds), dtype=np.float32) md_scenario[SD.METADATA][SD.METADRIVE_PROCESSED] = False md_scenario[SD.METADATA][SD.SDC_ID] = str(sdc_id) md_scenario[SD.METADATA]["dataset"] = "waymo" md_scenario[SD.METADATA]["scenario_id"] = scenario_pb2.scenario_id - # TODO Can we infer it? + # TODO LQY Can we infer it? # md_scenario[SD.METADATA]["source_file"] = str(file) md_scenario[SD.METADATA]["track_length"] = track_length @@ -425,5 +426,5 @@ def get_waymo_scenarios(waymo_data_direction): file_path = os.path.join(waymo_data_direction, file) if ("tfrecord" not in file_path) or (not os.path.isfile(file_path)): continue - scenarios += [s for s in tf.data.TFRecordDataset(file_path, compression_type="")] + scenarios += [s for s in tf.data.TFRecordDataset(file_path, compression_type="").as_numpy_iterator()] return scenarios diff --git a/scenarionet/examples/convert_nuplan.py b/scenarionet/examples/convert_nuplan.py index c1e1800..f7b73ba 100644 --- a/scenarionet/examples/convert_nuplan.py +++ b/scenarionet/examples/convert_nuplan.py @@ -5,7 +5,7 @@ MetaDrive. import os from scenarionet import SCENARIONET_DATASET_PATH -from scenarionet.converter.nuplan.utils import get_nuplan_scenarios, convert_nuplan_scenario, example_scenario_types +from scenarionet.converter.nuplan.utils import get_nuplan_scenarios, convert_nuplan_scenario, example_dataset_params from scenarionet.converter.utils import write_to_directory if __name__ == "__main__": @@ -14,7 +14,7 @@ if __name__ == "__main__": output_path = os.path.join(SCENARIONET_DATASET_PATH, dataset_name) version = 'v1.1' - scenarios = get_nuplan_scenarios(example_scenario_types) + scenarios = get_nuplan_scenarios(example_dataset_params) write_to_directory(convert_func=convert_nuplan_scenario, scenarios=scenarios,