From 18f85f6d27b7464aa58ff1d7760f26909cde621a Mon Sep 17 00:00:00 2001 From: QuanyiLi Date: Sat, 6 May 2023 17:52:20 +0100 Subject: [PATCH] remove set as well --- scenarionet/converter/utils.py | 20 +++++++++++--------- scenarionet/examples/convert_nuplan.py | 4 ++-- scenarionet/examples/convert_waymo.py | 4 ++-- 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/scenarionet/converter/utils.py b/scenarionet/converter/utils.py index 3b3db3d..af7a223 100644 --- a/scenarionet/converter/utils.py +++ b/scenarionet/converter/utils.py @@ -1,5 +1,4 @@ import ast -import json import copy import inspect import math @@ -45,12 +44,14 @@ def compute_angular_velocity(initial_heading, final_heading, dt): return angular_vel -def dict_recursive_remove_array(d): +def dict_recursive_remove_array_and_set(d): if isinstance(d, np.ndarray): return d.tolist() + if isinstance(d, set): + return tuple(d) if isinstance(d, dict): for k in d.keys(): - d[k] = dict_recursive_remove_array(d[k]) + d[k] = dict_recursive_remove_array_and_set(d[k]) return d @@ -141,7 +142,7 @@ def write_to_directory(convert_func, else: raise ValueError("Directory already exists! Abort") - summary_file = "dataset_summary.json" + summary_file = "dataset_summary.pkl" metadata_recorder = {} for scenario in tqdm.tqdm(scenarios): @@ -172,13 +173,14 @@ def write_to_directory(convert_func, with open(p, "wb") as f: pickle.dump(sd_scenario, f) + # store summary file, which is human-readable + summary_file = os.path.join(output_path, summary_file) + with open(summary_file, "wb") as file: + pickle.dump(dict_recursive_remove_array_and_set(metadata_recorder), file) + print("Summary is saved at: {}".format(summary_file)) + # rename and save if delay_remove is not None: assert delay_remove == save_path shutil.rmtree(delay_remove) os.rename(output_path, save_path) - summary_file = os.path.join(save_path, summary_file) - with open(summary_file, "wb") as file: - json.dump(dict_recursive_remove_array(metadata_recorder), file) - print("Summary is saved at: {}".format(summary_file)) - diff --git a/scenarionet/examples/convert_nuplan.py b/scenarionet/examples/convert_nuplan.py index 77938f4..7a7f3d2 100644 --- a/scenarionet/examples/convert_nuplan.py +++ b/scenarionet/examples/convert_nuplan.py @@ -11,7 +11,7 @@ import tqdm from metadrive.scenario.scenario_description import ScenarioDescription from scenarionet.converter.nuplan.utils import get_nuplan_scenarios, convert_nuplan_scenario -from scenarionet.converter.utils import dict_recursive_remove_array +from scenarionet.converter.utils import dict_recursive_remove_array_and_set def convert_nuplan(dataset_params, output_path, worker_index=None, force_overwrite=False): @@ -55,7 +55,7 @@ def convert_nuplan(dataset_params, output_path, worker_index=None, force_overwri os.rename(output_path, save_path) summary_file = os.path.join(save_path, summary_file) with open(summary_file, "wb") as file: - pickle.dump(dict_recursive_remove_array(metadata_recorder), file) + pickle.dump(dict_recursive_remove_array_and_set(metadata_recorder), file) print("Summary is saved at: {}".format(summary_file)) if delay_remove is not None: assert delay_remove == save_path, delay_remove + " vs. " + save_path diff --git a/scenarionet/examples/convert_waymo.py b/scenarionet/examples/convert_waymo.py index 6037d1b..c6a0809 100644 --- a/scenarionet/examples/convert_waymo.py +++ b/scenarionet/examples/convert_waymo.py @@ -13,7 +13,7 @@ import pickle import numpy as np -from scenarionet.converter.utils import dict_recursive_remove_array, get_agent_summary, get_number_summary +from scenarionet.converter.utils import dict_recursive_remove_array_and_set, get_agent_summary, get_number_summary try: import tensorflow as tf @@ -152,7 +152,7 @@ def convert_waymo(file_list, input_path, output_path, worker_index=None): summary_file = os.path.join(output_path, summary_file) with open(summary_file, "wb") as file: - pickle.dump(dict_recursive_remove_array(metadata_recorder), file) + pickle.dump(dict_recursive_remove_array_and_set(metadata_recorder), file) print("Summary is saved at: {}".format(summary_file))