move writing summary to metadrive

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"])
+            summary_dict[track_id] = SD.get_object_summary(state_dict=track["state"], id=track_id, type=track["type"])
-        sd_scenario[SD.METADATA]["object_summary"] = summary_dict
+        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)

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