-
Notifications
You must be signed in to change notification settings - Fork 373
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add nuScenes-based lidar examples (#4407)
### What Adds two new examples based on the nuScenes dataset. One example visualizing most of the data (lidar, radar, camera, and 3D bounding boxes) in the nuScenes dataset, and a more minimal example that only visualizes the lidar data in the lidar reference frame. Lidar example ![lidar](https://github.com/rerun-io/rerun/assets/9785832/6e3f863e-c1a9-46ba-8d88-6144b9051497) nuScenes example ![nuscenes](https://github.com/rerun-io/rerun/assets/9785832/f247ab96-2141-47b5-bdbb-4faf1e515885) The download module is included in both examples for now; if we have a place so this could be shared across examples, that would be nice. Closes #2510. ### Checklist * [x] I have read and agree to [Contributor Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and the [Code of Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md) * [x] I've included a screenshot or gif (if applicable) * [x] I have tested [app.rerun.io](https://app.rerun.io/pr/4407) (if applicable) * [x] The PR title and labels are set such as to maximize their usefulness for the next release's CHANGELOG - [PR Build Summary](https://build.rerun.io/pr/4407) - [Docs preview](https://rerun.io/preview/72f9caff3681fe2cde30f7c0f2ec193ded29831b/docs) <!--DOCS-PREVIEW--> - [Examples preview](https://rerun.io/preview/72f9caff3681fe2cde30f7c0f2ec193ded29831b/examples) <!--EXAMPLES-PREVIEW--> - [Recent benchmark results](https://build.rerun.io/graphs/crates.html) - [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)
- Loading branch information
Showing
14 changed files
with
527 additions
and
0 deletions.
There are no files selected for viewing
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 |
---|---|---|
|
@@ -210,6 +210,7 @@ | |
"Nikhila", | ||
"nohash", | ||
"noqa", | ||
"nuScenes", | ||
"numpy", | ||
"nyud", | ||
"obbs", | ||
|
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
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 @@ | ||
dataset/** |
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,23 @@ | ||
--- | ||
title: Lidar | ||
python: https://github.com/rerun-io/rerun/blob/latest/examples/python/lidar/main.py?speculative-link | ||
tags: [lidar, 3D] | ||
description: "Visualize the lidar data from the nuScenes dataset." | ||
thumbnail: https://static.rerun.io/lidar/bcea9337044919c1524429bd26bc51a3c4db8ccb/480w.png | ||
thumbnail_dimensions: [480, 286] | ||
--- | ||
|
||
<picture> | ||
<img src="https://static.rerun.io/lidar/bcea9337044919c1524429bd26bc51a3c4db8ccb/full.png" alt=""> | ||
<source media="(max-width: 480px)" srcset="https://static.rerun.io/lidar/bcea9337044919c1524429bd26bc51a3c4db8ccb/480w.png"> | ||
<source media="(max-width: 768px)" srcset="https://static.rerun.io/lidar/bcea9337044919c1524429bd26bc51a3c4db8ccb/768w.png"> | ||
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/lidar/bcea9337044919c1524429bd26bc51a3c4db8ccb/1024w.png"> | ||
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/lidar/bcea9337044919c1524429bd26bc51a3c4db8ccb/1200w.png"> | ||
</picture> | ||
|
||
This example visualizes only the lidar data from the [nuScenes dataset](https://www.nuscenes.org/) using Rerun. For a moe extensive example including other sensors and annotations check out the [nuScenes example](https://www.rerun.io/examples/real-data/nuscenes?speculative-link). | ||
|
||
```bash | ||
pip install -r examples/python/lidar/requirements.txt | ||
python examples/python/lidar/main.py | ||
``` |
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 @@ | ||
../nuscenes/dataset |
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,65 @@ | ||
"""Module to download nuScenes minisplit.""" | ||
from __future__ import annotations | ||
|
||
import os | ||
import pathlib | ||
import tarfile | ||
|
||
import requests | ||
import tqdm | ||
|
||
MINISPLIT_SCENES = [ | ||
"scene-0061", | ||
"scene-0103", | ||
"scene-0553", | ||
"scene-0655", | ||
"scene-0757", | ||
"scene-0796", | ||
"scene-0916", | ||
"scene-1077", | ||
"scene-1094", | ||
"scene-1100", | ||
] | ||
MINISPLIT_URL = "https://www.nuscenes.org/data/v1.0-mini.tgz" | ||
|
||
|
||
def download_file(url: str, dst_file_path: pathlib.Path) -> None: | ||
"""Download file from url to dst_fpath.""" | ||
dst_file_path.parent.mkdir(parents=True, exist_ok=True) | ||
print(f"Downloading {url} to {dst_file_path}") | ||
response = requests.get(url, stream=True) | ||
with tqdm.tqdm.wrapattr( | ||
open(dst_file_path, "wb"), | ||
"write", | ||
miniters=1, | ||
total=int(response.headers.get("content-length", 0)), | ||
desc=f"Downloading {dst_file_path.name}", | ||
) as f: | ||
for chunk in response.iter_content(chunk_size=4096): | ||
f.write(chunk) | ||
|
||
|
||
def untar_file(tar_file_path: pathlib.Path, dst_path: pathlib.Path, keep_tar: bool = True) -> bool: | ||
"""Untar tar file at tar_file_path to dst.""" | ||
print(f"Untar file {tar_file_path}") | ||
try: | ||
with tarfile.open(tar_file_path, "r") as tf: | ||
tf.extractall(dst_path) | ||
except Exception as error: | ||
print(f"Error unzipping {tar_file_path}, error: {error}") | ||
return False | ||
if not keep_tar: | ||
os.remove(tar_file_path) | ||
return True | ||
|
||
|
||
def download_minisplit(root_dir: pathlib.Path) -> None: | ||
""" | ||
Download nuScenes minisplit. | ||
Adopted from https://colab.research.google.com/github/nutonomy/nuscenes-devkit/blob/master/python-sdk/tutorials/nuscenes_tutorial.ipynb | ||
""" | ||
zip_file_path = pathlib.Path("./v1.0-mini.tgz") | ||
if not zip_file_path.is_file(): | ||
download_file(MINISPLIT_URL, zip_file_path) | ||
untar_file(zip_file_path, root_dir, keep_tar=True) |
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,102 @@ | ||
#!/usr/bin/env python3 | ||
from __future__ import annotations | ||
|
||
import argparse | ||
import os | ||
import pathlib | ||
from typing import Final | ||
|
||
import matplotlib | ||
import numpy as np | ||
import rerun as rr | ||
from download_dataset import MINISPLIT_SCENES, download_minisplit | ||
from nuscenes import nuscenes | ||
|
||
EXAMPLE_DIR: Final = pathlib.Path(os.path.dirname(__file__)) | ||
DATASET_DIR: Final = EXAMPLE_DIR / "dataset" | ||
|
||
# currently need to calculate the color manually | ||
# see https://github.com/rerun-io/rerun/issues/4409 | ||
cmap = matplotlib.colormaps["turbo_r"] | ||
norm = matplotlib.colors.Normalize( | ||
vmin=3.0, | ||
vmax=75.0, | ||
) | ||
|
||
|
||
def ensure_scene_available(root_dir: pathlib.Path, dataset_version: str, scene_name: str) -> None: | ||
""" | ||
Ensure that the specified scene is available. | ||
Downloads minisplit into root_dir if scene_name is part of it and root_dir is empty. | ||
Raises ValueError if scene is not available and cannot be downloaded. | ||
""" | ||
try: | ||
nusc = nuscenes.NuScenes(version=dataset_version, dataroot=root_dir, verbose=True) | ||
except AssertionError: # dataset initialization failed | ||
if dataset_version == "v1.0-mini" and scene_name in MINISPLIT_SCENES: | ||
download_minisplit(root_dir) | ||
nusc = nuscenes.NuScenes(version=dataset_version, dataroot=root_dir, verbose=True) | ||
else: | ||
print(f"Could not find dataset at {root_dir} and could not automatically download specified scene.") | ||
exit() | ||
|
||
scene_names = [s["name"] for s in nusc.scene] | ||
if scene_name not in scene_names: | ||
raise ValueError(f"{scene_name=} not found in dataset") | ||
|
||
|
||
def log_nuscenes_lidar(root_dir: pathlib.Path, dataset_version: str, scene_name: str) -> None: | ||
nusc = nuscenes.NuScenes(version=dataset_version, dataroot=root_dir, verbose=True) | ||
|
||
scene = next(s for s in nusc.scene if s["name"] == scene_name) | ||
|
||
rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, timeless=True) | ||
|
||
first_sample = nusc.get("sample", scene["first_sample_token"]) | ||
current_lidar_token = first_sample["data"]["LIDAR_TOP"] | ||
while current_lidar_token != "": | ||
sample_data = nusc.get("sample_data", current_lidar_token) | ||
|
||
data_file_path = nusc.dataroot / sample_data["filename"] | ||
pointcloud = nuscenes.LidarPointCloud.from_file(str(data_file_path)) | ||
points = pointcloud.points[:3].T # shape after transposing: (num_points, 3) | ||
point_distances = np.linalg.norm(points, axis=1) | ||
point_colors = cmap(norm(point_distances)) | ||
|
||
# timestamps are in microseconds | ||
rr.set_time_seconds("timestamp", sample_data["timestamp"] * 1e-6) | ||
rr.log("world/lidar", rr.Points3D(points, colors=point_colors)) | ||
|
||
current_lidar_token = sample_data["next"] | ||
|
||
|
||
def main() -> None: | ||
parser = argparse.ArgumentParser(description="Visualizes lidar scans using the Rerun SDK.") | ||
parser.add_argument( | ||
"--root_dir", | ||
type=pathlib.Path, | ||
default=DATASET_DIR, | ||
help="Root directory of nuScenes dataset", | ||
) | ||
parser.add_argument( | ||
"--scene_name", | ||
type=str, | ||
default="scene-0061", | ||
help="Scene name to visualize (typically of form 'scene-xxxx')", | ||
) | ||
parser.add_argument("--dataset_version", type=str, default="v1.0-mini", help="Scene id to visualize") | ||
rr.script_add_args(parser) | ||
args = parser.parse_args() | ||
|
||
ensure_scene_available(args.root_dir, args.dataset_version, args.scene_name) | ||
|
||
rr.script_setup(args, "rerun_example_lidar") | ||
log_nuscenes_lidar(args.root_dir, args.dataset_version, args.scene_name) | ||
|
||
rr.script_teardown(args) | ||
|
||
|
||
if __name__ == "__main__": | ||
main() |
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,4 @@ | ||
matplotlib | ||
numpy | ||
nuscenes-devkit | ||
rerun-sdk |
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 @@ | ||
dataset/** |
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,24 @@ | ||
--- | ||
title: nuScenes | ||
python: https://github.com/rerun-io/rerun/blob/latest/examples/python/nuscenes/main.py?speculative-link | ||
tags: [lidar, 3D, 2D, object-detection, pinhole-camera] | ||
description: "Visualize the nuScenes dataset including lidar, radar, images, and bounding boxes." | ||
thumbnail: https://static.rerun.io/nuscenes/64a50a9d67cbb69ae872551989ee807b195f6b5d/480w.png | ||
thumbnail_dimensions: [480, 282] | ||
--- | ||
|
||
<picture> | ||
<img src="https://static.rerun.io/nuscenes/64a50a9d67cbb69ae872551989ee807b195f6b5d/full.png" alt=""> | ||
<source media="(max-width: 480px)" srcset="https://static.rerun.io/nuscenes/64a50a9d67cbb69ae872551989ee807b195f6b5d/480w.png"> | ||
<source media="(max-width: 768px)" srcset="https://static.rerun.io/nuscenes/64a50a9d67cbb69ae872551989ee807b195f6b5d/768w.png"> | ||
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/nuscenes/64a50a9d67cbb69ae872551989ee807b195f6b5d/1024w.png"> | ||
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/nuscenes/64a50a9d67cbb69ae872551989ee807b195f6b5d/1200w.png"> | ||
</picture> | ||
|
||
This example visualizes the [nuScenes dataset](https://www.nuscenes.org/) using Rerun. The dataset | ||
contains lidar data, radar data, color images, and labeled bounding boxes. | ||
|
||
```bash | ||
pip install -r examples/python/nuscenes/requirements.txt | ||
python examples/python/nuscenes/main.py | ||
``` |
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,65 @@ | ||
"""Module to download nuScenes minisplit.""" | ||
from __future__ import annotations | ||
|
||
import os | ||
import pathlib | ||
import tarfile | ||
|
||
import requests | ||
import tqdm | ||
|
||
MINISPLIT_SCENES = [ | ||
"scene-0061", | ||
"scene-0103", | ||
"scene-0553", | ||
"scene-0655", | ||
"scene-0757", | ||
"scene-0796", | ||
"scene-0916", | ||
"scene-1077", | ||
"scene-1094", | ||
"scene-1100", | ||
] | ||
MINISPLIT_URL = "https://www.nuscenes.org/data/v1.0-mini.tgz" | ||
|
||
|
||
def download_file(url: str, dst_file_path: pathlib.Path) -> None: | ||
"""Download file from url to dst_fpath.""" | ||
dst_file_path.parent.mkdir(parents=True, exist_ok=True) | ||
print(f"Downloading {url} to {dst_file_path}") | ||
response = requests.get(url, stream=True) | ||
with tqdm.tqdm.wrapattr( | ||
open(dst_file_path, "wb"), | ||
"write", | ||
miniters=1, | ||
total=int(response.headers.get("content-length", 0)), | ||
desc=f"Downloading {dst_file_path.name}", | ||
) as f: | ||
for chunk in response.iter_content(chunk_size=4096): | ||
f.write(chunk) | ||
|
||
|
||
def untar_file(tar_file_path: pathlib.Path, dst_path: pathlib.Path, keep_tar: bool = True) -> bool: | ||
"""Untar tar file at tar_file_path to dst.""" | ||
print(f"Untar file {tar_file_path}") | ||
try: | ||
with tarfile.open(tar_file_path, "r") as tf: | ||
tf.extractall(dst_path) | ||
except Exception as error: | ||
print(f"Error unzipping {tar_file_path}, error: {error}") | ||
return False | ||
if not keep_tar: | ||
os.remove(tar_file_path) | ||
return True | ||
|
||
|
||
def download_minisplit(root_dir: pathlib.Path) -> None: | ||
""" | ||
Download nuScenes minisplit. | ||
Adopted from https://colab.research.google.com/github/nutonomy/nuscenes-devkit/blob/master/python-sdk/tutorials/nuscenes_tutorial.ipynb | ||
""" | ||
zip_file_path = pathlib.Path("./v1.0-mini.tgz") | ||
if not zip_file_path.is_file(): | ||
download_file(MINISPLIT_URL, zip_file_path) | ||
untar_file(zip_file_path, root_dir, keep_tar=True) |
Oops, something went wrong.