diff --git a/scenarionet/converter/utils.py b/scenarionet/converter/utils.py index af7a223..365dfb1 100644 --- a/scenarionet/converter/utils.py +++ b/scenarionet/converter/utils.py @@ -9,7 +9,7 @@ from collections import defaultdict import numpy as np import tqdm -from metadrive.scenario import ScenarioDescription as SD, ScenarioDescription +from metadrive.scenario import ScenarioDescription as SD def nuplan_to_metadrive_vector(vector, nuplan_center=(0, 0)): @@ -60,56 +60,6 @@ def mph_to_kmh(speed_in_mph: float): return speed_in_kmh -def get_agent_summary(state_dict, id, type): - track = state_dict["position"] - valid_track = track[state_dict["valid"].astype(int), :2] - distance = float(sum(np.linalg.norm(valid_track[i] - valid_track[i + 1]) for i in range(valid_track.shape[0] - 1))) - valid_length = int(sum(state_dict["valid"])) - - continuous_valid_length = 0 - for v in state_dict["valid"]: - if v: - continuous_valid_length += 1 - if continuous_valid_length > 0 and not v: - break - - return { - "type": type, - "object_id": str(id), - "track_length": int(len(track)), - "distance": float(distance), - "valid_length": int(valid_length), - "continuous_valid_length": int(continuous_valid_length) - } - - -def get_number_summary(scenario): - number_summary_dict = {} - number_summary_dict["object"] = len(scenario[SD.TRACKS]) - number_summary_dict["dynamic_object_states"] = len(scenario[SD.DYNAMIC_MAP_STATES]) - number_summary_dict["map_features"] = len(scenario[SD.MAP_FEATURES]) - number_summary_dict["object_types"] = set(v["type"] for v in scenario[SD.TRACKS].values()) - - object_types_counter = defaultdict(int) - for v in scenario[SD.TRACKS].values(): - object_types_counter[v["type"]] += 1 - number_summary_dict["object_types_counter"] = dict(object_types_counter) - - # Number of different dynamic object states - dynamic_object_states_types = set() - dynamic_object_states_counter = defaultdict(int) - for v in scenario[SD.DYNAMIC_MAP_STATES].values(): - for step_state in v["state"]["object_state"]: - if step_state is None: - continue - dynamic_object_states_types.add(step_state) - dynamic_object_states_counter[step_state] += 1 - number_summary_dict["dynamic_object_states_types"] = dynamic_object_states_types - number_summary_dict["dynamic_object_states_counter"] = dict(dynamic_object_states_counter) - - return number_summary_dict - - def contains_explicit_return(f): return any(isinstance(node, ast.Return) for node in ast.walk(ast.parse(inspect.getsource(f)))) @@ -153,20 +103,20 @@ def write_to_directory(convert_func, # add agents summary summary_dict = {} ego_car_id = sd_scenario[SD.METADATA][SD.SDC_ID] - summary_dict[ego_car_id] = get_agent_summary( + summary_dict[ego_car_id] = SD.get_object_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 + summary_dict[track_id] = SD.get_object_summary(state_dict=track["state"], id=track_id, type=track["type"]) + sd_scenario[SD.METADATA][SD.SUMMARY.OBJECT_SUMMARY] = summary_dict # count some objects occurrence - sd_scenario[SD.METADATA]["number_summary"] = get_number_summary(sd_scenario) + sd_scenario[SD.METADATA][SD.SUMMARY.NUMBER_SUMMARY] = SD.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) + SD.sanity_check(sd_scenario, check_self_type=True) # dump p = os.path.join(output_path, export_file_name) diff --git a/scenarionet/examples/convert_waymo.py b/scenarionet/examples/convert_waymo.py index c6a0809..4962af8 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_and_set, get_agent_summary, get_number_summary +from scenarionet.converter.utils import dict_recursive_remove_array_and_set, get_object_summary, get_number_summary try: import tensorflow as tf @@ -123,11 +123,11 @@ def convert_waymo(file_list, input_path, output_path, worker_index=None): export_file_name = "sd_{}_{}.pkl".format(file, scenario.scenario_id) summary_dict = {} - summary_dict["sdc"] = get_agent_summary( + summary_dict["sdc"] = get_object_summary( state_dict=md_scenario.get_sdc_track()["state"], id=sdc_id, type=md_scenario.get_sdc_track()["type"] ) for track_id, track in md_scenario[SD.TRACKS].items(): - summary_dict[track_id] = get_agent_summary(state_dict=track["state"], id=track_id, type=track["type"]) + summary_dict[track_id] = get_object_summary(state_dict=track["state"], id=track_id, type=track["type"]) md_scenario[SD.METADATA]["object_summary"] = summary_dict # Count some objects occurrence