Add nuScenes-based lidar examples

docs/cspell.json

@@ -210,6 +210,7 @@
     "Nikhila",
     "nohash",
     "noqa",
+    "nuScenes",
     "numpy",
     "nyud",
     "obbs",

examples/manifest.yml

@@ -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

examples/python/lidar/.gitignore

@@ -0,0 +1 @@
+dataset/**

examples/python/lidar/README.md

@@ -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
+```

examples/python/lidar/dataset

@@ -0,0 +1 @@
+../nuscenes/dataset

examples/python/lidar/download_dataset.py

@@ -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)

examples/python/lidar/main.py

@@ -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()

examples/python/lidar/requirements.txt

@@ -0,0 +1,4 @@
+matplotlib
+numpy
+nuscenes-devkit
+rerun-sdk

examples/python/nuscenes/.gitignore

@@ -0,0 +1 @@
+dataset/**

examples/python/nuscenes/README.md

@@ -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
+```

examples/python/nuscenes/download_dataset.py

@@ -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)