huangfu/scenarionet
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

remove set as well

Browse Source
This commit is contained in:
QuanyiLi
2023-05-06 17:52:20 +01:00
parent 13749b6522
commit 18f85f6d27
3 changed files with 15 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
scenarionet/converter/utils.py
View File

@@ -1,5 +1,4 @@
import ast import ast
import json
import copy import copy
import inspect import inspect
import math import math
@@ -45,12 +44,14 @@ def compute_angular_velocity(initial_heading, final_heading, dt):
return angular_vel return angular_vel
def dict_recursive_remove_array(d): def dict_recursive_remove_array_and_set(d):
if isinstance(d, np.ndarray): if isinstance(d, np.ndarray):
return d.tolist() return d.tolist()
if isinstance(d, set):
return tuple(d)
if isinstance(d, dict): if isinstance(d, dict):
for k in d.keys(): 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 return d
@@ -141,7 +142,7 @@ def write_to_directory(convert_func,
else: else:
raise ValueError("Directory already exists! Abort") raise ValueError("Directory already exists! Abort")
summary_file = "dataset_summary.json" summary_file = "dataset_summary.pkl"
metadata_recorder = {} metadata_recorder = {}
for scenario in tqdm.tqdm(scenarios): for scenario in tqdm.tqdm(scenarios):
@@ -172,13 +173,14 @@ def write_to_directory(convert_func,
with open(p, "wb") as f: with open(p, "wb") as f:
pickle.dump(sd_scenario, 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 # rename and save
if delay_remove is not None: if delay_remove is not None:
assert delay_remove == save_path assert delay_remove == save_path
shutil.rmtree(delay_remove) shutil.rmtree(delay_remove)
os.rename(output_path, save_path) 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))

4
scenarionet/examples/convert_nuplan.py
View File

@@ -11,7 +11,7 @@ import tqdm
from metadrive.scenario.scenario_description import ScenarioDescription from metadrive.scenario.scenario_description import ScenarioDescription
from scenarionet.converter.nuplan.utils import get_nuplan_scenarios, convert_nuplan_scenario 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): 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) os.rename(output_path, save_path)
summary_file = os.path.join(save_path, summary_file) summary_file = os.path.join(save_path, summary_file)
with open(summary_file, "wb") as 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)) print("Summary is saved at: {}".format(summary_file))
if delay_remove is not None: if delay_remove is not None:
assert delay_remove == save_path, delay_remove + " vs. " + save_path assert delay_remove == save_path, delay_remove + " vs. " + save_path

4
scenarionet/examples/convert_waymo.py
View File

@@ -13,7 +13,7 @@ import pickle
import numpy as np 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: try:
import tensorflow as tf 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) summary_file = os.path.join(output_path, summary_file)
with open(summary_file, "wb") as 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)) print("Summary is saved at: {}".format(summary_file))