-
Notifications
You must be signed in to change notification settings - Fork 101
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
107 changed files
with
7,169 additions
and
4,735 deletions.
There are no files selected for viewing
136 changes: 136 additions & 0 deletions
136
...jupyter_execute/build/jupyter_execute/build/jupyter_execute/notebooks/data_demo.ipynb.txt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,136 @@ | ||
{ | ||
"cells": [ | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": { | ||
"id": "rNYRA6k8Qfyo" | ||
}, | ||
"source": [ | ||
"# Scenario Data Loading\n", | ||
"\n", | ||
"This tutorial demonstrates how to load scenario data from the Waymo Open Motion Dataset (WOMD) using the Waymax dataloader." | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": { | ||
"id": "MtgRcYqmtMwD" | ||
}, | ||
"outputs": [], | ||
"source": [ | ||
"%%capture\n", | ||
"import numpy as np\n", | ||
"import mediapy\n", | ||
"from tqdm import tqdm\n", | ||
"import dataclasses\n", | ||
"\n", | ||
"from waymax import config as _config\n", | ||
"from waymax import dataloader\n", | ||
"from waymax import datatypes\n", | ||
"from waymax import visualization" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": { | ||
"id": "0o2sAapxRMAT" | ||
}, | ||
"source": [ | ||
"\n", | ||
"We first create a dataset config, using the default configs provided in the `waymax.config` module. In particular, `config.WOD_1_1_0_TRAINING` is a pre-defined configuration that points to version 1.1.0 of the Waymo Open Dataset.\n", | ||
"\n", | ||
"The data config contains a number of options to configure how and where the dataset is loaded from. By default, the `WOD_1_1_0_TRAINING` loads up to 128 objects (e.g. vehicles, pedestrians) per scenario. Here, we can save memory and compute by loading only the first 32 objects stored in the scenario.\n", | ||
"\n", | ||
"We use the `dataloader.simulator_state_generator` function to create an iterator\n", | ||
"through Open Motion Dataset scenarios. Calling next on the iterator will retrieve the first scenario in the dataset.\n" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": { | ||
"id": "dkJwTuSLr0gh" | ||
}, | ||
"outputs": [], | ||
"source": [ | ||
"config = dataclasses.replace(_config.WOD_1_1_0_TRAINING, max_num_objects=32)\n", | ||
"data_iter = dataloader.simulator_state_generator(config=config)\n", | ||
"scenario = next(data_iter)" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": { | ||
"id": "q1xyeYpLR8J6" | ||
}, | ||
"source": [ | ||
"Next, we can plot the initial state of this scenario. We use a matplotlib-based visualization available in the `waymax.visualization` package." | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": { | ||
"id": "OY3-OOArsFcU" | ||
}, | ||
"outputs": [], | ||
"source": [ | ||
"# Using logged trajectory\n", | ||
"img = visualization.plot_simulator_state(scenario, use_log_traj=True)\n", | ||
"mediapy.show_image(img)" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": { | ||
"id": "H0Z15epRSC23" | ||
}, | ||
"source": [ | ||
"The Waymo Open Motion Dataset consists of 9-second trajectory snippets. We can visualize the entire logged trajectory as a video as follows:" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": { | ||
"id": "06SjvXdRrV3N" | ||
}, | ||
"outputs": [], | ||
"source": [ | ||
"imgs = []\n", | ||
"\n", | ||
"state = scenario\n", | ||
"for _ in range(scenario.remaining_timesteps):\n", | ||
" state = datatypes.update_state_by_log(state, num_steps=1)\n", | ||
" imgs.append(visualization.plot_simulator_state(state, use_log_traj=True))\n", | ||
"\n", | ||
"mediapy.show_video(imgs, fps=10)" | ||
] | ||
} | ||
], | ||
"metadata": { | ||
"colab": { | ||
"last_runtime": { | ||
"build_target": "", | ||
"kind": "local" | ||
}, | ||
"private_outputs": true, | ||
"provenance": [ | ||
{ | ||
"file_id": "14w5MbrMNLsOsLuD5kXy5-rrNO3ZgsHat", | ||
"timestamp": 1678404744504 | ||
} | ||
] | ||
}, | ||
"kernelspec": { | ||
"display_name": "Python 3", | ||
"name": "python3" | ||
}, | ||
"language_info": { | ||
"name": "python" | ||
} | ||
}, | ||
"nbformat": 4, | ||
"nbformat_minor": 0 | ||
} |
Oops, something went wrong.