wehub/rerun-io--rerun
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:05:14 +08:00
commit 2a547be7fe
7904 changed files with 1000926 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/README.md
+83
View File
@@ -0,0 +1,83 @@
# Rerun Python examples
The simplest example is [`minimal`](minimal/minimal.py). You may want to start there!
Read more about our examples at <https://www.rerun.io/examples>.
## Setup
First install the Rerun Python SDK with `pip install rerun-sdk` or via [`conda`](https://github.com/conda-forge/rerun-sdk-feedstock).
> Note: Make sure your SDK version matches the code in the examples.
For example, if your SDK version is `0.15.1`, check out the matching tag
for this repository by running `git checkout v0.15.1`.
## Installing the example
Each example is packaged as a regular Python package, with a `pyproject.toml` file specifying its required dependencies. To run an example, it must first be installed.
For example, to install dependencies and run the toy `minimal` example (which doesn't need to download any data) run:
```sh
pip install -e examples/python/minimal
```
**Note**: it is import to install example in editable mode, which is done using the `-e` flag (short for `--editable`).
Once installed, the example can be run as a regular Python module:
```shell
python -m minimal
```
Examples also declare console script, so they can also be run directly:
```shell
minimal
```
## Running the examples
By default, the examples spawn a Rerun Viewer and stream log data to it.
For most examples you can instead save the log data to an `.rrd` file using `plots --save data.rrd`. You can then view that `.rrd` file with `rerun data.rrd`.
(`rerun` is an alias for `python -m rerun`).
NOTE: `.rrd` files do not yet guarantee any backwards or forwards compatibility. One version of Rerun will likely not be able to open an `.rrd` file generated by another Rerun version.
## Running examples with Pixi
The Rerun project makes extensive use of [Pixi](https://pixi.sh/latest/) for various developer tasks, and Pixi can be used to run examples as well. For this, you need to install Pixi as per the installation instructions on their website.
### Running examples with Pixi from source
You can build Rerun from source, and install it in the uv environment managed by Pixi. Note that this requires a Rust toolchain to be installed on your system.
Before running examples, build the SDK:
```shell
pixi run py-build
```
Now you can run examples via uv:
```shell
pixi run uv run examples/python/minimal/minimal.py
```
### Installing and running examples
You can install all of the rerun python examples into the uv managed environment with:
```shell
pixi run py-build-examples
```
Then you can just examples based on their installed name, for example:
```shell
pixi run uv run plots
```
## Datasets
Some examples will download a small datasets before they run. They will do so the first time you run the example. The datasets will be added to a subdir called `dataset`, which is in the repo-wide `.gitignore`.
## Contributions welcome
Feel free to open a PR to add a new example!
If you add an example, please make sure to add the corresponding entry to the top-level `pixi.toml` file as well, otherwise it will not be picked up.
See [`CONTRIBUTING.md`](../../CONTRIBUTING.md) for details on how to contribute.
examples/python/_isolated/mypy.ini
+8
View File
@@ -0,0 +1,8 @@
# Base mypy config shared across all isolated examples.
# Per-example `[[tool.mypy.overrides]]` from each example's pyproject.toml are
# merged onto this base at lint time by scripts/ci/isolated_examples.py.
[mypy]
strict = true
enable_error_code = redundant-expr, truthy-bool, ignore-without-code
no_implicit_reexport = false
disallow_untyped_calls = false
examples/python/air_traffic_data/.gitignore
+1
View File
@@ -0,0 +1 @@
/data
examples/python/air_traffic_data/README.md
+49
View File
@@ -0,0 +1,49 @@
<!--[metadata]
title = "Air traffic data"
tags = ["2D", "3D", "map", "crs"]
description = "Display aircraft traffic data"
thumbnail = "https://static.rerun.io/air_traffic_data/348dd2def3a55fd0bf481a35a0765eeacfa20b6f/480w.png"
thumbnail_dimensions = [480, 480]
channel = "nightly"
-->
Display air traffic data kindly provided by [INVOLI](https://www.involi.com).
<picture>
<img src="https://static.rerun.io/air_traffic_data/4a68b46a404c4f9e3c082f57a8a8ed4bf5b9b236/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/air_traffic_data/4a68b46a404c4f9e3c082f57a8a8ed4bf5b9b236/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/air_traffic_data/4a68b46a404c4f9e3c082f57a8a8ed4bf5b9b236/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/air_traffic_data/4a68b46a404c4f9e3c082f57a8a8ed4bf5b9b236/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/air_traffic_data/4a68b46a404c4f9e3c082f57a8a8ed4bf5b9b236/1200w.png">
</picture>
This example demonstrates multiple aspects of the Rerun viewer:
- Use of the [map view](https://rerun.io/docs/reference/types/views/map_view).
- Use of [pyproj](https://pyproj4.github.io/pyproj/stable/) to transform geospatial data from one coordinate system to another.
- Use [GeoPandas](https://geopandas.org/en/stable/) to load geospatial data into a 3D view.
- Use [Polars]https://pola.rs) to batch data to be sent via [`rr.send_columns()`](https://rerun.io/docs/howto/logging-and-ingestion/send-columns) (use `--batch`).
## Run the code
To run this example, make sure you have the [required Python version](https://ref.rerun.io/docs/python/main/common#supported-python-versions), the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/air_traffic_data
```
To experiment with the provided example, simply execute the main Python script:
```bash
python -m air_traffic_data
```
If you wish to customize it, explore additional features, or save it use the CLI with the `--help` option for guidance:
```bash
python -m air_traffic_data --help
```
examples/python/air_traffic_data/air_traffic_data.py
+485
View File
@@ -0,0 +1,485 @@
from __future__ import annotations
import dataclasses
import io
import itertools
import json
import re
import typing
import zipfile
from argparse import ArgumentParser
from pathlib import Path
from typing import Any
import geopandas as gpd
import numpy as np
import numpy.typing as npt
import polars
import pyproj
import requests
from pyproj import CRS, Transformer
from pyproj.aoi import AreaOfInterest
from pyproj.database import query_utm_crs_info
from tqdm import tqdm
import rerun as rr
import rerun.blueprint as rrb
if typing.TYPE_CHECKING:
import shapely
DATA_DIR = Path(__file__).parent / "dataset"
MAP_DATA_DIR = DATA_DIR / "map_data"
if not DATA_DIR.exists():
DATA_DIR.mkdir()
INVOLI_DATASETS = {
"10min": "https://storage.googleapis.com/rerun-example-datasets/involi/involi_demo_set_1_10min.zip",
"2h": "https://storage.googleapis.com/rerun-example-datasets/involi/involi_demo_set_2_2h.zip",
}
def download_with_progress(url: str, what: str) -> io.BytesIO:
"""Download file with tqdm progress bar."""
chunk_size = 1024 * 1024
try:
resp = requests.get(url, stream=True, timeout=30)
resp.raise_for_status()
total_size = int(resp.headers.get("content-length", 0))
with tqdm(
desc=f"Downloading {what}…",
total=total_size,
unit="iB",
unit_scale=True,
unit_divisor=1024,
) as progress:
download_file = io.BytesIO()
for data in resp.iter_content(chunk_size):
download_file.write(data)
progress.update(len(data))
if download_file.tell() != total_size and total_size > 0:
raise ValueError("Download incomplete: size mismatch")
download_file.seek(0)
return download_file
except (requests.RequestException, ValueError) as e:
raise RuntimeError(f"Failed to download {what} from {url}: {e}") from e
def shapely_geom_to_numpy(geom: shapely.Geometry) -> list[npt.NDArray[np.float64]]:
"""Convert shapely objects to numpy array suitable for logging as line batches."""
if geom.geom_type == "Polygon":
return [np.array(geom.exterior.coords)] + [np.array(interior.coords) for interior in geom.interiors]
elif geom.geom_type == "MultiPolygon":
res = []
for poly in geom.geoms:
res.extend(shapely_geom_to_numpy(poly))
return res
else:
print(f"Warning: unknown Shapely object {geom}")
return []
def log_region_boundaries_for_country(
country_code: str,
level: int,
color: tuple[float, float, float],
crs: CRS,
) -> None:
"""Log some boundaries for the given country and level."""
def download_eu_map_data() -> None:
"""Download some basic EU map data."""
if MAP_DATA_DIR.exists():
return
EU_MAP_DATA_URL = "https://gisco-services.ec.europa.eu/distribution/v2/nuts/download/ref-nuts-2021-01m.json.zip"
zip_data = download_with_progress(EU_MAP_DATA_URL, "map data")
with zipfile.ZipFile(zip_data) as zip_ref:
zip_ref.extractall(MAP_DATA_DIR)
download_eu_map_data()
# cspell:disable-next-line
map_data = gpd.read_file(MAP_DATA_DIR / f"NUTS_RG_01M_2021_4326_LEVL_{level}.json").set_crs("epsg:4326").to_crs(crs)
for _i, row in map_data[map_data.CNTR_CODE == country_code].iterrows():
entity_path = f"region_boundaries/{country_code}/{level}/{row.NUTS_ID}"
lines = shapely_geom_to_numpy(row.geometry)
rr.log(entity_path + "/2D", rr.LineStrips2D(lines, colors=color), static=True)
rr.log(
entity_path + "/3D",
rr.LineStrips3D(
[np.hstack([line, np.zeros((len(line), 1))]) for line in lines],
colors=color,
),
static=True,
)
metadata = row.to_dict()
metadata.pop("geometry")
rr.log(entity_path, rr.AnyValues(**metadata), static=True)
@dataclasses.dataclass
class Measurement:
"""One measurement loaded from INVOLI data. Corresponds to an "aircraft" record."""
icao_id: str
latitude: float | None
longitude: float | None
barometric_altitude: float | None
wg84_altitude: float | None
course: float | None
ground_speed: float | None
vertical_speed: float | None
ground_status: str | None
timestamp: float
@classmethod
def from_dict(cls, data: dict[str, Any]) -> Measurement:
return cls(
icao_id=data["ids"]["icao"],
latitude=data.get("latitude"),
longitude=data.get("longitude"),
barometric_altitude=data.get("barometric_altitude"),
wg84_altitude=data.get("wg84_altitude"),
course=data.get("course"),
ground_speed=data.get("ground_speed"),
vertical_speed=data.get("vertical_speed"),
ground_status=data.get("ground_status"),
timestamp=data["timestamp"][0] + data["timestamp"][1] / 1e9,
)
def find_best_utm_crs(measurements: list[Measurement]) -> CRS:
"""Returns the best UTM coordinates reference system given a list of measurements."""
def get_area_of_interest(measurements: list[Measurement]) -> AreaOfInterest:
"""Compute the span of coordinates for all provided measurements."""
print("Computing area of interest…")
all_long_lat = [
(a.longitude, a.latitude) for a in measurements if a.latitude is not None and a.longitude is not None
]
return AreaOfInterest(
west_lon_degree=min(x[0] for x in all_long_lat),
south_lat_degree=min(x[1] for x in all_long_lat),
east_lon_degree=max(x[0] for x in all_long_lat),
north_lat_degree=max(x[1] for x in all_long_lat),
)
area_of_interest = get_area_of_interest(measurements)
utm_crs_list = query_utm_crs_info(
datum_name="WGS 84",
area_of_interest=area_of_interest,
)
return CRS.from_epsg(utm_crs_list[0].code)
def load_measurements(paths: list[Path]) -> list[Measurement]:
"""Load measurements from a bunch of json files."""
all_measurements = []
for path in tqdm(paths, "Loading measurements"):
data = json.loads(path.read_text())
for data_rec in data:
for aircraft in data_rec["records"]:
all_measurements.append(Measurement.from_dict(aircraft["aircraft"]))
return all_measurements
def get_paths_for_directory(directory: Path) -> list[Path]:
"""
Get a sorted list of JSON file by recursively walking the provided directory.
Note: technically, we don't need the list to be sorted as Rerun accepts out of order data. However, it comes at a
(small) performance cost and any (cheap) sorting on the logging end is always better.
"""
def atoi(text: str) -> int | str:
return int(text) if text.isdigit() else text
def natural_keys(path: Path) -> list[int | str]:
"""
Human sort.
alist.sort(key=natural_keys) sorts in human order
https://nedbatchelder.com/blog/200712/human_sorting.html
(See Toothy's implementation in the comments)
"""
return [atoi(c) for c in re.split(r"(\d+)", str(path))]
return sorted(directory.rglob("*.json"), key=natural_keys)
class Logger(typing.Protocol):
def process_measurement(self, measurement: Measurement) -> None:
pass
def flush(self) -> None:
pass
# ================================================================================================
# Simple logger
class MeasurementLogger:
"""Logger class that uses regular `rr.log` calls."""
def __init__(self, proj: pyproj.Transformer, raw: bool) -> None:
self._proj = proj
self._raw = raw
self._ignored_fields = [
"icao_id", # already the entity's path
"timestamp", # already the clock's value
]
def process_measurement(self, measurement: Measurement) -> None:
rr.set_time("unix_time", timestamp=measurement.timestamp)
if self._raw:
metadata = dataclasses.asdict(measurement)
else:
metadata = dataclasses.asdict(
measurement,
dict_factory=lambda x: {k: v for (k, v) in x if k not in self._ignored_fields and v is not None},
)
entity_path = f"aircraft/{measurement.icao_id}"
color = rr.components.Color.from_string(entity_path)
if (
measurement.latitude is not None
and measurement.longitude is not None
and measurement.barometric_altitude is not None
):
rr.log(
entity_path,
rr.Points3D(
[
self._proj.transform(
measurement.longitude,
measurement.latitude,
measurement.barometric_altitude,
),
],
colors=color,
),
rr.GeoPoints(lat_lon=[measurement.latitude, measurement.longitude]),
)
if len(metadata) > 0:
rr.log(entity_path, rr.AnyValues(**metadata))
if measurement.barometric_altitude is not None:
rr.log(
entity_path + "/barometric_altitude",
rr.Scalars(measurement.barometric_altitude),
rr.SeriesLines(colors=color),
)
def flush(self) -> None:
pass
# ================================================================================================
# Batch logger
class MeasurementBatchLogger:
"""Logger class that batches measurements and uses `rr.send_columns` calls."""
def __init__(self, proj: pyproj.Transformer, batch_size: int = 8192) -> None:
self._proj = proj
self._batch_size = batch_size
self._measurements: list[Measurement] = []
self._position_indicators: set[str] = set()
def process_measurement(self, measurement: Measurement) -> None:
self._measurements.append(measurement)
if len(self._measurements) >= 8192:
self.flush()
def flush(self) -> None:
# !!! the raw data is not sorted by timestamp, so we sort it here
df = polars.DataFrame(self._measurements).sort("timestamp")
self._measurements = []
for (icao_id,), group in df.group_by("icao_id"):
icao_id = str(icao_id)
# Note: this splitting in 3 different functions is due to the pattern of nulls in the raw data.
self.log_position_and_altitude(group, icao_id)
self.log_ground_status(group, icao_id)
self.log_metadata(group, icao_id)
def log_position_and_altitude(self, df: polars.DataFrame, icao_id: str) -> None:
entity_path = f"aircraft/{icao_id}"
df = df["timestamp", "latitude", "longitude", "barometric_altitude"].drop_nulls()
if df.height == 0:
return
if icao_id not in self._position_indicators:
color = rr.components.Color.from_string(entity_path)
rr.log(
entity_path,
rr.Points3D.from_fields(colors=color),
# TODO(cmc): That would be UB right now (and doesn't matter as long as we are on the untagged index).
# rr.GeoPoints.from_fields(colors=color),
static=True,
)
rr.log(entity_path + "/barometric_altitude", rr.SeriesLines.from_fields(colors=color), static=True)
self._position_indicators.add(icao_id)
timestamps = rr.TimeColumn("unix_time", timestamp=df["timestamp"].to_numpy())
pos = self._proj.transform(df["longitude"], df["latitude"], df["barometric_altitude"])
raw_coordinates = rr.AnyValues.columns(
latitude=df["latitude"].to_numpy(),
longitude=df["longitude"].to_numpy(),
barometric_altitude=df["barometric_altitude"].to_numpy(),
)
rr.send_columns(
entity_path,
[timestamps],
[
*rr.Points3D.columns(positions=np.vstack(pos).T),
*rr.GeoPoints.columns(positions=np.vstack((df["latitude"], df["longitude"])).T),
*raw_coordinates,
],
)
rr.send_columns(
entity_path + "/barometric_altitude",
[timestamps],
rr.Scalars.columns(scalars=df["barometric_altitude"].to_numpy()),
)
def log_ground_status(self, df: polars.DataFrame, icao_id: str) -> None:
entity_path = f"aircraft/{icao_id}"
df = df["timestamp", "ground_status"].drop_nulls()
if df.height == 0:
return
timestamps = rr.TimeColumn("unix_time", timestamp=df["timestamp"].to_numpy())
columns = rr.AnyValues.columns(ground_status=df["ground_status"].to_numpy())
rr.send_columns(entity_path, [timestamps], columns)
def log_metadata(self, df: polars.DataFrame, icao_id: str) -> None:
entity_path = f"aircraft/{icao_id}"
df = df["timestamp", "course", "ground_speed", "vertical_speed"].drop_nulls()
if df.height == 0:
return
metadata = rr.AnyValues.columns(
course=df["course"].to_numpy(),
ground_speed=df["ground_speed"].to_numpy(),
vertical_speed=df["vertical_speed"].to_numpy(),
)
rr.send_columns(
entity_path,
[rr.TimeColumn("unix_time", timestamp=df["timestamp"].to_numpy())],
metadata,
)
# ================================================================================================
def log_everything(paths: list[Path], raw: bool, batch: bool, batch_size: int) -> None:
measurements = load_measurements(paths)
utm_crs = find_best_utm_crs(measurements)
proj = Transformer.from_crs("EPSG:4326", utm_crs, always_xy=True)
for country_code, (level, color) in itertools.product(["DE", "CH"], [(0, (1, 0.5, 0.5))]):
log_region_boundaries_for_country(country_code, level, color, utm_crs)
# Exaggerate altitudes
rr.log("aircraft", rr.Transform3D(scale=[1, 1, 10]), static=True)
if batch:
logger: Logger = MeasurementBatchLogger(proj, batch_size)
else:
logger = MeasurementLogger(proj, raw)
for measurement in tqdm(measurements, "Logging measurements"):
if measurement.icao_id is None:
continue
logger.process_measurement(measurement)
logger.flush()
def main() -> None:
parser = ArgumentParser(description="Visualize INVOLI data")
parser.add_argument(
"--dataset",
choices=INVOLI_DATASETS.keys(),
default="2h",
help="Which dataset to automatically download and visualize",
)
parser.add_argument(
"--raw",
action="store_true",
help="If true, logs the raw data with all its issues (useful to stress edge cases in the viewer)",
)
parser.add_argument(
"--batch",
action="store_true",
default=True,
help="If true, use the batch logger function (rerun 0.18 required)",
)
parser.add_argument(
"--batch-size",
type=int,
default=8192,
help="Batch size for the batch logger",
)
parser.add_argument(
"--dir",
type=Path,
help="Use this directory of data instead of downloading a dataset",
)
rr.script_add_args(parser)
args = parser.parse_args()
if args.dir:
dataset_directory = args.dir
else:
dataset = args.dataset
dataset_ulr = INVOLI_DATASETS[dataset]
dataset_directory = DATA_DIR / dataset
if not dataset_directory.exists():
dataset_directory.mkdir()
zip_data = download_with_progress(dataset_ulr, f"dataset {dataset}")
with zipfile.ZipFile(zip_data) as zip_ref:
zip_ref.extractall(dataset_directory)
# TODO(ab): this blueprint would be massively improved by setting the 3D view's orbit point to FRA's coordinates.
blueprint = rrb.Vertical(
rrb.Horizontal(rrb.Spatial3DView(origin="/"), rrb.MapView(origin="/")),
rrb.TimeSeriesView(origin="/aircraft"),
row_shares=[3, 1],
)
rr.script_setup(args, "rerun_example_air_traffic_data", default_blueprint=blueprint)
paths = get_paths_for_directory(dataset_directory)
log_everything(paths, args.raw, args.batch, args.batch_size)
if __name__ == "__main__":
main()
examples/python/air_traffic_data/pyproject.toml
+22
View File
@@ -0,0 +1,22 @@
[project]
name = "air_traffic_data"
version = "0.1.0"
readme = "README.md"
dependencies = [
"geopandas",
"numpy",
"polars",
"pyproj",
"requests",
"rerun-sdk",
"shapely",
"tqdm",
]
[project.scripts]
air_traffic_data = "air_traffic_data:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/animated_urdf/README.md
+47
View File
@@ -0,0 +1,47 @@
<!--[metadata]
title = "URDF"
tags = ["3D", "Mesh", "URDF", "Animation"]
channel = "main"
include_in_manifest = true
thumbnail = "https://static.rerun.io/animated-urdf-thumbnail/02cd73ad1155db0a202392b1fd8f8036070ad888/480w.png"
thumbnail_dimensions = [480, 480]
-->
<picture>
<img src="https://static.rerun.io/animated_urdf/ebdefa158ab6f26f9dc1cb1924fce4b846fe8db2/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/animated_urdf/ebdefa158ab6f26f9dc1cb1924fce4b846fe8db2/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/animated_urdf/ebdefa158ab6f26f9dc1cb1924fce4b846fe8db2/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/animated_urdf/ebdefa158ab6f26f9dc1cb1924fce4b846fe8db2/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/animated_urdf/ebdefa158ab6f26f9dc1cb1924fce4b846fe8db2/1200w.png">
</picture>
An example of how to load and animate a URDF given some changing joint angles.
## Logging and visualizing with Rerun
This example demonstrates how to:
1. Load and log a URDF file as a static resource
2. Parse the URDF structure using `UrdfTree`
3. Animate joints by logging dynamic transforms
The key steps are:
```python
import rerun as rr
import rerun.urdf import UrdfTree
# Log the URDF file once, as a static resource
rec.log_file_from_path(urdf_path, static=True)
# Load the URDF tree structure into memory
urdf_tree = UrdfTree.from_file_path(urdf_path)
# Animate joints by logging transforms
for joint in urdf_tree.joints():
if joint.joint_type == "revolute":
# compute_transform gives you a complete transform that is ready to log,
# calculated from joint origin and the current angle and with the frame names set.
transform = joint.compute_transform(angle)
rec.log("transforms", transform)
```
examples/python/animated_urdf/animated_urdf.py
Executable
+140
View File
@@ -0,0 +1,140 @@
#!/usr/bin/env python3
"""
An example of how to load and animate a URDF given some changing joint angles.
Usage:
python -m animated_urdf
"""
from __future__ import annotations
import argparse
import math
from pathlib import Path
import rerun as rr
import rerun.blueprint as rrb
TIMELINE = "example_time"
def main() -> None:
parser = argparse.ArgumentParser(
description="An example of how to load and animate a URDF given some changing joint angles.",
)
rr.script_add_args(parser)
parser.add_argument(
"--dual",
action="store_true",
help="Load the same URDF twice with different frame prefixes (dual-arm demo).",
)
args = parser.parse_args()
urdf_path = Path(__file__).parent.parent.parent / "rust" / "animated_urdf" / "data" / "so100.urdf"
if args.dual:
run_dual(args, urdf_path)
else:
run_single(args, urdf_path)
def run_single(args: argparse.Namespace, urdf_path: Path) -> None:
duration = 0.0
rec = rr.script_setup(args, "rerun_example_animated_urdf")
rec.set_time(TIMELINE, duration=duration)
# Log the URDF file once
rec.log_file_from_path(urdf_path)
# Load the URDF tree structure into memory
urdf_tree = rr.urdf.UrdfTree.from_file_path(urdf_path)
# Hide the collision geometries by default in the viewer.
blueprint = rrb.Grid(
rrb.Spatial3DView(
name="Animated URDF", overrides={"so_arm100/collision_geometries": rrb.EntityBehavior(visible=False)}
)
)
rec.send_blueprint(blueprint)
for step in range(10000):
for joint_index, joint in enumerate(urdf_tree.joints()):
if joint.joint_type == "revolute":
# Usually this angle would come from a measurement, here we just fake something
dynamic_angle = _fake_angle(joint, step, joint_index, phase=0.0)
# Rerun loads the URDF transforms with child/parent frame relations.
# To move a joint, we just need to log a new transform between those frames.
# Here, we use the `compute_transform` method that automatically takes care
# of setting the frame names and calculating the full transform from the joint angle.
transform = joint.compute_transform(dynamic_angle, clamp=True)
rec.log("transforms", transform)
# We can also work with links from the URDF tree.
# Here, we change the color of the visual mesh entities of the "jaw" link based on the joint angle.
link = urdf_tree.get_joint_child(joint)
if link.name == "jaw":
for visual_path in urdf_tree.get_visual_geometry_paths(link):
normalized_angle = (dynamic_angle - joint.limit_lower) / (joint.limit_upper - joint.limit_lower)
rgba = [1.0 - normalized_angle, normalized_angle, 0, 0.5]
rec.log(visual_path, rr.Asset3D.from_fields(albedo_factor=rgba))
duration += 0.03
rec.set_time(TIMELINE, duration=duration)
rr.script_teardown(args)
def run_dual(args: argparse.Namespace, urdf_path: Path) -> None:
"""Load the same URDF twice with different frame prefixes (dual-arm demo)."""
rec = rr.script_setup(args, "rerun_example_animated_urdf")
# Load the same URDF twice with different prefixes.
# - entity_path_prefix separates geometry in the entity tree
# - frame_prefix makes frame IDs unique ("left/base", "right/base", …)
left = rr.urdf.UrdfTree.from_file_path(urdf_path, entity_path_prefix="left", frame_prefix="left/")
right = rr.urdf.UrdfTree.from_file_path(urdf_path, entity_path_prefix="right", frame_prefix="right/")
# Log both robots (geometry + static transforms with prefixed frame IDs).
left.log_urdf_to_recording()
right.log_urdf_to_recording()
# Offset the arms so they don't overlap.
rec.log("left", rr.Transform3D(translation=[-0.2, 0, 0]), static=True)
rec.log("right", rr.Transform3D(translation=[0.2, 0, 0]), static=True)
blueprint = rrb.Grid(
rrb.Spatial3DView(
name="Dual Arm",
overrides={
"left/so_arm100/collision_geometries": rrb.EntityBehavior(visible=False),
"right/so_arm100/collision_geometries": rrb.EntityBehavior(visible=False),
},
)
)
rec.send_blueprint(blueprint)
for step in range(10000):
rec.set_time("step", sequence=step)
for joint_index, joint in enumerate(left.joints()):
if joint.joint_type == "revolute":
angle = _fake_angle(joint, step, joint_index, phase=0.0)
rec.log("left/joint_transforms", joint.compute_transform(angle, clamp=True))
for joint_index, joint in enumerate(right.joints()):
if joint.joint_type == "revolute":
angle = _fake_angle(joint, step, joint_index, phase=2.0)
rec.log("right/joint_transforms", joint.compute_transform(angle, clamp=True))
rr.script_teardown(args)
def _fake_angle(joint: rr.urdf.UrdfJoint, step: int, joint_index: int, phase: float) -> float:
"""Generate a smooth oscillating angle within the joint's limits."""
sin_value = math.sin(step * (0.02 + joint_index / 100.0) + phase)
return joint.limit_lower + (sin_value + 1.0) / 3.0 * (joint.limit_upper - joint.limit_lower)
if __name__ == "__main__":
main()
examples/python/animated_urdf/pyproject.toml
+12
View File
@@ -0,0 +1,12 @@
[project]
name = "animated_urdf"
version = "0.1.0"
readme = "README.md"
dependencies = ["rerun-sdk"]
[project.scripts]
animated_urdf = "animated_urdf:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/annotation_gradio/README.md
+35
View File
@@ -0,0 +1,35 @@
<!--[metadata]
title = "Interactive 3D annotation app with Rerun and Gradio"
tags = ["2D", "3D", "Pinhole camera", "Time series", "SAM", "Segmentation"]
source = "https://github.com/rerun-io/annotation-example"
thumbnail = "https://static.rerun.io/square-thumbnail/931d51df4f2cfd84293c5c9212e5e84cecbd79b4/480w.png"
thumbnail_dimensions = [480, 344]
-->
https://vimeo.com/1078165216?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=2802:1790
## Background
This example showcases how to use Rerun with gradio to generate an annotation app. It consists of two different modes both of which leverage Segment Anything 2.
The first mode focuses on tracking an object in a monocular video stream. In addition to segmentation masks, it generates real-time depth maps and point clouds to provide full 3D spatial context, enabling users to visualize, inspect, and annotate the tracked object directly in three-dimensional space.
The second mode uses a multiview RGB-D video dataset. By obtaining segmentation masks from two synchronized and calibrated RGB-D views, the app triangulates these 2D masks to reconstruct a precise 3D mask of the chosen object. It then propagates this 3D mask across all camera views and through each frame of the videos, resulting in a fully tracked 3D object trajectory over time.
## Run the code
This is an external example. Check the [repository](https://github.com/rerun-io/annotation-example) for more information on how to run the code.
TLDR: make sure you have the [Pixi package manager](https://pixi.sh/latest/#installation) installed and run
```
git clone https://github.com/rerun-io/annotation-example
cd annotation-example
pixi run app
```
this will run the single view (monocular) app
```
pixi run multiview-app
```
will run the multiview rgb-d app
examples/python/any_scalar/README.md
+144
View File
@@ -0,0 +1,144 @@
<!--[metadata]
title = "Any scalar"
tags = ["Any scalar", "Plotting", "DynamicArchetype"]
thumbnail = "https://static.rerun.io/any_scalar_example_market/4076a99f7fd5912af93258aa0c6c775a96f8b8e7/480w.png"
thumbnail_dimensions = [480, 259]
channel = "release"
-->
<iframe width="560" height="315" src="https://www.youtube.com/embed/G9Xxf0sNYcQ?si=jfb-WrY9WrFGh6mB" title="Any Scalar Example" frameborder="0" allow="accelerometer; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" referrerpolicy="strict-origin-when-cross-origin" allowfullscreen></iframe>
*A 6-minute narrated walkthrough of using the Rerun UI to plot arbitrary scalar data from a dataset (MCAP) is available on [Youtube](https://www.youtube.com/embed/G9Xxf0sNYcQ?si=jfb-WrY9WrFGh6mB).*
## Overview
This example demonstrates how to visualize arbitrary data, even when it was not logged with specific Rerun-semantics. With the **"Any Scalar"** feature, you can log complex data structures (like dictionaries or structs) once and use **Selectors** in the Blueprint to "pick" which internal fields to plot.
**Key Benefits:**
* **Decoupled Logging:** You no longer need to log separate scalar entities for every value you want to graph.
* **Selective Visualization:** Use a single data stream to power multiple different views by targeting specific component fields (e.g., `.position` or `.close`).
---
## Run the code
To run this example, make sure you have the [required Python version](https://ref.rerun.io/docs/python/main/common#supported-python-versions), the Rerun repository checked out and the latest SDK installed:
```sh
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```sh
pip install -e examples/python/any_scalar
```
To experiment with the provided example, simply execute the main Python script:
```sh
python -m any_scalar --demo robotics # Simulated PID control
python -m any_scalar --demo market # Real-time stock performance
```
If you wish to explore additional features, use the CLI with the `--help` option for guidance:
```sh
python -m any_scalar --help
```
---
## Guided demos
### 1. Robotics: PID controller telemetry
**Goal:** Visualize a control loop's internal state without logging separate scalars for every field.
In `robotics_demo.py`, we simulate a joint controller. Instead of logging `error`, `effort`, and `position` as individual Rerun entities, we log a single **Telemetry struct** per time step.
**Tutorial highlights:**
- **Decoupled logging:** We log one `ControllerTelemetry` object. Later, in the Blueprint, we "pick" which parts to see.
- **Visualizer mapping:** Notice how the `Error` field is plotted twice: once as a **Line** (to see trends) and once as **Points** (to see individual sample timing).
- **Step interpolation:** The `Effort` signal uses `StepAfter` interpolation, which accurately reflects how a digital controller holds its output constant between steps.
- **Boolean plotting:** The `is_stable` flag is visualized as a step-function scalar (0/1).
What you should see when running `python -m any_scalar --demo robotics`:
<picture>
<img src="https://static.rerun.io/any_scalar_example_robotics/f665e28b471f5b7b575c14e0a7fe11b10f636b88/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/any_scalar_example_robotics/f665e28b471f5b7b575c14e0a7fe11b10f636b88/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/any_scalar_example_robotics/f665e28b471f5b7b575c14e0a7fe11b10f636b88/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/any_scalar_example_robotics/f665e28b471f5b7b575c14e0a7fe11b10f636b88/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/any_scalar_example_robotics/f665e28b471f5b7b575c14e0a7fe11b10f636b88/1200w.png">
</picture>
### 2. Market data: relative performance
**Goal:** Compare multiple live data streams (tickers) using a centralized selector.
In `market_demo.py`, we fetch real stock data. We log the raw prices and a "normalized" % change field.
**Tutorial highlights:**
- **Dynamic normalization:** We log the price relative to the morning opening.
- **Dynamic archetype:** We log the stock data as a dictionary using `rerun.DynamicArchetype`.
- **Selectors:** We use [`jq`](https://jqlang.org/)-style selectors, like `.prices.normalized` and `.prices.close`, to power different parts of the dashboard.
What you should see when running `python -m any_scalar --demo market`:
<picture>
<img src="https://static.rerun.io/any_scalar_example_market/4076a99f7fd5912af93258aa0c6c775a96f8b8e7/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/any_scalar_example_market/4076a99f7fd5912af93258aa0c6c775a96f8b8e7/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/any_scalar_example_market/4076a99f7fd5912af93258aa0c6c775a96f8b8e7/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/any_scalar_example_market/4076a99f7fd5912af93258aa0c6c775a96f8b8e7/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/any_scalar_example_market/4076a99f7fd5912af93258aa0c6c775a96f8b8e7/1200w.png">
</picture>
### 3. Load datasets directly into the viewer
**Goal:** Plot values from dataset files without writing any code.
Because Rerun can now plot **Any Scalar**, you can drag an `.mcap` or `.rrd` file into the viewer and create a `Time Series` view. Use the UI in the viewer to drill into nested ROS messages or telemetry logs and start plotting immediately.
> [!TIP]
> **Watch the video at the top of this page** to see a step-by-step walkthrough of how to use the UI to plot any field from an MCAP/RRD file.
---
## The "Magic": component mapping
### What is "Any Scalar"?
Traditionally, to plot a graph, you had to log data specifically as one of Rerun's `Scalar` archetypes. With **Any Scalar**, you can log complex blobs (Dictionaries, TypedDicts, Arrow Structs) and Rerun will let you "map" internal fields to visualizers.
### Selectors (jq-style)
Rerun uses a path syntax inspired by [`jq`](https://jqlang.org/) to reach into your data:
- `.state.position` -> reaches into the `state` dict and finds `position`.
- `.prices.normalized` -> pulls the calculated performance from the market tick.
### Benefits
1. **Developer velocity:** Log your entire state object once; decide what to plot later in the UI.
2. **Smaller files:** Less metadata overhead than logging 50 separate entities.
3. **Flexibility:** Change what you are visualizing in the Blueprint without restarting your simulation or re-running your data pipeline.
---
## Resources
- [Customize views](https://rerun.io/docs/concepts/visualization/customize-views)
- [Plot any scalar](https://rerun.io/docs/howto/visualization/plot-any-scalar)
- [Component Mappings Guide](https://rerun.io/docs/howto/visualization/component-mappings)
---
## Used Rerun types
[`DynamicArchetype`](https://ref.rerun.io/docs/python/stable/common/custom_data/#rerun.dynamic_archetype.DynamicArchetype), [`SeriesLines`](https://rerun.io/docs/reference/types/archetypes/series_lines), [`SeriesPoints`](https://rerun.io/docs/reference/types/archetypes/series_points)
examples/python/any_scalar/any_scalar/__init__.py
View File
examples/python/any_scalar/any_scalar/__main__.py
Executable
+46
View File
@@ -0,0 +1,46 @@
#!/usr/bin/env python3
"""
A single entry point to run the Rerun any_scalar_example demos.
You can run either the robotics demo or the market demo:
python -m any_scalar_example --demo robotics
python -m any_scalar_example --demo market
"""
from __future__ import annotations
import argparse
import rerun as rr
from any_scalar import market_demo, robotics_demo
def main() -> None:
parser = argparse.ArgumentParser(description="Rerun Any Scalar Example")
parser.add_argument(
"--demo",
type=str,
default="robotics",
choices=["robotics", "market"],
help="Which demo to run (default: robotics)",
)
# Add standard Rerun arguments
rr.script_add_args(parser)
args = parser.parse_args()
# Initialize Rerun SDK
rr.script_setup(args, "rerun_example_any_scalar")
if args.demo == "robotics":
robotics_demo.run_robotics_simulation()
rr.send_blueprint(robotics_demo.generate_blueprint())
elif args.demo == "market":
market_demo.run_market_demo()
rr.send_blueprint(market_demo.generate_blueprint())
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/any_scalar/any_scalar/market_demo.py
Executable
+161
View File
@@ -0,0 +1,161 @@
#!/usr/bin/env python3
"""
Market Data Comparison Demo
Shows features for "Any Scalar" visualization:
- Normalizing multiple real-time tickers (NVDA, AAPL, MSFT, GOOGL, AMD, INTC).
- Comparing relative performance using Blueprint Selectors.
"""
from __future__ import annotations
import yfinance as yf
import rerun as rr
import rerun.blueprint as rrb
from rerun.blueprint.datatypes import ComponentSourceKind, VisualizerComponentMapping
def log_market_data(tickers: list[str]) -> None:
"""Fetch, normalize, and log real market data for multiple tickers."""
print(f"Fetching market data for {tickers}…")
# Fetch 7 days (max for 1m interval) to ensure we have data even on weekends/holidays
df = yf.download(tickers, period="7d", interval="1m", progress=False)
if df.empty:
print(f"Warning: No data found for {tickers}.")
return
# Find the most recent date that has at least 100 data points
# This ensures we do not show a nearly empty chart if the market just opened.
daily_counts = df.groupby(df.index.date).size()
valid_dates = daily_counts[daily_counts >= 100].index
if len(valid_dates) == 0:
print("Warning: No day found with at least 100 data points. Using the most recent day with any data.")
last_date = df.index.max().date()
else:
last_date = max(valid_dates)
df = df[df.index.date == last_date]
for ticker in tickers:
try:
if len(tickers) > 1:
ticker_data = df.xs(ticker, axis=1, level=1).dropna()
else:
ticker_data = df.dropna()
if ticker_data.empty:
continue
timestamps = ticker_data.index
# Normalization: Scale relative to the first Close price (%)
baseline = float(ticker_data.iloc[0]["Close"])
# Nested Any Scalar: prices.close, prices.normalized, details.volume
market_ticks = [
[
{
"prices": {
"close": float(row["Close"]),
"normalized": (float(row["Close"]) / baseline - 1.0) * 100.0,
},
"details": {"volume": float(row["Volume"])},
}
]
for _, row in ticker_data.iterrows()
]
rr.send_columns(
f"market/{ticker}",
indexes=[rr.TimeColumn("market_time", timestamp=timestamps)],
columns=[*rr.DynamicArchetype.columns(archetype="MarketTelemetry", components={"data": market_ticks})],
)
except Exception as e:
print(f"Error processing {ticker}: {e}")
def run_market_demo() -> None:
"""Run the market data demo and log data."""
tickers = ["NVDA", "AAPL", "MSFT", "GOOGL", "AMD", "INTC"]
log_market_data(tickers)
def generate_blueprint() -> rrb.Blueprint:
"""Generate the blueprint for the market demo."""
tickers = ["NVDA", "AAPL", "MSFT", "GOOGL", "AMD", "INTC"]
colors = {
"NVDA": [118, 185, 0], # NVIDIA Green
"AAPL": [85, 85, 85], # Apple Gray
"MSFT": [0, 164, 239], # Microsoft Blue
"GOOGL": [219, 68, 55], # Google Red
"AMD": [237, 28, 36], # AMD Red
"INTC": [0, 104, 181], # Intel Blue
}
return rrb.Blueprint(
rrb.Vertical(
rrb.Horizontal(
rrb.TimeSeriesView(
name="Market Relative Performance (%)",
origin="/market",
overrides={
f"market/{ticker}": [
rr.SeriesLines(names=f"{ticker} (Rel)", colors=colors.get(ticker)).visualizer(
mappings=[
VisualizerComponentMapping(
target="Scalars:scalars",
source_kind=ComponentSourceKind.SourceComponent,
source_component="MarketTelemetry:data",
selector=".prices.normalized",
)
]
)
]
for ticker in tickers
},
axis_x=rrb.TimeAxis(
view_range=rrb.TimeRange(
start=rrb.TimeRangeBoundary.infinite(), end=rrb.TimeRangeBoundary.infinite()
)
),
),
rrb.TimeSeriesView(
name="Market Close",
origin="/market",
overrides={
f"market/{ticker}": [
rr.SeriesLines(names=f"{ticker} (Close)", colors=colors.get(ticker)).visualizer(
mappings=[
VisualizerComponentMapping(
target="Scalars:scalars",
source_kind=ComponentSourceKind.SourceComponent,
source_component="MarketTelemetry:data",
selector=".prices.close",
)
]
)
]
for ticker in tickers
},
axis_x=rrb.TimeAxis(
view_range=rrb.TimeRange(
start=rrb.TimeRangeBoundary.infinite(), end=rrb.TimeRangeBoundary.infinite()
)
),
),
),
rrb.DataframeView(name="Market Inspector", origin="/market"),
)
)
def main() -> None:
rr.init("rerun_example_any_scalar_market", spawn=True)
run_market_demo()
rr.send_blueprint(generate_blueprint())
if __name__ == "__main__":
main()
examples/python/any_scalar/any_scalar/robotics_demo.py
Executable
+189
View File
@@ -0,0 +1,189 @@
#!/usr/bin/env python3
"""
Robotics PID Controller Demo
Shows features for "Any Scalar" visualization:
- Multiple visualizers (Lines + Points) on the same telemetry field.
- Step interpolation for discrete control effort and state flags.
- Boolean scalar plotting.
"""
from __future__ import annotations
import numpy as np
import rerun as rr
import rerun.blueprint as rrb
from rerun.blueprint.datatypes import ComponentSourceKind, VisualizerComponentMapping
def simulate_robot_controller() -> None:
"""Simulate a 1-DOF joint with a PID controller tracking a trajectory."""
print("Simulating robot PID controller…")
# PID Constants
Kp, Ki, Kd = 12.0, 1.5, 0.8
dt = 0.05
steps = 400
# State
position = 0.0
velocity = 0.0
integral = 0.0
prev_error = 0.0
telemetry_data = []
for i in range(steps):
t = i * dt
# Trajectory
setpoint = 2.0 * np.sin(1.0 * t) + 0.5 * np.cos(3.0 * t)
# Controller
error = setpoint - position
integral += error * dt
derivative = (error - prev_error) / dt
effort = Kp * error + Ki * integral + Kd * derivative
prev_error = error
# Physics
acceleration = effort - 0.5 * velocity
velocity += acceleration * dt
position += velocity * dt
# Deeply nested telemetry: state, control, status
# This showcases Rerun's powerful jq-style selectors for hierarchical data.
telemetry_data.append([
{
"state": {
"setpoint": float(setpoint),
"position": float(position),
},
"control": {
"error": float(error),
"effort": float(effort),
},
"status": {"is_stable": bool(abs(error) < 0.1)},
}
])
# Log telemetry using send_columns
rr.send_columns(
"robot/joint_0",
indexes=[rr.TimeColumn("robot_step", sequence=np.arange(steps))],
columns=[*rr.DynamicArchetype.columns(archetype="ControllerTelemetry", components={"data": telemetry_data})],
)
def run_robotics_simulation() -> None:
"""Run the robotics simulation and log data."""
simulate_robot_controller()
def generate_blueprint() -> rrb.Blueprint:
"""Generate the blueprint for the robotics demo."""
return rrb.Blueprint(
rrb.Vertical(
# View 1: Main Control Performance (Multiple Visualizers + Nested Selectors)
rrb.TimeSeriesView(
name="Control Performance",
origin="/robot/joint_0",
overrides={
"robot/joint_0": [
rr.SeriesLines(names="Setpoint", colors=[100, 100, 255]).visualizer(
mappings=[
VisualizerComponentMapping(
target="Scalars:scalars",
source_kind=ComponentSourceKind.SourceComponent,
source_component="ControllerTelemetry:data",
selector=".state.setpoint",
)
]
),
rr.SeriesLines(names="Position", colors=[255, 100, 0]).visualizer(
mappings=[
VisualizerComponentMapping(
target="Scalars:scalars",
source_kind=ComponentSourceKind.SourceComponent,
source_component="ControllerTelemetry:data",
selector=".state.position",
)
]
),
# Multiple visualizers on the SAME field (.control.error)
rr.SeriesLines(names="Error (Line)", colors=[255, 0, 0]).visualizer(
mappings=[
VisualizerComponentMapping(
target="Scalars:scalars",
source_kind=ComponentSourceKind.SourceComponent,
source_component="ControllerTelemetry:data",
selector=".control.error",
)
]
),
rr.SeriesPoints(names="Error (Dots)", colors=[255, 0, 0]).visualizer(
mappings=[
VisualizerComponentMapping(
target="Scalars:scalars",
source_kind=ComponentSourceKind.SourceComponent,
source_component="ControllerTelemetry:data",
selector=".control.error",
)
]
),
]
},
),
# View 2: Internals (Step Interpolation & Nested Booleans)
rrb.TimeSeriesView(
name="Controller Internals",
origin="/robot/joint_0",
overrides={
"robot/joint_0": [
rr.SeriesLines(names="Effort", colors=[0, 200, 200], interpolation_mode="StepAfter").visualizer(
mappings=[
VisualizerComponentMapping(
target="Scalars:scalars",
source_kind=ComponentSourceKind.SourceComponent,
source_component="ControllerTelemetry:data",
selector=".control.effort",
)
]
),
rr.SeriesLines(
names="Stability Flag", colors=[0, 255, 0], interpolation_mode="StepAfter"
).visualizer(
mappings=[
VisualizerComponentMapping(
target="Scalars:scalars",
source_kind=ComponentSourceKind.SourceComponent,
source_component="ControllerTelemetry:data",
selector=".status.is_stable",
)
]
),
]
},
),
# View 3: Dataframe Inspector
rrb.DataframeView(
name="Telemetry Inspector",
origin="/robot/joint_0",
query=rrb.archetypes.DataframeQuery(
auto_scroll=True,
),
),
)
)
def main() -> None:
rr.init("rerun_example_any_scalar_robotics", spawn=True)
run_robotics_simulation()
rr.send_blueprint(generate_blueprint())
if __name__ == "__main__":
main()
examples/python/any_scalar/pyproject.toml
+12
View File
@@ -0,0 +1,12 @@
[project]
name = "any_scalar"
version = "0.1.0"
readme = "README.md"
dependencies = ["psutil", "pyarrow", "rerun-sdk", "yfinance"]
[project.scripts]
any_scalar = "any_scalar.__main__:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/arflow/README.md
+52
View File
@@ -0,0 +1,52 @@
<!--[metadata]
title = "ARFlow: a framework for simplifying AR experimentation workflow"
source = "https://github.com/cake-lab/ARFlow"
tags = ["3D", "Augmented reality", "Spatial computing", "Integration"]
thumbnail = "https://static.rerun.io/arflow/a6b509af10a42b3c7ad3909d44e972a3cb1a9c41/480w.png"
thumbnail_dimensions = [480, 480]
-->
This is an external project that uses Rerun as a core component.
## External project presentation
[Paper](https://dl.acm.org/doi/10.1145/3638550.3643617) | [BibTeX](#citation) | [Project Page](https://cake.wpi.edu/ARFlow/) | [Video](https://youtu.be/mml8YrCgfTk)
ARFlow is designed to lower the barrier for AR researchers to evaluate ideas in hours instead of weeks or months, following:
- **Efficient AR experiment data collection** from various data sources, including camera, depth sensors, and IMU sensors with an efficient thin mobile client.
- Flexible AR runtime data **management** with **real-time visualization** (powered by Rerun).
- **Easy integration** with existing AR research projects without breaking the experimentation logic.
Watch our demo video:
[![Demo video](https://img.youtube.com/vi/mml8YrCgfTk/maxresdefault.jpg)](https://youtu.be/mml8YrCgfTk)
## Get started
Please see [the original project repo](https://github.com/cake-lab/ARFlow/blob/main/README.md), and refer to the individual [server](https://github.com/cake-lab/ARFlow/blob/090995a066e8394fc7358a889c655fa3020d20d4/python/README.md) and [client](https://github.com/cake-lab/ARFlow/tree/090995a066e8394fc7358a889c655fa3020d20d4/unity) installation guides.
## Citation
Please add the following citation in your publication if you used our code for your research project.
```bibtex
@inproceedings{zhao2024arflow,
author = {Zhao, Yiqin and Guo, Tian},
title = {Demo: ARFlow: A Framework for Simplifying AR Experimentation Workflow},
year = {2024},
isbn = {9798400704970},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://dl.acm.org/doi/10.1145/3638550.3643617},
doi = {10.1145/3638550.3643617},
abstract = {The recent advancement in computer vision and XR hardware has ignited the community's interest in AR systems research. Similar to traditional systems research, the evaluation of AR systems involves capturing real-world data with AR hardware and iteratively evaluating the targeted system designs [1]. However, it is challenging to conduct scalable and reproducible AR experimentation [2] due to two key reasons. First, there is a lack of integrated framework support in real-world data capturing, which makes it a time-consuming process. Second, AR data often exhibits characteristics, including temporal and spatial variations, and is in a multi-modal format, which makes it difficult to conduct controlled evaluations.},
booktitle = {Proceedings of the 25th International Workshop on Mobile Computing Systems and Applications},
pages = {154},
numpages = {1},
location = {<conf-loc>, <city>San Diego</city>, <state>CA</state>, <country>USA</country>, </conf-loc>},
series = {HOTMOBILE '24}
}
```
examples/python/arkit_scenes/.gitignore
+1
View File
@@ -0,0 +1 @@
dataset/**
examples/python/arkit_scenes/README.md
+132
View File
@@ -0,0 +1,132 @@
<!--[metadata]
title = "ARKit scenes"
tags = ["2D", "3D", "Depth", "Mesh", "Object detection", "Pinhole camera", "Blueprint"]
thumbnail = "https://static.rerun.io/arkit-scenes/6d920eaa42fb86cfd264d47180ecbecbb6dd3e09/480w.png"
thumbnail_dimensions = [480, 480]
channel = "main"
include_in_manifest = true
-->
This example visualizes the [ARKitScenes dataset](https://github.com/apple/ARKitScenes/) using Rerun.
The dataset contains color images, depth images, the reconstructed mesh, and labeled bounding boxes around furniture.
<picture data-inline-viewer="examples/arkit_scenes">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/1200w.png">
<img src="https://static.rerun.io/arkit_scenes/fb9ec9e8d965369d39d51b17fc7fc5bae6be10cc/full.png" alt="ARKit Scenes screenshot">
</picture>
## Used Rerun types
[`Image`](https://www.rerun.io/docs/reference/types/archetypes/image),
[`DepthImage`](https://www.rerun.io/docs/reference/types/archetypes/depth_image), [`Transform3D`](https://www.rerun.io/docs/reference/types/archetypes/transform3d),
[`Pinhole`](https://www.rerun.io/docs/reference/types/archetypes/pinhole), [`Mesh3D`](https://www.rerun.io/docs/reference/types/archetypes/mesh3d),
[`Boxes3D`](https://www.rerun.io/docs/reference/types/archetypes/boxes3d),
[`TextDocument`](https://www.rerun.io/docs/reference/types/archetypes/text_document)
## Background
The ARKitScenes dataset, captured using Apple's ARKit technology, encompasses a diverse array of indoor scenes, offering color and depth images, reconstructed 3D meshes, and labeled bounding boxes around objects like furniture. It's a valuable resource for researchers and developers in computer vision and augmented reality, enabling advancements in object recognition, depth estimation, and spatial understanding.
## Logging and visualizing with Rerun
This visualization through Rerun highlights the dataset's potential in developing immersive AR experiences and enhancing machine learning models for real-world applications while showcasing Reruns visualization capabilities.
### Logging a moving RGB-D camera
To log a moving RGB-D camera, we log four key components: the camera's intrinsics via a pinhole camera model, its pose or extrinsics, along with the color and depth images. The camera intrinsics, which define the camera's lens properties, and the pose, detailing its position and orientation, are logged to create a comprehensive 3D to 2D mapping. Both the RGB and depth images are then logged as child entities, capturing the visual and depth aspects of the scene, respectively. This approach ensures a detailed recording of the camera's viewpoint and the scene it captures, all stored compactly under the same entity path for simplicity.
```python
rr.log("world/camera_lowres", rr.Transform3D(transform=camera_from_world))
rr.log("world/camera_lowres", rr.Pinhole(image_from_camera=intrinsic, resolution=[w, h]))
rr.log(f"{entity_id}/rgb", rr.Image(rgb).compress(jpeg_quality=95))
rr.log(f"{entity_id}/depth", rr.DepthImage(depth, meter=1000))
```
### Ground-truth mesh
The mesh is logged as an [rr.Mesh3D archetype](https://www.rerun.io/docs/reference/types/archetypes/mesh3d).
In this case the mesh is composed of mesh vertices, indices (i.e., which vertices belong to the same face), and vertex
colors.
```python
rr.log(
"world/mesh",
rr.Mesh3D(
vertex_positions=mesh.vertices,
vertex_colors=mesh.visual.vertex_colors,
triangle_indices=mesh.faces,
),
static=True,
)
```
Here, the mesh is logged to the world/mesh entity and is marked as static, since it does not change in the context of this visualization.
### Logging 3D bounding boxes
Here we loop through the data and add bounding boxes to all the items found.
```python
for i, label_info in enumerate(annotation["data"]):
rr.log(
f"world/annotations/box-{uid}-{label}",
rr.Boxes3D(
half_sizes=half_size,
centers=centroid,
labels=label,
colors=colors[i],
),
rr.InstancePoses3D(mat3x3=mat3x3),
static=True,
)
```
### Setting up the default blueprint
This example benefits at lot from having a custom blueprint defined. This happens with the following code:
```python
primary_camera_entity = HIGHRES_ENTITY_PATH if args.include_highres else LOWRES_POSED_ENTITY_PATH
blueprint = rrb.Horizontal(
rrb.Spatial3DView(name="3D"),
rrb.Vertical(
rrb.Tabs(
rrb.Spatial2DView(
name="RGB",
origin=primary_camera_entity,
contents=["$origin/rgb", "/world/annotations/**"],
),
rrb.Spatial2DView(
name="Depth",
origin=primary_camera_entity,
contents=["$origin/depth", "/world/annotations/**"],
),
name="2D",
),
rrb.TextDocumentView(name="Readme"),
),
)
rr.script_setup(args, "rerun_example_arkit_scenes", default_blueprint=blueprint)
```
In particular, we want to reproject 3D annotations onto the 2D camera views. To configure such a view, two things are necessary:
- The view origin must be set to the entity that contains the pinhole transforms. In this example, the entity path is stored in the `primary_camera_entity` variable.
- The view contents must explicitly include the annotations, which are not logged in the subtree defined by the origin. This is done using the `contents` argument, here set to `["$origin/depth", "/world/annotations/**"]`.
## Run the code
To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/arkit_scenes
```
To run this example use
```bash
python -m arkit_scenes
```
examples/python/arkit_scenes/arkit_scenes/__init__.py
View File
examples/python/arkit_scenes/arkit_scenes/__main__.py
Executable
+322
View File
@@ -0,0 +1,322 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
import os
from typing import TYPE_CHECKING, Any
import cv2
import numpy as np
import trimesh
from scipy.spatial.transform import Rotation as R
from tqdm import tqdm
import rerun as rr # pip install rerun-sdk
import rerun.blueprint as rrb
from .download_dataset import AVAILABLE_RECORDINGS, ensure_recording_available
if TYPE_CHECKING:
from pathlib import Path
DESCRIPTION = """
# ARKitScenes
This example visualizes the [ARKitScenes dataset](https://github.com/apple/ARKitScenes/) using Rerun. The dataset
contains color images, depth images, the reconstructed mesh, and labeled bounding boxes around furniture.
The full source code for this example is available
[on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/arkit_scenes).
""".strip()
Color = tuple[float, float, float, float]
# hack for now since dataset does not provide orientation information, only known after initial visual inspection
ORIENTATION = {
"48458663": "landscape",
"42444949": "portrait",
"41069046": "portrait",
"41125722": "portrait",
"41125763": "portrait",
"42446167": "portrait",
}
assert len(ORIENTATION) == len(AVAILABLE_RECORDINGS)
assert set(ORIENTATION.keys()) == set(AVAILABLE_RECORDINGS)
LOWRES_POSED_ENTITY_PATH = "world/camera_lowres"
HIGHRES_ENTITY_PATH = "world/camera_highres"
def load_json(js_path: Path) -> dict[str, Any]:
with open(js_path, encoding="utf8") as f:
json_data: dict[str, Any] = json.load(f)
return json_data
def log_annotated_bboxes(annotation: dict[str, Any]) -> None:
"""
Logs annotated oriented bounding boxes to Rerun.
annotation json file
| |-- label: object name of bounding box
| |-- axesLengths[x, y, z]: size of the origin bounding-box before transforming
| |-- centroid[]: the translation matrix (1,3) of bounding-box
| |-- normalizedAxes[]: the rotation matrix (3,3) of bounding-box
"""
for label_info in annotation["data"]:
uid = label_info["uid"]
label = label_info["label"]
half_size = 0.5 * np.array(label_info["segments"]["obbAligned"]["axesLengths"]).reshape(-1, 3)[0]
centroid = np.array(label_info["segments"]["obbAligned"]["centroid"]).reshape(-1, 3)[0]
mat3x3 = np.array(label_info["segments"]["obbAligned"]["normalizedAxes"]).reshape(3, 3).T
rr.log(
f"world/annotations/box-{uid}-{label}",
rr.Boxes3D(
half_sizes=half_size,
labels=label,
),
rr.InstancePoses3D(translations=centroid, mat3x3=mat3x3),
static=True,
)
def log_camera(
intri_path: Path,
frame_id: str,
poses_from_traj: dict[str, rr.Transform3D],
entity_id: str,
) -> None:
"""Logs camera transform and 3D bounding boxes in the image frame."""
w, h, fx, fy, cx, cy = np.loadtxt(intri_path)
intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
camera_from_world = poses_from_traj[frame_id]
# clear previous centroid labels
rr.log(f"{entity_id}/bbox-2D-segments", rr.Clear(recursive=True))
# pathlib makes it easy to get the parent, but log methods requires a string
rr.log(entity_id, camera_from_world)
rr.log(entity_id, rr.Pinhole(image_from_camera=intrinsic, resolution=[w, h]))
def read_camera_from_world(traj_string: str) -> tuple[str, rr.Transform3D]:
"""
Reads out camera_from_world transform from trajectory string.
Args:
----
traj_string:
A space-delimited file where each line represents a camera position at a particular timestamp.
The file has seven columns:
* Column 1: timestamp
* Columns 2-4: rotation (axis-angle representation in radians)
* Columns 5-7: translation (usually in meters)
Returns
-------
timestamp: float
timestamp in seconds
camera_from_world: tuple of two numpy arrays
A tuple containing a translation vector and a quaternion that represent the camera_from_world transform
Raises
------
AssertionError:
If the input string does not contain 7 tokens.
"""
tokens = traj_string.split() # Split the input string into tokens
assert len(tokens) == 7, f"Input string must have 7 tokens, but found {len(tokens)}."
ts: str = tokens[0] # Extract timestamp from the first token
# Extract rotation from the second to fourth tokens
angle_axis = [float(tokens[1]), float(tokens[2]), float(tokens[3])]
rotation = R.from_rotvec(np.asarray(angle_axis))
# Extract translation from the fifth to seventh tokens
translation = np.asarray([float(tokens[4]), float(tokens[5]), float(tokens[6])])
# Create tuple in format log_transform3d expects
camera_from_world = rr.Transform3D(
translation=translation,
rotation=rr.Quaternion(xyzw=rotation.as_quat()),
relation=rr.TransformRelation.ChildFromParent,
)
return (ts, camera_from_world)
def find_closest_frame_id(target_id: str, frame_ids: dict[str, Any]) -> str:
"""Finds the closest frame id to the target id."""
target_value = float(target_id)
closest_id = min(frame_ids.keys(), key=lambda x: abs(float(x) - target_value))
return closest_id
def log_arkit(recording_path: Path, include_highres: bool) -> None:
"""
Logs ARKit recording data using Rerun.
Args:
----
recording_path (Path):
The path to the ARKit recording.
include_highres (bool):
Whether to include high resolution data.
Returns
-------
None
"""
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
video_id = recording_path.stem
lowres_image_dir = recording_path / "lowres_wide"
image_dir = recording_path / "wide"
lowres_depth_dir = recording_path / "lowres_depth"
depth_dir = recording_path / "highres_depth"
lowres_intrinsics_dir = recording_path / "lowres_wide_intrinsics"
intrinsics_dir = recording_path / "wide_intrinsics"
traj_path = recording_path / "lowres_wide.traj"
# frame_ids are indexed by timestamps, you can see more info here
# https://github.com/apple/ARKitScenes/blob/main/threedod/README.md#data-organization-and-format-of-input-data
depth_filenames = [x.name for x in sorted(lowres_depth_dir.iterdir())]
lowres_frame_ids = [x.split(".png")[0].split("_")[1] for x in depth_filenames]
lowres_frame_ids.sort()
# dict of timestamp to pose which is a tuple of translation and quaternion
camera_from_world_dict = {}
with open(traj_path, encoding="utf-8") as f:
trajectory = f.readlines()
for line in trajectory:
timestamp, camera_from_world = read_camera_from_world(line)
# round timestamp to 3 decimal places as seen in the original repo here
# https://github.com/apple/ARKitScenes/blob/e2e975128a0a9695ea56fa215fe76b4295241538/threedod/benchmark_scripts/utils/tenFpsDataLoader.py#L247
timestamp = f"{round(float(timestamp), 3):.3f}"
camera_from_world_dict[timestamp] = camera_from_world
rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
ply_path = recording_path / f"{recording_path.stem}_3dod_mesh.ply"
print(f"Loading {ply_path}…")
assert os.path.isfile(ply_path), f"Failed to find {ply_path}"
mesh = trimesh.load(str(ply_path))
rr.log(
"world/mesh",
rr.Mesh3D(
vertex_positions=mesh.vertices, # type: ignore[attr-defined]
vertex_colors=mesh.visual.vertex_colors, # type: ignore[attr-defined]
triangle_indices=mesh.faces, # type: ignore[attr-defined]
face_rendering="Front", # We want to hide the front facing faces, but the dataset uses mostly clockwise winding order which is the opposite of what Rerun assumes (CCW).
),
static=True,
)
# load the obb annotations and log them in the world frame
bbox_annotations_path = recording_path / f"{recording_path.stem}_3dod_annotation.json"
annotation = load_json(bbox_annotations_path)
log_annotated_bboxes(annotation)
print("Processing frames…")
for frame_timestamp in tqdm(lowres_frame_ids):
# frame_id is equivalent to timestamp
rr.set_time("time", duration=float(frame_timestamp))
# load the lowres image and depth
bgr = cv2.imread(f"{lowres_image_dir}/{video_id}_{frame_timestamp}.png")
depth = cv2.imread(f"{lowres_depth_dir}/{video_id}_{frame_timestamp}.png", cv2.IMREAD_ANYDEPTH)
high_res_exists: bool = (image_dir / f"{video_id}_{frame_timestamp}.png").exists() and include_highres
# Log the camera transforms:
if frame_timestamp in camera_from_world_dict:
lowres_intri_path = lowres_intrinsics_dir / f"{video_id}_{frame_timestamp}.pincam"
log_camera(
lowres_intri_path,
frame_timestamp,
camera_from_world_dict,
LOWRES_POSED_ENTITY_PATH,
)
rr.log(f"{LOWRES_POSED_ENTITY_PATH}/bgr", rr.Image(bgr, color_model="BGR").compress(jpeg_quality=95))
rr.log(f"{LOWRES_POSED_ENTITY_PATH}/depth", rr.DepthImage(depth, meter=1000))
# log the high res camera
if high_res_exists:
rr.set_time("time high resolution", duration=float(frame_timestamp))
# only low res camera has a trajectory, high res does not so need to find the closest low res frame id
closest_lowres_frame_id = find_closest_frame_id(frame_timestamp, camera_from_world_dict)
highres_intri_path = intrinsics_dir / f"{video_id}_{frame_timestamp}.pincam"
log_camera(
highres_intri_path,
closest_lowres_frame_id,
camera_from_world_dict,
HIGHRES_ENTITY_PATH,
)
# load the highres image and depth if they exist
highres_bgr = cv2.imread(f"{image_dir}/{video_id}_{frame_timestamp}.png")
highres_depth = cv2.imread(f"{depth_dir}/{video_id}_{frame_timestamp}.png", cv2.IMREAD_ANYDEPTH)
rr.log(f"{HIGHRES_ENTITY_PATH}/bgr", rr.Image(highres_bgr, color_model="BGR").compress(jpeg_quality=75))
rr.log(f"{HIGHRES_ENTITY_PATH}/depth", rr.DepthImage(highres_depth, meter=1000))
def main() -> None:
parser = argparse.ArgumentParser(description="Visualizes the ARKitScenes dataset using the Rerun SDK.")
parser.add_argument(
"--video-id",
type=str,
choices=AVAILABLE_RECORDINGS,
default=AVAILABLE_RECORDINGS[0],
help="Video ID of the ARKitScenes Dataset",
)
parser.add_argument(
"--include-highres",
action="store_true",
help="Include the high resolution camera and depth images",
)
rr.script_add_args(parser)
args = parser.parse_args()
primary_camera_entity = HIGHRES_ENTITY_PATH if args.include_highres else LOWRES_POSED_ENTITY_PATH
blueprint = rrb.Horizontal(
rrb.Spatial3DView(name="3D"),
rrb.Vertical(
rrb.Tabs(
# Note that we re-project the annotations into the 2D views:
# For this to work, the origin of the 2D views has to be a pinhole camera,
# this way the viewer knows how to project the 3D annotations into the 2D views.
rrb.Spatial2DView(
name="BGR",
origin=primary_camera_entity,
contents=["$origin/bgr", "/world/annotations/**"],
),
rrb.Spatial2DView(
name="Depth",
origin=primary_camera_entity,
contents=["$origin/depth", "/world/annotations/**"],
),
name="2D",
),
rrb.TextDocumentView(name="Readme"),
row_shares=[2, 1],
),
)
rr.script_setup(args, "rerun_example_arkit_scenes", default_blueprint=blueprint)
recording_path = ensure_recording_available(args.video_id, args.include_highres)
log_arkit(recording_path, args.include_highres)
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/arkit_scenes/arkit_scenes/download_dataset.py
+324
View File
@@ -0,0 +1,324 @@
# Copied from https://github.com/apple/ARKitScenes/blob/9ec0b99c3cd55e29fc0724e1229e2e6c2909ab45/download_data.py
# Licensing information: https://github.com/apple/ARKitScenes/blob/9ec0b99c3cd55e29fc0724e1229e2e6c2909ab45/LICENSE
from __future__ import annotations
import math
import os
import subprocess
import zipfile
from pathlib import Path
from typing import Final
import pandas as pd
ARkitscense_url = "https://docs-assets.developer.apple.com/ml-research/datasets/arkitscenes/v1"
TRAINING: Final = "Training"
VALIDATION: Final = "Validation"
HIGRES_DEPTH_ASSET_NAME: Final = "highres_depth"
POINT_CLOUDS_FOLDER: Final = "laser_scanner_point_clouds"
AVAILABLE_RECORDINGS: Final = ["48458663", "42444949", "41069046", "41125722", "41125763", "42446167"]
DATASET_DIR: Final = Path(__file__).parent.parent / "dataset"
default_raw_dataset_assets = [
"mov",
"annotation",
"mesh",
"confidence",
"highres_depth",
"lowres_depth",
"lowres_wide.traj",
"lowres_wide",
"lowres_wide_intrinsics",
"ultrawide",
"ultrawide_intrinsics",
"vga_wide",
"vga_wide_intrinsics",
]
missing_3dod_assets_video_ids = [
"47334522",
"47334523",
"42897421",
"45261582",
"47333152",
"47333155",
"48458535",
"48018733",
"47429677",
"48458541",
"42897848",
"47895482",
"47333960",
"47430089",
"42899148",
"42897612",
"42899153",
"42446164",
"48018149",
"47332198",
"47334515",
"45663223",
"45663226",
"45663227",
]
def raw_files(video_id: str, assets: list[str], metadata: pd.DataFrame) -> list[str]:
file_names = []
for asset in assets:
if HIGRES_DEPTH_ASSET_NAME == asset:
in_upsampling = metadata.loc[metadata["video_id"] == float(video_id), ["is_in_upsampling"]].iat[0, 0]
if not in_upsampling:
print(f"Skipping asset {asset} for video_id {video_id} - Video not in upsampling dataset")
continue # highres_depth asset only available for video ids from upsampling dataset
if asset in [
"confidence",
"highres_depth",
"lowres_depth",
"lowres_wide",
"lowres_wide_intrinsics",
"ultrawide",
"ultrawide_intrinsics",
"wide",
"wide_intrinsics",
"vga_wide",
"vga_wide_intrinsics",
]:
file_names.append(asset + ".zip")
elif asset == "mov":
file_names.append(f"{video_id}.mov")
elif asset == "mesh":
if video_id not in missing_3dod_assets_video_ids:
file_names.append(f"{video_id}_3dod_mesh.ply")
elif asset == "annotation":
if video_id not in missing_3dod_assets_video_ids:
file_names.append(f"{video_id}_3dod_annotation.json")
elif asset == "lowres_wide.traj":
if video_id not in missing_3dod_assets_video_ids:
file_names.append("lowres_wide.traj")
else:
raise Exception(f"No asset = {asset} in raw dataset")
return file_names
def download_file(url: str, file_name: str, dst: Path) -> bool:
os.makedirs(dst, exist_ok=True)
filepath = os.path.join(dst, file_name)
if not os.path.isfile(filepath):
command = f"curl {url} -o {file_name}.tmp --fail"
print(f"Downloading file {filepath}")
try:
subprocess.check_call(command, shell=True, cwd=dst)
except Exception as error:
error_msg = f"Error downloading {url}, error: {error}"
print(error_msg)
raise RuntimeError(error_msg) from error
os.rename(filepath + ".tmp", filepath)
else:
pass # skipping download of existing file
return True
def unzip_file(file_name: str, dst: Path, keep_zip: bool = True) -> bool:
filepath = os.path.join(dst, file_name)
print(f"Unzipping zip file {filepath}")
try:
with zipfile.ZipFile(filepath, "r") as zip_ref:
zip_ref.extractall(dst)
except Exception as error:
error_msg = f"Error unzipping {filepath}, error: {error}"
print(error_msg)
raise RuntimeError(error_msg) from error
if not keep_zip:
os.remove(filepath)
return True
def download_laser_scanner_point_clouds_for_video(video_id: str, metadata: pd.DataFrame, download_dir: Path) -> None:
video_metadata = metadata.loc[metadata["video_id"] == float(video_id)]
visit_id: float = video_metadata["visit_id"].iat[0] # type: ignore[assignment]
has_laser_scanner_point_clouds = video_metadata["has_laser_scanner_point_clouds"].iat[0]
if not has_laser_scanner_point_clouds:
print(f"Warning: Laser scanner point clouds for video {video_id} are not available")
return
if math.isnan(visit_id) or not visit_id.is_integer():
print(f"Warning: Downloading laser scanner point clouds for video {video_id} failed - Bad visit id {visit_id}")
return
visit_id = int(visit_id)
laser_scanner_point_clouds_ids = laser_scanner_point_clouds_for_visit_id(visit_id, download_dir)
for point_cloud_id in laser_scanner_point_clouds_ids:
download_laser_scanner_point_clouds(point_cloud_id, visit_id, download_dir)
def laser_scanner_point_clouds_for_visit_id(visit_id: int, download_dir: Path) -> list[str]:
point_cloud_to_visit_id_mapping_filename = "laser_scanner_point_clouds_mapping.csv"
if not os.path.exists(point_cloud_to_visit_id_mapping_filename):
point_cloud_to_visit_id_mapping_url = (
f"{ARkitscense_url}/raw/laser_scanner_point_clouds/{point_cloud_to_visit_id_mapping_filename}"
)
download_file(
point_cloud_to_visit_id_mapping_url,
point_cloud_to_visit_id_mapping_filename,
download_dir,
)
point_cloud_to_visit_id_mapping_filepath = os.path.join(download_dir, point_cloud_to_visit_id_mapping_filename)
point_cloud_to_visit_id_mapping = pd.read_csv(point_cloud_to_visit_id_mapping_filepath)
point_cloud_ids = point_cloud_to_visit_id_mapping.loc[
point_cloud_to_visit_id_mapping["visit_id"] == visit_id,
["laser_scanner_point_clouds_id"],
]
point_cloud_ids_list = [scan_id[0] for scan_id in point_cloud_ids.values]
return point_cloud_ids_list
def download_laser_scanner_point_clouds(laser_scanner_point_cloud_id: str, visit_id: int, download_dir: Path) -> None:
laser_scanner_point_clouds_folder_path = download_dir / POINT_CLOUDS_FOLDER / str(visit_id)
os.makedirs(laser_scanner_point_clouds_folder_path, exist_ok=True)
for extension in [".ply", "_pose.txt"]:
filename = f"{laser_scanner_point_cloud_id}{extension}"
filepath = os.path.join(laser_scanner_point_clouds_folder_path, filename)
if os.path.exists(filepath):
return
file_url = f"{ARkitscense_url}/raw/laser_scanner_point_clouds/{visit_id}/{filename}"
download_file(file_url, filename, laser_scanner_point_clouds_folder_path)
def get_metadata(dataset: str, download_dir: Path) -> pd.DataFrame:
filename = "metadata.csv"
url = f"{ARkitscense_url}/threedod/{filename}" if "3dod" == dataset else f"{ARkitscense_url}/{dataset}/{filename}"
dst_folder = download_dir / dataset
dst_file = dst_folder / filename
download_file(url, filename, dst_folder)
metadata = pd.read_csv(dst_file)
return metadata
def download_data(
dataset: str,
video_ids: list[str],
dataset_splits: list[str],
download_dir: Path,
keep_zip: bool,
raw_dataset_assets: list[str] | None = None,
should_download_laser_scanner_point_cloud: bool = False,
) -> None:
"""
Downloads data from the specified dataset and video IDs to the given download directory.
Args:
----
dataset: the name of the dataset to download from (raw, 3dod, or upsampling)
video_ids: the list of video IDs to download data for
dataset_splits: the list of splits for each video ID (train, validation, or test)
download_dir: the directory to download data to
keep_zip: whether to keep the downloaded zip files after extracting them
raw_dataset_assets: a list of asset types to download from the raw dataset, if dataset is "raw"
should_download_laser_scanner_point_cloud: whether to download the laser scanner point cloud data, if available
Returns: None
"""
metadata = get_metadata(dataset, download_dir)
for video_id in sorted(set(video_ids)):
split = dataset_splits[video_ids.index(video_id)]
dst_dir = download_dir / dataset / split
if dataset == "raw":
url_prefix = ""
file_names = []
if not raw_dataset_assets:
print(f"Warning: No raw assets given for video id {video_id}")
else:
dst_dir = dst_dir / str(video_id)
url_prefix = f"{ARkitscense_url}/raw/{split}/{video_id}" + "/{}"
file_names = raw_files(video_id, raw_dataset_assets, metadata)
elif dataset == "3dod":
url_prefix = f"{ARkitscense_url}/threedod/{split}" + "/{}"
file_names = [
f"{video_id}.zip",
]
elif dataset == "upsampling":
url_prefix = f"{ARkitscense_url}/upsampling/{split}" + "/{}"
file_names = [
f"{video_id}.zip",
]
else:
raise Exception(f"No such dataset = {dataset}")
if should_download_laser_scanner_point_cloud and dataset == "raw":
# Point clouds only available for the raw dataset
download_laser_scanner_point_clouds_for_video(video_id, metadata, download_dir)
for file_name in file_names:
dst_path = os.path.join(dst_dir, file_name)
url = url_prefix.format(file_name)
if not file_name.endswith(".zip") or not os.path.isdir(dst_path[: -len(".zip")]):
download_file(url, dst_path, dst_dir)
else:
pass # skipping download of existing zip file
if file_name.endswith(".zip") and os.path.isfile(dst_path):
unzip_file(file_name, dst_dir, keep_zip)
def ensure_recording_downloaded(video_id: str, include_highres: bool) -> Path:
"""Only downloads from validation set."""
data_path = DATASET_DIR / "raw" / "Validation" / video_id
assets_to_download = [
"lowres_wide",
"lowres_depth",
"lowres_wide_intrinsics",
"lowres_wide.traj",
"annotation",
"mesh",
]
if include_highres:
assets_to_download.extend(["highres_depth", "wide", "wide_intrinsics"])
download_data(
dataset="raw",
video_ids=[video_id],
dataset_splits=[VALIDATION],
download_dir=DATASET_DIR,
keep_zip=False,
raw_dataset_assets=assets_to_download,
should_download_laser_scanner_point_cloud=False,
)
return data_path
def ensure_recording_available(video_id: str, include_highres: bool) -> Path:
"""
Returns the path to the recording for a given video_id.
Args:
----
video_id (str):
Identifier for the recording.
include_highres (bool):
Whether to include the high resolution recording.
Returns
-------
Path: Path object representing the path to the recording.
Raises
------
AssertionError:
If the recording path does not exist.
"""
recording_path = ensure_recording_downloaded(video_id, include_highres)
assert recording_path.exists(), f"Recording path {recording_path} does not exist."
return recording_path # Return the path to the recording
examples/python/arkit_scenes/pyproject.toml
+22
View File
@@ -0,0 +1,22 @@
[project]
name = "arkit_scenes"
version = "0.1.0"
# requires-python = "<3.12"
readme = "README.md"
dependencies = [
"matplotlib",
"numpy",
"opencv-python",
"pandas",
"rerun-sdk",
"scipy",
"tqdm",
"trimesh",
]
[project.scripts]
arkit_scenes = "arkit_scenes.__main__:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/blueprint/blueprint.py
Executable
+76
View File
@@ -0,0 +1,76 @@
#!/usr/bin/env python3
"""Example of using the blueprint APIs to configure Rerun."""
from __future__ import annotations
import argparse
import numpy as np
import rerun as rr # pip install rerun-sdk
import rerun.blueprint as rrb
def main() -> None:
parser = argparse.ArgumentParser(description="Different options for how we might use blueprint")
parser.add_argument("--skip-blueprint", action="store_true", help="Don't send the blueprint")
parser.add_argument("--auto-views", action="store_true", help="Automatically add views")
args = parser.parse_args()
if args.skip_blueprint:
blueprint = None
else:
# Create a blueprint which includes 2 additional views each only showing 1 of the two
# rectangles.
#
# If auto_views is True, the blueprint will automatically add one of the heuristic
# views, which will include the image and both rectangles.
blueprint = rrb.Blueprint(
rrb.Grid(
rrb.Spatial2DView(name="Rect 0", origin="/", contents=["image", "rect/0"]),
rrb.Spatial2DView(
name="Rect 1",
origin="/",
contents=["/**"],
defaults=[rr.Boxes2D.from_fields(radii=2)], # Default all rectangles to have a radius of 2
overrides={"rect/0": rr.Boxes2D.from_fields(radii=1)}, # Override the radius of rect/0 to be 1
),
),
rrb.BlueprintPanel(state="collapsed"),
rrb.SelectionPanel(state="collapsed"),
rrb.TimePanel(
state="collapsed",
timeline="custom",
time_selection=rrb.components.AbsoluteTimeRange(10, 25),
loop_mode="selection",
play_state="playing",
),
auto_views=args.auto_views,
)
rr.init("rerun_example_blueprint", spawn=True, default_blueprint=blueprint)
rr.set_time("custom", sequence=0)
img = np.zeros([128, 128, 3], dtype="uint8")
for i in range(8):
img[(i * 16) + 4 : (i * 16) + 12, :] = (0, 0, 200)
rr.log("image", rr.Image(img))
rr.set_time("custom", sequence=10)
rr.log(
"rect/0",
rr.Boxes2D(mins=[16, 16], sizes=[64, 64], labels="Rect0", colors=(255, 0, 0)),
)
rr.set_time("custom", sequence=20)
rr.log(
"rect/1",
rr.Boxes2D(mins=[48, 48], sizes=[64, 64], labels="Rect1", colors=(0, 255, 0)),
)
if __name__ == "__main__":
main()
examples/python/blueprint/pyproject.toml
+12
View File
@@ -0,0 +1,12 @@
[project]
name = "blueprint"
version = "0.1.0"
# requires-python = "<3.12"
dependencies = ["numpy", "rerun-sdk"]
[project.scripts]
blueprint = "blueprint:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/blueprint_stocks/README.md
+38
View File
@@ -0,0 +1,38 @@
<!--[metadata]
title = "Stock charts"
tags = ["Time series", "Blueprint"]
thumbnail = "https://static.rerun.io/blueprint_stocks/8bfe6f16963acdceb2debb9de9a206dc2eb9b280/480w.png"
thumbnail_dimensions = [480, 270]
-->
This example fetches the last 5 days of stock data for a few different stocks.
We show how Rerun blueprints can then be used to present many different views of the same data.
<picture>
<img src="https://static.rerun.io/blueprint_stocks/8bfe6f16963acdceb2debb9de9a206dc2eb9b280/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/blueprint_stocks/8bfe6f16963acdceb2debb9de9a206dc2eb9b280/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/blueprint_stocks/8bfe6f16963acdceb2debb9de9a206dc2eb9b280/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/blueprint_stocks/8bfe6f16963acdceb2debb9de9a206dc2eb9b280/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/blueprint_stocks/8bfe6f16963acdceb2debb9de9a206dc2eb9b280/1200w.png">
</picture>
```bash
pip install -e examples/python/blueprint_stocks
python -m blueprint_stocks
```
The different blueprints can be explored using the `--blueprint` flag. For example:
```
python -m blueprint_stocks --blueprint=one-stock
```
Available choices are:
- `auto`: Reset the blueprint to the auto layout used by the viewer.
- `one-stock`: Uses a filter to show only a single chart.
- `one-stock-with-info`: Uses a container to layout a chart and its info document
- `one-stock-no-peaks`: Uses a filter to additionally remove some of the data from the chart.
- `compare-two`: Adds data from multiple sources to a single chart.
- `grid`: Shows all the charts in a grid layout.
examples/python/blueprint_stocks/blueprint_stocks.py
Executable
+222
View File
@@ -0,0 +1,222 @@
#!/usr/bin/env python3
"""
A simple application that fetches stock data from Yahoo Finance and visualizes it using the Rerun SDK.
The main focus of this example is using blueprints to control how the data is displayed in the viewer.
"""
from __future__ import annotations
import argparse
import datetime as dt
from typing import Any
import humanize
import pytz
import yfinance as yf
import rerun as rr
import rerun.blueprint as rrb
################################################################################
# Helper functions to create blueprints
################################################################################
def auto_blueprint() -> rrb.BlueprintLike:
"""A blueprint enabling auto views, which matches the application default."""
return rrb.Blueprint(auto_views=True, auto_layout=True)
def one_stock(symbol: str) -> rrb.ContainerLike:
"""Create a blueprint showing a single stock."""
return rrb.TimeSeriesView(name=f"{symbol}", origin=f"/stocks/{symbol}")
def one_stock_with_info(symbol: str) -> rrb.ContainerLike:
"""Create a blueprint showing a single stock with its info arranged vertically."""
return rrb.Vertical(
rrb.TextDocumentView(name=f"{symbol}", origin=f"/stocks/{symbol}/info"),
rrb.TimeSeriesView(name=f"{symbol}", origin=f"/stocks/{symbol}"),
row_shares=[1, 4],
)
def compare_two(symbol1: str, symbol2: str, day: Any) -> rrb.ContainerLike:
"""Create a blueprint comparing 2 stocks for a single day."""
return rrb.TimeSeriesView(
name=f"{symbol1} vs {symbol2} ({day})",
contents=[
f"+ /stocks/{symbol1}/{day}",
f"+ /stocks/{symbol2}/{day}",
],
)
def one_stock_no_peaks(symbol: str) -> rrb.ContainerLike:
"""
Create a blueprint showing a single stock without annotated peaks.
This uses an exclusion pattern to hide the peaks.
"""
return rrb.TimeSeriesView(
name=f"{symbol}",
origin=f"/stocks/{symbol}",
contents=[
"+ $origin/**",
"- $origin/peaks/**",
],
)
def stock_grid(symbols: list[str], dates: list[Any]) -> rrb.ContainerLike:
"""Create a grid of stocks and their time series over all days."""
return rrb.Vertical(
contents=[
rrb.Horizontal(
contents=[rrb.TextDocumentView(name=f"{symbol}", origin=f"/stocks/{symbol}/info")]
+ [rrb.TimeSeriesView(name=f"{day}", origin=f"/stocks/{symbol}/{day}") for day in dates],
name=symbol,
)
for symbol in symbols
],
)
def hide_panels(viewport: rrb.ContainerLike) -> rrb.BlueprintLike:
"""Wrap a viewport in a blueprint that hides the time and selection panels."""
return rrb.Blueprint(
viewport,
rrb.TimePanel(state="collapsed"),
rrb.SelectionPanel(state="collapsed"),
)
################################################################################
# Helper functions for styling
################################################################################
brand_colors = {
"AAPL": 0xA2AAADFF,
"AMZN": 0xFF9900FF,
"GOOGL": 0x34A853FF,
"META": 0x0081FBFF,
"MSFT": 0xF14F21FF,
}
def style_plot(symbol: str) -> rr.SeriesLines:
return rr.SeriesLines(
colors=brand_colors[symbol],
names=symbol,
)
def style_peak(symbol: str) -> rr.SeriesPoints:
return rr.SeriesPoints(
colors=0xFF0000FF,
names=f"{symbol} (peak)",
markers="up",
)
################################################################################
# Main script
################################################################################
def main() -> None:
parser = argparse.ArgumentParser(description="Visualize stock data using the Rerun SDK")
parser.add_argument(
"--blueprint",
choices=["auto", "one-stock", "one-stock-with-info", "compare-two", "one-stock-no-peaks", "grid"],
default="grid",
help="Select the blueprint to use",
)
parser.add_argument(
"--show_panels",
action="store_true",
help="Show the time and selection panels",
)
rr.script_add_args(parser)
args = parser.parse_args()
et_timezone = pytz.timezone("America/New_York")
current_date = dt.datetime.now(et_timezone).date()
symbols = ["AAPL", "AMZN", "GOOGL", "META", "MSFT"]
dates = list(filter(lambda x: x.weekday() < 5, [current_date - dt.timedelta(days=i) for i in range(7, 0, -1)]))
if args.blueprint == "auto":
blueprint = auto_blueprint()
else:
if args.blueprint == "one-stock":
viewport = one_stock("AAPL")
elif args.blueprint == "one-stock-with-info":
viewport = one_stock_with_info("AMZN")
elif args.blueprint == "one-stock-no-peaks":
viewport = one_stock_no_peaks("GOOGL")
elif args.blueprint == "compare-two":
viewport = compare_two("META", "MSFT", dates[-1])
elif args.blueprint == "grid":
viewport = stock_grid(symbols, dates)
else:
raise ValueError(f"Unknown blueprint: {args.blueprint}")
if not args.show_panels:
blueprint = hide_panels(viewport)
else:
blueprint = viewport
rr.script_setup(args, "rerun_example_blueprint_stocks")
rr.send_blueprint(blueprint)
# In a future blueprint release, this can move into the blueprint as well
for symbol in symbols:
for day in dates:
rr.log(f"stocks/{symbol}/{day}", style_plot(symbol), static=True)
rr.log(f"stocks/{symbol}/peaks/{day}", style_peak(symbol), static=True)
for symbol in symbols:
stock = yf.Ticker(symbol)
name = stock.info["shortName"]
industry = stock.info["industry"]
marketCap = humanize.intword(stock.info["marketCap"])
revenue = humanize.intword(stock.info["totalRevenue"])
info_md = (
f"- **Name**: {name}\n"
f"- **Industry**: {industry}\n"
f"- **Market cap**: ${marketCap}\n"
f"- **Total Revenue**: ${revenue}\n"
)
rr.log(
f"stocks/{symbol}/info",
rr.TextDocument(info_md, media_type=rr.MediaType.MARKDOWN),
static=True,
)
for day in dates:
open_time = dt.datetime.combine(day, dt.time(9, 30))
close_time = dt.datetime.combine(day, dt.time(16, 00))
hist = stock.history(start=open_time, end=close_time, interval="5m")
if len(hist.index) == 0:
continue
hist.index = hist.index - et_timezone.localize(open_time)
peak = hist.High.idxmax()
for row in hist.itertuples():
rr.set_time("time", duration=row.Index)
rr.log(f"stocks/{symbol}/{day}", rr.Scalars(row.High))
if row.Index == peak:
rr.log(f"stocks/{symbol}/peaks/{day}", rr.Scalars(row.High))
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/blueprint_stocks/pyproject.toml
+16
View File
@@ -0,0 +1,16 @@
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "blueprint_stocks"
version = "0.1.0"
requires-python = ">=3.10"
readme = "README.md"
dependencies = ["humanize", "rerun-sdk", "yfinance"]
[project.scripts]
blueprint_stocks = "blueprint_stocks:main"
[tool.rerun-example]
#skip = true
examples/python/camera_video_stream/README.md
+58
View File
@@ -0,0 +1,58 @@
<!--[metadata]
title = "Compressed camera video stream"
tags = ["2D", "Image encoding", "Video", "Streaming"]
thumbnail = "https://static.rerun.io/camera_video_stream/b2f8f61eb62424aa942bdb5183e49246cf417e60/480w.png"
thumbnail_dimensions = [480, 300]
-->
This example uses [pyAV](https://pypi.org/project/av/) to fetch and encode a video stream to H.264 video and streams it live
to the Viewer using the [`VideoStream`](https://www.rerun.io/docs/reference/types/archetypes/video_stream) archetype
<img src="https://static.rerun.io/camera_video_stream/b2f8f61eb62424aa942bdb5183e49246cf417e60/480w.png">
## Run the code
```bash
pip install -e examples/python/camera_video_stream
python -m camera_video_stream
```
## Details
To learn more about video support in general check [the docs page on the topic](https://rerun.io/docs/concepts/logging-and-ingestion/video).
The example first sets up a video stream on an entity called `"video_stream"` with the H264 codec.
```py
rr.log("video_stream", rr.VideoStream(codec=rr.VideoCodec.H264), static=True)
```
H264 is a very well established codec and well suited for streaming.
Note that in `setup_output_stream` we specicifially configure `pyAV` to use low latency encoding.
⚠️ Latency is expected to be well below a second, but as of writing Rerun is not well
optimized for low latency video transmission as used for teleoperations.
If you need lower latency you should consider tweaking the encoding settings further and
configure Rerun's [Micro Batching](https://rerun.io/docs/reference/sdk/micro-batching).
For each frame we have to set a new timestamp.
```py
rr.set_time("time", duration=float(packet.pts * packet.time_base))
```
The time set here is the [_presentation timestamp_](https://en.wikipedia.org/wiki/Presentation_timestamp) (PTS) of the frame.
Note that unlike with unlike with [`VideoAsset`](https://www.rerun.io/docs/reference/types/archetypes/asset_video),
there's no need to log [`VideoFrameReference`](https://www.rerun.io/docs/reference/types/archetypes/video_frame_reference),
to map the video's PTS to the Rerun timeline, since the time at which video samples
are logged directly represents the PTS.
TODO(#10090): In the presence of H.264/H.265 b-frames, separate _decode timestamps_ (DTS) are needed. This is not yet supported.
The frame data, known as a frame-`sample` since this may contain data relevant for an arbitrary number of frames in the future,
is then logged with:
```py
rr.log("video_stream", rr.VideoStream.from_fields(sample=bytes(packet)))
```
It's best practice to log the codec only once as it never changes, but there is little harm in
doing so for every frame:
```py
rr.log("video_stream", rr.VideoStream(codec=rr.VideoCodec.H264, sample=bytes(packet)))
```
examples/python/camera_video_stream/camera_video_stream.py
Executable
+154
View File
@@ -0,0 +1,154 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import platform
import subprocess
import av
import rerun as rr
def setup_camera_input(video_device: str | None = None) -> av.container.InputContainer:
"""
Setup the camera input container.
Uses first available video device. Exact behavior is platform dependent.
"""
if platform.system() == "Darwin":
if video_device is None:
video_device = "0"
return av.open(
video_device,
format="avfoundation",
container_options={"framerate": "30"}, # `avfoundation` fails if the frame rate is not set.
)
elif platform.system() == "Windows":
if video_device is None:
# On windows we *have* to know the device name. Query ffmpeg for it.
list_devices_cmd = ["ffmpeg", "-f", "dshow", "-list_devices", "true", "-i", "dummy", "-hide_banner"]
devices_output = subprocess.check_output(list_devices_cmd, stderr=subprocess.STDOUT).decode()
# Extract video device names
video_devices = []
for line in devices_output.split("\n"):
# FFmpeg will return both `(audio)` and `(video)` devices
if "(video)" in line and "Alternative name" not in line:
device_name = line.split('"')[1]
if device_name not in video_devices:
video_devices.append(device_name)
if not video_devices:
raise RuntimeError("No video devices found")
# Use first available video device
video_device = video_devices[0]
print(f"Using video device: {video_device}")
return av.open(f"video={video_device}", format="dshow")
else:
if video_device is None:
video_device = "/dev/video0"
return av.open(video_device, format="v4l2")
def setup_output_stream(width: int, height: int, codec: str = "h264") -> av.video.VideoStream:
"""Setup the output stream which encodes the video stream to the specified codec."""
if codec == "h264":
output_container = av.open("/dev/null", "w", format="h264") # Use AnnexB H.264 stream.
output_stream = output_container.add_stream("libx264") # type: ignore[assignment]
elif codec == "av1":
output_container = av.open("/dev/null", "w", format="ivf") # Use IVF container for AV1 stream.
output_stream = output_container.add_stream("libaom-av1") # type: ignore[assignment]
else:
raise ValueError(f"Unsupported codec: {codec}")
# Type narrowing
assert isinstance(output_stream, av.video.stream.VideoStream)
output_stream.width = width
output_stream.height = height
# Configure for low latency.
if codec == "h264":
output_stream.codec_context.options = {
"tune": "zerolatency",
"preset": "veryfast",
}
elif codec == "av1":
output_stream.codec_context.options = {
"cpu-used": "8",
"usage": "realtime", # Optimize for realtime encoding
}
output_stream.max_b_frames = 0 # Avoid b-frames for lower latency.
return output_stream
def stream_video_to_rerun(
input: av.container.InputContainer, output: av.video.VideoStream, codec: str = "h264"
) -> None:
"""Streams the video continuously to Rerun."""
# Log codec only once as static data (it naturally never changes). This isn't strictly necessary, but good practice.
video_codec = rr.VideoCodec.H264 if codec == "h264" else rr.VideoCodec.AV1
rr.log("video_stream", rr.VideoStream(codec=video_codec), static=True)
while True:
try:
for frame in input.decode(video=0):
# By default all the frames that come from the camera are marked as I-frames.
# If we pass this just on as-is, then we get an encoded video stream that
# just consists entirely of I-frames, thus having very poor compression!
# Instead, we want the encoder to decide when to use P & I frames.
frame.pict_type = av.video.frame.PictureType.NONE
for packet in output.encode(frame):
if packet.pts is None:
continue
rr.set_time("time", duration=float(packet.pts * packet.time_base))
rr.log("video_stream", rr.VideoStream.from_fields(sample=bytes(packet)))
except av.BlockingIOError:
pass
def main() -> None:
parser = argparse.ArgumentParser(description="Streams compressed video from camera to Rerun.")
parser.add_argument(
"--video-device",
type=str,
help="Video device to use. If not provided, the first available video device will be used.",
)
parser.add_argument(
"--codec",
type=str,
choices=["h264", "av1"],
default="h264",
help="Video codec to use for encoding (default: h264).",
)
rr.script_add_args(parser)
args = parser.parse_args()
rr.script_setup(args, "rerun_example_video_stream_camera")
av.logging.set_level(av.logging.VERBOSE)
input_container = setup_camera_input(args.video_device)
output_stream = setup_output_stream(
input_container.streams.video[0].width, input_container.streams.video[0].height, args.codec
)
try:
stream_video_to_rerun(input_container, output_stream, args.codec)
except KeyboardInterrupt:
print("Recording stopped by user")
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/camera_video_stream/pyproject.toml
+12
View File
@@ -0,0 +1,12 @@
[project]
name = "camera_video_stream"
version = "0.1.0"
readme = "README.md"
dependencies = ["rerun-sdk", "av>=14.2.0"]
[project.scripts]
camera_video_stream = "camera_video_stream:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/clock/README.md
+76
View File
@@ -0,0 +1,76 @@
<!--[metadata]
title = "Clock"
tags = ["3D", "API example"]
thumbnail = "https://static.rerun.io/clock/8c49e25f5cac4d6a1d7d0490b14cf6881bdb707b/480w.png"
thumbnail_dimensions = [480, 480]
-->
<picture>
<source media="(max-width: 480px)" srcset="https://static.rerun.io/clock/05e69dc20c9a28005f1ffe7f0f2ac9eeaa95ba3b/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/clock/05e69dc20c9a28005f1ffe7f0f2ac9eeaa95ba3b/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/clock/05e69dc20c9a28005f1ffe7f0f2ac9eeaa95ba3b/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/clock/05e69dc20c9a28005f1ffe7f0f2ac9eeaa95ba3b/1200w.png">
<img src="https://static.rerun.io/clock/05e69dc20c9a28005f1ffe7f0f2ac9eeaa95ba3b/full.png" alt="Clock example screenshot">
</picture>
An example visualizing an analog clock with hour, minute and seconds hands using Rerun Arrow3D primitives.
## Used Rerun types
[`Boxes3D`](https://www.rerun.io/docs/reference/types/archetypes/boxes3d), [`Points3D`](https://www.rerun.io/docs/reference/types/archetypes/points3d), [`Arrows3D`](https://www.rerun.io/docs/reference/types/archetypes/arrows3d)
## Logging and visualizing with Rerun
The visualizations in this example were created with the following Rerun code:
The clock's frame is logged as a 3D box using [`Boxes3D`](https://www.rerun.io/docs/reference/types/archetypes/boxes3d) archetype.
```python
rr.log(
"world/frame",
rr.Boxes3D(half_sizes=[LENGTH_S, LENGTH_S, 1.0], centers=[0.0, 0.0, 0.0]),
static=True,
)
```
Then, the positions and colors of points and arrows representing the hands of a clock for seconds, minutes, and hours are logged in each simulation time.
It first sets the simulation time using [`timelines`](https://www.rerun.io/docs/concepts/logging-and-ingestion/timelines), calculates the data for each hand, and logs it using [`Points3D`](https://www.rerun.io/docs/reference/types/archetypes/points3d) and [`Arrows3D`](https://www.rerun.io/docs/reference/types/archetypes/arrows3d) archetypes.
This enables the visualization of the clock's movement over time.
```python
for step in range(steps):
rr.set_time("sim_time", duration=t_secs)
# … calculating seconds …
rr.log("world/seconds_pt", rr.Points3D(positions=point_s, colors=color_s))
rr.log("world/seconds_hand", rr.Arrows3D(vectors=point_s, colors=color_s, radii=WIDTH_S))
# … calculating minutes …
rr.log("world/minutes_pt", rr.Points3D(positions=point_m, colors=color_m))
rr.log("world/minutes_hand", rr.Arrows3D(vectors=point_m, colors=color_m, radii=WIDTH_M))
# … calculating hours …
rr.log("world/hours_pt", rr.Points3D(positions=point_h, colors=color_h))
rr.log("world/hours_hand", rr.Arrows3D(vectors=point_h, colors=color_h, radii=WIDTH_H))
```
## Run the code
To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/clock
```
To experiment with the provided example, simply execute the main Python script:
```bash
python -m clock # run the example
```
If you wish to customize it, explore additional features, or save it use the CLI with the `--help` option for guidance:
```bash
python -m clock --help
```
examples/python/clock/clock.py
Executable
+81
View File
@@ -0,0 +1,81 @@
#!/usr/bin/env python3
"""
An example showing usage of `log_arrow`.
An analog clock is built with Rerun Arrow3D primitives.
"""
from __future__ import annotations
import argparse
import math
from typing import Final
import numpy as np
import rerun as rr # pip install rerun-sdk
LENGTH_S: Final = 20.0
LENGTH_M: Final = 10.0
LENGTH_H: Final = 4.0
WIDTH_S: Final = 0.125
WIDTH_M: Final = 0.2
WIDTH_H: Final = 0.3
def log_clock(steps: int) -> None:
def rotate(angle: float, len: float) -> tuple[float, float, float]:
return (
len * math.sin(angle),
len * math.cos(angle),
0.0,
)
rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, static=True)
rr.log(
"world/frame",
rr.Boxes3D(half_sizes=[LENGTH_S, LENGTH_S, 1.0], centers=[0.0, 0.0, 0.0]),
static=True,
)
for step in range(steps):
t_secs = step
rr.set_time("sim_time", duration=t_secs)
scaled_s = (t_secs % 60) / 60.0
point_s = np.array(rotate(math.tau * scaled_s, LENGTH_S))
color_s = (int(255 - (scaled_s * 255)), int(scaled_s * 255), 0, 128)
rr.log("world/seconds_pt", rr.Points3D(positions=point_s, colors=color_s))
rr.log("world/seconds_hand", rr.Arrows3D(vectors=point_s, colors=color_s, radii=WIDTH_S))
scaled_m = (t_secs % 3600) / 3600.0
point_m = np.array(rotate(math.tau * scaled_m, LENGTH_M))
color_m = (int(255 - (scaled_m * 255)), int(scaled_m * 255), 128, 128)
rr.log("world/minutes_pt", rr.Points3D(positions=point_m, colors=color_m))
rr.log("world/minutes_hand", rr.Arrows3D(vectors=point_m, colors=color_m, radii=WIDTH_M))
scaled_h = (t_secs % 43200) / 43200.0
point_h = np.array(rotate(math.tau * scaled_h, LENGTH_H))
color_h = (int(255 - (scaled_h * 255)), int(scaled_h * 255), 255, 255)
rr.log("world/hours_pt", rr.Points3D(positions=point_h, colors=color_h))
rr.log("world/hours_hand", rr.Arrows3D(vectors=point_h, colors=color_h, radii=WIDTH_H))
def main() -> None:
parser = argparse.ArgumentParser(
description="An example visualizing an analog clock is built with Rerun Arrow3D primitives.",
)
parser.add_argument("--steps", type=int, default=10_000, help="The number of time steps to log")
rr.script_add_args(parser)
args = parser.parse_args()
rr.script_setup(args, "rerun_example_clock")
log_clock(args.steps)
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/clock/pyproject.toml
+16
View File
@@ -0,0 +1,16 @@
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "clock"
version = "0.1.0"
readme = "README.md"
dependencies = ["numpy", "rerun-sdk"]
[project.scripts]
clock = "clock:main"
[tool.rerun-example]
#skip = false
extra-args = ["--steps=200"]
examples/python/controlnet/README.md
+92
View File
@@ -0,0 +1,92 @@
<!--[metadata]
title = "ControlNet"
tags = ["ControlNet", "Canny", "Hugging Face", "Stable diffusion", "Tensor", "Text"]
thumbnail = "https://static.rerun.io/controlnet/2e984b27dd8120fb89d4e805df9da506ea6d9138/480w.png"
thumbnail_dimensions = [480, 480]
-->
Use [Hugging Face's ControlNet](https://huggingface.co/docs/diffusers/using-diffusers/controlnet#controlnet) to generate an image from text, conditioned on detected edges from another image.
<picture>
<source media="(max-width: 480px)" srcset="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/1200w.png">
<img src="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/full.png" alt="">
</picture>
## Used Rerun types
[`Image`](https://www.rerun.io/docs/reference/types/archetypes/image), [`Tensor`](https://www.rerun.io/docs/reference/types/archetypes/tensor), [`TextDocument`](https://www.rerun.io/docs/reference/types/archetypes/text_document)
## Background
[Hugging Face's ControlNet](https://huggingface.co/docs/diffusers/using-diffusers/controlnet#controlnet) allows to condition Stable Diffusion on various modalities. In this example we condition on edges detected by the Canny edge detector to keep our shape intact while generating an image with Stable Diffusion. This generates a new image based on the text prompt, but that still retains the structure of the control image. We visualize the whole generation process with Rerun.
https://vimeo.com/870289439?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=1440:1080
## Logging and visualizing with Rerun
The visualizations in this example were created with the following Rerun code.
### Images
```python
rr.log("input/raw", rr.Image(image), static=True)
rr.log("input/canny", rr.Image(canny_image), static=True)
```
The input image and control canny_image are marked as static and logged in rerun.
Static entities belong to all timelines (existing ones, and ones not yet created) and are shown leftmost in the time panel in the viewer. This is useful for entities that aren't part of normal data capture, but set the scene for how they are shown.
This designation ensures their constant availability across all timelines in Rerun, aiding in consistent comparison and documentation.
### Prompts
```python
rr.log("positive_prompt", rr.TextDocument(prompt), static=True)
rr.log("negative_prompt", rr.TextDocument(negative_prompt), static=True)
```
The positive and negative prompt used for generation is logged to Rerun.
### Custom diffusion step callback
We use a custom callback function for ControlNet that logs the output and the latent values at each timestep, which makes it possible for us to view all timesteps of the generation in Rerun.
```python
def controlnet_callback(
iteration: int, timestep: float, latents: torch.Tensor, pipeline: StableDiffusionXLControlNetPipeline
) -> None:
rr.set_time("iteration", sequence=iteration)
rr.set_time("timestep", duration=timestep)
rr.log("output", rr.Image(image))
rr.log("latent", rr.Tensor(latents.squeeze(), dim_names=["channel", "height", "width"]))
```
### Output image
```python
rr.log("output", rr.Image(images))
```
Finally we log the output image generated by ControlNet.
## Run the code
To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/controlnet
```
To experiment with the provided example, simply execute the main Python script:
```bash
python -m controlnet
```
You can specify your own image and prompts using
```bash
python -m controlnet [--img-path IMG_PATH] [--prompt PROMPT] [--negative-prompt NEGATIVE_PROMPT]
```
This example requires a machine with CUDA backend available for pytorch (GPU) to work.
examples/python/controlnet/controlnet.py
Executable
+158
View File
@@ -0,0 +1,158 @@
#!/usr/bin/env python3
"""
Example running ControlNet conditioned on Canny edges.
Based on <https://huggingface.co/docs/diffusers/using-diffusers/controlnet>.
"""
from __future__ import annotations
import argparse
import os
from typing import Any
import cv2
import numpy as np
import PIL.Image
import requests
import torch
from diffusers import (
AutoencoderKL,
ControlNetModel,
StableDiffusionXLControlNetPipeline,
)
import rerun as rr
import rerun.blueprint as rrb
RERUN_LOGO_URL = "https://storage.googleapis.com/rerun-example-datasets/controlnet/rerun-icon-1000.png"
def controlnet_callback(
pipe: StableDiffusionXLControlNetPipeline,
step_index: int,
timestep: float,
callback_kwargs: dict[str, Any],
) -> dict[str, Any]:
rr.set_time("iteration", sequence=step_index)
rr.set_time("timestep", duration=timestep)
latents = callback_kwargs["latents"]
image = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0] # type: ignore[attr-defined]
image = pipe.image_processor.postprocess(image, output_type="np").squeeze() # type: ignore[attr-defined]
rr.log("output", rr.Image(image))
rr.log("latent", rr.Tensor(latents.squeeze(), dim_names=["channel", "height", "width"]))
return callback_kwargs
def run_canny_controlnet(image_path: str, prompt: str, negative_prompt: str) -> None:
if not torch.cuda.is_available():
print("This example requires a torch with CUDA, but no CUDA device found. Aborting.")
return
if image_path.startswith(("http://", "https://")):
pil_image = PIL.Image.open(requests.get(image_path, stream=True).content)
elif os.path.isfile(image_path):
pil_image = PIL.Image.open(image_path)
else:
raise ValueError(f"Invalid image_path: {image_path}")
image = np.array(pil_image)
if image.shape[2] == 4: # RGBA image
rgb_image = image[..., :3] # RGBA to RGB
rgb_image[image[..., 3] < 200] = 0.0 # reduces artifacts for transparent parts
else:
rgb_image = image
low_threshold = 100.0
high_threshold = 200.0
canny_data = cv2.Canny(rgb_image, low_threshold, high_threshold)
canny_data = canny_data[:, :, None]
# cv2.dilate(kjgk
canny_data = np.concatenate([canny_data, canny_data, canny_data], axis=2)
canny_image = PIL.Image.fromarray(canny_data)
rr.log("input/raw", rr.Image(image), static=True)
rr.log("input/canny", rr.Image(canny_image), static=True)
controlnet = ControlNetModel.from_pretrained(
"diffusers/controlnet-canny-sdxl-1.0",
torch_dtype=torch.float16,
use_safetensors=True,
)
vae = AutoencoderKL.from_pretrained(
"madebyollin/sdxl-vae-fp16-fix",
torch_dtype=torch.float16,
use_safetensors=True,
)
pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
controlnet=controlnet,
vae=vae,
torch_dtype=torch.float16,
use_safetensors=True,
)
pipeline.enable_model_cpu_offload()
rr.log("positive_prompt", rr.TextDocument(prompt), static=True)
rr.log("negative_prompt", rr.TextDocument(negative_prompt), static=True)
images = pipeline(
prompt,
negative_prompt=negative_prompt,
image=canny_image, # add batch dimension
controlnet_conditioning_scale=0.5,
callback_on_step_end=controlnet_callback,
).images[0]
rr.log("output", rr.Image(images))
def main() -> None:
parser = argparse.ArgumentParser(description="Use Canny-conditioned ControlNet to generate image.")
parser.add_argument(
"--img-path",
type=str,
help="Path to image used as input for Canny edge detector.",
default=RERUN_LOGO_URL,
)
parser.add_argument(
"--prompt",
type=str,
help="Prompt used as input for ControlNet.",
default="aerial view, a futuristic research complex in a bright foggy jungle, hard lighting",
)
parser.add_argument(
"--negative-prompt",
type=str,
help="Negative prompt used as input for ControlNet.",
default="low quality, bad quality, sketches",
)
rr.script_add_args(parser)
args = parser.parse_args()
rr.script_setup(
args,
"rerun_example_controlnet",
default_blueprint=rrb.Horizontal(
rrb.Grid(
rrb.Spatial2DView(origin="input/raw"),
rrb.Spatial2DView(origin="input/canny"),
rrb.Vertical(
rrb.TextDocumentView(origin="positive_prompt"),
rrb.TextDocumentView(origin="negative_prompt"),
),
rrb.TensorView(origin="latent"),
),
rrb.Spatial2DView(origin="output"),
),
)
run_canny_controlnet(args.img_path, args.prompt, args.negative_prompt)
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/controlnet/pyproject.toml
+22
View File
@@ -0,0 +1,22 @@
[project]
name = "controlnet"
version = "0.1.0"
requires-python = ">=3.10"
readme = "README.md"
dependencies = [
"accelerate",
"opencv-python",
"pillow",
"diffusers<0.39",
"numpy",
"torch", # this will use the version defined in the uv workspace
"transformers",
"rerun-sdk",
]
[project.scripts]
controlnet = "controlnet:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/dataframe_query/README.md
+16
View File
@@ -0,0 +1,16 @@
This example will query for the first 10 rows of data in your recording of choice,
and display the results as a table in your terminal.
You can use one of your recordings, or grab one from our hosted examples, e.g.:
```bash
curl 'https://app.rerun.io/version/latest/examples/dna.rrd' -o - > /tmp/dna.rrd
```
The results can be filtered further by specifying an entity filter expression:
```bash
python dataframe_query.py my_recording.rrd /helix/structure/**\
```
```bash
python dataframe_query.py <path_to_rrd> [entity_path_filter]
```
examples/python/dataframe_query/dataframe_query.py
Executable
+53
View File
@@ -0,0 +1,53 @@
#!/usr/bin/env python3
"""Demonstrates basic usage of the dataframe APIs."""
from __future__ import annotations
import argparse
import rerun as rr
DESCRIPTION = """
Usage: python dataframe_query.py <path_to_rrd> [entity_path_filter]
This example will query for the first 10 rows of data in your recording of choice,
and display the results as a table in your terminal.
You can use one of your recordings, or grab one from our hosted examples, e.g.:
curl 'https://app.rerun.io/version/latest/examples/dna.rrd' -o - > /tmp/dna.rrd
The results can be filtered further by specifying an entity filter expression:
{bin_name} my_recording.rrd /helix/structure/**
""".strip()
def query(path_to_rrd: str, entity_path_filter: str) -> None:
with rr.server.Server(datasets={"recording": [path_to_rrd]}) as server:
dataset = server.client().get_dataset("recording")
# Query the data
view = dataset.filter_contents([entity_path_filter])
df = view.reader(index="log_time")
# Convert to pandas and show first 10 rows
table = df.to_pandas()
print(table.head(10))
def main() -> None:
parser = argparse.ArgumentParser(description=DESCRIPTION)
parser.add_argument("path_to_rrd", type=str, help="Path to the .rrd file")
parser.add_argument(
"entity_path_filter",
type=str,
nargs="?",
default="/**",
help="Optional entity path filter expression",
)
args = parser.parse_args()
query(args.path_to_rrd, args.entity_path_filter)
if __name__ == "__main__":
main()
examples/python/dataframe_query/pyproject.toml
+13
View File
@@ -0,0 +1,13 @@
[project]
name = "dataframe_query"
version = "0.1.0"
# requires-python = "<3.12"
readme = "README.md"
dependencies = ["rerun-sdk"]
[project.scripts]
dataframe_query = "dataframe_query:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/dataloader/.gitignore
+3
View File
@@ -0,0 +1,3 @@
data/
act_checkpoint/
.venv/
examples/python/dataloader/README.md
+77
View File
@@ -0,0 +1,77 @@
Train a [LeRobot](https://github.com/huggingface/lerobot) ACT policy using Rerun's experimental PyTorch dataloader, streaming trajectory data directly from a Rerun catalog.
For an explanation of the dataloader API and how the example fits together, see the [Train PyTorch models with the Rerun dataloader](https://rerun.io/docs/howto/train) how-to guide.
## Run the code
### 1. Install dependencies
This example has its own `uv` project, separate from the workspace `.venv`, because LeRobot requires
Python >=3.12 while the workspace supports older versions.
**Standalone** (sparse-checkout of just this directory, no local Rerun build):
```bash
uv sync --no-sources --no-dev
```
**Monorepo dev** (full repo checkout, editable local `rerun-sdk`):
```bash
cd examples/python/dataloader
RERUN_ALLOW_MISSING_BIN=1 uv sync
uv pip install ../../../rerun_py/rerun_dev_fixup
```
Then either `source .venv/bin/activate` or prefix subsequent commands with `uv run`.
### 2. Start a local Rerun server
In a separate terminal:
```bash
rerun server
```
This serves a Rerun server at `rerun+http://127.0.0.1:51234` (the default used by the scripts).
### 3. Prepare and register the dataset
Downloads a LeRobot dataset from HuggingFace, splits it into per-episode RRDs, and registers them as a dataset in the catalog:
```bash
uv run python prepare_dataset.py
```
Pass `--repo-id user/other_lerobot_ds` to use a different dataset, or `--catalog-url ""` to skip registration and only write local RRDs.
### 4. Train
```bash
uv run python train.py
```
The script streams batches from the catalog, trains an ACT policy for a few epochs, and saves a checkpoint to `act_checkpoint/`.
It accepts a few CLI flags (run `uv run python train.py --help` for the full list):
```bash
uv run python train.py \
--catalog-url rerun+http://127.0.0.1:51234 \
--dataset rerun_so101-pick-and-place \
--num-segments 3 \
--epochs 5 \
--batch-size 8 \
--num-workers 8 \
--lr 1e-5 \
--checkpoint-dir act_checkpoint \
--dataset-style iterable # or "map"
```
Pass `--num-segments 0` to train on all segments in the dataset.
### Training with traces
```sh
TELEMETRY_ENABLED=true OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://localhost:4317 uv run python train.py
```
examples/python/dataloader/prepare_dataset.py
+152
View File
@@ -0,0 +1,152 @@
"""Download a LeRobot dataset from HuggingFace Hub and prepare it for the dataloader.
This script:
1. Downloads a LeRobot dataset from HuggingFace Hub.
2. Loads it into Rerun via the built-in LeRobot importer (`log_file_from_path`).
3. Splits the resulting archive into one RRD per episode.
4. Registers the per-episode RRDs to a catalog server instance.
"""
from __future__ import annotations
import argparse
import re
from pathlib import Path
from huggingface_hub import snapshot_download
import rerun as rr
DEFAULT_REPO_ID = "rerun/so101-pick-and-place"
DEFAULT_OUTPUT_DIR = Path(__file__).resolve().parent / "data"
APPLICATION_ID = "lerobot"
_EPISODE_RE = re.compile(r"^(episode_)(\d+)$")
def _zero_pad_episode_id(rec_id: str, width: int = 5) -> str:
"""Turn `episode_1` into `episode_00001` so segments sort lexicographically."""
m = _EPISODE_RE.match(rec_id)
if m:
return f"{m.group(1)}{int(m.group(2)):0{width}d}"
return rec_id
def download_dataset(repo_id: str, dest: Path) -> Path:
"""Download a LeRobot dataset from HuggingFace Hub into *dest* and return its path."""
print(f"Downloading {repo_id} to {dest} …")
local_dir = snapshot_download(repo_id=repo_id, repo_type="dataset", local_dir=dest)
return Path(local_dir)
def lerobot_to_combined_rrd(dataset_dir: Path, combined_rrd: Path) -> None:
"""Use Rerun's built-in LeRobot importer to turn the dataset into a single RRD."""
print(f"Converting {dataset_dir} -> {combined_rrd}")
with rr.RecordingStream(APPLICATION_ID) as rec:
rec.save(str(combined_rrd))
rec.log_file_from_path(str(dataset_dir))
def split_into_episode_rrds(combined_rrd: Path, rrd_dir: Path) -> list[Path]:
"""Split a combined RRD archive into one RRD per episode.
Returns the paths of the written per-episode RRDs.
"""
rrd_dir.mkdir(parents=True, exist_ok=True)
reader = rr.experimental.RrdReader(str(combined_rrd))
recordings = reader.recordings()
print(f"Archive contains {len(recordings)} recordings")
episode_paths: list[Path] = []
for entry in recordings:
store = reader.store(store=entry)
# Skip metadata-only recordings (e.g. the "root" recording that only carries properties).
if not store.schema().entity_paths():
continue
episode_id = _zero_pad_episode_id(entry.recording_id)
rrd_path = rrd_dir / f"{episode_id}.rrd"
with rr.RecordingStream(APPLICATION_ID, recording_id=episode_id, send_properties=False) as rec:
rec.save(str(rrd_path))
rec.send_chunks(store)
episode_paths.append(rrd_path)
print(f" wrote {rrd_path} ({rrd_path.stat().st_size / (1024 * 1024):.1f} MB)")
return episode_paths
def register_to_catalog(
rrd_paths: list[Path],
*,
catalog_url: str,
dataset_name: str,
) -> None:
"""Register per-episode RRDs to a catalog server instance.
Uses absolute file:// URIs so the catalog can read the RRDs directly from the local filesystem.
"""
print(f"\nRegistering {len(rrd_paths)} episodes to {catalog_url} as dataset '{dataset_name}' …")
client = rr.catalog.CatalogClient(catalog_url)
dataset = client.create_dataset(dataset_name, exist_ok=True)
uris = [f"file://{p.resolve()}" for p in rrd_paths]
on_duplicate = rr.catalog.OnDuplicateSegmentLayer(rr.catalog.OnDuplicateSegmentLayer.REPLACE)
dataset.register(uris, on_duplicate=on_duplicate).wait()
print(" registration done")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument(
"--repo-id",
default=DEFAULT_REPO_ID,
help=f"HuggingFace dataset repo id (default: {DEFAULT_REPO_ID}).",
)
parser.add_argument(
"--output-dir",
type=Path,
default=DEFAULT_OUTPUT_DIR,
help=f"Directory to store downloaded dataset and output RRDs (default: {DEFAULT_OUTPUT_DIR}).",
)
parser.add_argument(
"--catalog-url",
default="rerun+http://127.0.0.1:51234",
help="Rerun catalog URL to register episodes with. Pass an empty string to skip registration.",
)
parser.add_argument(
"--dataset-name",
default=None,
help="Name of the dataset to create/use in the catalog (default: derived from --repo-id).",
)
parser.add_argument(
"--keep-combined",
action="store_true",
help="Keep the intermediate combined RRD after splitting (useful for debugging).",
)
args = parser.parse_args()
output_dir = args.output_dir
repo_slug = args.repo_id.replace("/", "_")
dataset_dir = output_dir / "lerobot" / repo_slug
rrd_dir = output_dir / "rrds" / repo_slug
combined_rrd = output_dir / f"{repo_slug}_combined.rrd"
download_dataset(args.repo_id, dataset_dir)
lerobot_to_combined_rrd(dataset_dir, combined_rrd)
episode_paths = split_into_episode_rrds(combined_rrd, rrd_dir)
if not args.keep_combined:
combined_rrd.unlink()
print(f"\nWrote {len(episode_paths)} per-episode RRDs to {rrd_dir}")
if args.catalog_url:
dataset_name = args.dataset_name or repo_slug
register_to_catalog(episode_paths, catalog_url=args.catalog_url, dataset_name=dataset_name)
if __name__ == "__main__":
main()
examples/python/dataloader/pyproject.toml
+45
View File
@@ -0,0 +1,45 @@
[project]
name = "dataloader"
version = "0.1.0"
readme = "README.md"
requires-python = ">=3.12,<3.13"
dependencies = [
"rerun-sdk[dataloader,catalog,tracing]",
"huggingface-hub>=1.0",
"lerobot[dataset]==0.6.0",
]
[dependency-groups]
dev = ["mypy==1.19.1"]
[tool.rerun-example]
# Picked up by scripts/ci/isolated_examples.py and the `py-lint-isolated-examples` pixi task.
isolated = true
[tool.uv]
# The example is flat scripts, not a wheel — skip project build, just sync deps.
package = false
# pyarrow 24.0.0 segfaults rerun on import and ships an incomplete py.typed that breaks mypy.
# Isolated examples don't inherit the workspace root's constraint-dependencies, so pin it here.
# pyarrow arrives transitively (via rerun-sdk), hence a constraint rather than a direct dependency.
constraint-dependencies = ["pyarrow>=23.0.1,<24"]
# Default `uv sync` uses the in-repo editable rerun-sdk (monorepo dev mode).
# After syncing, also run `uv pip install ../../../rerun_py/rerun_dev_fixup` to
# install the .pth shim that makes `import rerun` resolve to the editable source tree.
#
# Standalone users (e.g. sparse-checkout of just this example) run instead:
# uv sync --no-sources --no-dev
# That ignores the path source below and resolves `rerun-sdk` from PyPI.
# rerun-dev-fixup is intentionally absent from this file: uv 0.7.x resolves all
# dependency groups and extras unconditionally, so any path-only package here would
# block standalone `--no-sources` resolution.
[tool.uv.sources]
rerun-sdk = { path = "../../../rerun_py", editable = true }
# Merged onto the shared base at `../_isolated/mypy.ini` by
# scripts/ci/isolated_examples.py — list the untyped third-party libs this
# example actually imports.
[[tool.mypy.overrides]]
module = ["lerobot.*", "torch.*", "torchvision.*", "diffusers.*", "accelerate.*"]
ignore_missing_imports = true
examples/python/dataloader/train.py
+286
View File
@@ -0,0 +1,286 @@
"""Train a LeRobot ACT policy using the Rerun dataloader.
Demonstrates how to stream robot trajectory data from Rerun's catalog
into an imitation learning policy (Action Chunking Transformers).
The Rerun dataloader's Field.window feature fetches future action chunks in a single query per batch.
"""
from __future__ import annotations
import argparse
import time
from pathlib import Path
from typing import cast
import torch
import torch.nn.functional as F
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.policies.act.modeling_act import ACTPolicy
from torch.utils.data import DataLoader
from rerun._tracing import tracing_scope, with_tracing
from rerun.catalog import CatalogClient
from rerun.experimental.dataloader import (
DataSource,
Field,
NumericDecoder,
RerunIterableDataset,
RerunMapDataset,
VideoFrameDecoder,
)
CHECKPOINT_DIR = Path(__file__).resolve().parent / "act_checkpoint"
IMAGE_H = 32
IMAGE_W = 128
CAMERAS = ("laptop", "phone", "side")
IMAGE_KEYS = tuple(f"observation.images.{cam}" for cam in CAMERAS)
CHUNK_SIZE = 50
EPOCHS = 5
BATCH_SIZE = 8
LR = 1e-5
NUM_WORKERS = 4
FETCH_SIZE = 256
class CollateFn:
"""Picklable collate callable for PyTorch DataLoader multiprocessing."""
def __init__(self, chunk_size: int, state_dim: int) -> None:
self.chunk_size = chunk_size
self.state_dim = state_dim
@with_tracing("CollateFn")
def __call__(self, samples: list[dict[str, torch.Tensor | None]]) -> dict[str, torch.Tensor]:
# `VideoFrameDecoder` returns `None` when a target precedes the first keyframe; filter those out.
complete: list[dict[str, torch.Tensor]] = [
cast("dict[str, torch.Tensor]", s) for s in samples if all(s[f"image_{cam}"] is not None for cam in CAMERAS)
]
batch_size = len(complete)
states = torch.stack([s["state"] for s in complete]).float()
# Future action chunks: reshape windowed flat tensors
actions = torch.stack([s["action"].reshape(self.chunk_size, self.state_dim) for s in complete]).float()
batch: dict[str, torch.Tensor] = {
"observation.state": states,
"action": actions,
"action_is_pad": torch.zeros(batch_size, self.chunk_size, dtype=torch.bool),
}
# Per-camera images: (3, H, W) uint8 -> float in [0, 1], resized to (IMAGE_H, IMAGE_W)
for cam, key in zip(CAMERAS, IMAGE_KEYS):
imgs = torch.stack([s[f"image_{cam}"] for s in complete]).float() / 255.0
batch[key] = F.interpolate(imgs, size=(IMAGE_H, IMAGE_W), mode="bilinear", align_corners=False)
return batch
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument(
"--catalog-url",
default="rerun+http://127.0.0.1:51234",
help="Rerun catalog URL",
)
parser.add_argument(
"--dataset",
default="rerun_so101-pick-and-place",
help="Dataset name in the catalog",
)
parser.add_argument("--num-segments", type=int, default=3, help="Number of segments to use (0 for all)")
parser.add_argument("--epochs", type=int, default=EPOCHS, help="Number of training epochs")
parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="Training batch size")
parser.add_argument("--num-workers", type=int, default=NUM_WORKERS, help="DataLoader worker processes")
parser.add_argument(
"--fetch-size",
type=int,
default=FETCH_SIZE,
help="Samples fetched per server query for the iterable dataset",
)
parser.add_argument("--lr", type=float, default=LR, help="Learning rate")
parser.add_argument(
"--dataset-style",
choices=("iterable", "map"),
default="iterable",
help="Which Rerun dataset class to use: 'iterable' (RerunIterableDataset, internal shuffling) "
"or 'map' (RerunMapDataset, random access via DataLoader samplers).",
)
parser.add_argument(
"--checkpoint-dir",
type=Path,
default=CHECKPOINT_DIR,
help="Directory to save the trained policy checkpoint",
)
return parser.parse_args()
@with_tracing("main")
def main() -> None:
args = parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
client = CatalogClient(args.catalog_url)
dataset_entry = client.get_dataset(args.dataset)
all_segments = dataset_entry.segment_ids()
segments = all_segments if args.num_segments == 0 else all_segments[: args.num_segments]
print(f"Using {len(segments)} segments")
source = DataSource(dataset_entry, segments=segments)
fields = {
"state": Field("/observation.state:Scalars:scalars", decode=NumericDecoder()),
"action": Field(
"/action:Scalars:scalars",
decode=NumericDecoder(),
window=(1, CHUNK_SIZE),
),
"image_laptop": Field(
"/observation.images.laptop:VideoStream:sample",
decode=VideoFrameDecoder(codec="av1", keyframe_interval=2),
),
"image_phone": Field(
"/observation.images.phone:VideoStream:sample",
decode=VideoFrameDecoder(codec="av1", keyframe_interval=2),
),
"image_side": Field(
"/observation.images.side:VideoStream:sample",
decode=VideoFrameDecoder(codec="av1", keyframe_interval=2),
),
}
ds: RerunIterableDataset | RerunMapDataset
if args.dataset_style == "map":
ds = RerunMapDataset(source=source, index="frame_index", fields=fields)
else:
ds = RerunIterableDataset(source=source, index="frame_index", fields=fields, fetch_size=args.fetch_size)
print(f"Using {args.dataset_style} dataset with {len(ds)} samples (after window trimming)")
# IterableDataset doesn't support indexing, so probe shape via iteration.
state_tensor = next(iter(ds))["state"]
assert state_tensor is not None # NumericDecoder never returns None
state_dim = state_tensor.shape[0]
action_dim = state_dim
print(f"Dimensions: {state_dim=}, {action_dim=}")
config = ACTConfig(
chunk_size=CHUNK_SIZE,
n_action_steps=CHUNK_SIZE,
use_vae=True,
kl_weight=10.0,
dim_model=256,
n_heads=8,
dim_feedforward=1024,
n_encoder_layers=4,
n_decoder_layers=1,
latent_dim=32,
n_vae_encoder_layers=4,
dropout=0.1,
vision_backbone="resnet18",
pretrained_backbone_weights=None,
normalization_mapping={
"STATE": NormalizationMode.MEAN_STD,
"VISUAL": NormalizationMode.MEAN_STD,
"ACTION": NormalizationMode.MEAN_STD,
},
input_features={
"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(state_dim,)),
**{key: PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMAGE_H, IMAGE_W)) for key in IMAGE_KEYS},
},
output_features={
"action": PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,)),
},
)
policy = ACTPolicy(config)
policy.train()
policy.to(device)
print(f"ACT policy created ({sum(p.numel() for p in policy.parameters()):,} parameters, device={device})")
optimizer = torch.optim.AdamW(
policy.get_optim_params(),
lr=args.lr,
weight_decay=1e-4,
)
collate_fn = CollateFn(CHUNK_SIZE, state_dim)
# For the map-style dataset, shuffling is driven by the DataLoader's default RandomSampler.
# Swap in `sampler=DistributedSampler(ds)` (and call `sampler.set_epoch(epoch)` each epoch)
# for multi-node training, or plug in any other PyTorch sampler.
loader = DataLoader(
ds,
batch_size=args.batch_size,
shuffle=isinstance(ds, RerunMapDataset),
num_workers=args.num_workers,
collate_fn=collate_fn,
persistent_workers=True,
prefetch_factor=8,
)
num_batches = len(loader)
print(f"\nTraining for {args.epochs} epochs, {num_batches} batches/epoch, batch_size={args.batch_size}\n")
for epoch in range(args.epochs):
with tracing_scope(f"epoch {epoch}"):
if isinstance(ds, RerunIterableDataset):
ds.set_epoch(epoch)
total_loss = 0.0
total_l1 = 0.0
total_kld = 0.0
n = 0
t_last_print = time.perf_counter()
data_sum = 0.0
model_sum = 0.0
t_data_start = time.perf_counter()
for batch in loader:
data_time = time.perf_counter() - t_data_start
t_model_start = time.perf_counter()
batch = {k: v.to(device) for k, v in batch.items()}
loss, loss_dict = policy.forward(batch)
optimizer.zero_grad()
loss.backward()
optimizer.step()
model_time = time.perf_counter() - t_model_start
total_loss += loss.item()
total_l1 += loss_dict["l1_loss"]
total_kld += loss_dict.get("kld_loss", 0.0)
data_sum += data_time
model_sum += model_time
n += 1
if n % 10 == 0 or n == 1:
now = time.perf_counter()
since_last = now - t_last_print
t_last_print = now
print(
f" epoch {epoch + 1}/{args.epochs} batch {n}/{num_batches}"
f" loss={loss.item():.4f}"
f" data={data_sum:.1f}s model={model_sum:.1f}s"
f" since_last={since_last:.1f}s",
flush=True,
)
data_sum = 0.0
model_sum = 0.0
t_data_start = time.perf_counter()
avg_loss = total_loss / max(n, 1)
avg_l1 = total_l1 / max(n, 1)
avg_kld = total_kld / max(n, 1)
print(f"Epoch {epoch + 1}/{args.epochs} loss={avg_loss:.4f} l1={avg_l1:.4f} kld={avg_kld:.4f}")
with tracing_scope("save_pretrained"):
policy.save_pretrained(str(args.checkpoint_dir))
print(f"\nSaved checkpoint to {args.checkpoint_dir}")
if __name__ == "__main__":
main()
examples/python/dataloader/uv.lock
Generated
+1490
View File
File diff suppressed because it is too large Load Diff
examples/python/depth_compare/README.md
+23
View File
@@ -0,0 +1,23 @@
<!--[metadata]
title = "Depth compare"
tags = ["2D", "3D", "HuggingFace", "Depth", "Pinhole camera"]
source = "https://github.com/pablovela5620/monoprior"
thumbnail = "https://static.rerun.io/depth-compare-thumbnail/270fdb11a2b6e85c69d83d52d475a5ec644de9df/480w.png"
thumbnail_dimensions = [480, 258]
-->
https://vimeo.com/990128518?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=2802:1790
## Background
Depth Compare allows for easy comparison between different depth models, both metric and scale + shift invariant. There has been a recent flurry of monocular depth estimation networks, and often the only method of comparison is a 2D depth image. This does not provide the full scope of how well a depth network performs in reality. Depth Compare allows for 3D back-projection of the depth image into world space, along with outputting the point cloud to give a better understanding.
## Run the code
This is an external example. Check the [repository](https://github.com/pablovela5620/monoprior) for more information.
You can try the example on Rerun's HuggingFace space [here](https://huggingface.co/spaces/pablovela5620/depth-compare).
You can also run things locally by cloning the above repo and running:
```
pixi run app
```
examples/python/detect_and_track_objects/.gitignore
+2
View File
@@ -0,0 +1,2 @@
cache/
dataset/
examples/python/detect_and_track_objects/README.md
+172
View File
@@ -0,0 +1,172 @@
<!--[metadata]
title = "Detect and track objects"
tags = ["2D", "Hugging face", "Object detection", "Object tracking", "OpenCV"]
thumbnail = "https://static.rerun.io/detect-and-track-objects/63d7684ab1504c86a5375cb5db0fc515af433e08/480w.png"
thumbnail_dimensions = [480, 480]
channel = "release"
include_in_manifest = true
allow_warnings = true # TODO(emilk): torch produces a warning because of `transformers` (I think?). We should fix that, if we can.
-->
Visualize object detection and segmentation using the [Huggingface's Transformers](https://huggingface.co/docs/transformers/index) and optical flow tracking from OpenCV.
<picture data-inline-viewer="examples/detect_and_track_objects">
<img src="https://static.rerun.io/detact_and_track_objects/ce1939b8f2d22b36c4ca8b36dc0441e106b51da5/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/detact_and_track_objects/ce1939b8f2d22b36c4ca8b36dc0441e106b51da5/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/detact_and_track_objects/ce1939b8f2d22b36c4ca8b36dc0441e106b51da5/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/detact_and_track_objects/ce1939b8f2d22b36c4ca8b36dc0441e106b51da5/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/detact_and_track_objects/ce1939b8f2d22b36c4ca8b36dc0441e106b51da5/1200w.png">
</picture>
## Used Rerun types
[`Image`](https://www.rerun.io/docs/reference/types/archetypes/image), [`AssetVideo`](https://www.rerun.io/docs/reference/types/archetypes/asset_video), [`VideoFrameReference`](https://rerun.io/docs/reference/types/archetypes/video_frame_reference), [`SegmentationImage`](https://www.rerun.io/docs/reference/types/archetypes/segmentation_image), [`AnnotationContext`](https://www.rerun.io/docs/reference/types/archetypes/annotation_context), [`Boxes2D`](https://www.rerun.io/docs/reference/types/archetypes/boxes2d), [`TextLog`](https://www.rerun.io/docs/reference/types/archetypes/text_log)
## Background
In this example, optical flow tracking from OpenCV is employed for tracking objects across frames.
Additionally, the example showcases basic object detection and segmentation on a video using the Huggingface transformers library.
## Logging and visualizing with Rerun
The visualizations in this example were created with the following Rerun code.
### Timelines
For each processed video frame, all data sent to Rerun is associated with the [`timelines`](https://www.rerun.io/docs/concepts/logging-and-ingestion/timelines) `frame_idx`.
```python
rr.set_time("frame", sequence=frame_idx)
```
### Video
The input video is logged as a static [`AssetVideo`](https://www.rerun.io/docs/reference/types/archetypes/asset_video) to the `video` entity.
```python
video_asset = rr.AssetVideo(path=video_path)
frame_timestamps_ns = video_asset.read_frame_timestamps_nanos()
rr.log("video", video_asset, static=True)
```
Each frame is processed and the timestamp is logged to the `frame` timeline using a [`VideoFrameReference`](https://www.rerun.io/docs/reference/types/archetypes/video_frame_reference).
```python
rr.log("video", rr.VideoFrameReference(nanoseconds=frame_timestamps_ns[frame_idx]))
```
Since the detection and segmentation model operates on smaller images the resized images are logged to the separate `segmentation/rgb_scaled` entity.
This allows us to subsequently visualize the segmentation mask on top of the video.
```python
rr.log("segmentation/rgb_scaled", rr.Image(rgb_scaled).compress(jpeg_quality=85))
```
### Segmentations
The segmentation results is logged through a combination of two archetypes.
The segmentation image itself is logged as an
[`SegmentationImage`](https://www.rerun.io/docs/reference/types/archetypes/segmentation_image) and
contains the id for each pixel. It is logged to the `segmentation` entity.
```python
rr.log("segmentation", rr.SegmentationImage(mask))
```
The color and label for each class is determined by the
[`AnnotationContext`](https://www.rerun.io/docs/reference/types/archetypes/annotation_context) which is
logged to the root entity using `rr.log("/", …, static=True)` as it should apply to the whole sequence and all
entities that have a class id.
```python
class_descriptions = [rr.AnnotationInfo(id=cat["id"], color=cat["color"], label=cat["name"]) for cat in coco_categories]
rr.log("/", rr.AnnotationContext(class_descriptions), static=True)
```
### Detections
The detections and tracked bounding boxes are visualized by logging the [`Boxes2D`](https://www.rerun.io/docs/reference/types/archetypes/boxes2d) to Rerun.
#### Detections
```python
rr.log(
"segmentation/detections/things",
rr.Boxes2D(
array=thing_boxes,
array_format=rr.Box2DFormat.XYXY,
class_ids=thing_class_ids,
),
)
```
```python
rr.log(
f"image/tracked/{self.tracking_id}",
rr.Boxes2D(
array=self.tracked.bbox_xywh,
array_format=rr.Box2DFormat.XYWH,
class_ids=self.tracked.class_id,
),
)
```
#### Tracked bounding boxes
```python
rr.log(
"segmentation/detections/background",
rr.Boxes2D(
array=background_boxes,
array_format=rr.Box2DFormat.XYXY,
class_ids=background_class_ids,
),
)
```
The color and label of the bounding boxes is determined by their class id, relying on the same
[`AnnotationContext`](https://www.rerun.io/docs/reference/types/archetypes/annotation_context) as the
segmentation images. This ensures that a bounding box and a segmentation image with the same class id will also have the
same color.
Note that it is also possible to log multiple annotation contexts should different colors and / or labels be desired.
The annotation context is resolved by seeking up the entity hierarchy.
### Text log
Rerun integrates with the [Python logging module](https://docs.python.org/3/library/logging.html).
Through the [`TextLog`](https://www.rerun.io/docs/reference/types/archetypes/text_log#textlogintegration) text at different importance level can be logged. After an initial setup that is described on the
[`TextLog`](https://www.rerun.io/docs/reference/types/archetypes/text_log#textlogintegration), statements
such as `logging.info("…")`, `logging.debug("…")`, etc. will show up in the Rerun viewer.
```python
def setup_logging() -> None:
logger = logging.getLogger()
rerun_handler = rr.LoggingHandler("logs")
rerun_handler.setLevel(-1)
logger.addHandler(rerun_handler)
def main() -> None:
# … existing code …
setup_logging() # setup logging
track_objects(video_path, max_frame_count=args.max_frame) # start tracking
```
In the Viewer you can adjust the filter level and look at the messages time-synchronized with respect to other logged data.
## Run the code
To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/detect_and_track_objects
```
To experiment with the provided example, simply execute the main Python script:
```bash
python -m detect_and_track_objects # run the example
```
If you wish to customize it for various videos, adjust the maximum frames, explore additional features, or save it use the CLI with the `--help` option for guidance:
```bash
python -m detect_and_track_objects --help
```
examples/python/detect_and_track_objects/detect_and_track_objects.py
Executable
+506
View File
@@ -0,0 +1,506 @@
#!/usr/bin/env python3
"""Example applying simple object detection and tracking on a video."""
from __future__ import annotations
import argparse
import json
import logging
import os
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final
import cv2
import numpy as np
import numpy.typing as npt
import requests
from PIL import Image
import rerun as rr # pip install rerun-sdk
DESCRIPTION = """
# Detect and track objects
This is a more elaborate example applying simple object detection and segmentation on a video using the Huggingface
`transformers` library. Tracking across frames is performed using optical flow from OpenCV. The results are
visualized using Rerun.
The full source code for this example is available
[on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/detect_and_track_objects).
""".strip()
EXAMPLE_DIR: Final = Path(os.path.dirname(__file__))
DATASET_DIR: Final = EXAMPLE_DIR / "dataset" / "tracking_sequences"
DATASET_URL_BASE: Final = "https://storage.googleapis.com/rerun-example-datasets/tracking_sequences"
CACHE_DIR: Final = EXAMPLE_DIR / "cache"
# panoptic_coco_categories.json comes from:
# https://github.com/cocodataset/panopticapi/blob/7bb4655548f98f3fedc07bf37e9040a992b054b0/panoptic_coco_categories.json
# License: https://github.com/cocodataset/panopticapi/blob/7bb4655548f98f3fedc07bf37e9040a992b054b0/license.txt
COCO_CATEGORIES_PATH = EXAMPLE_DIR / "panoptic_coco_categories.json"
DOWNSCALE_FACTOR = 2
DETECTION_SCORE_THRESHOLD = 0.8
os.environ["HF_HOME"] = str(CACHE_DIR.absolute())
from transformers import (
DetrForSegmentation,
DetrImageProcessor,
)
if TYPE_CHECKING:
from collections.abc import Sequence
@dataclass
class Detection:
"""Information about a detected object."""
class_id: int
bbox_xywh: list[float]
image_width: int
image_height: int
def scaled_to_fit_image(self, target_image: npt.NDArray[Any]) -> Detection:
"""Rescales detection to fit to target image."""
target_height, target_width = target_image.shape[:2]
return self.scaled_to_fit_size(target_width=target_width, target_height=target_height)
def scaled_to_fit_size(self, target_width: int, target_height: int) -> Detection:
"""Rescales detection to fit to target image with given size."""
if target_height == self.image_height and target_width == self.image_width:
return self
width_scale = target_width / self.image_width
height_scale = target_height / self.image_height
target_bbox = [
self.bbox_xywh[0] * width_scale,
self.bbox_xywh[1] * height_scale,
self.bbox_xywh[2] * width_scale,
self.bbox_xywh[3] * height_scale,
]
return Detection(self.class_id, target_bbox, target_width, target_height)
class Detector:
"""Detects objects to track."""
def __init__(self, coco_categories: list[dict[str, Any]]) -> None:
logging.info("Initializing neural net for detection and segmentation.")
self.feature_extractor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50-panoptic")
self.model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")
self.is_thing_from_id: dict[int, bool] = {cat["id"]: bool(cat["isthing"]) for cat in coco_categories}
def detect_objects_to_track(self, rgb: cv2.typing.MatLike, frame_idx: int) -> list[Detection]:
logging.info("Looking for things to track on frame %d", frame_idx)
logging.debug("Preprocess image for detection network")
pil_im_small = Image.fromarray(rgb)
inputs = self.feature_extractor(images=pil_im_small, return_tensors="pt")
_, _, scaled_height, scaled_width = inputs["pixel_values"].shape
scaled_size = (scaled_width, scaled_height)
rgb_scaled = cv2.resize(rgb, scaled_size)
rr.log("segmentation/rgb_scaled", rr.Image(rgb_scaled).compress(jpeg_quality=85))
logging.debug("Pass image to detection network")
outputs = self.model(**inputs)
logging.debug("Extracting detections and segmentations from network output")
processed_sizes = [(scaled_height, scaled_width)]
segmentation_mask = self.feature_extractor.post_process_semantic_segmentation(outputs, processed_sizes)[0]
detections = self.feature_extractor.post_process_object_detection(
outputs,
threshold=0.8,
target_sizes=processed_sizes,
)[0]
mask = segmentation_mask.detach().cpu().numpy().astype(np.uint8)
rr.log("segmentation", rr.SegmentationImage(mask))
boxes = detections["boxes"].detach().cpu().numpy()
class_ids = detections["labels"].detach().cpu().numpy()
things = [self.is_thing_from_id[id] for id in class_ids]
self.log_detections(boxes, class_ids, things)
objects_to_track: list[Detection] = []
for idx, (class_id, is_thing) in enumerate(zip(class_ids, things, strict=False)):
if is_thing:
x_min, y_min, x_max, y_max = boxes[idx, :]
bbox_xywh = [x_min, y_min, x_max - x_min, y_max - y_min]
objects_to_track.append(
Detection(
class_id=class_id,
bbox_xywh=bbox_xywh,
image_width=scaled_width,
image_height=scaled_height,
),
)
return objects_to_track
def log_detections(self, boxes: npt.NDArray[np.float32], class_ids: list[int], things: list[bool]) -> None:
things_np = np.array(things)
class_ids_np = np.array(class_ids, dtype=np.uint16)
thing_boxes = boxes[things_np, :]
thing_class_ids = class_ids_np[things_np]
rr.log(
"segmentation/detections/things",
rr.Boxes2D(
array=thing_boxes,
array_format=rr.Box2DFormat.XYXY,
class_ids=thing_class_ids,
),
)
background_boxes = boxes[~things_np, :]
background_class_ids = class_ids[~things_np]
rr.log(
"segmentation/detections/background",
rr.Boxes2D(
array=background_boxes,
array_format=rr.Box2DFormat.XYXY,
class_ids=background_class_ids,
),
)
class Tracker:
"""
Each instance takes care of tracking a single object using optical flow.
The factory class method `create_new_tracker` is used to give unique tracking id's per instance.
"""
next_tracking_id = 0
MAX_TIMES_UNDETECTED = 2
def __init__(self, tracking_id: int, detection: Detection, bgr: cv2.typing.MatLike) -> None:
self.tracking_id = tracking_id
self.tracked = detection.scaled_to_fit_image(bgr)
self.num_recent_undetected_frames = 0
# Store the previous frame and points for optical flow tracking
self.prev_gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
self.is_active = True
self.prev_points: npt.NDArray[np.float32] = np.array([]) # Will be initialized below
self._init_tracking_points()
self.log_tracked()
def _init_tracking_points(self) -> None:
"""Initialize corner points within the bounding box for tracking."""
x, y, w, h = [int(v) for v in self.tracked.bbox_xywh]
# Create a grid of points within the bounding box
points = []
grid_size = 5
for i in range(grid_size):
for j in range(grid_size):
px = x + (w * (i + 1)) // (grid_size + 1)
py = y + (h * (j + 1)) // (grid_size + 1)
points.append([[px, py]])
self.prev_points = np.array(points, dtype=np.float32)
@classmethod
def create_new_tracker(cls, detection: Detection, bgr: cv2.typing.MatLike) -> Tracker:
new_tracker = cls(cls.next_tracking_id, detection, bgr)
cls.next_tracking_id += 1
return new_tracker
def update(self, bgr: cv2.typing.MatLike) -> None:
if not self.is_tracking:
return
curr_gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
# Calculate optical flow
next_points, status, _error = cv2.calcOpticalFlowPyrLK( # type: ignore[call-overload]
self.prev_gray,
curr_gray,
self.prev_points,
None,
winSize=(15, 15),
maxLevel=2,
)
if next_points is None or status is None:
logging.info("Optical flow failed for tracker with id #%d", self.tracking_id)
self.is_active = False
self.log_tracked()
return
# Filter good points
status_mask = status.flatten() == 1
good_new = next_points[status_mask].reshape(-1, 2)
good_old = self.prev_points[status_mask].reshape(-1, 2)
if len(good_new) < 3:
logging.info("Too few points tracked for tracker with id #%d", self.tracking_id)
self.is_active = False
self.log_tracked()
return
# Calculate displacement to adjust bbox
displacement_x = np.median(good_new[:, 0] - good_old[:, 0])
displacement_y = np.median(good_new[:, 1] - good_old[:, 1])
x, y, w, h = self.tracked.bbox_xywh
new_x = x + displacement_x
new_y = y + displacement_y
self.tracked.bbox_xywh = clip_bbox_to_image(
bbox_xywh=[new_x, new_y, w, h],
image_width=self.tracked.image_width,
image_height=self.tracked.image_height,
)
# Update for next iteration
self.prev_gray = curr_gray.copy()
self.prev_points = good_new.reshape(-1, 1, 2)
self.log_tracked()
def log_tracked(self) -> None:
if self.is_tracking:
rr.log(
f"video/tracked/{self.tracking_id}",
rr.Boxes2D(
array=self.tracked.bbox_xywh,
array_format=rr.Box2DFormat.XYWH,
class_ids=self.tracked.class_id,
),
)
else:
rr.log(f"video/tracked/{self.tracking_id}", rr.Boxes2D.cleared())
def update_with_detection(self, detection: Detection, bgr: cv2.typing.MatLike) -> None:
self.num_recent_undetected_frames = 0
self.tracked = detection.scaled_to_fit_image(bgr)
self.prev_gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
self.is_active = True
self._init_tracking_points()
self.log_tracked()
def set_not_detected_in_frame(self) -> None:
self.num_recent_undetected_frames += 1
if self.num_recent_undetected_frames >= Tracker.MAX_TIMES_UNDETECTED:
logging.info(
"Dropping tracker with id #%d after not being detected %d times",
self.tracking_id,
self.num_recent_undetected_frames,
)
self.is_active = False
self.log_tracked()
@property
def is_tracking(self) -> bool:
return self.is_active
def match_score(self, other: Detection) -> float:
"""Returns bbox IoU if classes match, otherwise 0."""
if self.tracked.class_id != other.class_id:
return 0.0
if not self.is_tracking:
return 0.0
other = other.scaled_to_fit_size(target_width=self.tracked.image_width, target_height=self.tracked.image_height)
tracked_bbox = self.tracked.bbox_xywh
other_bbox = other.bbox_xywh
return box_iou(tracked_bbox, other_bbox)
def box_iou(first: list[float], second: list[float]) -> float:
"""Calculate Intersection over Union (IoU) between two 2D rectangles in XYWH format."""
left = max(first[0], second[0])
right = min(first[0] + first[2], second[0] + second[2])
top = min(first[1] + first[3], second[1] + second[3])
bottom = max(first[1], second[1])
overlap_width = max(0.0, right - left)
overlap_height = max(0.0, top - bottom)
intersection_area = overlap_width * overlap_height
tracked_area = first[2] * first[3]
other_area = second[2] * second[3]
union_area = tracked_area + other_area - intersection_area
return intersection_area / union_area
def clip_bbox_to_image(bbox_xywh: list[float], image_width: int, image_height: int) -> list[float]:
x_min = max(0, bbox_xywh[0])
y_min = max(0, bbox_xywh[1])
x_max = min(image_width - 1, bbox_xywh[0] + bbox_xywh[2])
y_max = min(image_height - 1, bbox_xywh[1] + bbox_xywh[3])
return [x_min, y_min, x_max - x_min, y_max - y_min]
def update_trackers_with_detections(
trackers: list[Tracker],
detections: Sequence[Detection],
label_strs: Sequence[str],
bgr: cv2.typing.MatLike,
) -> list[Tracker]:
"""
Tries to match detections to existing trackers and updates the trackers if they match.
Any detections that don't match existing trackers will generate new trackers.
Returns the new set of trackers.
"""
non_updated_trackers = list(trackers) # shallow copy
updated_trackers: list[Tracker] = []
logging.debug("Updating %d trackers with %d new detections", len(trackers), len(detections))
for detection in detections:
top_match_score = 0.0
best_match_idx = -1
if non_updated_trackers:
scores = [tracker.match_score(detection) for tracker in non_updated_trackers]
best_match_idx = int(np.argmax(scores))
top_match_score = scores[best_match_idx]
if top_match_score > 0.0 and best_match_idx >= 0:
best_tracker = non_updated_trackers.pop(best_match_idx)
best_tracker.update_with_detection(detection, bgr)
updated_trackers.append(best_tracker)
else:
updated_trackers.append(Tracker.create_new_tracker(detection, bgr))
logging.info(
"Tracking newly detected %s with tracking id #%d",
label_strs[detection.class_id],
Tracker.next_tracking_id,
)
logging.debug("Updating %d trackers without matching detections", len(non_updated_trackers))
for tracker in non_updated_trackers:
tracker.set_not_detected_in_frame()
tracker.update(bgr)
if tracker.is_tracking:
updated_trackers.append(tracker)
logging.info("Tracking %d objects after updating with %d new detections", len(updated_trackers), len(detections))
return updated_trackers
def track_objects(video_path: str, *, max_frame_count: int | None) -> None:
with open(COCO_CATEGORIES_PATH, encoding="utf8") as f:
coco_categories = json.load(f)
class_descriptions = [
rr.AnnotationInfo(id=cat["id"], color=cat["color"], label=cat["name"]) for cat in coco_categories
]
rr.log("/", rr.AnnotationContext(class_descriptions), static=True)
logging.info("Initializing detector…")
# This call has a tendency to hard exit on failure (no exceptions):
detector = Detector(coco_categories=coco_categories)
logging.info("Detector initialized.")
video_asset = rr.AssetVideo(path=video_path)
frame_timestamps_ns = video_asset.read_frame_timestamps_nanos()
rr.log("video", video_asset, static=True)
logging.info("Loading input video: %s", str(video_path))
cap = cv2.VideoCapture(video_path)
frame_idx = 0
label_strs = [cat["name"] or str(cat["id"]) for cat in coco_categories]
trackers: list[Tracker] = []
while cap.isOpened():
if max_frame_count is not None and frame_idx >= max_frame_count:
break
ret, bgr = cap.read()
rr.set_time("frame", sequence=frame_idx)
if not ret:
logging.info("End of video")
break
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
rr.log("video", rr.VideoFrameReference(nanoseconds=frame_timestamps_ns[frame_idx]))
if not trackers or frame_idx % 40 == 0:
detections = detector.detect_objects_to_track(rgb=rgb, frame_idx=frame_idx)
trackers = update_trackers_with_detections(trackers, detections, label_strs, bgr)
else:
if frame_idx % 10 == 0:
logging.debug("Running tracking update step for frame %d", frame_idx)
for tracker in trackers:
tracker.update(bgr)
trackers = [tracker for tracker in trackers if tracker.is_tracking]
frame_idx += 1
def get_downloaded_path(dataset_dir: Path, video_name: str) -> str:
video_file_name = f"{video_name}.mp4"
destination_path = dataset_dir / video_file_name
if destination_path.exists():
logging.info("%s already exists. No need to download", destination_path)
return str(destination_path)
source_path = f"{DATASET_URL_BASE}/{video_file_name}"
logging.info("Downloading video from %s to %s", source_path, destination_path)
os.makedirs(dataset_dir.absolute(), exist_ok=True)
with requests.get(source_path, stream=True) as req:
req.raise_for_status()
with open(destination_path, "wb") as f:
f.writelines(req.iter_content(chunk_size=8192))
return str(destination_path)
def setup_logging() -> None:
logger = logging.getLogger()
rerun_handler = rr.LoggingHandler("logs")
rerun_handler.setLevel(-1)
logger.addHandler(rerun_handler)
def main() -> None:
# Ensure the logging gets written to stderr:
logging.getLogger().addHandler(logging.StreamHandler())
logging.getLogger().setLevel("DEBUG")
parser = argparse.ArgumentParser(description="Example applying simple object detection and tracking on a video.")
parser.add_argument(
"--video",
type=str,
default="horses",
choices=["horses", "driving", "boats"],
help="The example video to run on.",
)
parser.add_argument("--dataset-dir", type=Path, default=DATASET_DIR, help="Directory to save example videos to.")
parser.add_argument("--video-path", type=str, default="", help="Full path to video to run on. Overrides `--video`.")
parser.add_argument(
"--max-frame",
type=int,
help="Stop after processing this many frames. If not specified, will run until interrupted.",
)
rr.script_add_args(parser)
args = parser.parse_args()
rr.script_setup(args, "rerun_example_detect_and_track_objects")
setup_logging()
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
video_path: str = args.video_path
if not video_path:
video_path = get_downloaded_path(args.dataset_dir, args.video)
track_objects(video_path, max_frame_count=args.max_frame)
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/detect_and_track_objects/panoptic_coco_categories.json
+1465
View File
File diff suppressed because it is too large Load Diff
examples/python/detect_and_track_objects/pyproject.toml
+25
View File
@@ -0,0 +1,25 @@
[project]
name = "detect_and_track_objects"
version = "0.1.0"
# requires-python = "<3.12"
readme = "README.md"
dependencies = [
"numpy",
"opencv-python>4.9",
"pillow",
"requests>=2.31,<3",
"rerun-sdk",
"timm==1.0.19",
"torch", # this will use the version defined in the uv workspace
"transformers>=4.55.0",
]
[project.scripts]
detect_and_track_objects = "detect_and_track_objects:main"
[tool.rerun-example]
extra-args = "--max-frame=10"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/dicom_mri/.gitignore
+1
View File
@@ -0,0 +1 @@
dataset/**
examples/python/dicom_mri/README.md
+63
View File
@@ -0,0 +1,63 @@
<!--[metadata]
title = "Dicom MRI"
tags = ["Tensor", "MRI", "DICOM"]
thumbnail = "https://static.rerun.io/dicom-mri/d5a434f92504e8dda8af6c7f4eded2a9d662c991/480w.png"
thumbnail_dimensions = [480, 480]
channel = "main"
include_in_manifest = true
-->
Visualize a [DICOM](https://en.wikipedia.org/wiki/DICOM) MRI scan. This demonstrates the flexible tensor slicing capabilities of the Rerun viewer.
<picture data-inline-viewer="examples/dicom_mri">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/dicom_mri/e39f34a1b1ddd101545007f43a61783e1d2e5f8e/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/dicom_mri/e39f34a1b1ddd101545007f43a61783e1d2e5f8e/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/dicom_mri/e39f34a1b1ddd101545007f43a61783e1d2e5f8e/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/dicom_mri/e39f34a1b1ddd101545007f43a61783e1d2e5f8e/1200w.png">
<img src="https://static.rerun.io/dicom_mri/e39f34a1b1ddd101545007f43a61783e1d2e5f8e/full.png" alt="">
</picture>
## Used Rerun types
[`Tensor`](https://www.rerun.io/docs/reference/types/archetypes/tensor), [`TextDocument`](https://www.rerun.io/docs/reference/types/archetypes/text_document)
## Background
Digital Imaging and Communications in Medicine (DICOM) serves as a technical standard for the digital storage and transmission of medical images. In this instance, an MRI scan is visualized using Rerun.
## Logging and visualizing with Rerun
The visualizations in this example were created with just the following line.
```python
rr.log("tensor", rr.Tensor(voxels_volume_u16, dim_names=["right", "back", "up"]))
```
A `numpy.array` named `voxels_volume_u16` representing volumetric MRI intensities with a shape of `(512, 512, 512)`.
To visualize this data effectively in Rerun, we can log the `numpy.array` as [`Tensor`](https://www.rerun.io/docs/reference/types/archetypes/tensor) to the `tensor` entity.
In the Rerun Viewer you can also inspect the data in detail. The `dim_names` provided in the above call to `rr.log` help to
give semantic meaning to each axis. After selecting the tensor view, you can adjust various settings in the Blueprint
settings on the right-hand side. For example, you can adjust the color map, the brightness, which dimensions to show as
an image and which to select from, and more.
## Run the code
To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/dicom_mri
```
To experiment with the provided example, simply execute the main Python script:
```bash
python -m dicom_mri # run the example
```
If you wish to customize it, explore additional features, or save it, use the CLI with the `--help` option for guidance:
```bash
python -m dicom_mri --help
```
examples/python/dicom_mri/dicom_mri.py
Executable
+92
View File
@@ -0,0 +1,92 @@
#!/usr/bin/env python3
"""Example using MRI scan data in the DICOM format."""
from __future__ import annotations
import argparse
import io
import os
import zipfile
from pathlib import Path
from typing import TYPE_CHECKING, Final
import dicom_numpy
import numpy as np
import numpy.typing as npt
import pydicom as dicom
import requests
import rerun as rr # pip install rerun-sdk
if TYPE_CHECKING:
from collections.abc import Iterable
DESCRIPTION = """
# Dicom MRI
This example visualizes an MRI scan using Rerun.
The visualization of the data consists of just the following line
```python
rr.log("tensor", rr.Tensor(voxels_volume_u16, dim_names=["right", "back", "up"]))
```
The full source code for this example is available
[on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/dicom_mri).
"""
DATASET_DIR: Final = Path(os.path.dirname(__file__)) / "dataset"
DATASET_URL: Final = "https://storage.googleapis.com/rerun-example-datasets/dicom.zip"
def extract_voxel_data(
dicom_files: Iterable[Path],
) -> tuple[npt.NDArray[np.int16], npt.NDArray[np.float32]]:
slices = [dicom.read_file(f) for f in dicom_files] # type: ignore[misc]
voxel_ndarray, ijk_to_xyz = dicom_numpy.combine_slices(slices)
return voxel_ndarray, ijk_to_xyz
def list_dicom_files(dir: Path) -> Iterable[Path]:
for path, _, files in os.walk(dir):
for f in files:
if f.endswith(".dcm"):
yield Path(path) / f
def read_and_log_dicom_dataset(dicom_files: Iterable[Path]) -> None:
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
voxels_volume, _ = extract_voxel_data(dicom_files)
# the data is i16, but in range [0, 536].
voxels_volume_u16: npt.NDArray[np.uint16] = np.require(voxels_volume, np.uint16)
rr.log("tensor", rr.Tensor(voxels_volume_u16, dim_names=["right", "back", "up"]))
def ensure_dataset_downloaded() -> Iterable[Path]:
dicom_files = list(list_dicom_files(DATASET_DIR))
if dicom_files:
return dicom_files
print("downloading dataset…")
os.makedirs(DATASET_DIR.absolute(), exist_ok=True)
resp = requests.get(DATASET_URL, stream=True)
z = zipfile.ZipFile(io.BytesIO(resp.content))
z.extractall(DATASET_DIR.absolute())
return list_dicom_files(DATASET_DIR)
def main() -> None:
parser = argparse.ArgumentParser(description="Example using MRI scan data in the DICOM format.")
rr.script_add_args(parser)
args = parser.parse_args()
rr.script_setup(args, "rerun_example_dicom_mri")
dicom_files = ensure_dataset_downloaded()
read_and_log_dicom_dataset(dicom_files)
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/dicom_mri/pyproject.toml
+20
View File
@@ -0,0 +1,20 @@
[project]
name = "dicom_mri"
version = "0.1.0"
# requires-python = "<3.12"
readme = "README.md"
dependencies = [
"dicom_numpy==0.6.2",
"numpy",
"pydicom==2.4.5",
"requests>=2.31,<3",
"rerun-sdk",
"types-requests>=2.31,<3",
]
[project.scripts]
dicom_mri = "dicom_mri:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/differentiable_blocks_world/README.md
+43
View File
@@ -0,0 +1,43 @@
<!--[metadata]
title = "Differentiable blocks world: qualitative 3D decomposition by rendering primitives"
source = "https://github.com/rerun-io/differentiable-blocksworld"
tags = ["3D", "Mesh", "Pinhole camera", "Paper walkthrough"]
thumbnail = "https://static.rerun.io/differentiable-blocks/42f3a5481162a0e75f1c52ef1a12d4fedb35389e/480w.png"
thumbnail_dimensions = [480, 480]
-->
This example is a visual walkthrough of the paper "Differentiable Block Worlds".
All the visualizations were created by editing the original source code to log data with the Rerun SDK.
## Visual paper walkthrough
Finding a textured mesh decomposition from a collection of posed images is a very challenging optimization problem. "Differentiable Block Worlds" by Tom Monnier et al. shows impressive results using differentiable rendering. Here we visualize how this optimization works using the Rerun SDK.
https://vimeo.com/865326948?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=10000:7309
In "Differentiable Blocks World: Qualitative 3D Decomposition by Rendering Primitives" the authors describe an optimization of a background icosphere, a ground plane, and multiple superquadrics. The goal is to find the shapes and textures that best explain the observations.
<picture>
<source media="(max-width: 480px)" srcset="https://static.rerun.io/dbw-overview/83fe4a19b65b2c9a5c0e10aef00e4a82026e2b46/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/dbw-overview/83fe4a19b65b2c9a5c0e10aef00e4a82026e2b46/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/dbw-overview/83fe4a19b65b2c9a5c0e10aef00e4a82026e2b46/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/dbw-overview/83fe4a19b65b2c9a5c0e10aef00e4a82026e2b46/1200w.png">
<img src="https://static.rerun.io/dbw-overview/83fe4a19b65b2c9a5c0e10aef00e4a82026e2b46/full.png" alt="">
</picture>
The optimization is initialized with an initial set of superquadrics ("blocks"), a ground plane, and a sphere for the background. From here, the optimization can only reduce the number of blocks, not add additional ones.
https://vimeo.com/865327350?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=10000:6497
A key difference to other differentiable renderers is the addition of transparency handling. Each mesh has an opacity associated with it that is optimized. When the opacity becomes lower than a threshold the mesh is discarded in the visualization. This allows to optimize the number of meshes.
https://vimeo.com/865327387?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=10000:7037
To stabilize the optimization and avoid local minima, a 3-stage optimization is employed:
1. the texture resolution is reduced by a factor of 8,
2. the full resolution texture is optimized, and
3. transparency-based optimization is deactivated, only optimizing the opaque meshes from here.
https://vimeo.com/865329177?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=10000:8845
Make sure to read the [paper](https://arxiv.org/abs/2307.05473) by Tom Monnier, Jake Austin, Angjoo Kanazawa, Alexei A. Efros, Mathieu Aubry. Interesting study of how to approach such a difficult optimization problem.
examples/python/dna/README.md
+22
View File
@@ -0,0 +1,22 @@
<!--[metadata]
title = "Helix"
tags = ["3D", "API example"]
thumbnail = "https://static.rerun.io/dna/40d9744af3f0e21d3b174054f0a935662a574ce0/480w.png"
thumbnail_dimensions = [480, 480]
channel = "main"
include_in_manifest = true
-->
Simple example of logging point and line primitives to draw a 3D helix.
<picture data-inline-viewer="examples/dna">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/helix/f4c375546fa9d24f7cd3a1a715ebf75b2978817a/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/helix/f4c375546fa9d24f7cd3a1a715ebf75b2978817a/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/helix/f4c375546fa9d24f7cd3a1a715ebf75b2978817a/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/helix/f4c375546fa9d24f7cd3a1a715ebf75b2978817a/1200w.png">
<img src="https://static.rerun.io/helix/f4c375546fa9d24f7cd3a1a715ebf75b2978817a/full.png" alt="">
</picture>
```bash
python examples/python/dna/dna.py
```
examples/python/dna/dna.py
Executable
+95
View File
@@ -0,0 +1,95 @@
#!/usr/bin/env python3
"""The example from our Getting Started page."""
from __future__ import annotations
import argparse
from math import tau
import numpy as np
import rerun as rr # pip install rerun-sdk
from rerun import blueprint as rrb
from rerun.utilities import bounce_lerp, build_color_spiral
DESCRIPTION = """
# DNA
This is a minimal example that logs synthetic 3D data in the shape of a double helix. The underlying data is generated
using numpy and visualized using Rerun.
The full source code for this example is available
[on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/dna).
""".strip()
def log_data() -> None:
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
rr.set_time("stable_time", duration=0)
NUM_POINTS = 100
# points and colors are both np.array((NUM_POINTS, 3))
points1, colors1 = build_color_spiral(NUM_POINTS)
points2, colors2 = build_color_spiral(NUM_POINTS, angular_offset=tau * 0.5)
rr.log("helix/structure/left", rr.Points3D(points1, colors=colors1, radii=0.08), static=True)
rr.log("helix/structure/right", rr.Points3D(points2, colors=colors2, radii=0.08), static=True)
rr.log(
"helix/structure/scaffolding",
rr.LineStrips3D(np.stack((points1, points2), axis=1), colors=[128, 128, 128]),
static=True,
)
time_offsets = np.random.rand(NUM_POINTS)
for i in range(400):
time = i * 0.01
rr.set_time("stable_time", duration=time)
times = np.repeat(time, NUM_POINTS) + time_offsets
beads = [bounce_lerp(points1[n], points2[n], times[n]) for n in range(NUM_POINTS)]
colors = [[int(bounce_lerp(80, 230, times[n] * 2))] for n in range(NUM_POINTS)]
rr.log(
"helix/structure/scaffolding/beads",
rr.Points3D(beads, radii=0.06, colors=np.repeat(colors, 3, axis=-1)),
)
rr.log(
"helix/structure",
rr.Transform3D(rotation=rr.RotationAxisAngle(axis=[0, 0, 1], radians=time / 4.0 * tau)),
)
def main() -> None:
parser = argparse.ArgumentParser(description="Logs rich data using the Rerun SDK.")
rr.script_add_args(parser)
parser.add_argument(
"--blueprint",
action="store_true",
help="Logs a blueprint that enables a cursor-relative time range on the beads.",
)
args = parser.parse_args()
rr.script_setup(args, "rerun_example_dna_abacus")
log_data()
if args.blueprint:
blueprint = rrb.Blueprint(
rrb.Spatial3DView(
origin="/",
overrides={
"helix/structure/scaffolding/beads": rrb.VisibleTimeRanges(
timeline="stable_time",
start=rrb.TimeRangeBoundary.cursor_relative(seconds=-0.3),
end=rrb.TimeRangeBoundary.cursor_relative(seconds=0.3),
),
},
),
)
rr.send_blueprint(blueprint)
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/dna/pyproject.toml
+13
View File
@@ -0,0 +1,13 @@
[project]
name = "dna"
version = "0.1.0"
# requires-python = "<3.12"
readme = "README.md"
dependencies = ["numpy", "rerun-sdk", "scipy"]
[project.scripts]
dna = "dna:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/droid_dataset/README.md
+19
View File
@@ -0,0 +1,19 @@
<!--[metadata]
title = "DROID"
source = "https://github.com/rerun-io/python-example-droid-dataset"
tags = ["2D", "3D", "Depth", "Pinhole camera", "Blueprint"]
thumbnail = "https://static.rerun.io/droid-screenshot-thumbnail.png/87462a48d06a3c2fcc9225ac86df617048c262e6/480w.png"
thumbnail_dimensions = [480, 480]
-->
This example visualizes the [DROID dataset](https://droid-dataset.github.io/).
https://vimeo.com/940717540?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=1920:1080
## Background
DROID is a robot manipulation dataset collected over a diverse set of environments that can be used to train robotic manipulation policies. The dataset contains sensor data from each joint of the robot as well as video footage from 3 stereo cameras located on wrist and on either side of the arm. In each frame it also stores the action taken which is the target position and velocity given by the teleoperator.
## Run the code
You can find the build instructions here: [python-example-droid-dataset](https://github.com/rerun-io/python-example-droid-dataset)
examples/python/droid_semantic_search/.gitignore
+10
View File
@@ -0,0 +1,10 @@
# Local vector index built by `ingest.py` (LanceDB / Qdrant backends)
droid_lancedb/
droid_qdrant/
# DROID episodes downloaded by `prepare_dataset.py` (Hugging Face path)
data/
# Optimized (keyframe-fixed) episode copies written by `prepare_dataset.py`
optimized/
# uv-managed virtualenv for this isolated example
.venv/
__pycache__/
examples/python/droid_semantic_search/README.md
+162
View File
@@ -0,0 +1,162 @@
<!--[metadata]
title = "DROID semantic frame search"
tags = ["Robotics", "Semantic search", "Embeddings", "SigLIP", "LanceDB", "Qdrant", "Rerun Hub"]
thumbnail = "https://static.rerun.io/droid_semantic_search/fd265e110c9d05c9fdd018c68dd83561aa836c52/480w.png"
thumbnail_dimensions = [480, 320]
include_in_manifest = true
-->
Find moments in robot demonstrations by describing them in plain language, and jump straight to the matching frame in the Rerun viewer.
<picture>
<source media="(max-width: 480px)" srcset="https://static.rerun.io/droid_semantic_search/fd265e110c9d05c9fdd018c68dd83561aa836c52/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/droid_semantic_search/fd265e110c9d05c9fdd018c68dd83561aa836c52/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/droid_semantic_search/fd265e110c9d05c9fdd018c68dd83561aa836c52/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/droid_semantic_search/fd265e110c9d05c9fdd018c68dd83561aa836c52/1200w.png">
<img src="https://static.rerun.io/droid_semantic_search/fd265e110c9d05c9fdd018c68dd83561aa836c52/full.png" alt="DROID semantic frame search screenshot">
</picture>
This pulls frame embeddings from a Rerun dataset, indexes them in an external vector store, and resolves text queries back into deep-links that open the relevant frames in the viewer.
It ships two interchangeable backends — [LanceDB](https://lancedb.com) (default) and [Qdrant](https://qdrant.tech), both fully local on disk and selected with `--backend`.
They're worked examples of one pattern — embeddings out of Rerun, search in your vector database of choice, deep-links back into the viewer — that applies to any vector store.
## What it shows off
- Schema introspection and DataFusion content-filtered reads.
- The experimental video dataloader (`DataSource` / `Field` / `VideoFrameDecoder`) for streaming and decoding H.264 frames on demand.
- SigLIP-2 embeddings, with text and image features sharing one vector space.
- `segment_url` deep-links that focus the viewer on a specific frame.
## How it works
Three scripts:
1. `prepare_dataset.py` — registers a few DROID episodes to your local catalog as a dataset (using the episodes bundled in the repo, or downloading them from the Hugging Face Hub when those aren't present). This is the data the other two scripts read.
2. `ingest.py` — builds the index. For each camera it **auto-detects** the embedding source:
- **Read path:** if the dataset already has a `/camera/{role}/embedding` column (DROID registered with `--create-embeddings`), it reads those vectors directly via a DataFusion query — no model, no video decoding.
- **Compute path:** otherwise it streams the `VideoStream` through the dataloader, decodes frames, and embeds them with SigLIP-2.
Either way it writes `(segment_id, camera, timestamp_ms, vector)` rows into the selected vector store: LanceDB by default, or Qdrant with `--backend qdrant`.
3. `search.py` — embeds your example image or text prompt with the SigLIP-2 encoder, runs a vector search over that store (pass the same `--backend`), prints the ranked matches, and opens the best one in the viewer.
Both scripts reach the store only through a small `VectorStore` interface in `vector_store.py`, so picking a backend is one `--backend` choice and adding one is a single subclass.
Each backend keeps its index in its own directory (`./droid_lancedb` vs `./droid_qdrant`, override with `--db-path`), so the `--backend` you pass to `search.py` must match the one `ingest.py` wrote.
## Run the code
### 1. Install dependencies
This example has its own `uv` project, separate from the workspace `.venv`, because it needs the experimental `rerun-sdk[dataloader]` extras plus heavy ML deps (`transformers`, `lancedb`).
**Standalone** (sparse-checkout of just this directory, no local Rerun build):
```bash
uv sync --no-sources --no-dev
```
**Monorepo dev** (full repo checkout, editable local `rerun-sdk`):
```bash
cd examples/python/droid_semantic_search
RERUN_ALLOW_MISSING_BIN=1 uv sync
uv pip install ../../../rerun_py/rerun_dev_fixup
```
The second command installs the `.pth` shim that points `import rerun` (and the `rerun` CLI) at the in-repo editable source tree.
It's a separate `uv pip install` rather than a dev-group dependency because uv resolves all dependency groups unconditionally, so a path-only package in `pyproject.toml` would break the standalone `--no-sources` resolution above.
Then either `source .venv/bin/activate` or prefix subsequent commands with `uv run`.
### 2. Start a local Rerun server
This example reads data from a local open-source Rerun catalog server.
In a separate terminal, start one:
```bash
rerun server
```
This serves a catalog at `rerun+http://127.0.0.1:51234` — the default the scripts use.
Leave it running; the steps below connect to it.
### 3. Register a dataset
Registers a few DROID episodes to the catalog as `droid:sample`:
```bash
uv run python prepare_dataset.py
```
By default it auto-selects the source:
- **Monorepo checkout** — it uses the episodes bundled in the repo at `tests/assets/rrd/sample_5` (via git-LFS), so there's nothing to download. If those are un-pulled LFS pointers, run `git lfs install && git lfs pull` first.
- **Standalone checkout** — when the bundled episodes aren't present, it downloads a few from [`rerun/droid_sample`](https://huggingface.co/datasets/rerun/droid_sample) on the Hugging Face Hub into `./data`.
Useful flags:
- `--source bundled|huggingface` to force a source (default `auto`).
- `--num-episodes N` to register more (or fewer) episodes; `0` for all (the full Hub dataset is ~3.3 GB).
- `--dataset-name` to register under a different name (pass the same name to `ingest.py`/`search.py`).
- `--no-optimize` to register episodes as-is (see the video decode yield note below).
- `--catalog-url ""` to skip registration.
### 4. Build the search index
Index a handful of segments (the exterior camera works well for scene-level search):
```bash
uv run python ingest.py --num-segments 15 --cameras ext1
```
These sample episodes ship video only (no pre-computed embeddings), so `ingest.py` takes the **compute path**:
the first run downloads the SigLIP-2 model (a few hundred MB) and then decodes and embeds frames — the slow step.
Start with a small `--num-segments` to keep it quick; raise it once you've seen it work.
### 5. Search
Search by text and open the best hit in the viewer:
```bash
uv run python search.py "an open drawer full of tools" --top-k 5
```
Or search by example image instead of text (same vector space):
```bash
uv run python search.py --image ./query.jpg --top-k 5
```
Both scripts default to LanceDB; pass `--backend qdrant` to use [Qdrant](https://qdrant.tech) instead.
Give the same `--backend` to `ingest.py` and `search.py`, since each writes to its own directory:
```bash
uv run python ingest.py --backend qdrant --num-segments 15 --cameras ext1
uv run python search.py --backend qdrant "an open drawer full of tools" --top-k 5
```
Queries that discriminate well on DROID describe concrete, visible objects/scenes, e.g. `"a pink flower"`, `"a cardboard box"`, `"a white plastic bag"`, `"a robot arm over an empty table"`.
Run `ingest.py --help` and `search.py --help` for the full flag list — index multiple cameras, change the sampling rate, widen the time selection around a hit, and more.
## Scope and extension points
- **Text or image queries.** Search by a text prompt or, with `--image <path>`, by an example frame. SigLIP-2 puts text and image features in one space, so image-to-image search reuses the exact same index and ranking — only the query encoder differs.
- **Bring your own vector store.** Rerun supplies the embeddings (read from the dataset, or computed from platform-hosted video) and resolves matches back into viewer deep-links; search itself runs in an external store. This example ships two backends, LanceDB and Qdrant, behind the small `VectorStore` interface in `vector_store.py` — the same write-then-query flow drops onto any vector database, so adding a third is a single subclass.
## Notes and gotchas
- **SigLIP text tokenization.** SigLIP is trained with a fixed 64-token sequence and *must* be tokenized with `padding="max_length", max_length=64`. With dynamic padding the text embeddings are malformed and text→image retrieval collapses onto a single "hub" frame that wins every query.
- **Video decode yield.** DROID doesn't log the `VideoStream:is_keyframe` markers the decoder needs to seek, so without them only ~25 % of sampled frames decode.
`prepare_dataset.py` derives the markers up front (via `optimize`), which brings yield to ~100 %.
Optimized copies land in `./optimized`; the originals are untouched.
Skip it with `--no-optimize` if you'd rather register the raw episodes.
## Files
- `prepare_dataset.py` — register DROID sample episodes (bundled or downloaded) to the catalog.
- `ingest.py` — build the vector index from the dataset.
- `search.py` — query the index and open results in the viewer.
- `vector_store.py` — the `VectorStore` interface, with LanceDB and Qdrant backends.
- `embeddings.py` — SigLIP-2 helpers (adapted from the DROID loader's `embedding_util.py`).
examples/python/droid_semantic_search/embeddings.py
+121
View File
@@ -0,0 +1,121 @@
"""SigLIP-2 embedding helpers shared by `ingest.py` and `search.py`.
These are trimmed copies of the helpers in the DROID loader
(`dataplatform/examples/droid/droid-loader/src/droid_loader/embedding_util.py`),
with the `Timer` instrumentation removed so this example stays standalone and
doesn't pull in the `droid_loader` package. The model is the same one the
loader uses to populate `/camera/{role}/embedding`, so query embeddings land in
the same vector space as any pre-computed frame embeddings.
"""
from __future__ import annotations
import os
from pathlib import Path
from typing import TYPE_CHECKING, Any
import torch
# Disable HF `tokenizers` (Rust) parallelism *before* importing transformers. Otherwise,
# once the SigLIP tokenizer has been used and the process later forks (e.g. a DataLoader
# worker), tokenizers prints "the current process just got forked, after parallelism has
# already been used". We only tokenize tiny queries, so parallelism buys nothing here.
# `setdefault` lets a caller still override via the real environment variable.
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
from transformers import AutoModel, AutoProcessor # imported after TOKENIZERS_PARALLELISM is set (above)
if TYPE_CHECKING:
from PIL.Image import Image
# The concrete SigLIP-2 model/processor that `from_pretrained` returns.
EmbeddingModel = Any
EmbeddingProcessor = Any
# Dual image/text encoder; image and text features share one space, so text
# queries retrieve image frames directly. 768-dim, L2-normalized output.
EMBEDDING_MODEL = "google/siglip2-base-patch16-224"
def _resolve_device(device: str | torch.device | None) -> torch.device:
if device is not None:
return torch.device(device)
if torch.cuda.is_available():
return torch.device("cuda")
if torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
def load_embedding_model(
cache_dir: str | Path | None = None,
use_fast: bool = True,
) -> tuple[EmbeddingModel, EmbeddingProcessor]:
"""Load the SigLIP-2 model and its processor."""
print(f"Loading model '{EMBEDDING_MODEL}'")
model = AutoModel.from_pretrained(EMBEDDING_MODEL, cache_dir=cache_dir)
processor = AutoProcessor.from_pretrained(EMBEDDING_MODEL, cache_dir=cache_dir, use_fast=use_fast)
return model, processor
def get_text_embeddings(
text: str | list[str],
model: EmbeddingModel,
processor: EmbeddingProcessor,
device: str | torch.device | None = None,
) -> torch.Tensor:
"""Embed one or more strings into the SigLIP-2 space.
Returns an L2-normalized `[N, 768]` CPU tensor (one row per input string).
"""
if isinstance(text, str):
text = [text]
if not text:
raise ValueError("Input 'text' must be a non-empty string or list of strings.")
device = _resolve_device(device)
model = model.to(device)
model.eval()
# SigLIP is trained with a fixed 64-token sequence; it MUST be tokenized with
# padding="max_length" (max_length=64). With dynamic padding="True" the text
# embeddings are malformed and text->image retrieval collapses onto a hub image.
inputs = processor(text=text, return_tensors="pt", padding="max_length", max_length=64, truncation=True).to(device)
with torch.inference_mode():
# transformers 5.x: get_text_features returns the full encoder output, not a
# bare tensor — the embedding is its `pooler_output` (`[N, 768]`).
features = model.get_text_features(**inputs).pooler_output
normalized: torch.Tensor = features / features.norm(p=2, dim=-1, keepdim=True)
return normalized.cpu()
def compute_image_embeddings(
images: list[Image],
model: EmbeddingModel,
processor: EmbeddingProcessor,
device: str | torch.device | None = None,
batch_size: int = 64,
) -> torch.Tensor:
"""Embed a list of PIL images into the SigLIP-2 space.
Returns an L2-normalized `[len(images), 768]` CPU tensor.
"""
if not images:
raise ValueError("Input 'images' list cannot be empty.")
device = _resolve_device(device)
model = model.to(device)
model.eval()
all_embeddings: list[torch.Tensor] = []
with torch.inference_mode():
for start in range(0, len(images), batch_size):
batch = images[start : start + batch_size]
inputs = processor(images=batch, return_tensors="pt").to(device)
# transformers 5.x: get_image_features returns the full encoder output, not a
# bare tensor — the embedding is its `pooler_output` (`[N, 768]`).
features = model.get_image_features(**inputs).pooler_output
normalized = features / features.norm(p=2, dim=1, keepdim=True)
all_embeddings.append(normalized.cpu())
return torch.cat(all_embeddings, dim=0)
examples/python/droid_semantic_search/ingest.py
+330
View File
@@ -0,0 +1,330 @@
"""Build a local vector index of DROID camera frames.
For each requested camera the script auto-detects how to get embeddings:
* **Read path** — if the dataset already has a `/camera/{role}/embedding` column
(DROID registered with `--create-embeddings`), read it straight out of the
catalog with a DataFusion query. No video decoding, no model needed.
* **Compute path** — otherwise stream the H.264 `VideoStream` via the
experimental dataloader, decode frames, and embed them with SigLIP-2.
Either way we end up with a columnar `(segment_id, camera, timestamp_ms, vector)`
Arrow table, which we write to a local vector store (LanceDB or Qdrant, see
`--backend`) and index for ANN search.
Run inside the rerun SDK venv, e.g.:
pixi run uv run ../droid_semantic_search/ingest.py --num-segments 5 --cameras ext1
"""
from __future__ import annotations
import argparse
import itertools
from collections.abc import Iterator
from typing import Any
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
from vector_store import BACKENDS, DEFAULT_PATHS, open_store
from rerun.catalog import CatalogClient
# DROID camera roles, and the timeline everything is logged on.
ALL_CAMERAS = ("wrist", "ext1", "ext2")
TIMELINE = "real_time"
# DROID is H.264, GOP size 64 at ~15 fps (see the droid-loader). These knobs are only a
# fallback for episodes registered with `--no-optimize` (no keyframe markers): the decoder
# then seeks by a fixed window, so we use 2x the GOP to make sure each window holds a keyframe.
DROID_CODEC = "h264"
DROID_KEYFRAME_INTERVAL = 128
DROID_FPS_ESTIMATE = 15.0
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--catalog-url", default="rerun+http://127.0.0.1:51234", help="Rerun catalog URL")
parser.add_argument("--dataset", default="droid:sample", help="Dataset name in the catalog")
parser.add_argument("--token", default=None, help="Auth token (if the catalog requires one)")
parser.add_argument(
"--cameras",
default="ext1",
help="Comma-separated camera roles, or 'all'. Exterior cams (ext1/ext2) give better scene-level matches.",
)
parser.add_argument("--num-segments", type=int, default=10, help="Number of segments to index (0 for all)")
parser.add_argument(
"--backend",
choices=BACKENDS,
default="lance",
help="Local vector store to write the index to.",
)
parser.add_argument(
"--db-path",
default=None,
help="Directory for the local vector DB (default: ./droid_lancedb or ./droid_qdrant per backend).",
)
parser.add_argument("--table", default="droid_frames", help="Table/collection name")
parser.add_argument(
"--rate-hz",
type=float,
default=2.0,
help="Compute-path only: frames per second to sample from each segment.",
)
parser.add_argument(
"--fetch-batch",
type=int,
default=32,
help="Compute-path only: samples fetched per server round-trip.",
)
parser.add_argument(
"--num-workers",
type=int,
default=0,
help="Compute-path only: DataLoader workers for fetching/decoding.",
)
return parser.parse_args()
def resolve_cameras(arg: str) -> list[str]:
if arg.strip().lower() == "all":
return list(ALL_CAMERAS)
roles = [c.strip() for c in arg.split(",") if c.strip()]
unknown = [r for r in roles if r not in ALL_CAMERAS]
if unknown:
raise SystemExit(f"Unknown camera role(s) {unknown}; expected a subset of {ALL_CAMERAS} or 'all'.")
return roles
def cameras_with_embeddings(schema: object, cameras: list[str]) -> set[str]:
"""Return the subset of *cameras* that already have an embedding entity in the dataset."""
present = {p.strip("/") for p in schema.entity_paths()} # type: ignore[attr-defined]
return {role for role in cameras if f"camera/{role}/embedding" in present}
def _is_list_type(t: pa.DataType) -> bool:
return bool(pa.types.is_list(t) or pa.types.is_large_list(t) or pa.types.is_fixed_size_list(t))
def _vector_dim(vectors: pa.Array) -> int:
"""Embedding dimensionality of a (variable- or fixed-size) list array."""
if pa.types.is_fixed_size_list(vectors.type):
return int(vectors.type.list_size)
return int(pc.max(pc.list_value_length(vectors)).as_py())
def _embedding_table(role: str, segment_ids: pa.Array, timestamps_ms: pa.Array, vectors: pa.Array) -> pa.Table:
"""Assemble the four index columns into one Arrow table.
Both ingest paths funnel through here, so they share a single schema and
`pa.concat_tables` can stitch their results together with no per-row work.
"""
return pa.table(
{
"segment_id": segment_ids,
"camera": pa.array([role] * len(segment_ids), pa.string()),
"timestamp_ms": timestamps_ms,
"vector": vectors,
},
)
def _find_embedding_column(schema: pa.Schema, role: str) -> str:
"""Locate the list-typed embedding column for *role* in a query result schema.
The DROID loader logs embeddings via `rr.AnyValues(embeddings=...)`, so the
exact column name (e.g. `/camera/ext1/embedding:embeddings`) is derived by
the platform — discover it rather than hardcoding.
"""
candidates: list[str] = [f.name for f in schema if _is_list_type(f.type) and "embedding" in f.name.lower()]
if not candidates:
raise RuntimeError(f"No list-typed embedding column for camera '{role}'. Columns: {schema.names}")
for name in candidates:
if role in name:
return name
return candidates[0]
def read_embedding_table(dataset: object, segments: list[str], role: str) -> pa.Table | None:
"""Read pre-computed embeddings for *role* out of the catalog.
The query result is already columnar Arrow, so we stay in Arrow the whole
way — filter out missing rows and normalize the column types without ever
materializing Python row objects.
"""
view = dataset.filter_segments(segments).filter_contents( # type: ignore[attr-defined]
[f"/camera/{role}/embedding", f"/camera/{role}/embedding/**"],
)
table = view.reader(index=TIMELINE).to_arrow_table()
emb_col = _find_embedding_column(table.schema, role)
if "rerun_segment_id" not in table.schema.names or TIMELINE not in table.schema.names:
raise RuntimeError(f"Expected 'rerun_segment_id' and '{TIMELINE}' columns, got {table.schema.names}")
# Columnar equivalent of the per-row `if vec is None or seg is None: continue`.
keep = pc.and_(pc.is_valid(table.column(emb_col)), pc.is_valid(table.column("rerun_segment_id")))
table = table.filter(keep)
if table.num_rows == 0:
print(f" [{role}] no pre-computed embeddings")
return None
vectors = table.column(emb_col).combine_chunks()
vectors = vectors.cast(pa.list_(pa.float32(), _vector_dim(vectors)))
segment_ids = table.column("rerun_segment_id").cast(pa.string())
# DROID's index timeline is nanosecond timestamps; LanceDB just wants an int.
timestamps_ms = table.column(TIMELINE).cast(pa.timestamp("ms")).cast(pa.int64())
out = _embedding_table(role, segment_ids, timestamps_ms, vectors)
print(f" [{role}] read {out.num_rows} pre-computed embeddings")
return out
def _identity_collate(batch: list[Any]) -> list[Any]:
"""Collate that leaves the list of per-sample dicts untouched (picklable for workers)."""
return batch
def compute_embedding_table(
dataset: object,
segments: list[str],
role: str,
*,
rate_hz: float,
fetch_batch: int,
num_workers: int,
) -> pa.Table | None:
"""Decode frames for *role* and embed them with SigLIP-2."""
# Heavy / optional deps are imported lazily so a pure read-path run stays light.
from embeddings import compute_image_embeddings, load_embedding_model
from PIL import Image
from torch.utils.data import DataLoader
from tqdm import tqdm
from rerun.experimental.dataloader import (
DataSource,
Field,
FixedRateSampling,
RerunMapDataset,
VideoFrameDecoder,
)
field_name = f"img_{role}"
source = DataSource(dataset, segments=segments) # type: ignore[arg-type]
fields = {
field_name: Field(
f"/camera/{role}:VideoStream:sample",
decode=VideoFrameDecoder(
codec=DROID_CODEC,
keyframe_interval=DROID_KEYFRAME_INTERVAL,
fps_estimate=DROID_FPS_ESTIMATE,
),
),
}
ds = RerunMapDataset(
source=source,
index=TIMELINE,
fields=fields,
timeline_sampling=FixedRateSampling(rate_hz=rate_hz),
)
total = len(ds)
print(f" [{role}] decoding ~{total} frames at {rate_hz} Hz …")
# The DataLoader earns its keep on the *fetch* side: `batch_size` batches the
# catalog round-trips, and `num_workers > 0` fans the CPU-bound video decode
# across worker processes. `shuffle=False` keeps indices in 0..N-1 order, which
# the running counter in `decoded_frames` relies on for the (segment, timestamp)
# pairing below.
loader = DataLoader(
ds,
batch_size=fetch_batch,
shuffle=False,
num_workers=num_workers,
collate_fn=_identity_collate,
)
def decoded_frames(pbar: tqdm[Any]) -> Iterator[tuple[Image.Image, str, int]]:
# The loader visits indices 0..N-1 in order, so a running counter pairs each
# decoded frame back to its (segment, timestamp) via `global_to_local`. The
# progress bar advances once per sample pulled from the loader (the slow,
# video-decoding step), including the ones we skip below.
global_idx = 0
for batch in loader:
for sample in batch:
tensor = sample[field_name]
seg_meta, idx_val = ds.sample_index.global_to_local(global_idx)
global_idx += 1
pbar.update(1)
if tensor is None: # target preceded the first keyframe; skip
continue
ts_ms = int(np.datetime64(idx_val).astype("datetime64[ms]").astype(np.int64)) # type: ignore[arg-type]
rgb = tensor.permute(1, 2, 0).cpu().numpy() # [C,H,W] uint8 -> [H,W,C]
yield Image.fromarray(rgb), seg_meta.segment_id, ts_ms
# Embed the decoded stream chunk-by-chunk so peak memory stays at ~embed_batch
# frames rather than the whole role. Each chunk becomes one small Arrow table;
# `pa.concat_tables` stitches them at the end with no per-row work.
embed_batch = 64
model, processor = load_embedding_model()
chunks: list[pa.Table] = []
with tqdm(total=total, desc=f"[{role}] decode+embed", unit="frame") as pbar:
for chunk in itertools.batched(decoded_frames(pbar), embed_batch): # type: ignore[attr-defined, unused-ignore]
frames = [frame for frame, _, _ in chunk]
segs = [seg for _, seg, _ in chunk]
timestamps = [ts_ms for _, _, ts_ms in chunk]
vectors = compute_image_embeddings(frames, model, processor, batch_size=embed_batch).numpy()
_, dim = vectors.shape
vector_col = pa.FixedSizeListArray.from_arrays(pa.array(vectors.reshape(-1), pa.float32()), dim)
chunks.append(
_embedding_table(role, pa.array(segs, pa.string()), pa.array(timestamps, pa.int64()), vector_col),
)
if not chunks:
print(f" [{role}] no frames decoded")
return None
out = pa.concat_tables(chunks)
print(f" [{role}] computed {out.num_rows} embeddings")
return out
def main() -> None:
args = parse_args()
cameras = resolve_cameras(args.cameras)
client = CatalogClient(args.catalog_url, token=args.token)
dataset = client.get_dataset(args.dataset)
all_segments = dataset.segment_ids()
segments = all_segments if args.num_segments == 0 else all_segments[: args.num_segments]
if not segments:
raise SystemExit(f"Dataset '{args.dataset}' has no segments.")
have_emb = cameras_with_embeddings(dataset.schema(), cameras)
print(f"Indexing {len(segments)} segment(s); cameras={cameras}; pre-computed embeddings for {sorted(have_emb)}")
tables: list[pa.Table] = []
for role in cameras:
if role in have_emb:
table = read_embedding_table(dataset, segments, role)
else:
table = compute_embedding_table(
dataset,
segments,
role,
rate_hz=args.rate_hz,
fetch_batch=args.fetch_batch,
num_workers=args.num_workers,
)
if table is not None:
tables.append(table)
if not tables:
raise SystemExit("No embeddings produced; nothing to index.")
db_path = args.db_path or DEFAULT_PATHS[args.backend]
open_store(args.backend, db_path, args.table).write(pa.concat_tables(tables))
if __name__ == "__main__":
main()
examples/python/droid_semantic_search/prepare_dataset.py
+222
View File
@@ -0,0 +1,222 @@
"""Register a few DROID episodes to a local Rerun catalog so the rest of the example has data to index.
Two sources, auto-selected (override with `--source`):
* **Bundled** — the `tests/assets/rrd/sample_5` episodes shipped in the Rerun repo (via git-LFS).
Used automatically in a monorepo checkout: no download, works offline.
* **Hugging Face** — a few episodes from the [`rerun/droid_sample`](https://huggingface.co/datasets/rerun/droid_sample)
dataset, downloaded into `./data`. Used when the bundled episodes aren't available
(e.g. a standalone sparse-checkout of just this example).
Either way the episodes are registered to the catalog as a dataset (default name `droid:sample`).
They carry H.264 `VideoStream`s but no pre-computed embeddings, so `ingest.py` will take its
(slower) compute path: decode frames and embed them with SigLIP-2.
Episodes are optimized first to derive keyframe markers, so the compute path can decode
frames (DROID doesn't log the markers the decoder needs). Pass `--no-optimize` to skip.
Run inside the rerun SDK venv, with a `rerun server` running in another terminal, e.g.:
uv run python prepare_dataset.py
"""
from __future__ import annotations
import argparse
from dataclasses import replace
from pathlib import Path
from tqdm import tqdm
import rerun as rr
from rerun.experimental import OptimizationProfile, RrdReader
# `tests/assets/rrd/sample_5`, relative to this file at `examples/python/droid_semantic_search/`.
BUNDLED_SAMPLE_DIR = Path(__file__).resolve().parents[3] / "tests" / "assets" / "rrd" / "sample_5"
DEFAULT_REPO_ID = "rerun/droid_sample"
DEFAULT_OUTPUT_DIR = Path(__file__).resolve().parent / "data"
DEFAULT_OPTIMIZED_DIR = Path(__file__).resolve().parent / "optimized"
DEFAULT_DATASET_NAME = "droid:sample"
DEFAULT_CATALOG_URL = "rerun+http://127.0.0.1:51234"
_LFS_POINTER_PREFIX = b"version https://git-lfs.github.com/spec/v1"
def _is_lfs_pointer(path: Path) -> bool:
"""True if *path* is an un-pulled git-LFS pointer file rather than the real RRD."""
with path.open("rb") as f:
return f.read(len(_LFS_POINTER_PREFIX)) == _LFS_POINTER_PREFIX
def bundled_episode_paths() -> list[Path]:
"""Return the bundled sample_5 RRDs (sorted), or an empty list if they're not available.
Raises if the directory exists but the files are un-pulled git-LFS pointers — that's a
recoverable monorepo setup issue worth surfacing rather than silently working around.
"""
if not BUNDLED_SAMPLE_DIR.is_dir():
return []
rrds = sorted(BUNDLED_SAMPLE_DIR.glob("*.rrd"))
if not rrds:
return []
if any(_is_lfs_pointer(p) for p in rrds):
raise SystemExit(
f"The bundled DROID episodes in {BUNDLED_SAMPLE_DIR} are un-pulled git-LFS pointers.\n"
"Fetch them with `git lfs install && git lfs pull`, or pass `--source huggingface` to download instead.",
)
return rrds
def download_episodes(repo_id: str, num_episodes: int, dest: Path) -> list[Path]:
"""Download the first *num_episodes* (0 for all) `.rrd` files from *repo_id* into *dest*.
`snapshot_download` renders its own per-file progress bars, so the user sees the download advance.
"""
from huggingface_hub import HfApi, snapshot_download
files = sorted(f for f in HfApi().list_repo_files(repo_id, repo_type="dataset") if f.endswith(".rrd"))
if not files:
raise SystemExit(f"No .rrd files found in '{repo_id}'.")
files = files if num_episodes == 0 else files[:num_episodes]
print(f"Downloading {len(files)} episode(s) from '{repo_id}' to {dest} …")
local_dir = snapshot_download(repo_id=repo_id, repo_type="dataset", allow_patterns=files, local_dir=dest)
return [Path(local_dir) / f for f in files]
def resolve_episodes(source: str, *, num_episodes: int, repo_id: str, output_dir: Path) -> list[Path]:
"""Pick the episode RRDs to register, per the requested *source* (`auto`/`bundled`/`huggingface`)."""
if source in ("auto", "bundled"):
bundled = bundled_episode_paths()
if bundled:
paths = bundled if num_episodes == 0 else bundled[:num_episodes]
print(f"Using {len(paths)} bundled episode(s) from {BUNDLED_SAMPLE_DIR}")
return paths
if source == "bundled":
raise SystemExit(f"No bundled episodes found at {BUNDLED_SAMPLE_DIR}; pass `--source huggingface`.")
print(f"Bundled episodes not found at {BUNDLED_SAMPLE_DIR}; downloading from the Hub instead.")
return download_episodes(repo_id, num_episodes, output_dir)
def optimize_episodes(rrd_paths: list[Path], dest_dir: Path) -> list[Path]:
"""Derive `VideoStream:is_keyframe` markers (DROID doesn't log them) so the decoder can seek.
Writes a fixed copy of each episode to *dest_dir* and returns the new paths. IDs are
preserved, so catalog segment IDs and `segment_url` links are unchanged.
"""
dest_dir.mkdir(parents=True, exist_ok=True)
profile = replace(OptimizationProfile.OBJECT_STORE, fix_keyframe=True)
optimized: list[Path] = []
for src in tqdm(rrd_paths, desc="Optimizing", unit="episode"):
reader = RrdReader(src)
recordings = reader.recordings()
if len(recordings) != 1:
raise SystemExit(f"Expected one recording in {src}, found {len(recordings)}.")
entry = recordings[0]
store = reader.stream(store=entry).collect(optimize=profile)
dst = dest_dir / src.name
store.write_rrd(dst, application_id=entry.application_id, recording_id=entry.recording_id)
optimized.append(dst)
return optimized
def register_to_catalog(rrd_paths: list[Path], *, catalog_url: str, dataset_name: str) -> None:
"""Register per-episode RRDs to a catalog server instance.
Uses absolute `file://` URIs so the catalog can read the RRDs directly from the local filesystem.
Streams `iter_results()` so a progress bar advances as each segment finishes, rather than blocking
silently on `wait()`.
"""
print(f"\nRegistering {len(rrd_paths)} episode(s) to {catalog_url} as dataset '{dataset_name}' …")
client = rr.catalog.CatalogClient(catalog_url)
dataset = client.create_dataset(dataset_name, exist_ok=True)
uris = [f"file://{p.resolve()}" for p in rrd_paths]
on_duplicate = rr.catalog.OnDuplicateSegmentLayer(rr.catalog.OnDuplicateSegmentLayer.REPLACE)
handle = dataset.register(uris, on_duplicate=on_duplicate)
failures: list[str] = []
for result in tqdm(handle.iter_results(), total=len(uris), desc="Registering", unit="segment"):
if result.is_error:
failures.append(f"{result.uri}: {result.error}")
if failures:
joined = "\n ".join(failures)
raise SystemExit(f"Failed to register {len(failures)} of {len(uris)} episode(s):\n {joined}")
print(" registration done")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument(
"--source",
choices=("auto", "bundled", "huggingface"),
default="auto",
help="Where to get episodes: 'bundled' (in-repo sample_5), 'huggingface' (download), "
"or 'auto' (bundled if available, else download). Default: auto.",
)
parser.add_argument(
"--repo-id",
default=DEFAULT_REPO_ID,
help=f"Hugging Face dataset repo id, for the download path (default: {DEFAULT_REPO_ID}).",
)
parser.add_argument(
"--num-episodes",
type=int,
default=5,
help="Number of episodes to register (0 for all). The full Hub dataset is ~3.3 GB.",
)
parser.add_argument(
"--output-dir",
type=Path,
default=DEFAULT_OUTPUT_DIR,
help=f"Directory to download episode RRDs into, for the download path (default: {DEFAULT_OUTPUT_DIR}).",
)
parser.add_argument(
"--optimize",
default=True,
action=argparse.BooleanOptionalAction,
help="Derive keyframe markers before registering, so ingest.py can decode frames "
"(~100%% yield vs ~25%%). Use --no-optimize to register episodes as-is.",
)
parser.add_argument(
"--optimized-dir",
type=Path,
default=DEFAULT_OPTIMIZED_DIR,
help=f"Directory to write optimized episode RRDs into (default: {DEFAULT_OPTIMIZED_DIR}).",
)
parser.add_argument(
"--catalog-url",
default=DEFAULT_CATALOG_URL,
help="Rerun catalog URL to register episodes with. Pass an empty string to skip registration.",
)
parser.add_argument(
"--dataset-name",
default=DEFAULT_DATASET_NAME,
help=f"Name of the dataset to create/use in the catalog (default: {DEFAULT_DATASET_NAME}).",
)
args = parser.parse_args()
rrd_paths = resolve_episodes(
args.source,
num_episodes=args.num_episodes,
repo_id=args.repo_id,
output_dir=args.output_dir,
)
if args.optimize:
print(f"\nOptimizing {len(rrd_paths)} episode(s) into {args.optimized_dir} …")
rrd_paths = optimize_episodes(rrd_paths, args.optimized_dir)
if args.catalog_url:
register_to_catalog(rrd_paths, catalog_url=args.catalog_url, dataset_name=args.dataset_name)
else:
print(f"Skipping registration (empty --catalog-url). Episodes: {[str(p) for p in rrd_paths]}")
if __name__ == "__main__":
main()
examples/python/droid_semantic_search/pyproject.toml
+66
View File
@@ -0,0 +1,66 @@
[project]
name = "droid_semantic_search"
version = "0.1.0"
readme = "README.md"
requires-python = ">=3.12,<3.13"
dependencies = [
# `catalog` brings datafusion + pandas; `dataloader` brings torch, torchvision, av, pillow.
"rerun-sdk[catalog,dataloader]",
"lancedb", # default local vector store + ANN index (--backend lance)
"qdrant-client", # alternative local vector store (--backend qdrant)
"transformers>=5.13.0", # SigLIP-2 model + processor
"pyarrow<24", # building the vector-index table
"huggingface-hub", # prepare_dataset.py: download DROID sample episodes from the Hub
"tqdm", # progress bars for download/register/ingest
]
[dependency-groups]
dev = ["mypy==1.19.1", "types-tqdm"]
[tool.rerun-example]
# Picked up by scripts/ci/isolated_examples.py and the `py-lint-isolated-examples` pixi task.
isolated = true
[tool.uv]
# The example is flat scripts, not a wheel — skip project build, just sync deps.
package = false
# Default `uv sync` uses the in-repo editable rerun-sdk (monorepo dev mode).
# After syncing, also run `uv pip install ../../../rerun_py/rerun_dev_fixup` to
# install the .pth shim that makes `import rerun` resolve to the editable source tree.
#
# Standalone users (e.g. sparse-checkout of just this example) run instead:
# uv sync --no-sources --no-dev
# That ignores the path source below and resolves `rerun-sdk` from PyPI.
# rerun-dev-fixup is intentionally absent from this file: uv 0.7.x resolves all
# dependency groups and extras unconditionally, so any path-only package here would
# block standalone `--no-sources` resolution.
[tool.uv.sources]
rerun-sdk = { path = "../../../rerun_py", editable = true }
# Merged onto the shared base at `../_isolated/mypy.ini` by
# scripts/ci/isolated_examples.py — list the untyped third-party libs this
# example actually imports.
[[tool.mypy.overrides]]
module = [
"lancedb.*",
"qdrant_client.*",
"torch.*",
"torchvision.*",
"av.*",
"PIL.*",
"huggingface_hub.*",
# pyarrow is pinned <24 (24.0.0 segfaults rerun on import); 23.x ships no py.typed,
# so all of pyarrow — including pyarrow.compute — is untyped to mypy.
"pyarrow.*",
]
ignore_missing_imports = true
# transformers ships a `py.typed` marker but with incomplete stubs, so type-checking its
# internals yields false positives we can't fix from here: model methods are wrapped in
# `_Wrapped` descriptors that mistype `self` (`AutoModel has no attribute "to"`,
# `AutoProcessor not callable`, etc.). Skip following it — we only need our own usage typed.
[[tool.mypy.overrides]]
module = ["transformers.*"]
follow_imports = "skip"
ignore_missing_imports = true
examples/python/droid_semantic_search/search.py
+128
View File
@@ -0,0 +1,128 @@
"""Query the DROID frame index with a text prompt or example image and open the best hit in Rerun.
Embeds the query with the SigLIP-2 text *or* image encoder (both share one
vector space, the same one the indexed frame embeddings live in), runs a cosine
nearest-neighbor search over the local vector store (LanceDB or Qdrant, see
`--backend`), prints the ranked matches, and mints a `segment_url` deep-link
that opens the top result focused on that frame in the Rerun viewer.
Run inside the rerun SDK venv, e.g.:
pixi run uv run ../droid_semantic_search/search.py "a robot gripper reaching for a cup"
pixi run uv run ../droid_semantic_search/search.py --image ./query.jpg
"""
from __future__ import annotations
import argparse
import webbrowser
from datetime import datetime, timedelta, timezone
from typing import cast
from embeddings import (
EmbeddingModel,
EmbeddingProcessor,
compute_image_embeddings,
get_text_embeddings,
load_embedding_model,
)
from vector_store import BACKENDS, DEFAULT_PATHS, open_store
import rerun as rr
from rerun.catalog import CatalogClient
TIMELINE = "real_time"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("query", nargs="?", help="Text prompt to search for")
parser.add_argument(
"--image", help="Path to an image to search by (image-to-image); mutually exclusive with the text query"
)
parser.add_argument("--catalog-url", default="rerun+http://127.0.0.1:51234", help="Rerun catalog URL")
parser.add_argument("--dataset", default="droid:sample", help="Dataset name (used to mint the viewer link)")
parser.add_argument(
"--login", action="store_true", help="Authenticate with the catalog via rr.login() before connecting"
)
parser.add_argument(
"--backend",
choices=BACKENDS,
default="lance",
help="Local vector store to query (must match what ingest.py wrote).",
)
parser.add_argument(
"--db-path",
default=None,
help="Directory of the local vector DB (default: ./droid_lancedb or ./droid_qdrant per backend).",
)
parser.add_argument("--table", default="droid_frames", help="Table/collection name")
parser.add_argument("--top-k", type=int, default=5, help="Number of matches to return")
parser.add_argument("--window-secs", type=float, default=2.0, help="Time window around the matched frame to show")
parser.add_argument(
"--open",
default=True,
action=argparse.BooleanOptionalAction,
help="Open the top hit in the viewer",
)
return parser.parse_args()
def embed_image_query(path: str, model: EmbeddingModel, processor: EmbeddingProcessor) -> list[float]:
"""Embed an image file for image-to-image search.
Wired to the `--image` flag: text and image features share one SigLIP-2
vector space, so a query image retrieves frames the same way a text prompt does.
"""
from PIL import Image
image = Image.open(path).convert("RGB")
vector = compute_image_embeddings([image], model, processor).numpy()[0]
return cast("list[float]", vector.tolist())
def viewer_url(dataset: object, segment_id: str, timestamp_ms: int, window_secs: float) -> str:
ts = datetime.fromtimestamp(timestamp_ms / 1000, tz=timezone.utc)
half = timedelta(seconds=window_secs / 2)
return dataset.segment_url(segment_id, timeline=TIMELINE, start=ts - half, end=ts + half) # type: ignore[attr-defined, no-any-return]
def main() -> None:
args = parse_args()
if bool(args.query) == bool(args.image):
raise SystemExit("Provide exactly one of: a text query or --image <path>.")
model, processor = load_embedding_model()
if args.image:
query_vec = embed_image_query(args.image, model, processor)
query_label = f"image {args.image}"
else:
query_vec = get_text_embeddings(args.query, model, processor).numpy()[0].tolist()
query_label = args.query
db_path = args.db_path or DEFAULT_PATHS[args.backend]
hits = open_store(args.backend, db_path, args.table).search(query_vec, args.top_k)
if not hits:
raise SystemExit("No matches found — is the index populated? Run ingest.py first.")
print(f'\nTop {len(hits)} matches for: "{query_label}"\n')
print(f"{'#':>2} {'sim':>5} {'camera':<6} {'timestamp (UTC)':<24} segment")
for rank, hit in enumerate(hits, start=1):
ts_iso = datetime.fromtimestamp(hit["timestamp_ms"] / 1000, tz=timezone.utc).isoformat(timespec="milliseconds")
print(f"{rank:>2} {hit['similarity']:>5.3f} {hit['camera']:<6} {ts_iso:<24} {hit['segment_id']}")
# Mint a deep-link into the viewer for the best match.
if args.login:
rr.login()
client = CatalogClient(args.catalog_url)
dataset = client.get_dataset(args.dataset)
best = hits[0]
url = viewer_url(dataset, best["segment_id"], best["timestamp_ms"], args.window_secs)
print(f"\nTop hit in viewer:\n{url}")
if args.open:
webbrowser.open(url)
if __name__ == "__main__":
main()
examples/python/droid_semantic_search/uv.lock
Generated
+1298
View File
File diff suppressed because it is too large Load Diff
examples/python/droid_semantic_search/vector_store.py
+161
View File
@@ -0,0 +1,161 @@
"""Pluggable local vector-store backends for the DROID frame index.
`ingest.py` and `search.py` talk to a vector store only through the small
`VectorStore` interface here, so the same code path works whether you pick
LanceDB (`--backend lance`) or Qdrant (`--backend qdrant`). Both run fully
locally on disk — no server to stand up — so the example stays one-command.
Adding another store (Pinecone, pgvector, Milvus, …) is a third subclass.
The index is written from the columnar `(segment_id, camera, timestamp_ms,
vector)` Arrow table that `ingest.py` assembles. Searches return hits carrying
those three metadata fields plus a `similarity` in `[-1, 1]` (cosine),
normalized here so callers never have to know which backend's distance/score
convention is in play.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
import pyarrow as pa
BACKENDS = ("lance", "qdrant")
# Per-backend default on-disk location, used when `--db-path` is omitted.
DEFAULT_PATHS = {
"lance": "./droid_lancedb",
"qdrant": "./droid_qdrant",
}
class VectorStore(ABC):
"""A local on-disk vector index over `(segment_id, camera, timestamp_ms, vector)` rows."""
@abstractmethod
def write(self, table: pa.Table) -> None:
"""(Over)write the index from *table*, replacing any existing table/collection.
*table* has columns `segment_id` (string), `camera` (string),
`timestamp_ms` (int64), and `vector` (fixed-size list of float32).
"""
@abstractmethod
def search(self, vector: list[float], top_k: int) -> list[dict[str, Any]]:
"""Cosine nearest-neighbor search.
Returns up to *top_k* hit dicts, each with `segment_id`, `camera`,
`timestamp_ms`, and a cosine `similarity` (higher is closer).
"""
def open_store(backend: str, path: str, table: str) -> VectorStore:
if backend == "lance":
return LanceStore(path, table)
if backend == "qdrant":
return QdrantStore(path, table)
raise ValueError(f"Unknown backend '{backend}'; expected one of {BACKENDS}.")
class LanceStore(VectorStore):
"""LanceDB backend. Returns cosine *distance*, which we flip to similarity."""
def __init__(self, path: str, table: str) -> None:
self._path = path
self._table = table
def write(self, table: pa.Table) -> None:
import lancedb
dim = table.schema.field("vector").type.list_size
db = lancedb.connect(self._path)
tbl = db.create_table(self._table, data=table, mode="overwrite")
print(f"Wrote {table.num_rows} rows ({dim}-dim) to LanceDB table '{self._table}' in {self._path}")
# An ANN index needs enough rows to train; small demo tables fall back to
# brute-force search, which is exact and plenty fast at this scale.
try:
tbl.create_index(metric="cosine", vector_column_name="vector")
print("Built ANN index (cosine).")
except Exception as exc:
print(f"Skipped ANN index ({exc}); brute-force cosine search will be used.")
def search(self, vector: list[float], top_k: int) -> list[dict[str, Any]]:
import lancedb
tbl = lancedb.connect(self._path).open_table(self._table)
hits = tbl.search(vector).metric("cosine").limit(top_k).to_list()
return [
{
"segment_id": h["segment_id"],
"camera": h["camera"],
"timestamp_ms": h["timestamp_ms"],
"similarity": 1.0 - float(h["_distance"]), # cosine distance -> similarity
}
for h in hits
]
class QdrantStore(VectorStore):
"""Qdrant backend in local (embedded) mode. Returns cosine *score* directly."""
def __init__(self, path: str, collection: str) -> None:
self._path = path
self._collection = collection
def write(self, table: pa.Table) -> None:
from qdrant_client import QdrantClient, models
dim = table.schema.field("vector").type.list_size
segment_ids = table.column("segment_id").to_pylist()
cameras = table.column("camera").to_pylist()
timestamps_ms = table.column("timestamp_ms").to_pylist()
vectors = table.column("vector").to_pylist()
client = QdrantClient(path=self._path)
# Mirror Lance's overwrite semantics: drop any prior collection first.
if client.collection_exists(self._collection):
client.delete_collection(self._collection)
client.create_collection(
collection_name=self._collection,
vectors_config=models.VectorParams(size=dim, distance=models.Distance.COSINE),
)
points = [
models.PointStruct(
id=i,
vector=vectors[i],
payload={
"segment_id": segment_ids[i],
"camera": cameras[i],
"timestamp_ms": timestamps_ms[i],
},
)
for i in range(table.num_rows)
]
client.upsert(collection_name=self._collection, points=points)
print(f"Wrote {table.num_rows} rows ({dim}-dim) to Qdrant collection '{self._collection}' in {self._path}")
def search(self, vector: list[float], top_k: int) -> list[dict[str, Any]]:
from qdrant_client import QdrantClient
client = QdrantClient(path=self._path)
result = client.query_points(
collection_name=self._collection,
query=vector,
limit=top_k,
with_payload=True,
)
return [
{
"segment_id": payload["segment_id"],
"camera": payload["camera"],
"timestamp_ms": payload["timestamp_ms"],
"similarity": float(point.score), # Qdrant cosine score is already a similarity
}
for point in result.points
if (payload := point.payload) is not None
]
examples/python/drone_lidar/.gitignore
+1
View File
@@ -0,0 +1 @@
/dataset
examples/python/drone_lidar/README.md
+40
View File
@@ -0,0 +1,40 @@
<!--[metadata]
title = "Drone LiDAR"
tags = ["3D", "drone", "Lidar"]
description = "Display drone-based LiDAR data"
-->
<picture>
<img src="https://static.rerun.io/drone_lidar/95c49d78abc01513d344c06e2d9a0c8b84376a0d/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/drone_lidar/95c49d78abc01513d344c06e2d9a0c8b84376a0d/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/drone_lidar/95c49d78abc01513d344c06e2d9a0c8b84376a0d/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/drone_lidar/95c49d78abc01513d344c06e2d9a0c8b84376a0d/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/drone_lidar/95c49d78abc01513d344c06e2d9a0c8b84376a0d/1200w.png">
</picture>
## Background
This example displays drone-based indoor LiDAR data loaded from a [`.las`](https://en.wikipedia.org/wiki/LAS_file_format) file. This dataset contains 18.7M points, acquired at 4013 distinct time points (~4650 points per time point). The point data is loaded using the [laspy](https://laspy.readthedocs.io/en/latest/) Python package, and then sent in one go to the viewer thanks to the [`rr.send_columns()`](https://ref.rerun.io/docs/python/0.18.2/common/columnar_api/#rerun.send_columns) API and its `.partition()` helper. Together, these APIs enable associating subgroups of points with each of their corresponding, non-repeating timestamps.
[Flyability](https://www.flyability.com) kindly provided the data for this example.
## Running
Install the example package:
```bash
pip install -e examples/python/drone_lidar
```
To experiment with the provided example, simply execute the main Python script:
```bash
python -m drone_lidar
```
If you wish to customize it, explore additional features, or save it, use the CLI with the `--help` option for guidance:
```bash
python -m drone_lidar --help
```
examples/python/drone_lidar/drone_lidar.py
+180
View File
@@ -0,0 +1,180 @@
from __future__ import annotations
import io
import typing
import zipfile
from argparse import ArgumentParser
from pathlib import Path
import laspy
import numpy as np
import numpy.typing as npt
import requests
from tqdm import tqdm
import rerun as rr
import rerun.blueprint as rrb
DATASET_DIR = Path(__file__).parent / "dataset"
if not DATASET_DIR.exists():
DATASET_DIR.mkdir()
LIDAR_DATA_FILE = DATASET_DIR / "livemap.las"
TRAJECTORY_DATA_FILE = DATASET_DIR / "livetraj.csv"
LIDAR_DATA_URL = "https://storage.googleapis.com/rerun-example-datasets/flyability/basement/livemap.las.zip"
TRAJECTORY_DATA_URL = "https://storage.googleapis.com/rerun-example-datasets/flyability/basement/livetraj.csv"
def download_with_progress(url: str, what: str) -> io.BytesIO:
"""Download a file with a tqdm progress bar."""
chunk_size = 1024 * 1024
resp = requests.get(url, stream=True)
total_size = int(resp.headers.get("content-length", 0))
with tqdm(
desc=f"Downloading {what}",
total=total_size,
unit="iB",
unit_scale=True,
unit_divisor=1024,
) as progress:
download_file = io.BytesIO()
for data in resp.iter_content(chunk_size):
download_file.write(data)
progress.update(len(data))
download_file.seek(0)
return download_file
def unzip_file_from_archive_with_progress(zip_data: typing.BinaryIO, file_name: str, dest_dir: Path) -> None:
"""Unzip the file named `file_name` from the zip archive contained in `zip_data` to `dest_dir`."""
with zipfile.ZipFile(zip_data, "r") as zip_ref:
file_info = zip_ref.getinfo(file_name)
total_size = file_info.file_size
with (
tqdm(
total=total_size,
desc=f"Extracting file {file_name}",
unit="iB",
unit_scale=True,
unit_divisor=1024,
) as progress,
zip_ref.open(file_name) as source,
open(dest_dir / file_name, "wb") as target,
):
for chunk in iter(lambda: source.read(1024 * 1024), b""):
target.write(chunk)
progress.update(len(chunk))
def download_dataset() -> None:
if not LIDAR_DATA_FILE.exists():
unzip_file_from_archive_with_progress(
download_with_progress(LIDAR_DATA_URL, LIDAR_DATA_FILE.name),
LIDAR_DATA_FILE.name,
LIDAR_DATA_FILE.parent,
)
if not TRAJECTORY_DATA_FILE.exists():
TRAJECTORY_DATA_FILE.write_bytes(
download_with_progress(TRAJECTORY_DATA_URL, TRAJECTORY_DATA_FILE.name).getvalue(),
)
# TODO(#7333): this utility should be included in the Rerun SDK
def compute_partitions(
times: npt.NDArray[np.float64],
) -> tuple[typing.Sequence[float], typing.Sequence[np.uintp]]:
"""
Compute partitions given possibly repeating times.
This function returns two arrays:
- Non-repeating times: a filtered version of `times` where repeated times are removed.
- Partitions: an array of integers where each element indicates the number of elements for the corresponding time
values in the original `times` array.
By construction, both arrays should have the same length, and the sum of all elements in `partitions` should be
equal to the length of `times`.
"""
change_indices = (np.argwhere(times != np.concatenate([times[1:], np.array([np.nan])])).T + 1).reshape(-1)
partitions = np.concatenate([[change_indices[0]], np.diff(change_indices)])
non_repeating_times = times[change_indices - 1]
assert np.sum(partitions) == len(times)
assert len(non_repeating_times) == len(partitions)
return non_repeating_times, partitions # type: ignore[return-value]
def log_lidar_data() -> None:
las_data = laspy.read(LIDAR_DATA_FILE)
# get positions and convert to meters
points = las_data.points
positions = np.column_stack((points.X / 1000.0, points.Y / 1000.0, points.Z / 1000.0))
times = las_data.gps_time
non_repeating_times, partitions = compute_partitions(times)
# log all positions at once using the computed partitions
rr.send_columns(
"/lidar",
[rr.TimeColumn("time", duration=non_repeating_times)],
rr.Points3D.columns(positions=positions).partition(partitions),
)
rr.log(
"/lidar",
# negative radii are interpreted in UI units (instead of scene units)
rr.Points3D.from_fields(colors=(128, 128, 255), radii=-0.1),
static=True,
)
def log_drone_trajectory() -> None:
data = np.genfromtxt(TRAJECTORY_DATA_FILE, delimiter=" ", skip_header=1)
timestamp = data[:, 0]
positions = data[:, 1:4]
rr.send_columns(
"/drone",
[rr.TimeColumn("time", duration=timestamp)],
rr.Points3D.columns(positions=positions),
)
rr.log(
"/drone",
rr.Points3D.from_fields(colors=(255, 0, 0), radii=0.5),
static=True,
)
def main() -> None:
parser = ArgumentParser(description="Visualize drone-based LiDAR data")
rr.script_add_args(parser)
args = parser.parse_args()
download_dataset()
blueprint = rrb.Spatial3DView(
origin="/",
time_ranges=[
rrb.VisibleTimeRange(
timeline="time",
start=rrb.TimeRangeBoundary.cursor_relative(seconds=-60.0),
end=rrb.TimeRangeBoundary.cursor_relative(),
),
],
)
rr.script_setup(args, "rerun_example_drone_lidar", default_blueprint=blueprint)
log_lidar_data()
log_drone_trajectory()
if __name__ == "__main__":
main()
examples/python/drone_lidar/pyproject.toml
+13
View File
@@ -0,0 +1,13 @@
[project]
name = "drone_lidar"
version = "0.1.0"
readme = "README.md"
dependencies = ["laspy", "numpy", "requests", "rerun-sdk", "tqdm"]
[project.scripts]
drone_lidar = "drone_lidar:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/egoexo_forge/README.md
+35
View File
@@ -0,0 +1,35 @@
<!--[metadata]
title = "EgoExo forge"
tags = ["3D", "HuggingFace", "Egocentric", "Exocentric", "manipulation"]
source = "https://github.com/rerun-io/egoexo-forge"
thumbnail = "https://static.rerun.io/egoexo_forge/629a093f1e2653711ad8fdd59c68b2318ca6bc6c/480w.png"
thumbnail_dimensions = [480, 436]
-->
https://vimeo.com/1134260310?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=2386:1634
A collection of datasets and tools for egocentric and exocentric human activity understanding, featuring hand-object interactions, manipulation tasks, and multi-view recordings.
## Background
EgoExo Forge provides a consistent labeling scheme and data layout across multiple egocentric and exocentric human datasets with varying sensor configurations and annotations.
The following datasets are supported:
* [Assembly101](https://assembly-101.github.io/): A procedural activity dataset with 4321 multi-view videos of people assembling and disassembling 101 take-apart toy vehicles, featuring rich variations in action ordering, mistakes, and corrections.
* [HO-Cap](https://irvlutd.github.io/HOCap/): A dataset for 3D reconstruction and pose tracking of hands and objects in videos, featuring humans interacting with objects for various tasks including pick-and-place actions and handovers.
* [EgoDex](https://arxiv.org/abs/2505.11709): The largest and most diverse dataset of dexterous human manipulation with 829 hours of egocentric video and paired 3D hand tracking, covering 194 different tabletop tasks with everyday household objects.
## Run the code
This is an external example. Check the [repository](https://github.com/rerun-io/egoexo-forge) for more information.
You can try the example on a HuggingFace space [here](https://pablovela5620-egoexo-forge-viewer.hf.space/).
Or locally, make sure you have the [Pixi package manager](https://pixi.sh/latest/#installation) installed and run
```sh
git clone https://github.com/rerun-io/egoexo-forge.git
cd egoexo-forge
pixi run app
```
examples/python/external_importer/README.md
+30
View File
@@ -0,0 +1,30 @@
---
title: External importer example
python: https://github.com/rerun-io/rerun/tree/latest/examples/python/external_importer/rerun-importer-python-file.py
rust: https://github.com/rerun-io/rerun/tree/latest/examples/rust/external_importer/src/main.rs
cpp: https://github.com/rerun-io/rerun/tree/latest/examples/cpp/external_importer/main.cpp
thumbnail: https://static.rerun.io/external_data_loader_py/6c5609f5dd7d1de373c81babe19221b72d616da3/480w.png
thumbnail_dimensions: [480, 302]
---
<picture>
<img src="https://static.rerun.io/external_data_loader_py/6c5609f5dd7d1de373c81babe19221b72d616da3/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/external_data_loader_py/6c5609f5dd7d1de373c81babe19221b72d616da3/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/external_data_loader_py/6c5609f5dd7d1de373c81babe19221b72d616da3/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/external_data_loader_py/6c5609f5dd7d1de373c81babe19221b72d616da3/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/external_data_loader_py/6c5609f5dd7d1de373c81babe19221b72d616da3/1200w.png">
</picture>
This is an example executable importer plugin for the Rerun Viewer.
It will log Python source code files as markdown documents.
On Linux & Mac you can simply copy it in your $PATH as `rerun-importer-python-file.py`, then open a Python source file with Rerun (`rerun file.py`).
Make sure the file has a shebang (`#!/usr/bin/env python3`) and is executable (`chmod +x`).
On Windows you have to install the script as an executable first and then put the executable under %PATH%.
One way to do this is to use `pyinstaller`: `pyinstaller .\examples\python\external_importer\rerun-importer-python-file.py -n rerun-importer-python-file --onefile`
Consider using the [`send_columns`](https://ref.rerun.io/docs/python/stable/common/columnar_api/#rerun.send_columns) API for importers that ingest time series data from a file.
This can be much more efficient that the stateful `log` API as it allows bundling
component data over time into a single call consuming a continuous block of memory.
examples/python/external_importer/rerun-importer-python-file.py
Executable
+114
View File
@@ -0,0 +1,114 @@
#!/usr/bin/env python3
"""Example of an executable importer plugin for the Rerun Viewer."""
from __future__ import annotations
import argparse
import os
import sys
import rerun as rr # pip install rerun-sdk
# The Rerun Viewer will always pass these two pieces of information:
# 1. The path to be loaded, as a positional arg.
# 2. A shared recording ID, via the `--recording-id` flag.
#
# It is up to you whether you make use of that shared recording ID or not.
# If you use it, the data will end up in the same recording as all other plugins interested in
# that file, otherwise you can just create a dedicated recording for it. Or both.
#
# Check out `re_importer::ImporterSettings` documentation for an exhaustive listing of
# the available CLI parameters.
parser = argparse.ArgumentParser(
description="""
This is an example executable importer plugin for the Rerun Viewer.
Any executable on your `$PATH` with a name that starts with `rerun-importer-` will be
treated as an external importer.
This particular one will log Python source code files as markdown documents, and return a
special exit code to indicate that it doesn't support anything else.
To try it out, copy it in your $PATH as `rerun-importer-python-file`, then open a Python source
file with Rerun (`rerun file.py`).
""",
)
parser.add_argument("filepath", type=str)
parser.add_argument("--application-id", type=str, help="optional recommended ID for the application")
parser.add_argument("--recording-id", type=str, help="optional recommended ID for the recording")
parser.add_argument("--entity-path-prefix", type=str, help="optional prefix for all entity paths")
parser.add_argument("--static", action="store_true", default=False, help="optionally mark data to be logged as static")
parser.add_argument(
"--time_sequence",
type=str,
action="append",
help="optional sequences to log at (e.g. `--time_sequence sim_frame=42`)",
)
parser.add_argument(
"--time_duration_nanos",
type=str,
action="append",
help="optional duration(s) (in nanoseconds) to log at (e.g. `--time_duration_nanos sim_time=123`) (repeatable)",
)
parser.add_argument(
"--time_timestamp_nanos",
type=str,
action="append",
help="optional timestamp(s) (in nanoseconds since epochj) to log at (e.g. `--time_timestamp_nanos sim_time=1709203426123456789`) (repeatable)",
)
args = parser.parse_args()
def main() -> None:
is_file = os.path.isfile(args.filepath)
is_python_file = os.path.splitext(args.filepath)[1].lower() == ".py"
# Inform the Rerun Viewer that we do not support that kind of file.
if not is_file or not is_python_file:
sys.exit(rr.EXTERNAL_IMPORTER_INCOMPATIBLE_EXIT_CODE)
app_id = "rerun_example_external_importer"
if args.application_id is not None:
app_id = args.application_id
rr.init(app_id, recording_id=args.recording_id)
# The most important part of this: log to standard output so the Rerun Viewer can ingest it!
rr.stdout()
set_time_from_args()
if args.entity_path_prefix:
entity_path = f"{args.entity_path_prefix}/{args.filepath}"
else:
entity_path = args.filepath
with open(args.filepath, encoding="utf8") as file:
body = file.read()
text = f"""## Some Python code\n```python\n{body}\n```\n"""
rr.log(entity_path, rr.TextDocument(text, media_type=rr.MediaType.MARKDOWN), static=args.static)
def set_time_from_args() -> None:
if not args.static and args.time is not None:
for time_str in args.time_sequence:
parts = time_str.split("=")
if len(parts) != 2:
continue
timeline_name, sequence = parts
rr.set_time(timeline_name, sequence=int(sequence))
for time_str in args.time_duration_nanos:
parts = time_str.split("=")
if len(parts) != 2:
continue
timeline_name, nanos = parts
rr.set_time(timeline_name, duration=1e-9 * int(nanos))
for time_str in args.time_timestamp_nanos:
parts = time_str.split("=")
if len(parts) != 2:
continue
timeline_name, nanos = parts
rr.set_time(timeline_name, timestamp=1e-9 * int(nanos))
if __name__ == "__main__":
main()
examples/python/eye_control/README.md
+65
View File
@@ -0,0 +1,65 @@
<!--[metadata]
title = "Eye control"
tags = ["Eye control", "3D", "Pinhole camera"]
source = "https://github.com/rerun-io/eye_control_example"
thumbnail = "https://static.rerun.io/eye_control_example/01288e2cd92ec68715258e281104701fc8908c37/480w.png"
thumbnail_dimensions = [480, 306]
-->
<video width="100%" autoplay loop muted controls>
<source src="https://static.rerun.io/d75d85d16b70cf6e7863367cd716c10e5725f852_eye_control_example_kth_rpl.mp4" type="video/mp4" />
</video>
<video width="100%" autoplay loop muted controls>
<source src="https://static.rerun.io/7ef9a59523051d45c25de02e6a844e7179e205d2_eye_control_example_ntu_viral.mp4" type="video/mp4" />
</video>
## Used Rerun types
[`EyeControls3D`](https://ref.rerun.io/docs/python/0.27.3/common/blueprint_archetypes/#rerun.blueprint.archetypes.EyeControls3D),
[`Transform3D`](https://rerun.io/docs/reference/types/archetypes/transform3d), [`Points3D`](https://rerun.io/docs/reference/types/archetypes/points3d), [`Pinhole`](https://rerun.io/docs/reference/types/archetypes/pinhole), [`EncodedImage`](https://rerun.io/docs/reference/types/archetypes/encoded_image), [`Image`](https://rerun.io/docs/reference/types/archetypes/image), [`LineStrips3D`](https://rerun.io/docs/reference/types/archetypes/line_strips3d), [`Scalars`](https://rerun.io/docs/reference/types/archetypes/scalars), [`TextDocument`](https://rerun.io/docs/reference/types/archetypes/text_document)
## Background
This example demonstrates how to programmatically configure and control the 3D view camera using the Rerun Blueprint API. By defining camera states in Python, you can precisely tailor your workspace to highlight the most relevant aspects of your data.
In this example, we define several specialized perspectives:
* Top-down overview: a global scene perspective for general spatial awareness.
* Comparative close-up: A focused view designed to analyze trajectory deviations between different localization methods.
* 3rd-person follow: dynamic camera that tracks the ego vehicle as it moves through the environment.
Finally, we demonstrate how to control the camera at runtime, enabling the creation of cinematic visualizations or automated data storytelling for presentations and datasets.
## Useful resources
Below you will find a collection of useful Rerun resources for this example:
* [Blueprints](https://rerun.io/docs/concepts/visualization/blueprints)
* [Building blueprints programmatically](https://rerun.io/docs/howto/build-a-blueprint-programmatically)
## Run the code
This is an external example. Check the [repository](https://github.com/rerun-io/eye_control_example) for more information.
To run this example, make sure you have the [Pixi](https://pixi.sh/latest/#installation) package manager installed.
### KTH RPL (indoor handheld dataset)
```sh
pixi run kth_rpl
```
You can type:
```sh
pixi run kth_rpl -h
```
to see all available commands. For example, you can set the voxel size used for downsampling, where the dataset is located, and for how long to sleep in-between frames.
### NTU VIRAL (drone dataset)
```sh
pixi run ntu_viral
```
examples/python/face_tracking/.gitignore
+2
View File
@@ -0,0 +1,2 @@
dataset/
model/
examples/python/face_tracking/README.md
+183
View File
@@ -0,0 +1,183 @@
<!--[metadata]
title = "Face tracking"
tags = ["2D", "3D", "Camera", "Face tracking", "Live", "MediaPipe", "Time series"]
thumbnail = "https://static.rerun.io/face-tracking/f798733b72c703ee82cc946df39f32fa1145c23b/480w.png"
thumbnail_dimensions = [480, 480]
-->
Use the [MediaPipe](https://github.com/google-ai-edge/mediapipe) Face Detector and Landmarker solutions to detect and track a human face in image, video, and camera stream.
<picture>
<source media="(max-width: 480px)" srcset="https://static.rerun.io/mp_face/f5ee03278408bf8277789b637857d5a4fda7eba3/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/mp_face/f5ee03278408bf8277789b637857d5a4fda7eba3/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/mp_face/f5ee03278408bf8277789b637857d5a4fda7eba3/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/mp_face/f5ee03278408bf8277789b637857d5a4fda7eba3/1200w.png">
<img src="https://static.rerun.io/mp_face/f5ee03278408bf8277789b637857d5a4fda7eba3/full.png" alt="screenshot of the Rerun visualization of the MediaPipe Face Detector and Landmarker">
</picture>
## Used Rerun types
[`Image`](https://www.rerun.io/docs/reference/types/archetypes/image), [`Points2D`](https://www.rerun.io/docs/reference/types/archetypes/points2d), [`Points3D`](https://www.rerun.io/docs/reference/types/archetypes/points3d), [`Boxes2D`](https://www.rerun.io/docs/reference/types/archetypes/boxes2d), [`AnnotationContext`](https://www.rerun.io/docs/reference/types/archetypes/annotation_context), [`Scalars`](https://www.rerun.io/docs/reference/types/archetypes/scalars)
## Background
The face and face landmark detection technology aims to give the ability of the devices to interpret face movements and facial expressions as commands or inputs.
At the core of this technology, a pre-trained machine-learning model analyses the visual input, locates face and identifies face landmarks and blendshape scores (coefficients representing facial expression).
Human-Computer Interaction, Robotics, Gaming, and Augmented Reality are among the fields where this technology shows significant promise for applications.
In this example, the [MediaPipe](https://developers.google.com/mediapipe/) Face and Face Landmark Detection solutions were utilized to detect human face, detect face landmarks and identify facial expressions.
Rerun was employed to visualize the output of the Mediapipe solution over time to make it easy to analyze the behavior.
## Logging and visualizing with Rerun
The visualizations in this example were created with the following Rerun code.
### Timelines
For each processed video frame, all data sent to Rerun is associated with the two [`timelines`](https://www.rerun.io/docs/concepts/logging-and-ingestion/timelines) `time` and `frame_idx`.
```python
rr.set_time("time", duration=bgr_frame.time)
rr.set_time("frame_idx", sequence=bgr_frame.idx)
```
### Video
The input video is logged as a sequence of [`Image`](https://www.rerun.io/docs/reference/types/archetypes/image) objects to the 'Video' entity.
```python
rr.log("video/image", rr.Image(frame).compress(jpeg_quality=75))
```
### Face landmark points
Logging the face landmarks involves specifying connections between the points, extracting face landmark points and logging them to the Rerun SDK.
The 2D points are visualized over the video/image for a better understanding and visualization of the face.
The 3D points allows the creation of a 3D model of the face reconstruction for a more comprehensive representation of the face.
The 2D and 3D points are logged through a combination of two archetypes. First, a static
[`ClassDescription`](https://www.rerun.io/docs/reference/types/datatypes/class_description) is logged, that contains the information which maps keypoint ids to labels and how to connect
the keypoints. Defining these connections automatically renders lines between them.
Second, the actual keypoint positions are logged in 2D and 3D as [`Points2D`](https://www.rerun.io/docs/reference/types/archetypes/points2d) and [`Points3D`](https://www.rerun.io/docs/reference/types/archetypes/points3d) archetypes, respectively.
#### Label mapping and keypoint connections
An annotation context is logged with one class ID assigned per facial feature. The class description includes the connections between corresponding keypoints extracted from the MediaPipe face mesh solution.
A class ID array is generated to match the class IDs in the annotation context with keypoint indices (to be utilized as the class_ids argument to rr.log).
```python
# Initialize a list of facial feature classes from MediaPipe face mesh solution
classes = [
mp.solutions.face_mesh.FACEMESH_LIPS,
mp.solutions.face_mesh.FACEMESH_LEFT_EYE,
mp.solutions.face_mesh.FACEMESH_LEFT_IRIS,
mp.solutions.face_mesh.FACEMESH_LEFT_EYEBROW,
mp.solutions.face_mesh.FACEMESH_RIGHT_EYE,
mp.solutions.face_mesh.FACEMESH_RIGHT_EYEBROW,
mp.solutions.face_mesh.FACEMESH_RIGHT_IRIS,
mp.solutions.face_mesh.FACEMESH_FACE_OVAL,
mp.solutions.face_mesh.FACEMESH_NOSE,
]
# Initialize class descriptions and class IDs array
self._class_ids = [0] * mp.solutions.face_mesh.FACEMESH_NUM_LANDMARKS_WITH_IRISES
class_descriptions = []
# Loop through each facial feature class
for i, klass in enumerate(classes):
# MediaPipe only provides connections for class, not actual class per keypoint. So we have to extract the
# classes from the connections.
ids = set()
for connection in klass:
ids.add(connection[0])
ids.add(connection[1])
for id_ in ids:
self._class_ids[id_] = i
# Append class description with class ID and keypoint connections
class_descriptions.append(
rr.ClassDescription(
info=rr.AnnotationInfo(id=i),
keypoint_connections=klass,
)
)
# Log annotation context for video/landmarker and reconstruction entities
rr.log("video/landmarker", rr.AnnotationContext(class_descriptions), static=True)
rr.log("reconstruction", rr.AnnotationContext(class_descriptions), static=True)
rr.log("reconstruction", rr.ViewCoordinates.RDF, static=True) # properly align the 3D face in the viewer
```
With the below annotation, the keypoints will be connected with lines to enhance visibility in the `video/detector` entity.
```python
rr.log(
"video/detector",
rr.ClassDescription(
info=rr.AnnotationInfo(id=0), keypoint_connections=[(0, 1), (1, 2), (2, 0), (2, 3), (0, 4), (1, 5)]
),
static=True,
)
```
#### Bounding box
```python
rr.log(
f"video/detector/faces/{i}/bbox",
rr.Boxes2D(array=[bbox.origin_x, bbox.origin_y, bbox.width, bbox.height], array_format=rr.Box2DFormat.XYWH),
rr.AnyValues(index=index, score=score),
)
```
#### 2D points
```python
rr.log(f"video/detector/faces/{i}/keypoints", rr.Points2D(pts, radii=3, keypoint_ids=list(range(6))))
```
```python
rr.log(
f"video/landmarker/faces/{i}/landmarks",
rr.Points2D(pts, radii=3, keypoint_ids=keypoint_ids, class_ids=self._class_ids),
)
```
#### 3D points
```python
rr.log(
f"reconstruction/faces/{i}",
rr.Points3D(
[(lm.x, lm.y, lm.z) for lm in landmark],
keypoint_ids=keypoint_ids,
class_ids=self._class_ids,
),
)
```
### Scalars
Blendshapes are essentially predefined facial expressions or configurations that can be detected by the face landmark detection model. Each blendshape typically corresponds to a specific facial movement or expression, such as blinking, squinting, smiling, etc.
The blendshapes are logged along with their corresponding scores.
```python
for blendshape in blendshapes:
if blendshape.category_name in BLENDSHAPES_CATEGORIES:
rr.log(f"blendshapes/{i}/{blendshape.category_name}", rr.Scalars(blendshape.score))
```
## Run the code
To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/face_tracking
```
To experiment with the provided example, simply execute the main Python script:
```bash
python -m face_tracking # run the example
```
If you wish to customize it for various videos, adjust the maximum frames, explore additional features, or save it use the CLI with the `--help` option for guidance:
```bash
python -m face_tracking --help
```
examples/python/face_tracking/face_tracking.py
Executable
+518
View File
@@ -0,0 +1,518 @@
#!/usr/bin/env python3
"""Use the MediaPipe Face detection and Face landmark detection solutions to track human faces in images and videos."""
from __future__ import annotations
import argparse
import itertools
import logging
import math
import os
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final
import cv2
import mediapipe as mp
import numpy as np
import requests
import tqdm
from mediapipe.tasks.python import vision
import rerun as rr # pip install rerun-sdk
import rerun.blueprint as rrb
if TYPE_CHECKING:
from collections.abc import Iterable, Iterator
# If set, log everything as static.
#
# Generally, the Viewer accumulates data until its set memory budget at which point it will
# remove the oldest data from the recording (see https://rerun.io/docs/howto/visualization/limit-ram)
# By instead logging data as static, no data will be accumulated over time since previous
# data is overwritten.
# Naturally, the drawback of this is that there's no history of previous data sent to the viewer,
# as well as no timestamps, making the Viewer's timeline effectively inactive.
ALL_STATIC: bool = False
EXAMPLE_DIR: Final = Path(os.path.dirname(__file__))
DATASET_DIR: Final = EXAMPLE_DIR / "dataset"
MODEL_DIR: Final = EXAMPLE_DIR / "model"
SAMPLE_IMAGE_PATH = (DATASET_DIR / "image.jpg").resolve()
# from https://pixabay.com/photos/brother-sister-girl-family-boy-977170/
SAMPLE_IMAGE_URL = "https://i.imgur.com/Vu2Nqwb.jpg"
# uncomment blendshapes of interest
BLENDSHAPES_CATEGORIES = {
"_neutral",
"browDownLeft",
"browDownRight",
"browInnerUp",
"browOuterUpLeft",
"browOuterUpRight",
"cheekPuff",
"cheekSquintLeft",
"cheekSquintRight",
"eyeBlinkLeft",
"eyeBlinkRight",
"eyeLookDownLeft",
"eyeLookDownRight",
"eyeLookInLeft",
"eyeLookInRight",
"eyeLookOutLeft",
"eyeLookOutRight",
"eyeLookUpLeft",
"eyeLookUpRight",
"eyeSquintLeft",
"eyeSquintRight",
"eyeWideLeft",
"eyeWideRight",
"jawForward",
"jawLeft",
"jawOpen",
"jawRight",
"mouthClose",
"mouthDimpleLeft",
"mouthDimpleRight",
"mouthFrownLeft",
"mouthFrownRight",
"mouthFunnel",
"mouthLeft",
"mouthLowerDownLeft",
"mouthLowerDownRight",
"mouthPressLeft",
"mouthPressRight",
"mouthPucker",
"mouthRight",
"mouthRollLower",
"mouthRollUpper",
"mouthShrugLower",
"mouthShrugUpper",
"mouthSmileLeft",
"mouthSmileRight",
"mouthStretchLeft",
"mouthStretchRight",
"mouthUpperUpLeft",
"mouthUpperUpRight",
"noseSneerLeft",
"noseSneerRight",
}
class FaceDetectorLogger:
"""
Logger for the MediaPipe Face Detection solution.
<https://developers.google.com/mediapipe/solutions/vision/face_detector>
"""
MODEL_PATH: Final = (MODEL_DIR / "blaze_face_short_range.tflite").resolve()
MODEL_URL: Final = (
"https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/"
"blaze_face_short_range.tflite"
)
def __init__(self, video_mode: bool = False) -> None:
self._video_mode = video_mode
# download model if necessary
if not self.MODEL_PATH.exists():
download_file(self.MODEL_URL, self.MODEL_PATH)
self._base_options = mp.tasks.BaseOptions(
model_asset_path=str(self.MODEL_PATH),
)
self._options = vision.FaceDetectorOptions(
base_options=self._base_options,
running_mode=mp.tasks.vision.RunningMode.VIDEO if self._video_mode else mp.tasks.vision.RunningMode.IMAGE,
)
self._detector = vision.FaceDetector.create_from_options(self._options)
# With this annotation, the viewer will connect the keypoints with some lines to improve visibility.
rr.log(
"video/detector",
rr.ClassDescription(
info=rr.AnnotationInfo(id=0),
keypoint_connections=[(0, 1), (1, 2), (2, 0), (2, 3), (0, 4), (1, 5)],
),
static=True,
)
def detect_and_log(self, image: cv2.typing.MatLike, frame_time_nano: int) -> None:
height, width, _ = image.shape
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
detection_result = (
self._detector.detect_for_video(image, int(frame_time_nano / 1e6))
if self._video_mode
else self._detector.detect(image)
)
rr.log("video/detector/faces", rr.Clear(recursive=True), static=ALL_STATIC)
for i, detection in enumerate(detection_result.detections):
# log bounding box
bbox = detection.bounding_box
index, score = detection.categories[0].index, detection.categories[0].score
# log bounding box
rr.log(
f"video/detector/faces/{i}/bbox",
rr.Boxes2D(
array=[bbox.origin_x, bbox.origin_y, bbox.width, bbox.height],
array_format=rr.Box2DFormat.XYWH,
),
rr.AnyValues(index=index, score=score),
static=ALL_STATIC,
)
# MediaPipe's keypoints are normalized to [0, 1], so we need to scale them to get pixel coordinates.
pts = [
(math.floor(keypoint.x * width), math.floor(keypoint.y * height)) for keypoint in detection.keypoints
]
rr.log(
f"video/detector/faces/{i}/keypoints",
rr.Points2D(pts, radii=3, keypoint_ids=list(range(6))),
static=ALL_STATIC,
)
class FaceLandmarkerLogger:
"""
Logger for the MediaPipe Face Landmark Detection solution.
<https://developers.google.com/mediapipe/solutions/vision/face_landmarker>
"""
MODEL_PATH: Final = (MODEL_DIR / "face_landmarker.task").resolve()
MODEL_URL: Final = (
"https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/"
"face_landmarker.task"
)
def __init__(self, video_mode: bool = False, num_faces: int = 1) -> None:
self._video_mode = video_mode
# download model if necessary
if not self.MODEL_PATH.exists():
download_file(self.MODEL_URL, self.MODEL_PATH)
self._base_options = mp.tasks.BaseOptions(
model_asset_path=str(self.MODEL_PATH),
)
self._options = vision.FaceLandmarkerOptions(
base_options=self._base_options,
output_face_blendshapes=True,
num_faces=num_faces,
running_mode=mp.tasks.vision.RunningMode.VIDEO if self._video_mode else mp.tasks.vision.RunningMode.IMAGE,
)
self._detector = vision.FaceLandmarker.create_from_options(self._options)
# Extract classes from MediaPipe face mesh solution. The goal of this code is:
# 1) Log an annotation context with one class ID per facial feature. For each class ID, the class description
# contains the connections between corresponding keypoints (taken from the MediaPipe face mesh solution)
# 2) A class ID array matching the class IDs in the annotation context to keypoint indices (to be passed as
# the `class_ids` argument to `rr.log`).
classes = [
mp.solutions.face_mesh.FACEMESH_LIPS,
mp.solutions.face_mesh.FACEMESH_LEFT_EYE,
mp.solutions.face_mesh.FACEMESH_LEFT_IRIS,
mp.solutions.face_mesh.FACEMESH_LEFT_EYEBROW,
mp.solutions.face_mesh.FACEMESH_RIGHT_EYE,
mp.solutions.face_mesh.FACEMESH_RIGHT_EYEBROW,
mp.solutions.face_mesh.FACEMESH_RIGHT_IRIS,
mp.solutions.face_mesh.FACEMESH_FACE_OVAL,
mp.solutions.face_mesh.FACEMESH_NOSE,
]
self._class_ids = [0] * mp.solutions.face_mesh.FACEMESH_NUM_LANDMARKS_WITH_IRISES
class_descriptions = []
for i, klass in enumerate(classes):
# MediaPipe only provides connections for class, not actual class per keypoint. So we have to extract the
# classes from the connections.
ids = set()
for connection in klass:
ids.add(connection[0])
ids.add(connection[1])
for id_ in ids:
self._class_ids[id_] = i
class_descriptions.append(
rr.ClassDescription(
info=rr.AnnotationInfo(id=i),
keypoint_connections=klass,
),
)
rr.log("video/landmarker", rr.AnnotationContext(class_descriptions), static=True)
rr.log("reconstruction", rr.AnnotationContext(class_descriptions), static=True)
# properly align the 3D face in the viewer
rr.log("reconstruction", rr.ViewCoordinates.RDF, static=True)
def detect_and_log(self, image: cv2.typing.MatLike, frame_time_nano: int) -> None:
height, width, _ = image.shape
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
detection_result = (
self._detector.detect_for_video(image, int(frame_time_nano / 1e6))
if self._video_mode
else self._detector.detect(image)
)
def is_empty(i: Iterator[Any]) -> bool:
try:
next(i)
return False
except StopIteration:
return True
if is_empty(zip(detection_result.face_landmarks, detection_result.face_blendshapes, strict=False)):
rr.log("video/landmarker/faces", rr.Clear(recursive=True), static=ALL_STATIC)
rr.log("reconstruction/faces", rr.Clear(recursive=True), static=ALL_STATIC)
rr.log("blendshapes", rr.Clear(recursive=True), static=ALL_STATIC)
for i, (landmark, blendshapes) in enumerate(
zip(detection_result.face_landmarks, detection_result.face_blendshapes, strict=False),
):
if len(landmark) == 0 or len(blendshapes) == 0:
rr.log(
f"video/landmarker/faces/{i}/landmarks",
rr.Clear(recursive=True),
static=ALL_STATIC,
)
rr.log(
f"reconstruction/faces/{i}",
rr.Clear(recursive=True),
static=ALL_STATIC,
)
rr.log(f"blendshapes/{i}", rr.Clear(recursive=True), static=ALL_STATIC)
continue
# MediaPipe's keypoints are normalized to [0, 1], so we need to scale them to get pixel coordinates.
pts = [(math.floor(lm.x * width), math.floor(lm.y * height)) for lm in landmark]
keypoint_ids = list(range(len(landmark)))
rr.log(
f"video/landmarker/faces/{i}/landmarks",
rr.Points2D(pts, radii=3, keypoint_ids=keypoint_ids, class_ids=self._class_ids),
static=ALL_STATIC,
)
rr.log(
f"reconstruction/faces/{i}",
rr.Points3D(
[(lm.x, lm.y, lm.z) for lm in landmark],
keypoint_ids=keypoint_ids,
class_ids=self._class_ids,
),
static=ALL_STATIC,
)
for blendshape in blendshapes:
if blendshape.category_name in BLENDSHAPES_CATEGORIES:
# NOTE(cmc): That one we still log as temporal, otherwise it's really meh.
rr.log(
f"blendshapes/{i}/{blendshape.category_name}",
rr.Scalars(blendshape.score),
)
# ========================================================================================
# Main & CLI code
def download_file(url: str, path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
logging.info("Downloading %s to %s", url, path)
response = requests.get(url, stream=True)
with tqdm.tqdm.wrapattr(
open(path, "wb"),
"write",
miniters=1,
total=int(response.headers.get("content-length", 0)),
desc=f"Downloading {path.name}",
) as f:
for chunk in response.iter_content(chunk_size=4096):
f.write(chunk)
def resize_image(image: cv2.typing.MatLike, max_dim: int | None) -> cv2.typing.MatLike:
"""Resize an image if it is larger than max_dim."""
if max_dim is None:
return image
height, width, _ = image.shape
scale = max_dim / max(height, width)
if scale < 1:
image = cv2.resize(image, (0, 0), fx=scale, fy=scale)
return image
def run_from_video_capture(vid: int | str, max_dim: int | None, max_frame_count: int | None, num_faces: int) -> None:
"""
Run the face detector on a video stream.
Parameters
----------
vid:
The video stream to run the detector on. Use 0 for the default camera or a path to a video file.
max_dim:
The maximum dimension of the image. If the image is larger, it will be scaled down.
max_frame_count:
The maximum number of frames to process. If None, process all frames.
num_faces:
The number of faces to track. If set to 1, temporal smoothing will be applied.
"""
cap = cv2.VideoCapture(vid)
fps = cap.get(cv2.CAP_PROP_FPS)
detector = FaceDetectorLogger(video_mode=True)
landmarker = FaceLandmarkerLogger(video_mode=True, num_faces=num_faces)
print("Capturing video stream. Press ctrl-c to stop.")
try:
it: Iterable[int] = itertools.count() if max_frame_count is None else range(max_frame_count)
for frame_idx in tqdm.tqdm(it, desc="Processing frames"):
# Capture frame-by-frame
ret, frame = cap.read()
if not ret:
break
# OpenCV sometimes returns a blank frame, so we skip it
if np.all(frame == 0):
continue
frame = resize_image(frame, max_dim)
# get frame time
frame_time_nano = int(cap.get(cv2.CAP_PROP_POS_MSEC) * 1e6)
if frame_time_nano == 0:
# On some platforms it always returns zero, so we compute from the frame counter and fps
frame_time_nano = int(frame_idx * 1000 / fps * 1e6)
# log data
rr.set_time("frame_nr", sequence=frame_idx)
rr.set_time("frame_time", duration=1e-9 * frame_time_nano)
detector.detect_and_log(frame, frame_time_nano)
landmarker.detect_and_log(frame, frame_time_nano)
rr.log(
"video/image",
rr.Image(frame, color_model="BGR"),
static=ALL_STATIC,
)
except KeyboardInterrupt:
pass
# When everything done, release the capture
cap.release()
cv2.destroyAllWindows()
def run_from_sample_image(path: Path, max_dim: int | None, num_faces: int) -> None:
"""Run the face detector on a single image."""
image = cv2.imread(str(path))
image = resize_image(image, max_dim)
logger = FaceDetectorLogger(video_mode=False)
landmarker = FaceLandmarkerLogger(video_mode=False, num_faces=num_faces)
logger.detect_and_log(image, 0)
landmarker.detect_and_log(image, 0)
rr.log(
"video/image",
rr.Image(image, color_model="BGR"),
static=ALL_STATIC,
)
def main() -> None:
logging.getLogger().addHandler(logging.StreamHandler())
logging.getLogger().setLevel("INFO")
parser = argparse.ArgumentParser(description="Uses the MediaPipe Face Detection to track a human pose in video.")
parser.add_argument(
"--demo-image",
action="store_true",
help="Run on a demo image automatically downloaded",
)
parser.add_argument(
"--image",
type=Path,
help="Run on the provided image",
)
parser.add_argument("--video", type=Path, help="Run on the provided video file.")
parser.add_argument(
"--camera",
type=int,
default=0,
help="Run from the camera stream (parameter is the camera ID, usually 0)",
)
parser.add_argument(
"--max-frame",
type=int,
help="Stop after processing this many frames. If not specified, will run until interrupted.",
)
parser.add_argument(
"--max-dim",
type=int,
help="Resize the image such as its maximum dimension is not larger than this value.",
)
parser.add_argument(
"--num-faces",
type=int,
default=1,
help=(
"Max number of faces detected by the landmark model (temporal smoothing is applied only for a value of 1)."
),
)
parser.add_argument("--static", action="store_true", help="If set, logs everything as static")
rr.script_add_args(parser)
args, unknown = parser.parse_known_args()
for arg in unknown:
logging.warning(f"unknown arg: {arg}")
rr.script_setup(
args,
"rerun_example_mp_face_detection",
default_blueprint=rrb.Horizontal(
rrb.Spatial3DView(origin="reconstruction"),
rrb.Vertical(
rrb.Spatial2DView(origin="video"),
rrb.TimeSeriesView(
origin="blendshapes",
# Enable only certain blend shapes by default. More can be added in the viewer ui
contents=[
"+ blendshapes/0/eyeBlinkLeft",
"+ blendshapes/0/eyeBlinkRight",
"+ blendshapes/0/jawOpen",
"+ blendshapes/0/mouthSmileLeft",
"+ blendshapes/0/mouthSmileRight",
],
),
),
),
)
global ALL_STATIC
ALL_STATIC = args.static
if args.demo_image:
if not SAMPLE_IMAGE_PATH.exists():
download_file(SAMPLE_IMAGE_URL, SAMPLE_IMAGE_PATH)
run_from_sample_image(SAMPLE_IMAGE_PATH, args.max_dim, args.num_faces)
elif args.image is not None:
run_from_sample_image(args.image, args.max_dim, args.num_faces)
elif args.video is not None:
run_from_video_capture(str(args.video), args.max_dim, args.max_frame, args.num_faces)
else:
run_from_video_capture(args.camera, args.max_dim, args.max_frame, args.num_faces)
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/face_tracking/pyproject.toml
+24
View File
@@ -0,0 +1,24 @@
[project]
name = "face_tracking"
version = "0.1.0"
requires-python = ">=3.10,<3.12" # TODO(ab): relax when mediapipe supports 3.12
readme = "README.md"
dependencies = [
"mediapipe==0.10.11 ; sys_platform != 'darwin'",
"mediapipe==0.10.9 ; sys_platform == 'darwin'", # https://github.com/google/mediapipe/issues/5188
"numpy",
"opencv-python>4.6",
"requests",
"rerun-sdk",
"tqdm",
]
[project.scripts]
face_tracking = "face_tracking:main"
[tool.rerun-example]
extra-args = "--maxframe=30"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/gesture_detection/.gitignore
+2
View File
@@ -0,0 +1,2 @@
dataset/hand_gestures
/model
examples/python/gesture_detection/README.md
+134
View File
@@ -0,0 +1,134 @@
<!--[metadata]
title = "Hand tracking and gesture recognition"
tags = ["MediaPipe", "Keypoint detection", "2D", "3D"]
thumbnail = "https://static.rerun.io/hand-tracking-and-gesture-recognition/56d097e347af2a4b7c4649c7d994cc038c02c2f4/480w.png"
thumbnail_dimensions = [480, 480]
-->
Use the [MediaPipe](https://github.com/google-ai-edge/mediapipe/) Hand Landmark and Gesture Detection solutions to
track hands and recognize gestures in images, video, and camera stream.
<picture>
<img src="https://static.rerun.io/gesture_detection/2a5a3ec83962623063297fd95de57062372d5db0/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/gesture_detection/2a5a3ec83962623063297fd95de57062372d5db0/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/gesture_detection/2a5a3ec83962623063297fd95de57062372d5db0/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/gesture_detection/2a5a3ec83962623063297fd95de57062372d5db0/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/gesture_detection/2a5a3ec83962623063297fd95de57062372d5db0/1200w.png">
</picture>
## Used Rerun types
[`Image`](https://www.rerun.io/docs/reference/types/archetypes/image), [`Points2D`](https://www.rerun.io/docs/reference/types/archetypes/points2d), [`Points3D`](https://www.rerun.io/docs/reference/types/archetypes/points3d), [`LineStrips2D`](https://www.rerun.io/docs/reference/types/archetypes/line_strips2d), [`ClassDescription`](https://www.rerun.io/docs/reference/types/datatypes/class_description), [`AnnotationContext`](https://www.rerun.io/docs/reference/types/archetypes/annotation_context), [`TextDocument`](https://www.rerun.io/docs/reference/types/archetypes/text_document)
## Background
The hand tracking and gesture recognition technology aims to give the ability of the devices to interpret hand movements and gestures as commands or inputs.
At the core of this technology, a pre-trained machine-learning model analyses the visual input and identifies hand landmarks and hand gestures.
The real applications of such technology vary, as hand movements and gestures can be used to control smart devices.
Human-Computer Interaction, Robotics, Gaming, and Augmented Reality are a few of the fields where the potential applications of this technology appear most promising.
In this example, the [MediaPipe](https://developers.google.com/mediapipe/) Gesture and Hand Landmark Detection solutions were utilized to detect and track hand landmarks and recognize gestures.
Rerun was employed to visualize the output of the Mediapipe solution over time to make it easy to analyze the behavior.
## Logging and visualizing with Rerun
The visualizations in this example were created with the following Rerun code.
### Timelines
For each processed video frame, all data sent to Rerun is associated with the two [`timelines`](https://www.rerun.io/docs/concepts/logging-and-ingestion/timelines) `time` and `frame_idx`.
```python
rr.set_time("frame_nr", sequence=frame_idx)
rr.set_time("frame_time", duration=1e-9 * frame_time_nano)
```
### Video
The input video is logged as a sequence of [`Image`](https://www.rerun.io/docs/reference/types/archetypes/image) objects to the `Media/Video` entity.
```python
rr.log("Media/Video", rr.Image(frame).compress(jpeg_quality=75))
```
### Hand landmark points
Logging the hand landmarks involves specifying connections between the points, extracting pose landmark points and logging them to the Rerun SDK.
The 2D points are visualized over the video and at a separate entity.
Meanwhile, the 3D points allows the creation of a 3D model of the hand for a more comprehensive representation of the hand landmarks.
The 2D and 3D points are logged through a combination of two archetypes.
For the 2D points, the Points2D and LineStrips2D archetypes are utilized. These archetypes help visualize the points and connect them with lines, respectively.
As for the 3D points, the logging process involves two steps. First, a static [`ClassDescription`](https://www.rerun.io/docs/reference/types/datatypes/class_description) is logged, that contains the information which maps keypoint ids to labels and how to connect
the keypoints. Defining these connections automatically renders lines between them. Mediapipe provides the `HAND_CONNECTIONS` variable which contains the list of `(from, to)` landmark indices that define the connections.
Second, the actual keypoint positions are logged in 3D [`Points3D`](https://www.rerun.io/docs/reference/types/archetypes/points3d) archetype.
#### Label mapping and keypoint connections
```python
rr.log(
"/",
rr.AnnotationContext(
rr.ClassDescription(
info=rr.AnnotationInfo(id=0, label="Hand3D"),
keypoint_connections=mp.solutions.hands.HAND_CONNECTIONS,
)
),
static=True,
)
rr.log("Hand3D", rr.ViewCoordinates.LEFT_HAND_Y_DOWN, static=True)
```
#### 2D points
```python
# Log points to the image and Hand entity
for log_key in ["Media/Points", "Hand/Points"]:
rr.log(log_key, rr.Points2D(points, radii=10, colors=[255, 0, 0]))
# Log connections to the image and Hand entity [128, 128, 128]
for log_key in ["Media/Connections", "Hand/Connections"]:
rr.log(log_key, rr.LineStrips2D(np.stack((points1, points2), axis=1), colors=[255, 165, 0]))
```
#### 3D points
```python
rr.log(
"Hand3D/Points",
rr.Points3D(
landmark_positions_3d,
radii=20,
class_ids=0,
keypoint_ids=[i for i in range(len(landmark_positions_3d))],
),
)
```
### Detection
To showcase gesture recognition, an image of the corresponding gesture emoji is displayed within a `TextDocument` under the `Detection` entity.
```python
# Log the detection by using the appropriate image
rr.log(
"Detection",
rr.TextDocument(f"![Image]({GESTURE_URL + gesture_pic})".strip(), media_type=rr.MediaType.MARKDOWN),
)
```
## Run the code
To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/gesture_detection
```
To experiment with the provided example, simply execute the main Python script:
```bash
python -m gesture_detection # run the example
```
If you wish to customize it for various videos, adjust the maximum frames, explore additional features, or save it use the CLI with the `--help` option for guidance:
```bash
$ python -m gesture_detection --help
```
examples/python/gesture_detection/gesture_detection.py
Executable
+346
View File
@@ -0,0 +1,346 @@
#!/usr/bin/env python3
"""Use the MediaPipe Gesture detection and Gesture landmark detection solutions to track hands and recognize gestures in images and videos."""
from __future__ import annotations
import argparse
import itertools
import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING, Final
import cv2
import mediapipe as mp
import numpy as np
import requests
import tqdm
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import rerun as rr # pip install rerun-sdk
import rerun.blueprint as rrb
if TYPE_CHECKING:
from collections.abc import Iterable
from mediapipe.tasks.python.components.containers import NormalizedLandmark
EXAMPLE_DIR: Final = Path(os.path.dirname(__file__))
DATASET_DIR: Final = EXAMPLE_DIR / "dataset" / "hand_gestures"
SAMPLE_IMAGE_PATH = EXAMPLE_DIR / "dataset" / "hand_gestures" / "victory.jpg"
# More samples: 'thumbs_down.jpg', 'victory.jpg', 'pointing_up.jpg', 'thumbs_up.jpg'
SAMPLE_IMAGE_URL = "https://storage.googleapis.com/mediapipe-tasks/gesture_recognizer/victory.jpg"
SAMPLE_VIDEO_PATH = EXAMPLE_DIR / "dataset" / "hand_gestures" / "peace.mp4"
SAMPLE_VIDEO_URL = "https://storage.googleapis.com/rerun-example-datasets/hand_gestures/peace.mp4"
# Emojis from https://github.com/googlefonts/noto-emoji/tree/main
GESTURE_URL = (
"https://raw.githubusercontent.com/googlefonts/noto-emoji/9cde38ef5ee6f090ce23f9035e494cb390a2b051/png/128/"
)
# Mapping of gesture categories to corresponding emojis
GESTURE_PICTURES = {
"None": "emoji_u2754.png",
"Closed_Fist": "emoji_u270a.png",
"Open_Palm": "emoji_u270b.png",
"Pointing_Up": "emoji_u261d.png",
"Thumb_Down": "emoji_u1f44e.png",
"Thumb_Up": "emoji_u1f44d.png",
"Victory": "emoji_u270c.png",
"ILoveYou": "emoji_u1f91f.png",
}
class GestureDetectorLogger:
"""
Logger for the MediaPipe Gesture Detection solution.
This class provides logging and utility functions for handling gesture recognition.
For more information on MediaPipe Gesture Detection:
<https://developers.google.com/mediapipe/solutions/vision/gesture_recognizer>
"""
# URL to the pre-trained MediaPipe Gesture Detection model
MODEL_DIR: Final = EXAMPLE_DIR / "model"
MODEL_PATH: Final = (MODEL_DIR / "gesture_recognizer.task").resolve()
MODEL_URL: Final = "https://storage.googleapis.com/mediapipe-models/gesture_recognizer/gesture_recognizer/float16/latest/gesture_recognizer.task"
def __init__(self, video_mode: bool = False) -> None:
self._video_mode = video_mode
if not self.MODEL_PATH.exists():
download_file(self.MODEL_URL, self.MODEL_PATH)
base_options = python.BaseOptions(model_asset_path=str(self.MODEL_PATH))
options = vision.GestureRecognizerOptions(
base_options=base_options,
running_mode=mp.tasks.vision.RunningMode.VIDEO if self._video_mode else mp.tasks.vision.RunningMode.IMAGE,
)
self.recognizer = vision.GestureRecognizer.create_from_options(options)
rr.log(
"/",
rr.AnnotationContext(
rr.ClassDescription(
info=rr.AnnotationInfo(id=0, label="Hand3D"),
keypoint_connections=mp.solutions.hands.HAND_CONNECTIONS,
),
),
static=True,
)
rr.log("hand3d", rr.ViewCoordinates.LEFT_HAND_Y_DOWN, static=True)
@staticmethod
def convert_landmarks_to_image_coordinates(
hand_landmarks: list[list[NormalizedLandmark]],
width: int,
height: int,
) -> list[tuple[int, int]]:
return [(int(lm.x * width), int(lm.y * height)) for hand_landmark in hand_landmarks for lm in hand_landmark]
@staticmethod
def convert_landmarks_to_3d(hand_landmarks: list[list[NormalizedLandmark]]) -> list[tuple[float, float, float]]:
return [(lm.x, lm.y, lm.z) for hand_landmark in hand_landmarks for lm in hand_landmark]
def detect_and_log(self, image: cv2.typing.MatLike, frame_time_nano: int) -> None:
# Recognize gestures in the image
height, width, _ = image.shape
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
recognition_result = (
self.recognizer.recognize_for_video(image, int(frame_time_nano / 1e6))
if self._video_mode
else self.recognizer.recognize(image)
)
for log_key in ["hand2d/points", "hand2d/connections", "hand3d/points"]:
rr.log(log_key, rr.Clear(recursive=True))
for gesture in recognition_result.gestures:
# Get the top gesture from the recognition result
gesture_category = gesture[0].category_name if recognition_result.gestures else "None"
self.present_detected_gesture(gesture_category) # Log the detected gesture
if recognition_result.hand_landmarks:
hand_landmarks = recognition_result.hand_landmarks
landmark_positions_3d = self.convert_landmarks_to_3d(hand_landmarks)
if landmark_positions_3d is not None:
rr.log(
"hand3d/points",
rr.Points3D(
landmark_positions_3d,
radii=20,
class_ids=0,
keypoint_ids=list(range(len(landmark_positions_3d))),
),
)
# Convert normalized coordinates to image coordinates
points = self.convert_landmarks_to_image_coordinates(hand_landmarks, width, height)
# Log points to the image and Hand Entity
rr.log("hand2d/points", rr.Points2D(points, radii=10, colors=[255, 0, 0]))
# Obtain hand connections from MediaPipe
mp_hands_connections = mp.solutions.hands.HAND_CONNECTIONS
points1 = [points[connection[0]] for connection in mp_hands_connections]
points2 = [points[connection[1]] for connection in mp_hands_connections]
# Log connections to the image and Hand Entity [128, 128, 128]
rr.log("hand2d/connections", rr.LineStrips2D(np.stack((points1, points2), axis=1), colors=[255, 165, 0]))
def present_detected_gesture(self, category: str) -> None:
# Get the corresponding ulr of the picture for the detected gesture category
gesture_pic = GESTURE_PICTURES.get(
category,
"emoji_u2754.png", # default
)
# Log the detection by using the appropriate image
rr.log(
"detection",
rr.TextDocument(f"![Image]({GESTURE_URL + gesture_pic})".strip(), media_type=rr.MediaType.MARKDOWN),
)
def download_file(url: str, path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
logging.info("Downloading %s to %s", url, path)
response = requests.get(url, stream=True)
with tqdm.tqdm.wrapattr(
open(path, "wb"),
"write",
miniters=1,
total=int(response.headers.get("content-length", 0)),
desc=f"Downloading {path.name}",
) as f:
for chunk in response.iter_content(chunk_size=4096):
f.write(chunk)
def resize_image(image: cv2.typing.MatLike, max_dim: int | None) -> cv2.typing.MatLike:
"""Resize an image if it is larger than max_dim."""
if max_dim is None:
return image
height, width, _ = image.shape
scale = max_dim / max(height, width)
if scale < 1:
image = cv2.resize(image, (0, 0), fx=scale, fy=scale)
return image
def run_from_sample_image(path: Path | str) -> None:
"""Run the gesture recognition on a single image."""
image = cv2.imread(str(path))
# image = resize_image(image, max_dim)
rr.log("media/image", rr.Image(image, color_model="BGR"))
detect_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
logger = GestureDetectorLogger(video_mode=False)
logger.detect_and_log(detect_image, 0)
def run_from_video_capture(vid: int | str, max_frame_count: int | None) -> None:
"""
Run the detector on a video stream.
Parameters
----------
vid:
The video stream to run the detector on. Use 0/1 for the default camera or a path to a video file.
max_frame_count:
The maximum number of frames to process. If None, process all frames.
"""
cap = cv2.VideoCapture(vid)
fps = cap.get(cv2.CAP_PROP_FPS)
detector = GestureDetectorLogger(video_mode=True)
try:
it: Iterable[int] = itertools.count() if max_frame_count is None else range(max_frame_count)
for frame_idx in tqdm.tqdm(it, desc="Processing frames"):
# Capture frame-by-frame
ret, frame = cap.read()
if not ret:
break
# OpenCV sometimes returns a blank frame, so we skip it
if np.all(frame == 0):
continue
# frame = resize_image(frame, max_dim)
# get frame time
frame_time_nano = int(cap.get(cv2.CAP_PROP_POS_MSEC) * 1e6)
if frame_time_nano == 0:
# On some platforms it always returns zero, so we compute from the frame counter and fps
frame_time_nano = int(frame_idx * 1000 / fps * 1e6)
# log data
rr.set_time("frame_nr", sequence=frame_idx)
rr.set_time("frame_time", duration=1e-9 * frame_time_nano)
detector.detect_and_log(frame, frame_time_nano)
rr.log("media/video", rr.Image(frame, color_model="BGR").compress(jpeg_quality=75))
except KeyboardInterrupt:
pass
# When everything done, release the capture
cap.release()
cv2.destroyAllWindows()
def main() -> None:
# Ensure the logging gets written to stderr
logging.getLogger().addHandler(logging.StreamHandler())
logging.getLogger().setLevel("INFO")
# Set up argument parser with description
parser = argparse.ArgumentParser(
description="Uses the MediaPipe Gesture Recognition to track a hand and recognize gestures in image or video.",
)
parser.add_argument(
"--demo-image",
action="store_true",
help="Run on a demo image automatically downloaded",
)
parser.add_argument("--demo-video", action="store_true", help="Run on a demo image automatically downloaded.")
parser.add_argument(
"--image",
type=Path,
help="Run on the provided image",
)
parser.add_argument("--video", type=Path, help="Run on the provided video file.")
parser.add_argument(
"--camera",
type=int,
default=0,
help="Run from the camera stream (parameter is the camera ID, usually 0; or maybe 1 on mac)",
)
parser.add_argument(
"--max-frame",
type=int,
help="Stop after processing this many frames. If not specified, will run until interrupted.",
)
# Add Rerun specific arguments
rr.script_add_args(parser)
# Parse command line arguments
args, unknown = parser.parse_known_args()
for arg in unknown: # Log any unknown arguments
logging.warning(f"unknown arg: {arg}")
# Set up Rerun with script name
rr.script_setup(
args,
"rerun_example_mp_gesture_recognition",
default_blueprint=rrb.Horizontal(
rrb.Spatial2DView(name="Input & Hand", contents=["media/**", "hand2d/**"]),
rrb.Vertical(
rrb.Tabs(
rrb.Spatial3DView(name="Hand 3D", origin="hand3d"),
rrb.Spatial2DView(name="Hand 2D", origin="hand2d"),
),
rrb.TextDocumentView(name="Detection", origin="detection"),
row_shares=[3, 2],
),
column_shares=[3, 1],
),
)
# Choose the appropriate run mode based on provided arguments
if args.demo_image:
if not SAMPLE_IMAGE_PATH.exists():
download_file(SAMPLE_IMAGE_URL, SAMPLE_IMAGE_PATH)
run_from_sample_image(SAMPLE_IMAGE_PATH)
elif args.demo_video:
if not SAMPLE_VIDEO_PATH.exists():
download_file(SAMPLE_VIDEO_URL, SAMPLE_VIDEO_PATH)
run_from_video_capture(str(SAMPLE_VIDEO_PATH), args.max_frame)
elif args.image:
run_from_sample_image(args.image)
elif args.video:
run_from_video_capture(args.video, args.max_frame)
elif args.camera:
run_from_video_capture(int(args.camera), args.max_frame)
else:
if not SAMPLE_VIDEO_PATH.exists():
download_file(SAMPLE_VIDEO_URL, SAMPLE_VIDEO_PATH)
run_from_video_capture(str(SAMPLE_VIDEO_PATH), args.max_frame)
# Tear down Rerun script
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/gesture_detection/pyproject.toml
+21
View File
@@ -0,0 +1,21 @@
[project]
name = "gesture_detection"
version = "0.1.0"
requires-python = "<3.12" # TODO(ab): relax when mediapipe supports 3.12
readme = "README.md"
dependencies = [
"mediapipe==0.10.11 ; sys_platform != 'darwin'",
"mediapipe==0.10.9 ; sys_platform == 'darwin'", # https://github.com/google/mediapipe/issues/5188
"numpy",
"opencv-python>4.9",
"requests>=2.31,<3",
"rerun-sdk",
"tqdm",
]
[project.scripts]
gesture_detection = "gesture_detection:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/glomap/README.md
+17
View File
@@ -0,0 +1,17 @@
<!--[metadata]
title = "GLOMAP"
tags = ["3D", "Point cloud", "GLOMAP"]
source = "https://github.com/rerun-io/glomap"
thumbnail = "https://static.rerun.io/glomap-thumbnail/d0052fe6083b1e1d5698792ed0efec4d9cfd2f75/480w.png"
thumbnail_dimensions = [480, 342]
-->
https://vimeo.com/995792438?loop=1&autopause=0&background=1&muted=1&ratio=2802:1790
## Background
GLOMAP is a general purpose global structure-from-motion pipeline for image-based sparse reconstruction. As compared to COLMAP it provides a much more efficient and scalable reconstruction process, typically 1-2 orders of magnitude faster, with on-par or superior reconstruction quality. In the video we see it's global positioning step where it performs a joint global triangulation and camera position estimation starting from a randomly initialized state.
## Run the code
This is an external example. Check the [repository](https://github.com/rerun-io/glomap) for more information on how to run the code.
examples/python/graph_lattice/README.md
+28
View File
@@ -0,0 +1,28 @@
<!--[metadata]
title = "Graph lattice"
tags = ["Graph", "Layout"]
thumbnail = "https://static.rerun.io/graph_lattice/f9169da9c3f35b7260c9d74cd5be5fe710aec6a8/480w.png"
thumbnail_dimensions = [480, 269]
-->
This example shows different attributes that you can associate with nodes in a graph.
Since no explicit positions are passed for the nodes, Rerun will layout the graph automatically.
<picture>
<img src="https://static.rerun.io/graph_lattice/f9169da9c3f35b7260c9d74cd5be5fe710aec6a8/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/graph_lattice/f9169da9c3f35b7260c9d74cd5be5fe710aec6a8/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/graph_lattice/f9169da9c3f35b7260c9d74cd5be5fe710aec6a8/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/graph_lattice/f9169da9c3f35b7260c9d74cd5be5fe710aec6a8/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/graph_lattice/f9169da9c3f35b7260c9d74cd5be5fe710aec6a8/1200w.png">
</picture>
## Used Rerun types
[`GraphNodes`](https://www.rerun.io/docs/reference/types/archetypes/graph_nodes),
[`GraphEdges`](https://www.rerun.io/docs/reference/types/archetypes/graph_edges)
## Run the code
```bash
pip install -e examples/python/graph_lattice
python -m graph_lattice
```
examples/python/graph_lattice/graph_lattice.py
Executable
+75
View File
@@ -0,0 +1,75 @@
#!/usr/bin/env python3
"""Examples of logging graph data to Rerun and performing a force-based layout."""
from __future__ import annotations
import argparse
import itertools
import rerun as rr
NUM_NODES = 10
DESCRIPTION = """
# Graph Lattice
This is a minimal example that logs a graph (node-link diagram) that represents a lattice.
In this example, the node positions — and therefore the graph layout — are computed by Rerun internally.
The full source code for this example is available
[on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/graph_lattice).
""".strip()
def log_data() -> None:
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
coordinates = itertools.product(range(NUM_NODES), range(NUM_NODES))
nodes, colors = zip(
*[
(
str(i),
rr.components.Color([round((x / (NUM_NODES - 1)) * 255), round((y / (NUM_NODES - 1)) * 255), 0]),
)
for i, (x, y) in enumerate(coordinates)
],
strict=False,
)
rr.log(
"/lattice",
rr.GraphNodes(
nodes,
colors=colors,
labels=[f"({x}, {y})" for x, y in itertools.product(range(NUM_NODES), range(NUM_NODES))],
),
static=True,
)
edges = []
for x, y in itertools.product(range(NUM_NODES), range(NUM_NODES)):
if y > 0:
source = (y - 1) * NUM_NODES + x
target = y * NUM_NODES + x
edges.append((str(source), str(target)))
if x > 0:
source = y * NUM_NODES + (x - 1)
target = y * NUM_NODES + x
edges.append((str(source), str(target)))
rr.log("/lattice", rr.GraphEdges(edges, graph_type="directed"), static=True)
def main() -> None:
parser = argparse.ArgumentParser(description="Logs a graph lattice using the Rerun SDK.")
rr.script_add_args(parser)
args = parser.parse_args()
rr.script_setup(args, "rerun_example_graph_lattice")
log_data()
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/graph_lattice/pyproject.toml
+12
View File
@@ -0,0 +1,12 @@
[project]
name = "graph_lattice"
version = "0.1.0"
readme = "README.md"
dependencies = ["rerun-sdk"]
[project.scripts]
graph_lattice = "graph_lattice:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/graphs/README.md
+52
View File
@@ -0,0 +1,52 @@
<!--[metadata]
title = "Graphs"
tags = ["Graph", "Layout", "Node-link diagrams", "Bubble charts"]
thumbnail = "https://static.rerun.io/graphs/c1070214bed5e50c9e7d452835f32759b991383e/480w.png"
thumbnail_dimensions = [480, 480]
channel = "main"
include_in_manifest = true
-->
This example shows different types of graphs (and layouts) that you can visualize using Rerun.
<picture>
<img src="https://static.rerun.io/graphs/c1070214bed5e50c9e7d452835f32759b991383e/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/graphs/c1070214bed5e50c9e7d452835f32759b991383e/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/graphs/c1070214bed5e50c9e7d452835f32759b991383e/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/graphs/c1070214bed5e50c9e7d452835f32759b991383e/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/graphs/c1070214bed5e50c9e7d452835f32759b991383e/1200w.png">
</picture>
Rerun ships with an integrated engine to produce [force-based layouts](https://en.wikipedia.org/wiki/Force-directed_graph_drawing) to visualize graphs.
Force-directed layout approaches have to advantage that they are flexible and can therefore be used to create different kinds of visualizations.
This example shows different types of layouts:
* Regular force-directed layouts of node-link diagrams
* Bubble charts, which are based on packing circles
## Used Rerun types
[`GraphNodes`](https://www.rerun.io/docs/reference/types/archetypes/graph_nodes),
[`GraphEdges`](https://www.rerun.io/docs/reference/types/archetypes/graph_edges)
## Force-based layouts
To compute the graph layouts, Rerun implements a physics simulation that is very similar to [`d3-force`](https://d3js.org/d3-force). In particular, we implement the following forces:
* Centering force, which shifts the center of mass of the entire graph.
* Collision radius force, which resolves collisions between nodes in the graph, taking their radius into account.
* Many-Body force, which can be used to create attraction or repulsion between nodes.
* Link force, which acts like a spring between two connected nodes.
* Position force, which pull all nodes towards a given position, similar to gravity.
If you want to learn more about these forces, we recommend looking at the [D3 documentation](https://d3js.org/d3-force) as well.
Our implementation of the physics simulation is called _Fjädra_. You can find it on [GitHub](https://github.com/grtlr/fjadra) and on [`crates.io`](https://crates.io/crates/fjadra).
## Run the code
```bash
pip install -e examples/python/graphs
python -m graphs
```
examples/python/graphs/graphs.py
Executable
+206
View File
@@ -0,0 +1,206 @@
#!/usr/bin/env python3
"""Examples of logging graph data to Rerun and performing force-based layouts."""
from __future__ import annotations
import argparse
import itertools
import random
import numpy as np
import rerun as rr
import rerun.blueprint as rrb
from rerun.blueprint.archetypes.force_collision_radius import ForceCollisionRadius
from rerun.blueprint.archetypes.force_link import ForceLink
from rerun.blueprint.archetypes.force_many_body import ForceManyBody
color_scheme = [
[228, 26, 28, 255], # Red
[55, 126, 184, 255], # Blue
[77, 175, 74, 255], # Green
[152, 78, 163, 255], # Purple
[255, 127, 0, 255], # Orange
[255, 255, 51, 255], # Yellow
[166, 86, 40, 255], # Brown
[247, 129, 191, 255], # Pink
[153, 153, 153, 255], # Gray
]
DESCRIPTION = """
# Graphs
This example shows various graph visualizations that you can create using Rerun.
In this example, the node positions — and therefore the graph layout — are computed by Rerun internally using a force-based layout algorithm.
You can modify how these graphs look by changing the parameters of the force-based layout algorithm in the selection panel.
The full source code for this example is available
[on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/graphs).
""".strip()
# We want reproducible results
random.seed(42)
def log_lattice(num_nodes: int) -> None:
coordinates = itertools.product(range(num_nodes), range(num_nodes))
nodes, colors = zip(
*[
(
str(i),
rr.components.Color([round((x / (num_nodes - 1)) * 255), round((y / (num_nodes - 1)) * 255), 0, 255]),
)
for i, (x, y) in enumerate(coordinates)
],
strict=False,
)
rr.log(
"lattice",
rr.GraphNodes(
nodes,
colors=colors,
labels=[f"({x}, {y})" for x, y in itertools.product(range(num_nodes), range(num_nodes))],
),
static=True,
)
edges = []
for x, y in itertools.product(range(num_nodes), range(num_nodes)):
if y > 0:
source = (y - 1) * num_nodes + x
target = y * num_nodes + x
edges.append((str(source), str(target)))
if x > 0:
source = y * num_nodes + (x - 1)
target = y * num_nodes + x
edges.append((str(source), str(target)))
rr.log("lattice", rr.GraphEdges(edges, graph_type="directed"), static=True)
def log_trees() -> None:
nodes = ["root"]
radii = [42]
colors = [[81, 81, 81, 255]]
edges = []
# Randomly add nodes and edges to the graph
for i in range(50):
existing = random.choice(nodes)
new_node = str(i)
nodes.append(new_node)
radii.append(random.randint(10, 50))
colors.append(random.choice(color_scheme))
edges.append((existing, new_node))
rr.set_time("frame", sequence=i)
rr.log(
"node_link",
rr.GraphNodes(nodes, labels=nodes, radii=radii, colors=colors),
rr.GraphEdges(edges, graph_type=rr.GraphType.Directed),
)
rr.log(
"bubble_chart",
rr.GraphNodes(nodes, labels=nodes, radii=radii, colors=colors),
)
def log_markov_chain() -> None:
transition_matrix = np.array([
[0.8, 0.1, 0.1], # Transitions from sunny
[0.3, 0.4, 0.3], # Transitions from rainy
[0.2, 0.3, 0.5], # Transitions from cloudy
])
state_names = ["sunny", "rainy", "cloudy"]
# For this example, we use hardcoded positions.
positions = [[0, 0], [150, 150], [300, 0]]
inactive_color = [153, 153, 153, 255] # Gray
active_colors = [
[255, 127, 0, 255], # Orange
[55, 126, 184, 255], # Blue
[152, 78, 163, 255], # Purple
]
edges = [
(state_names[i], state_names[j])
for i in range(len(state_names))
for j in range(len(state_names))
if transition_matrix[i][j] > 0
]
edges.append(("start", "sunny"))
# We start in state "sunny"
state = "sunny"
for i in range(50):
current_state_index = state_names.index(state)
next_state_index = np.random.choice(range(len(state_names)), p=transition_matrix[current_state_index])
state = state_names[next_state_index]
colors = [inactive_color] * len(state_names)
colors[next_state_index] = active_colors[next_state_index]
rr.set_time("frame", sequence=i)
rr.log(
"markov_chain",
rr.GraphNodes(state_names, labels=state_names, colors=colors, positions=positions),
rr.GraphEdges(edges, graph_type="directed"),
)
def log_blueprint() -> None:
rr.send_blueprint(
rrb.Blueprint(
rrb.Grid(
rrb.GraphView(
origin="node_link",
name="Node-link diagram",
force_link=ForceLink(distance=60),
force_many_body=ForceManyBody(strength=-60),
),
rrb.GraphView(
origin="bubble_chart",
name="Bubble chart",
force_link=ForceLink(enabled=False),
force_many_body=ForceManyBody(enabled=False),
force_collision_radius=ForceCollisionRadius(enabled=True),
defaults=[rr.GraphNodes.from_fields(show_labels=False)],
),
rrb.GraphView(
origin="lattice",
name="Lattice",
force_link=ForceLink(distance=60),
force_many_body=ForceManyBody(strength=-60),
defaults=[rr.GraphNodes.from_fields(show_labels=False, radii=10)],
),
rrb.Horizontal(
rrb.GraphView(
origin="markov_chain",
name="Markov Chain",
# We don't need any forces for this graph, because the nodes have fixed positions.
),
rrb.TextDocumentView(origin="description", name="Description"),
),
),
),
)
def main() -> None:
parser = argparse.ArgumentParser(description="Logs various graphs using the Rerun SDK.")
rr.script_add_args(parser)
args = parser.parse_args()
rr.script_setup(args, "rerun_example_graphs")
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
log_trees()
log_lattice(10)
log_markov_chain()
log_blueprint()
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/graphs/pyproject.toml
+12
View File
@@ -0,0 +1,12 @@
[project]
name = "graphs"
version = "0.1.0"
readme = "README.md"
dependencies = ["rerun-sdk"]
[project.scripts]
graphs = "graphs:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
examples/python/hloc_glomap/README.md
+24
View File
@@ -0,0 +1,24 @@
<!--[metadata]
title = "Hierarchical-Localization and GLOMAP"
tags = ["2D", "3D", "COLMAP", "Pinhole camera", "Time series", "GLOMAP", ]
source = "https://github.com/pablovela5620/hloc-glomap"
thumbnail = "https://static.rerun.io/thumbnail/6a5b887927834a6ea3db9474ce1e843ecd28b3cf/480w.png"
thumbnail_dimensions = [480, 300]
-->
https://vimeo.com/1037241347?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=2802:1790
## Background
This examples allows use of the Hierarchical-Localization (hloc) repo and GLOMAP for easy and fast Structure-from-Motion with deep learned features and matchers. The Hierarchical-Localization repo (hloc for short) is a modular toolbox for state-of-the-art 6-DoF visual localization. It implements Hierarchical Localization, leveraging image retrieval and feature matching, and is fast, accurate, and scalable. This codebase combines and makes easily accessible years of research on image matching and Structure-from-Motion. GLOMAP is a general purpose global structure-from-motion pipeline for image-based sparse reconstruction. As compared to COLMAP it provides a much more efficient and scalable reconstruction process, typically 1-2 orders of magnitude faster, with on-par or superior reconstruction quality.
## Run the code
This is an external example. Check the [repository](https://github.com/pablovela5620/hloc-glomap) for more information on how to run the code.
TLDR: make sure you have the [Pixi package manager](https://pixi.sh/latest/#installation) installed and run
```
git clone https://github.com/pablovela5620/hloc-glomap.git
cd hloc-glomap
pixi run app
```
examples/python/human_pose_tracking/.gitignore
+2
View File
@@ -0,0 +1,2 @@
dataset/
model/
examples/python/human_pose_tracking/README.md
+149
View File
@@ -0,0 +1,149 @@
<!--[metadata]
title = "Human pose tracking"
tags = ["MediaPipe", "Keypoint detection", "2D", "3D"]
thumbnail = "https://static.rerun.io/human-pose-tracking/5d62a38b48bed1467698d4dc95c1f9fba786d254/480w.png"
thumbnail_dimensions = [480, 480]
-->
Use the [MediaPipe Pose Landmark Detection](https://developers.google.com/mediapipe/solutions/vision/pose_landmarker) solution to detect and track a human pose in video.
<picture data-inline-viewer="examples/human_pose_tracking">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/human_pose_tracking/37d47fe7e3476513f9f58c38da515e2cd4a093f9/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/human_pose_tracking/37d47fe7e3476513f9f58c38da515e2cd4a093f9/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/human_pose_tracking/37d47fe7e3476513f9f58c38da515e2cd4a093f9/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/human_pose_tracking/37d47fe7e3476513f9f58c38da515e2cd4a093f9/1200w.png">
<img src="https://static.rerun.io/human_pose_tracking/37d47fe7e3476513f9f58c38da515e2cd4a093f9/full.png" alt="">
</picture>
## Used Rerun types
[`Image`](https://www.rerun.io/docs/reference/types/archetypes/image), [`Points2D`](https://www.rerun.io/docs/reference/types/archetypes/points2d), [`Points3D`](https://www.rerun.io/docs/reference/types/archetypes/points3d), [`ClassDescription`](https://www.rerun.io/docs/reference/types/datatypes/class_description), [`AnnotationContext`](https://www.rerun.io/docs/reference/types/archetypes/annotation_context), [`SegmentationImage`](https://www.rerun.io/docs/reference/types/archetypes/segmentation_image)
## Background
Human pose tracking is a task in computer vision that focuses on identifying key body locations, analyzing posture, and categorizing movements.
At the heart of this technology is a pre-trained machine-learning model to assess the visual input and recognize landmarks on the body in both image coordinates and 3D world coordinates.
The use cases and applications of this technology include but are not limited to Human-Computer Interaction, Sports Analysis, Gaming, Virtual Reality, Augmented Reality, Health, etc.
In this example, the [MediaPipe Pose Landmark Detection](https://developers.google.com/mediapipe/solutions/vision/pose_landmarker) solution was utilized to detect and track human pose landmarks and produces segmentation masks for humans.
Rerun was employed to visualize the output of the Mediapipe solution over time to make it easy to analyze the behavior.
## Logging and visualizing with Rerun
The visualizations in this example were created with the following Rerun code.
### Timelines
For each processed video frame, all data sent to Rerun is associated with the two [`timelines`](https://www.rerun.io/docs/concepts/logging-and-ingestion/timelines) `time` and `frame_idx`.
```python
rr.set_time("time", duration=bgr_frame.time)
rr.set_time("frame_idx", sequence=bgr_frame.idx)
```
### Video
The input video is logged as a sequence of
[`Image`](https://www.rerun.io/docs/reference/types/archetypes/image) objects to the 'Video' entity.
```python
rr.log("video/rgb", rr.Image(rgb).compress(jpeg_quality=75))
```
### Segmentation mask
The segmentation result is logged through a combination of two archetypes. The segmentation
image itself is logged as a
[`SegmentationImage`](https://www.rerun.io/docs/reference/types/archetypes/segmentation_image) and
contains the id for each pixel. The color is determined by the
[`AnnotationContext`](https://www.rerun.io/docs/reference/types/archetypes/annotation_context) which is
logged with `static=True` as it should apply to the whole sequence.
#### Label mapping
```python
rr.log(
"video/mask",
rr.AnnotationContext([
rr.AnnotationInfo(id=0, label="Background"),
rr.AnnotationInfo(id=1, label="Person", color=(0, 0, 0)),
]),
static=True,
)
```
#### Segmentation image
```python
rr.log("video/mask", rr.SegmentationImage(binary_segmentation_mask.astype(np.uint8)))
```
### Body pose points
Logging the body pose as a skeleton involves specifying the connectivity of its keypoints (i.e., pose landmarks), extracting the pose landmarks, and logging them as points to Rerun. In this example, both the 2D and 3D estimates from Mediapipe are visualized.
The skeletons are logged through a combination of two archetypes. First, a static
[`ClassDescription`](https://www.rerun.io/docs/reference/types/datatypes/class_description) is logged, that contains the information which maps keypoint ids to labels and how to connect
the keypoints. By defining these connections Rerun will automatically add lines between them. Mediapipe provides the `POSE_CONNECTIONS` variable which contains the list of `(from, to)` landmark indices that define the connections. Second, the actual keypoint positions are logged in 2D
and 3D as [`Points2D`](https://www.rerun.io/docs/reference/types/archetypes/points2d) and
[`Points3D`](https://www.rerun.io/docs/reference/types/archetypes/points3d) archetypes, respectively.
#### Label mapping and keypoint connections
```python
rr.log(
"/",
rr.AnnotationContext(
rr.ClassDescription(
info=rr.AnnotationInfo(id=1, label="Person"),
keypoint_annotations=[rr.AnnotationInfo(id=lm.value, label=lm.name) for lm in mp_pose.PoseLandmark],
keypoint_connections=mp_pose.POSE_CONNECTIONS,
)
),
static=True,
)
```
#### 2D points
```python
rr.log("video/pose/points", rr.Points2D(landmark_positions_2d, class_ids=1, keypoint_ids=mp_pose.PoseLandmark))
```
#### 3D points
```python
rr.log(
"person/pose/points",
rr.Points3D(landmark_positions_3d, class_ids=1, keypoint_ids=mp_pose.PoseLandmark),
)
```
## Run the code
To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
```bash
pip install --upgrade rerun-sdk # install the latest Rerun SDK
git clone git@github.com:rerun-io/rerun.git # Clone the repository
cd rerun
git checkout latest # Check out the commit matching the latest SDK release
```
Install the necessary libraries specified in the requirements file:
```bash
pip install -e examples/python/human_pose_tracking
```
To experiment with the provided example, simply execute the main Python script:
```bash
python -m human_pose_tracking # run the example
```
If you wish to customize it for various videos, adjust the maximum frames, or explore additional features, use the CLI with the `--help` option for guidance:
```bash
python -m human_pose_tracking --help
```
examples/python/human_pose_tracking/human_pose_tracking.py
Executable
+263
View File
@@ -0,0 +1,263 @@
#!/usr/bin/env python3
"""Use the MediaPipe Pose solution to detect and track a human pose in video."""
from __future__ import annotations
import argparse
import logging
import os
from contextlib import closing
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final
import cv2
import mediapipe as mp
import mediapipe.python.solutions.pose as mp_pose
import numpy as np
import numpy.typing as npt
import requests
import rerun as rr # pip install rerun-sdk
import rerun.blueprint as rrb
if TYPE_CHECKING:
from collections.abc import Iterator
DESCRIPTION = """
# Human pose tracking
This example uses Rerun to visualize the output of [MediaPipe](https://developers.google.com/mediapipe)-based tracking
of a human pose in 2D and 3D.
The full source code for this example is available
[on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/human_pose_tracking).
""".strip()
EXAMPLE_DIR: Final = Path(os.path.dirname(__file__))
DATASET_DIR: Final = EXAMPLE_DIR / "dataset" / "pose_movement"
MODEL_DIR: Final = EXAMPLE_DIR / "model" / "pose_movement"
DATASET_URL_BASE: Final = "https://storage.googleapis.com/rerun-example-datasets/pose_movement"
MODEL_URL_TEMPLATE: Final = "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_{model_name}/float16/latest/pose_landmarker_{model_name}.task"
def track_pose(video_path: str, model_path: str, *, max_frame_count: int | None) -> None:
options = mp.tasks.vision.PoseLandmarkerOptions(
base_options=mp.tasks.BaseOptions(
model_asset_path=model_path,
),
running_mode=mp.tasks.vision.RunningMode.VIDEO,
output_segmentation_masks=True,
)
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
rr.log(
"/",
rr.AnnotationContext(
rr.ClassDescription(
info=rr.AnnotationInfo(id=1, label="Person"),
keypoint_annotations=[rr.AnnotationInfo(id=lm.value, label=lm.name) for lm in mp_pose.PoseLandmark],
keypoint_connections=mp_pose.POSE_CONNECTIONS,
),
),
static=True,
)
# Use a separate annotation context for the segmentation mask.
rr.log(
"video/mask",
rr.AnnotationContext([
rr.AnnotationInfo(id=0, label="Background"),
rr.AnnotationInfo(id=1, label="Person", color=(0, 0, 0)),
]),
static=True,
)
rr.log("person", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, static=True)
pose_landmarker = mp.tasks.vision.PoseLandmarker.create_from_options(options)
with closing(VideoSource(video_path)) as video_source:
for idx, bgr_frame in enumerate(video_source.stream_bgr()):
if max_frame_count is not None and idx >= max_frame_count:
break
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=bgr_frame.data)
rr.set_time("time", duration=bgr_frame.time)
rr.set_time("frame_idx", sequence=bgr_frame.idx)
results = pose_landmarker.detect_for_video(mp_image, int(bgr_frame.time * 1000))
h, w, _ = bgr_frame.data.shape
landmark_positions_2d = read_landmark_positions_2d(results, w, h)
rr.log("video/bgr", rr.Image(bgr_frame.data, color_model="BGR").compress(jpeg_quality=75))
if landmark_positions_2d is not None:
rr.log(
"video/pose/points",
rr.Points2D(landmark_positions_2d, class_ids=1, keypoint_ids=mp_pose.PoseLandmark),
)
landmark_positions_3d = read_landmark_positions_3d(results)
if landmark_positions_3d is not None:
rr.log(
"person/pose/points",
rr.Points3D(landmark_positions_3d, class_ids=1, keypoint_ids=mp_pose.PoseLandmark),
)
if results.segmentation_masks is not None:
segmentation_mask = results.segmentation_masks[0].numpy_view()
binary_segmentation_mask = segmentation_mask > 0.5
rr.log("video/mask", rr.SegmentationImage(binary_segmentation_mask.astype(np.uint8)))
def read_landmark_positions_2d(
results: Any,
image_width: int,
image_height: int,
) -> npt.NDArray[np.float32] | None:
if results.pose_landmarks is None or len(results.pose_landmarks) == 0:
return None
else:
pose_landmarks = results.pose_landmarks[0]
normalized_landmarks = [pose_landmarks[lm] for lm in mp_pose.PoseLandmark]
return np.array([(image_width * lm.x, image_height * lm.y) for lm in normalized_landmarks])
def read_landmark_positions_3d(
results: Any,
) -> npt.NDArray[np.float32] | None:
if results.pose_landmarks is None or len(results.pose_landmarks) == 0:
return None
else:
pose_landmarks = results.pose_landmarks[0]
landmarks = [pose_landmarks[lm] for lm in mp_pose.PoseLandmark]
return np.array([(lm.x, lm.y, lm.z) for lm in landmarks])
@dataclass
class VideoFrame:
data: cv2.typing.MatLike
time: float
idx: int
class VideoSource:
def __init__(self, path: str) -> None:
self.capture = cv2.VideoCapture(path)
if not self.capture.isOpened():
logging.error("Couldn't open video at %s", path)
def close(self) -> None:
self.capture.release()
def stream_bgr(self) -> Iterator[VideoFrame]:
while self.capture.isOpened():
idx = int(self.capture.get(cv2.CAP_PROP_POS_FRAMES))
is_open, bgr = self.capture.read()
time_ms = self.capture.get(cv2.CAP_PROP_POS_MSEC)
if not is_open:
break
yield VideoFrame(data=bgr, time=time_ms * 1e-3, idx=idx)
def get_downloaded_video_path(dataset_dir: Path, video_name: str) -> str:
video_file_name = f"{video_name}.mp4"
destination_path = dataset_dir / video_file_name
if destination_path.exists():
logging.info("%s already exists. No need to download", destination_path)
return str(destination_path)
source_path = f"{DATASET_URL_BASE}/{video_file_name}"
logging.info("Downloading video from %s to %s", source_path, destination_path)
os.makedirs(dataset_dir.absolute(), exist_ok=True)
download(source_path, destination_path)
return str(destination_path)
def get_downloaded_model_path(model_dir: Path, model_name: str) -> str:
model_file_name = f"{model_name}.task"
destination_path = model_dir / model_file_name
if destination_path.exists():
logging.info("%s already exists. No need to download", destination_path)
return str(destination_path)
model_url = MODEL_URL_TEMPLATE.format(model_name=model_name)
logging.info("Downloading model from %s to %s", model_url, destination_path)
download(model_url, destination_path)
return str(destination_path)
def download(url: str, destination_path: Path) -> None:
os.makedirs(destination_path.parent, exist_ok=True)
with requests.get(url, stream=True) as req:
req.raise_for_status()
with open(destination_path, "wb") as f:
f.writelines(req.iter_content(chunk_size=8192))
def main() -> None:
# Ensure the logging gets written to stderr:
logging.getLogger().addHandler(logging.StreamHandler())
logging.getLogger().setLevel("INFO")
parser = argparse.ArgumentParser(description="Uses the MediaPipe Pose solution to track a human pose in video.")
parser.add_argument(
"--video",
type=str,
default="backflip",
choices=["backflip", "soccer"],
help="The example video to run on.",
)
parser.add_argument("--dataset-dir", type=Path, default=DATASET_DIR, help="Directory to save example videos to.")
parser.add_argument("--video-path", type=str, default="", help="Full path to video to run on. Overrides `--video`.")
parser.add_argument(
"--model",
type=str,
default="heavy",
choices=["lite", "full", "heavy"],
help="The mediapipe model to use (see https://developers.google.com/mediapipe/solutions/vision/pose_landmarker).",
)
parser.add_argument("--model-dir", type=Path, default=MODEL_DIR, help="Directory to save downloaded model to.")
parser.add_argument("--model-path", type=str, default="", help="Full path of mediapipe model. Overrides `--model`.")
parser.add_argument(
"--max-frame",
type=int,
help="Stop after processing this many frames. If not specified, will run until interrupted.",
)
rr.script_add_args(parser)
args = parser.parse_args()
rr.script_setup(
args,
"rerun_example_human_pose_tracking",
default_blueprint=rrb.Horizontal(
rrb.Vertical(
rrb.Spatial2DView(origin="video", name="Result"),
rrb.Spatial3DView(origin="person", name="3D pose"),
),
rrb.Vertical(
rrb.Spatial2DView(origin="video/bgr", name="Raw video"),
rrb.TextDocumentView(origin="description", name="Description"),
row_shares=[2, 3],
),
column_shares=[3, 2],
),
)
video_path = args.video_path # type: str
if not video_path:
video_path = get_downloaded_video_path(args.dataset_dir, args.video)
model_path = args.model_path # type: str
if not args.model_path:
model_path = get_downloaded_model_path(args.model_dir, args.model)
track_pose(video_path, model_path, max_frame_count=args.max_frame)
rr.script_teardown(args)
if __name__ == "__main__":
main()
examples/python/human_pose_tracking/pyproject.toml
+25
View File
@@ -0,0 +1,25 @@
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "human_pose_tracking"
version = "0.1.0"
requires-python = "<3.12" # TODO(ab): relax when mediapipe supports 3.12
readme = "README.md"
dependencies = [
"mediapipe==0.10.11 ; sys_platform != 'darwin'",
"mediapipe==0.10.9 ; sys_platform == 'darwin'", # https://github.com/google/mediapipe/issues/5188
"numpy",
"opencv-python>4.6", # Avoid opencv-4.6 since it rotates images incorrectly (https://github.com/opencv/opencv/issues/22088)
"requests>=2.31,<3",
"rerun-sdk",
]
[project.scripts]
human_pose_tracking = "human_pose_tracking:main"
[tool.rerun-example]
skip = false
extra-args = "--max-fram=e10"
examples/python/imu_signals/.gitignore
+2
View File
@@ -0,0 +1,2 @@
dataset/*
examples/python/imu_signals/README.md
+69
View File
@@ -0,0 +1,69 @@
<!--[metadata]
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"
include_in_manifest = true
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).
<picture>
<img src="https://static.rerun.io/imu_signals/1184ab6e2df3275b8b7a574d7f0e42b1aed8343a/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/imu_signals/1184ab6e2df3275b8b7a574d7f0e42b1aed8343a/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/imu_signals/1184ab6e2df3275b8b7a574d7f0e42b1aed8343a/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/imu_signals/1184ab6e2df3275b8b7a574d7f0e42b1aed8343a/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/imu_signals/1184ab6e2df3275b8b7a574d7f0e42b1aed8343a/1200w.png">
</picture>
## Background
This example shows how to log multi-dimensional signals efficiently using the [`rr.send_columns()`](https://ref.rerun.io/docs/python/0.22.1/common/columnar_api/#rerun.send_columns) API.
The API automatically selects the right partition sizes, making it simple to log scalar signals like this:
```py
# Load IMU data from CSV into a dataframe
imu_data = pd.read_csv(
cwd / 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"])
# Extract gyroscope data (x, y, z axes) and log it to a single entity.
gyro = imu_data[["gyro.x", "gyro.y", "gyro.z"]]
rr.send_columns("/gyroscope", indexes=[times], columns=rr.Scalars.columns(scalars=gyro))
# Extract accelerometer data (x, y, z axes) and log it to a single entity.
accel = imu_data[["accel.x", "accel.y", "accel.z"]]
rr.send_columns("/accelerometer", indexes=[times], columns=rr.Scalars.columns(scalars=accel))
```
## Running
Install the example package:
```bash
pip install -e examples/python/imu_signals
```
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/)

Some files were not shown because too many files have changed in this diff Show More