From 730efa86a84fa3f8ba41c331dad0f66c5b2a1267 Mon Sep 17 00:00:00 2001 From: QuanyiLi Date: Sun, 7 May 2023 23:18:45 +0100 Subject: [PATCH] format --- scenarionet/builder/utils.py | 7 +-- scenarionet/converter/nuplan/utils.py | 4 +- scenarionet/converter/pg/utils.py | 16 +++-- scenarionet/converter/utils.py | 2 +- .../examples/combine_dataset_and_run.py | 18 +++--- .../local_test/_test_combine_dataset_local.py | 10 ++-- .../tests/local_test/_test_filter_local.py | 40 ++++++------- .../tests/script/generate_test_dataset.py | 3 +- scenarionet/tests/script/run_env.py | 8 +-- scenarionet/tests/test_combine_dataset.py | 12 ++-- scenarionet/tests/test_filter.py | 59 ++++++++----------- scenarionet/verifier/utils.py | 27 ++++----- 12 files changed, 96 insertions(+), 110 deletions(-) diff --git a/scenarionet/builder/utils.py b/scenarionet/builder/utils.py index 93f4521..5e871cb 100644 --- a/scenarionet/builder/utils.py +++ b/scenarionet/builder/utils.py @@ -41,10 +41,9 @@ def try_generating_mapping(file_folder): return mapping -def combine_multiple_dataset(output_path, *dataset_paths, - force_overwrite=False, - try_generate_missing_file=True, - filters: List[Callable] = None): +def combine_multiple_dataset( + output_path, *dataset_paths, force_overwrite=False, try_generate_missing_file=True, filters: List[Callable] = None +): """ Combine multiple datasets. Each dataset should have a dataset_summary.pkl :param output_path: The path to store the output dataset diff --git a/scenarionet/converter/nuplan/utils.py b/scenarionet/converter/nuplan/utils.py index 9462fa2..4e1db4b 100644 --- a/scenarionet/converter/nuplan/utils.py +++ b/scenarionet/converter/nuplan/utils.py @@ -357,9 +357,9 @@ def extract_traffic(scenario: NuPlanScenario, center): type=MetaDriveType.UNSET, state=dict( position=np.zeros(shape=(episode_len, 3)), - heading=np.zeros(shape=(episode_len,)), + heading=np.zeros(shape=(episode_len, )), velocity=np.zeros(shape=(episode_len, 2)), - valid=np.zeros(shape=(episode_len,)), + valid=np.zeros(shape=(episode_len, )), length=np.zeros(shape=(episode_len, 1)), width=np.zeros(shape=(episode_len, 1)), height=np.zeros(shape=(episode_len, 1)) diff --git a/scenarionet/converter/pg/utils.py b/scenarionet/converter/pg/utils.py index d5e66d0..f59f600 100644 --- a/scenarionet/converter/pg/utils.py +++ b/scenarionet/converter/pg/utils.py @@ -17,10 +17,14 @@ def convert_pg_scenario(scenario_index, version, env): def get_pg_scenarios(num_scenarios, policy, start_seed=0): - env = MetaDriveEnv(dict(start_seed=start_seed, - num_scenarios=num_scenarios, - traffic_density=0.2, - agent_policy=policy, - crash_vehicle_done=False, - map=2)) + env = MetaDriveEnv( + dict( + start_seed=start_seed, + num_scenarios=num_scenarios, + traffic_density=0.2, + agent_policy=policy, + crash_vehicle_done=False, + map=2 + ) + ) return [i for i in range(num_scenarios)], env diff --git a/scenarionet/converter/utils.py b/scenarionet/converter/utils.py index b3a30b0..428a7d1 100644 --- a/scenarionet/converter/utils.py +++ b/scenarionet/converter/utils.py @@ -67,7 +67,7 @@ def contains_explicit_return(f): def write_to_directory( - convert_func, scenarios, output_path, dataset_version, dataset_name, force_overwrite=False, **kwargs + convert_func, scenarios, output_path, dataset_version, dataset_name, force_overwrite=False, **kwargs ): """ Convert a batch of scenarios. diff --git a/scenarionet/examples/combine_dataset_and_run.py b/scenarionet/examples/combine_dataset_and_run.py index bf8beb9..04b881d 100644 --- a/scenarionet/examples/combine_dataset_and_run.py +++ b/scenarionet/examples/combine_dataset_and_run.py @@ -8,10 +8,12 @@ from scenarionet import SCENARIONET_DATASET_PATH from scenarionet.builder.utils import combine_multiple_dataset if __name__ == '__main__': - dataset_paths = [os.path.join(SCENARIONET_DATASET_PATH, "nuscenes"), - os.path.join(SCENARIONET_DATASET_PATH, "nuplan"), - os.path.join(SCENARIONET_DATASET_PATH, "waymo"), - os.path.join(SCENARIONET_DATASET_PATH, "pg")] + dataset_paths = [ + os.path.join(SCENARIONET_DATASET_PATH, "nuscenes"), + os.path.join(SCENARIONET_DATASET_PATH, "nuplan"), + os.path.join(SCENARIONET_DATASET_PATH, "waymo"), + os.path.join(SCENARIONET_DATASET_PATH, "pg") + ] combine_path = os.path.join(SCENARIONET_DATASET_PATH, "combined_dataset") combine_multiple_dataset(combine_path, *dataset_paths, force_overwrite=True, try_generate_missing_file=True) @@ -47,11 +49,9 @@ if __name__ == '__main__': c_lane = env.vehicle.lane long, lat, = c_lane.local_coordinates(env.vehicle.position) # if env.config["use_render"]: - env.render( - text={ - "seed": env.engine.global_seed + env.config["start_scenario_index"], - } - ) + env.render(text={ + "seed": env.engine.global_seed + env.config["start_scenario_index"], + }) if d: if info["arrive_dest"]: diff --git a/scenarionet/tests/local_test/_test_combine_dataset_local.py b/scenarionet/tests/local_test/_test_combine_dataset_local.py index 000953e..7045d14 100644 --- a/scenarionet/tests/local_test/_test_combine_dataset_local.py +++ b/scenarionet/tests/local_test/_test_combine_dataset_local.py @@ -6,10 +6,12 @@ from scenarionet.verifier.utils import verify_loading_into_metadrive def _test_combine_dataset(): - dataset_paths = [os.path.join(SCENARIONET_DATASET_PATH, "nuscenes"), - os.path.join(SCENARIONET_DATASET_PATH, "nuplan"), - os.path.join(SCENARIONET_DATASET_PATH, "waymo"), - os.path.join(SCENARIONET_DATASET_PATH, "pg")] + dataset_paths = [ + os.path.join(SCENARIONET_DATASET_PATH, "nuscenes"), + os.path.join(SCENARIONET_DATASET_PATH, "nuplan"), + os.path.join(SCENARIONET_DATASET_PATH, "waymo"), + os.path.join(SCENARIONET_DATASET_PATH, "pg") + ] combine_path = os.path.join(SCENARIONET_DATASET_PATH, "combined_dataset") combine_multiple_dataset(combine_path, *dataset_paths, force_overwrite=True, try_generate_missing_file=True) diff --git a/scenarionet/tests/local_test/_test_filter_local.py b/scenarionet/tests/local_test/_test_filter_local.py index 6e83a1d..1f02a4c 100644 --- a/scenarionet/tests/local_test/_test_filter_local.py +++ b/scenarionet/tests/local_test/_test_filter_local.py @@ -22,40 +22,34 @@ def test_filter_dataset(): # ========================= test 1 ========================= # nuscenes data has no light # light_condition = ScenarioFilter.make(ScenarioFilter.has_traffic_light) - sdc_driving_condition = ScenarioFilter.make(ScenarioFilter.sdc_moving_dist, - target_dist=30, - condition="greater") - summary, mapping = combine_multiple_dataset(output_path, - *dataset_paths, - force_overwrite=True, - try_generate_missing_file=True, - filters=[sdc_driving_condition] - ) + sdc_driving_condition = ScenarioFilter.make(ScenarioFilter.sdc_moving_dist, target_dist=30, condition="greater") + summary, mapping = combine_multiple_dataset( + output_path, + *dataset_paths, + force_overwrite=True, + try_generate_missing_file=True, + filters=[sdc_driving_condition] + ) assert len(summary) > 0 # ========================= test 2 ========================= - num_condition = ScenarioFilter.make(ScenarioFilter.object_number, - number_threshold=50, - object_type=MetaDriveType.PEDESTRIAN, - condition="greater") + num_condition = ScenarioFilter.make( + ScenarioFilter.object_number, number_threshold=50, object_type=MetaDriveType.PEDESTRIAN, condition="greater" + ) - summary, mapping = combine_multiple_dataset(output_path, - *dataset_paths, - force_overwrite=True, - try_generate_missing_file=True, - filters=[num_condition]) + summary, mapping = combine_multiple_dataset( + output_path, *dataset_paths, force_overwrite=True, try_generate_missing_file=True, filters=[num_condition] + ) assert len(summary) > 0 # ========================= test 3 ========================= traffic_light = ScenarioFilter.make(ScenarioFilter.has_traffic_light) - summary, mapping = combine_multiple_dataset(output_path, - *dataset_paths, - force_overwrite=True, - try_generate_missing_file=True, - filters=[traffic_light]) + summary, mapping = combine_multiple_dataset( + output_path, *dataset_paths, force_overwrite=True, try_generate_missing_file=True, filters=[traffic_light] + ) assert len(summary) > 0 diff --git a/scenarionet/tests/script/generate_test_dataset.py b/scenarionet/tests/script/generate_test_dataset.py index 71e50cb..e0f08b9 100644 --- a/scenarionet/tests/script/generate_test_dataset.py +++ b/scenarionet/tests/script/generate_test_dataset.py @@ -25,4 +25,5 @@ if __name__ == "__main__": dataset_version=version, dataset_name=dataset_name, force_overwrite=force_overwrite, - nuscenes=nusc) + nuscenes=nusc + ) diff --git a/scenarionet/tests/script/run_env.py b/scenarionet/tests/script/run_env.py index 090620d..94a89ef 100644 --- a/scenarionet/tests/script/run_env.py +++ b/scenarionet/tests/script/run_env.py @@ -49,11 +49,9 @@ if __name__ == '__main__': c_lane = env.vehicle.lane long, lat, = c_lane.local_coordinates(env.vehicle.position) if env.config["use_render"]: - env.render( - text={ - "seed": env.engine.global_seed + env.config["start_scenario_index"], - } - ) + env.render(text={ + "seed": env.engine.global_seed + env.config["start_scenario_index"], + }) if d and info["arrive_dest"]: print("seed:{}, success".format(env.engine.global_random_seed)) diff --git a/scenarionet/tests/test_combine_dataset.py b/scenarionet/tests/test_combine_dataset.py index 8f24603..b14bbb8 100644 --- a/scenarionet/tests/test_combine_dataset.py +++ b/scenarionet/tests/test_combine_dataset.py @@ -12,20 +12,16 @@ def test_combine_multiple_dataset(): dataset_paths = [original_dataset_path + "_{}".format(i) for i in range(5)] output_path = os.path.join(SCENARIONET_PACKAGE_PATH, "tests", "combine") - combine_multiple_dataset(output_path, - *dataset_paths, - force_overwrite=True, - try_generate_missing_file=True) + combine_multiple_dataset(output_path, *dataset_paths, force_overwrite=True, try_generate_missing_file=True) dataset_paths.append(output_path) for dataset_path in dataset_paths: summary, sorted_scenarios, mapping = read_dataset_summary(dataset_path) for scenario_file in sorted_scenarios: read_scenario(os.path.join(dataset_path, mapping[scenario_file], scenario_file)) num_worker = 4 if len(summary) > 4 else 1 - success, result = verify_loading_into_metadrive(dataset_path, - result_save_dir="test_dataset", - steps_to_run=1000, - num_workers=num_worker) + success, result = verify_loading_into_metadrive( + dataset_path, result_save_dir="test_dataset", steps_to_run=1000, num_workers=num_worker + ) assert success diff --git a/scenarionet/tests/test_filter.py b/scenarionet/tests/test_filter.py index db04858..3cda91b 100644 --- a/scenarionet/tests/test_filter.py +++ b/scenarionet/tests/test_filter.py @@ -18,15 +18,15 @@ def test_filter_dataset(): # ========================= test 1 ========================= # nuscenes data has no light # light_condition = ScenarioFilter.make(ScenarioFilter.has_traffic_light) - sdc_driving_condition = ScenarioFilter.make(ScenarioFilter.sdc_moving_dist, - target_dist=30, - condition="smaller") + sdc_driving_condition = ScenarioFilter.make(ScenarioFilter.sdc_moving_dist, target_dist=30, condition="smaller") answer = ['scene-0553', 'scene-0757', 'scene-1100'] - summary, mapping = combine_multiple_dataset(output_path, - *dataset_paths, - force_overwrite=True, - try_generate_missing_file=True, - filters=[sdc_driving_condition]) + summary, mapping = combine_multiple_dataset( + output_path, + *dataset_paths, + force_overwrite=True, + try_generate_missing_file=True, + filters=[sdc_driving_condition] + ) assert len(answer) == len(summary) for a in answer: in_ = False @@ -36,42 +36,35 @@ def test_filter_dataset(): break assert in_ - sdc_driving_condition = ScenarioFilter.make(ScenarioFilter.sdc_moving_dist, - target_dist=5, - condition="greater") - summary, mapping = combine_multiple_dataset(output_path, - *dataset_paths, - force_overwrite=True, - try_generate_missing_file=True, - filters=[sdc_driving_condition]) + sdc_driving_condition = ScenarioFilter.make(ScenarioFilter.sdc_moving_dist, target_dist=5, condition="greater") + summary, mapping = combine_multiple_dataset( + output_path, + *dataset_paths, + force_overwrite=True, + try_generate_missing_file=True, + filters=[sdc_driving_condition] + ) assert len(summary) == 8 # ========================= test 2 ========================= - num_condition = ScenarioFilter.make(ScenarioFilter.object_number, - number_threshold=50, - object_type=MetaDriveType.PEDESTRIAN, - condition="greater") + num_condition = ScenarioFilter.make( + ScenarioFilter.object_number, number_threshold=50, object_type=MetaDriveType.PEDESTRIAN, condition="greater" + ) answer = ['sd_nuscenes_v1.0-mini_scene-0061.pkl', 'sd_nuscenes_v1.0-mini_scene-1094.pkl'] - summary, mapping = combine_multiple_dataset(output_path, - *dataset_paths, - force_overwrite=True, - try_generate_missing_file=True, - filters=[num_condition]) + summary, mapping = combine_multiple_dataset( + output_path, *dataset_paths, force_overwrite=True, try_generate_missing_file=True, filters=[num_condition] + ) assert len(answer) == len(summary) for a in answer: assert a in summary - num_condition = ScenarioFilter.make(ScenarioFilter.object_number, - number_threshold=50, - condition="greater") + num_condition = ScenarioFilter.make(ScenarioFilter.object_number, number_threshold=50, condition="greater") - summary, mapping = combine_multiple_dataset(output_path, - *dataset_paths, - force_overwrite=True, - try_generate_missing_file=True, - filters=[num_condition]) + summary, mapping = combine_multiple_dataset( + output_path, *dataset_paths, force_overwrite=True, try_generate_missing_file=True, filters=[num_condition] + ) assert len(summary) > 0 diff --git a/scenarionet/verifier/utils.py b/scenarionet/verifier/utils.py index 4e63f0b..54fb23c 100644 --- a/scenarionet/verifier/utils.py +++ b/scenarionet/verifier/utils.py @@ -13,14 +13,12 @@ from functools import partial def verify_loading_into_metadrive(dataset_path, result_save_dir, steps_to_run=1000, num_workers=8): if result_save_dir is not None: - assert os.path.exists(result_save_dir) and os.path.isdir( - result_save_dir), "Argument result_save_dir must be an existing dir" + assert os.path.exists(result_save_dir + ) and os.path.isdir(result_save_dir), "Argument result_save_dir must be an existing dir" num_scenario = get_number_of_scenarios(dataset_path) argument_list = [] - func = partial(loading_wrapper, - dataset_path=dataset_path, - steps_to_run=steps_to_run) + func = partial(loading_wrapper, dataset_path=dataset_path, steps_to_run=steps_to_run) num_scenario_each_worker = int(num_scenario // num_workers) for i in range(num_workers): @@ -50,8 +48,11 @@ def verify_loading_into_metadrive(dataset_path, result_save_dir, steps_to_run=10 def loading_into_metadrive(start_scenario_index, num_scenario, dataset_path, steps_to_run): - print("================ Begin Scenario Loading Verification for scenario {}-{} ================ \n".format( - start_scenario_index, num_scenario + start_scenario_index)) + print( + "================ Begin Scenario Loading Verification for scenario {}-{} ================ \n".format( + start_scenario_index, num_scenario + start_scenario_index + ) + ) success = True env = ScenarioEnv( { @@ -81,9 +82,10 @@ def loading_into_metadrive(start_scenario_index, num_scenario, dataset_path, ste file_path = os.path.join(dataset_path, env.engine.data_manager.mapping[file_name], file_name) error_file = {"scenario_index": scenario_index, "file_path": file_path, "error": str(e)} error_files.append(error_file) - logger.warning("\n Scenario Error, " - "scenario_index: {}, file_path: {}.\n Error message: {}".format(scenario_index, file_path, - str(e))) + logger.warning( + "\n Scenario Error, " + "scenario_index: {}, file_path: {}.\n Error message: {}".format(scenario_index, file_path, str(e)) + ) success = False finally: env.close() @@ -92,7 +94,4 @@ def loading_into_metadrive(start_scenario_index, num_scenario, dataset_path, ste def loading_wrapper(arglist, dataset_path, steps_to_run): assert len(arglist) == 2, "Too much arguments!" - return loading_into_metadrive(arglist[0], - arglist[1], - dataset_path=dataset_path, - steps_to_run=steps_to_run) + return loading_into_metadrive(arglist[0], arglist[1], dataset_path=dataset_path, steps_to_run=steps_to_run)