src.dataset.waymo_dataset module

src.dataset.waymo_dataset module#

class src.dataset.waymo_dataset.WayMoDataset(*args: Any, **kwargs: Any)[source]#

Bases: BaseDataset

Waymo Open Motion Dataset loader.

Handles autonomous-driving scenes containing both pedestrians and vehicles.

raw_fps = 10#
classmethod load_data(args: Namespace, data_path: str, with_shape=False) → WayMoDataset[source]#

Load a single Waymo scene segment.

Read the processed csv data, rename coordinate columns, map type labels, filter out trajectories that are too short and unrelated vehicles that are too far from pedestrians, and process the map image by inverting colors, resizing, and projecting it into world coordinates.

Parameters:

  • args (Namespace) – Global arguments.

  • data_path (str) – Path to data.csv.gz.

  • with_shape (bool, optional) – Whether to load object-shape information. This is not fully implemented yet.

Returns:

Initialized dataset instance.

Return type:

WayMoDataset

classmethod load_data_batch(args: Namespace, data_path: str, show_tqdm=True, total=200) → List[WayMoDataset][source]#

Batch-load Waymo datasets.

Supports sorting scenes by pedestrian count via summary.csv and selecting the top total scenes.

Parameters:

  • args (Namespace) – Global arguments.

  • data_path (str) – Dataset root directory.

  • total (int, optional) – Maximum number of scenes to load.

Returns:

Dataset list.

Return type:

List[WayMoDataset]

static filter_vehicle_trajectories(df_data, distance_threshold=5.0)[source]#

Filter out vehicle trajectories whose distance to every pedestrian trajectory exceeds the threshold, reducing irrelevant data.

Parameters:

  • df_data (pd.DataFrame) – Raw data.

  • distance_threshold (float) – Distance threshold in meters.

Returns:

Filtered data.

Return type:

pd.DataFrame

static filter_short_trajectories(df_data, distance_threshold=3)[source]#

Filter out short trajectories whose displacement from start to end is below the threshold.

Parameters:

  • df_data (pd.DataFrame) – Raw data.

  • distance_threshold (float) – Minimum displacement threshold in meters.

Returns:

Filtered data.

Return type:

pd.DataFrame

Contents