diff --git a/scenarionet/__init__.py b/scenarionet/__init__.py index ce2e85a..9eb856c 100644 --- a/scenarionet/__init__.py +++ b/scenarionet/__init__.py @@ -2,3 +2,4 @@ import os SCENARIONET_PACKAGE_PATH = os.path.dirname(os.path.abspath(__file__)) SCENARIONET_REPO_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SCENARIONET_DATASET_PATH = os.path.join(SCENARIONET_REPO_PATH, "dataset") diff --git a/scenarionet/builder/conditions.py b/scenarionet/builder/conditions.py new file mode 100644 index 0000000..333cb95 --- /dev/null +++ b/scenarionet/builder/conditions.py @@ -0,0 +1,29 @@ +import numpy as np + + +def validate_sdc_track(sdc_state): + """ + This function filters the scenario based on SDC information. + + Rule 1: Filter out if the trajectory length < 10 + + Rule 2: Filter out if the whole trajectory last < 5s, assuming sampling frequency = 10Hz. + """ + valid_array = sdc_state["valid"] + sdc_trajectory = sdc_state["position"][valid_array, :2] + sdc_track_length = [ + np.linalg.norm(sdc_trajectory[i] - sdc_trajectory[i + 1]) for i in range(sdc_trajectory.shape[0] - 1) + ] + sdc_track_length = sum(sdc_track_length) + + # Rule 1 + if sdc_track_length < 10: + return False + + print("sdc_track_length: ", sdc_track_length) + + # Rule 2 + if valid_array.sum() < 50: + return False + + return True diff --git a/scenarionet/converter/scripts/__init__.py b/scenarionet/converter/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/scenarionet/converter/nuplan/convert_nuplan.py b/scenarionet/converter/scripts/convert_nuplan.py similarity index 100% rename from scenarionet/converter/nuplan/convert_nuplan.py rename to scenarionet/converter/scripts/convert_nuplan.py diff --git a/scenarionet/converter/nuscenes/convert_nuscenes.py b/scenarionet/converter/scripts/convert_nuscenes.py similarity index 99% rename from scenarionet/converter/nuscenes/convert_nuscenes.py rename to scenarionet/converter/scripts/convert_nuscenes.py index 10386cb..38005a2 100644 --- a/scenarionet/converter/nuscenes/convert_nuscenes.py +++ b/scenarionet/converter/scripts/convert_nuscenes.py @@ -55,6 +55,7 @@ def convert_nuscenes(version, dataroot, output_path, worker_index=None, verbose= with open(p, "wb") as f: pickle.dump(sd_scene, f) metadata_recorder[export_file_name] = copy.deepcopy(sd_scene[ScenarioDescription.METADATA]) + # rename and save if delay_remove is not None: shutil.rmtree(delay_remove) diff --git a/scenarionet/converter/waymo/convert_waymo.py b/scenarionet/converter/scripts/convert_waymo.py similarity index 69% rename from scenarionet/converter/waymo/convert_waymo.py rename to scenarionet/converter/scripts/convert_waymo.py index 042ab73..6037d1b 100644 --- a/scenarionet/converter/waymo/convert_waymo.py +++ b/scenarionet/converter/scripts/convert_waymo.py @@ -10,11 +10,10 @@ import argparse import copy import os import pickle -from collections import defaultdict import numpy as np -from scenarionet.converter.utils import dict_recursive_remove_array +from scenarionet.converter.utils import dict_recursive_remove_array, get_agent_summary, get_number_summary try: import tensorflow as tf @@ -39,84 +38,6 @@ from scenarionet.converter.waymo.utils import extract_tracks, extract_dynamic_ma import sys -def validate_sdc_track(sdc_state): - """ - This function filters the scenario based on SDC information. - - Rule 1: Filter out if the trajectory length < 10 - - Rule 2: Filter out if the whole trajectory last < 5s, assuming sampling frequency = 10Hz. - """ - valid_array = sdc_state["valid"] - sdc_trajectory = sdc_state["position"][valid_array, :2] - sdc_track_length = [ - np.linalg.norm(sdc_trajectory[i] - sdc_trajectory[i + 1]) for i in range(sdc_trajectory.shape[0] - 1) - ] - sdc_track_length = sum(sdc_track_length) - - # Rule 1 - if sdc_track_length < 10: - return False - - print("sdc_track_length: ", sdc_track_length) - - # Rule 2 - if valid_array.sum() < 50: - return False - - return True - - -def _get_agent_summary(state_dict, id, type): - track = state_dict["position"] - valid_track = track[state_dict["valid"], :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 convert_waymo(file_list, input_path, output_path, worker_index=None): scenario = scenario_pb2.Scenario() @@ -202,15 +123,15 @@ 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_agent_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_agent_summary(state_dict=track["state"], id=track_id, type=track["type"]) md_scenario[SD.METADATA]["object_summary"] = summary_dict # Count some objects occurrence - md_scenario[SD.METADATA]["number_summary"] = _get_number_summary(md_scenario) + md_scenario[SD.METADATA]["number_summary"] = get_number_summary(md_scenario) metadata_recorder[export_file_name] = copy.deepcopy(md_scenario[SD.METADATA]) diff --git a/scenarionet/converter/utils.py b/scenarionet/converter/utils.py index 253f2ae..0de7c80 100644 --- a/scenarionet/converter/utils.py +++ b/scenarionet/converter/utils.py @@ -1,15 +1,13 @@ import math +from collections import defaultdict import numpy as np +from metadrive.scenario import ScenarioDescription as SD def nuplan_to_metadrive_vector(vector, nuplan_center=(0, 0)): "All vec in nuplan should be centered in (0,0) to avoid numerical explosion" vector = np.array(vector) - # if len(vector.shape) == 1: - # vector[1] *= -1 - # else: - # vector[:, 1] *= -1 vector -= np.asarray(nuplan_center) return vector @@ -47,6 +45,57 @@ def dict_recursive_remove_array(d): d[k] = dict_recursive_remove_array(d[k]) return d + def mph_to_kmh(speed_in_mph: float): speed_in_kmh = speed_in_mph * 1.609344 - return speed_in_kmh \ No newline at end of file + return speed_in_kmh + + +def get_agent_summary(state_dict, id, type): + track = state_dict["position"] + valid_track = track[state_dict["valid"], :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