Skip to content

Commit

Permalink
Add nuScenes-based lidar examples (#4407)
Browse files Browse the repository at this point in the history
### 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
roym899 authored Dec 7, 2023
1 parent 3ab8d0a commit f6781a9
Show file tree
Hide file tree
Showing 14 changed files with 527 additions and 0 deletions.
1 change: 1 addition & 0 deletions docs/cspell.json
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,7 @@
"Nikhila",
"nohash",
"noqa",
"nuScenes",
"numpy",
"nyud",
"obbs",
Expand Down
6 changes: 6 additions & 0 deletions examples/manifest.yml
Original file line number Diff line number Diff line change
Expand Up @@ -72,12 +72,18 @@ root:
- name: human-pose-tracking
python: python/human_pose_tracking

- name: lidar
python: python/lidar

- name: live-camera-edge-detection
python: python/live_camera_edge_detection

- name: live-depth-sensor
python: python/live_depth_sensor

- name: nuscenes
python: python/nuscenes

- name: objectron
python: python/objectron
rust: rust/objectron
Expand Down
1 change: 1 addition & 0 deletions examples/python/lidar/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
dataset/**
23 changes: 23 additions & 0 deletions examples/python/lidar/README.md
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
```
1 change: 1 addition & 0 deletions examples/python/lidar/dataset
65 changes: 65 additions & 0 deletions examples/python/lidar/download_dataset.py
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)
102 changes: 102 additions & 0 deletions examples/python/lidar/main.py
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()
4 changes: 4 additions & 0 deletions examples/python/lidar/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
matplotlib
numpy
nuscenes-devkit
rerun-sdk
1 change: 1 addition & 0 deletions examples/python/nuscenes/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
dataset/**
24 changes: 24 additions & 0 deletions examples/python/nuscenes/README.md
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
```
65 changes: 65 additions & 0 deletions examples/python/nuscenes/download_dataset.py
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)
Loading

0 comments on commit f6781a9

Please sign in to comment.