Add IMU example to manifest (#9319)

oxkitsune · web-flow · commit 20cec414e700 · 2025-03-21T13:04:35.000+01:00
### Related Closes #9235 ### What Adds the IMU example to the nightly manifest - [x] Based on #9321 - [x] Upload the TUM VI dataset to a GCP bucket to make it available more reliably
diff --git a/examples/python/imu_signals/.gitignore b/examples/python/imu_signals/.gitignore
@@ -1,2 +1,2 @@
-dataset-corridor4_512_16*
+dataset/*
 
diff --git a/examples/python/imu_signals/README.md b/examples/python/imu_signals/README.md
@@ -2,7 +2,10 @@
 title = "IMU signals"
 tags = ["Plots"]
 description = "Log multi dimensional signals under a single entity."
+thumbnail = "https://static.rerun.io/imu_signals/64f773d238a0456a0f233abeea7e521cfb871b67/480w.jpg"
 thumbnail_dimensions = [480, 480]
+channel = "main"
+build_args = ["--seconds=10"]
 -->
 
 This example demonstrates how to log multi dimensional signals with the Rerun SDK, using the [TUM VI Benchmark](https://cvg.cit.tum.de/data/datasets/visual-inertial-dataset).
@@ -55,3 +58,11 @@ To experiment with the provided example, simply execute the main Python script:
 ```bash
 python -m imu_signals
 ```
+
+## Attribution
+
+This example uses a scene from the **TUM VI Benchmark dataset**, originally provided by [Technical University of Munich (TUM)](https://cvg.cit.tum.de/data/datasets/visual-inertial-dataset).
+The dataset is licensed under **Creative Commons Attribution 4.0 (CC BY 4.0)**.
+
+- Original dataset: [TUM VI Benchmark](https://cvg.cit.tum.de/data/datasets/visual-inertial-dataset)
+- License details: [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/)
diff --git a/examples/python/imu_signals/imu_signals.py b/examples/python/imu_signals/imu_signals.py
@@ -1,99 +1,119 @@
 from __future__ import annotations
 
+import argparse
 import os
-import pathlib
 import tarfile
+from pathlib import Path
 
+import numpy as np
 import pandas as pd
 import requests
 import rerun as rr
 from rerun import blueprint as rrb
+from tqdm.auto import tqdm
 
-cwd = pathlib.Path(__file__).parent.resolve()
+DATA_DIR = Path(__file__).parent / "dataset"
 
-DATASET_URL = "https://vision.in.tum.de/tumvi/exported/euroc/512_16/dataset-corridor4_512_16.tar"
+DATASET_URL = "https://storage.googleapis.com/rerun-example-datasets/imu_signals/tum_vi_corridor4_512_16.tar"
 DATASET_NAME = "dataset-corridor4_512_16"
 XYZ_AXIS_NAMES = ["x", "y", "z"]
 XYZ_AXIS_COLORS = [[(231, 76, 60), (39, 174, 96), (52, 120, 219)]]
 
 
 def main() -> None:
-    dataset_path = cwd / DATASET_NAME
+    dataset_path = DATA_DIR / DATASET_NAME
     if not dataset_path.exists():
-        _download_dataset(cwd)
+        _download_dataset(DATA_DIR)
+
+    parser = argparse.ArgumentParser(description="Visualizes the TUM Visual-Inertial dataset using the Rerun SDK.")
+    parser.add_argument(
+        "--seconds",
+        type=float,
+        default=float("inf"),
+        help="If specified, limits the number of seconds logged",
+    )
+    rr.script_add_args(parser)
+    args = parser.parse_args()
+
+    blueprint = rrb.Horizontal(
+        rrb.Vertical(
+            rrb.TimeSeriesView(
+                origin="gyroscope",
+                name="Gyroscope",
+                overrides={
+                    # TODO(#9022): Pluralize series line type.
+                    "/gyroscope": rr.SeriesLine.from_fields(name=XYZ_AXIS_NAMES, color=XYZ_AXIS_COLORS),  # type: ignore[arg-type]
+                },
+            ),
+            rrb.TimeSeriesView(
+                origin="accelerometer",
+                name="Accelerometer",
+                overrides={
+                    # TODO(#9022): Pluralize series line type.
+                    "/accelerometer": rr.SeriesLine.from_fields(name=XYZ_AXIS_NAMES, color=XYZ_AXIS_COLORS),  # type: ignore[arg-type]
+                },
+            ),
+        ),
+        rrb.Spatial3DView(origin="/", name="World position"),
+        column_shares=[0.45, 0.55],
+    )
+
+    rr.script_setup(args, "rerun_example_imu_signals", default_blueprint=blueprint)
 
-    _setup_rerun()
-    _log_imu_data()
-    _log_image_data()
-    _log_gt_imu()
+    _log_imu_data(args.seconds)
+    _log_image_data(args.seconds)
+    _log_gt_imu(args.seconds)
 
 
-def _download_dataset(root: pathlib.Path, dataset_url: str = DATASET_URL) -> None:
+def _download_dataset(root: Path, dataset_url: str = DATASET_URL) -> None:
     os.makedirs(root, exist_ok=True)
-    tar_path = os.path.join(root, "dataset-corridor4_512_16.tar")
-    print("Downloading dataset...")
-    with requests.get(dataset_url, stream=True) as r:
-        r.raise_for_status()
-        with open(tar_path, "wb") as f:
-            for chunk in r.iter_content(chunk_size=8192):
-                if chunk:
-                    f.write(chunk)
+    tar_path = os.path.join(root, f"{DATASET_NAME}.tar")
+    response = requests.get(dataset_url, stream=True)
+
+    total_size = int(response.headers.get("content-length", 0))
+    block_size = 1024
+
+    with tqdm(desc="Downloading dataset", total=total_size, unit="B", unit_scale=True) as pb:
+        with open(tar_path, "wb") as file:
+            for data in response.iter_content(chunk_size=block_size):
+                pb.update(len(data))
+                file.write(data)
+
+    if total_size != 0 and pb.n != total_size:
+        raise RuntimeError("Failed to download complete dataset!")
+
     print("Extracting dataset...")
     with tarfile.open(tar_path, "r:") as tar:
         tar.extractall(path=root)
     os.remove(tar_path)
 
 
-def _setup_rerun() -> None:
-    rr.init("rerun_example_imu_data", spawn=True)
-
-    rr.send_blueprint(
-        rrb.Horizontal(
-            rrb.Vertical(
-                rrb.TimeSeriesView(
-                    origin="gyroscope",
-                    name="Gyroscope",
-                    overrides={
-                        # TODO(#9022): Pluralize series line type.
-                        "/gyroscope": rr.SeriesLine.from_fields(name=XYZ_AXIS_NAMES, color=XYZ_AXIS_COLORS),  # type: ignore[arg-type]
-                    },
-                ),
-                rrb.TimeSeriesView(
-                    origin="accelerometer",
-                    name="Accelerometer",
-                    overrides={
-                        # TODO(#9022): Pluralize series line type.
-                        "/accelerometer": rr.SeriesLine.from_fields(name=XYZ_AXIS_NAMES, color=XYZ_AXIS_COLORS),  # type: ignore[arg-type]
-                    },
-                ),
-            ),
-            rrb.Spatial3DView(origin="/", name="World position"),
-            column_shares=[0.45, 0.55],
-        ),
-    )
-
-
-def _log_imu_data() -> None:
+def _log_imu_data(max_time_sec: float) -> None:
     imu_data = pd.read_csv(
-        cwd / DATASET_NAME / "dso/imu.txt",
+        DATA_DIR / DATASET_NAME / "dso/imu.txt",
         sep=" ",
         header=0,
         names=["timestamp", "gyro.x", "gyro.y", "gyro.z", "accel.x", "accel.y", "accel.z"],
         comment="#",
     )
 
-    times = rr.TimeColumn("timestamp", timestamp=imu_data["timestamp"])
+    timestamps = imu_data["timestamp"].to_numpy()
+    max_time_ns = imu_data["timestamp"][0] + max_time_sec * 1e9
+    selected = imu_data[imu_data["timestamp"] <= max_time_ns]
+
+    timestamps = selected["timestamp"].astype("datetime64[ns]")
+    times = rr.TimeColumn("timestamp", timestamp=timestamps)
 
-    gyro = imu_data[["gyro.x", "gyro.y", "gyro.z"]]
+    gyro = selected[["gyro.x", "gyro.y", "gyro.z"]].to_numpy()
     rr.send_columns("/gyroscope", indexes=[times], columns=rr.Scalar.columns(scalar=gyro))
 
-    accel = imu_data[["accel.x", "accel.y", "accel.z"]]
+    accel = selected[["accel.x", "accel.y", "accel.z"]]
     rr.send_columns("/accelerometer", indexes=[times], columns=rr.Scalar.columns(scalar=accel))
 
 
-def _log_image_data() -> None:
+def _log_image_data(max_time_sec: float) -> None:
     times = pd.read_csv(
-        cwd / DATASET_NAME / "dso/cam0/times.txt",
+        DATA_DIR / DATASET_NAME / "dso/cam0/times.txt",
         sep=" ",
         header=0,
         names=["filename", "timestamp", "exposure_time"],
@@ -103,35 +123,48 @@ def _log_image_data() -> None:
 
     rr.set_time("timestamp", timestamp=times["timestamp"][0])
     rr.log(
-        "/cam0",
+        "/world",
+        rr.Transform3D(rotation_axis_angle=rr.RotationAxisAngle(axis=(1, 0, 0), angle=-np.pi / 2)),
+        static=True,
+    )
+    rr.log(
+        "/world/cam0",
         rr.Pinhole(
             focal_length=(0.373 * 512, 0.373 * 512),
             resolution=(512, 512),
-            camera_xyz=rr.components.ViewCoordinates.FLU,
             image_plane_distance=0.4,
         ),
         static=True,
     )
 
+    max_time_sec = times["timestamp"][0] + max_time_sec
     for _, (filename, timestamp, _) in times.iterrows():
-        image_path = cwd / DATASET_NAME / "dso/cam0/images" / f"{filename}.png"
+        if timestamp > max_time_sec:
+            break
+
+        image_path = DATA_DIR / DATASET_NAME / "dso/cam0/images" / f"{filename}.png"
         rr.set_time("timestamp", timestamp=timestamp)
-        rr.log("/cam0/image", rr.ImageEncoded(path=image_path))
+        rr.log("/world/cam0/image", rr.EncodedImage(path=image_path))
 
 
-def _log_gt_imu() -> None:
+def _log_gt_imu(max_time_sec: float) -> None:
     gt_imu = pd.read_csv(
-        cwd / DATASET_NAME / "dso/gt_imu.csv",
+        DATA_DIR / DATASET_NAME / "dso/gt_imu.csv",
         sep=",",
         header=0,
         names=["timestamp", "t.x", "t.y", "t.z", "q.w", "q.x", "q.y", "q.z"],
         comment="#",
     )
 
-    times = rr.TimeColumn("timestamp", timestamp=gt_imu["timestamp"])
+    timestamps = gt_imu["timestamp"].to_numpy()
+    max_time_ns = gt_imu["timestamp"][0] + max_time_sec * 1e9
+    selected = gt_imu[gt_imu["timestamp"] <= max_time_ns]
 
-    translations = gt_imu[["t.x", "t.y", "t.z"]]
-    quaternions = gt_imu[
+    timestamps = selected["timestamp"].astype("datetime64[ns]")
+    times = rr.TimeColumn("timestamp", timestamp=timestamps)
+
+    translations = selected[["t.x", "t.y", "t.z"]]
+    quaternions = selected[
         [
             "q.x",
             "q.y",
@@ -140,9 +173,12 @@ def _log_gt_imu() -> None:
         ]
     ]
     rr.send_columns(
-        "/cam0",
+        "/world/cam0",
         indexes=[times],
-        columns=rr.Transform3D.columns(translation=translations, quaternion=quaternions),
+        columns=rr.Transform3D.columns(
+            translation=translations,
+            quaternion=quaternions,
+        ),
     )
 
 
diff --git a/examples/python/imu_signals/pyproject.toml b/examples/python/imu_signals/pyproject.toml
@@ -1,16 +1,17 @@
-[build-system]
-requires = ["hatchling"]
-build-backend = "hatchling.build"
-
 [project]
 name = "imu_signals"
+description = "Log multi dimensional signals under a single entity."
 version = "0.1.0"
 requires-python = "<3.12"
-dependencies = ["numpy", "requests>=2.31,<3", "rerun-sdk", "pandas"]
-
+dependencies = ["numpy", "requests>=2.31,<3", "rerun-sdk", "pandas", "tqdm"]
 
 [project.scripts]
 imu_signals = "imu_signals:main"
 
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+
 [tool.rerun-example]
 skip = false
diff --git a/pixi.lock b/pixi.lock
diff --git a/pixi.toml b/pixi.toml

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`		`-dataset-corridor4_512_16*`
	`1`	`+dataset/*`
`2`	`2`