chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,2 @@
|
||||
input_data/meshes/tablelegs.obj filter=lfs diff=lfs merge=lfs -text
|
||||
input_data/meshes/tabletop.obj filter=lfs diff=lfs merge=lfs -text
|
||||
@@ -0,0 +1 @@
|
||||
output/*.rrd
|
||||
@@ -0,0 +1,176 @@
|
||||
<!--[metadata]
|
||||
title = "Robot data preprocessing example"
|
||||
tags = ["API example"]
|
||||
thumbnail = "https://static.rerun.io/robot_postprocessing_thumb/ae27d24c3f530e71ed15ce47745eda56312ad014/480w.png"
|
||||
thumbnail_dimensions = [480, 299]
|
||||
-->
|
||||
|
||||
This example demonstrates how Rerun's [chunk processing API](https://rerun.io/docs/concepts/logging-and-ingestion/chunk-processing-api) can be used to assemble a robot recording from multiple file sources, including preprocessing to modify or augment the data.
|
||||
|
||||
<picture>
|
||||
<img src="https://static.rerun.io/robot_postprocessing_thumb/ae27d24c3f530e71ed15ce47745eda56312ad014/full.png" alt="">
|
||||
<source media="(max-width: 480px)" srcset="https://static.rerun.io/robot_postprocessing_thumb/ae27d24c3f530e71ed15ce47745eda56312ad014/480w.png">
|
||||
<source media="(max-width: 768px)" srcset="https://static.rerun.io/robot_postprocessing_thumb/ae27d24c3f530e71ed15ce47745eda56312ad014/768w.png">
|
||||
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/robot_postprocessing_thumb/ae27d24c3f530e71ed15ce47745eda56312ad014/1024w.png">
|
||||
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/robot_postprocessing_thumb/ae27d24c3f530e71ed15ce47745eda56312ad014/1200w.png">
|
||||
</picture>
|
||||
|
||||
## Introduction
|
||||
|
||||
### Input data
|
||||
|
||||
While the example uses simulated data, it's intentionally designed to cover real-world challenges that should sound familiar to most roboticists:
|
||||
|
||||
- incomplete data requiring preprocessing
|
||||
- custom data types
|
||||
- bugs in the recorded data
|
||||
- data spread across multiple files in different formats
|
||||
|
||||
Specifically, we use a recording of a dual-robot-arm setup, consisting of:
|
||||
|
||||
| `episode.mcap` | `offsets.json` | URDF files |
|
||||
| --- | --- | --- |
|
||||
| Base recording (videos, sensors, …). | Static world offsets for each robot. | Robot & scene models as [URDF](https://en.wikipedia.org/wiki/URDF).
|
||||
| Some cameras have wrong parameters.<br>No dynamic 3D transforms were recorded,<br>only joint states in a custom Protobuf schema. | Saved outside of base recording. | `robot.urdf`, `scene.urdf`, mesh data |
|
||||
|
||||
### Goals
|
||||
|
||||
Our task is to handle and process all the different data sources:
|
||||
|
||||
- read, convert and fix MCAP data
|
||||
- compute 3D transforms using MCAP joint states and URDF
|
||||
- handle URDFs
|
||||
- add `scene.urdf` and 2x `robot.urdf`
|
||||
- modify visual meshes with a custom color & transparency per robot
|
||||
- add static transforms from JSON
|
||||
|
||||
…and merge them into one coherent recording.
|
||||
|
||||
## Processing pipeline
|
||||
|
||||
Solving such a task in an elegant way requires a non-trivial amount of engineering, but Rerun's [chunk processing API](https://rerun.io/docs/concepts/logging-and-ingestion/chunk-processing-api) gives us all the tools to properly structure the pipeline:
|
||||
|
||||
<!-- Figma: https://www.figma.com/board/xOvrUjklsfPH8OB3GUDnax/Michael-s-scratchpad-%F0%9F%91%A8%F0%9F%8F%BB%E2%80%8D%F0%9F%8E%A8?node-id=0-1&t=3uvGqaikPNKr80iH-1 -->
|
||||
|
||||
<picture>
|
||||
<img src="https://static.rerun.io/robot_postprocessing/722231a7e1523c45f22a2fa4162a9e960df88f08/full.png" alt="">
|
||||
<source media="(max-width: 480px)" srcset="https://static.rerun.io/robot_postprocessing/722231a7e1523c45f22a2fa4162a9e960df88f08/480w.png">
|
||||
<source media="(max-width: 768px)" srcset="https://static.rerun.io/robot_postprocessing/722231a7e1523c45f22a2fa4162a9e960df88f08/768w.png">
|
||||
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/robot_postprocessing/722231a7e1523c45f22a2fa4162a9e960df88f08/1024w.png">
|
||||
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/robot_postprocessing/722231a7e1523c45f22a2fa4162a9e960df88f08/1200w.png">
|
||||
</picture>
|
||||
|
||||
The example code implements this pipeline and contains several explanatory comments.
|
||||
We recommend reading the concept explanations below, before going through the `main()` function of [`robot_data_preprocessing.py`](robot_data_preprocessing.py) to understand the code structure.
|
||||
|
||||
> ℹ️ Note that we create two separate RRD files in this example.
|
||||
For the Rerun Viewer or Catalog, both *physical* files form one [*logical* recording](https://rerun.io/docs/concepts/logging-and-ingestion/recordings#logical-vs-physical-recordings) since they specify the same recording ID.
|
||||
|
||||
### Chunk streams
|
||||
|
||||
[Chunks](https://rerun.io/docs/concepts/logging-and-ingestion/chunks) are the core datastructure of Rerun.
|
||||
In this example, we use [chunk _streams_](https://rerun.io/docs/concepts/logging-and-ingestion/chunk-processing-api) as the "glue" of our pipeline.
|
||||
|
||||
In a nutshell, `LazyChunkStream`s allow us to define how `Chunk`s get routed through filtering, transformation and output steps.
|
||||
As the name suggests, these streams are lazily evaluated.
|
||||
We use an expressive Python API to define the pipeline, but the final execution happens in a multithreaded, GIL-free execution engine written in Rust for maximum efficiency.
|
||||
|
||||
In this example, we use the following sources that can emit `LazyChunkStream`s:
|
||||
* `McapReader.stream()` for the MCAP recording
|
||||
* `UrdfTree.stream()` for the URDF models
|
||||
* manually constructed `LazyChunkStream` for the custom JSON file, using [`Chunk.from_columns(…)`](https://rerun.io/docs/concepts/logging-and-ingestion/chunks#sending-actual-chunks-sendchunks)
|
||||
|
||||
### Lenses
|
||||
|
||||
[Lenses](https://rerun.io/docs/concepts/query-and-transform/lenses) allow us to modify the chunks' components via [`MutateLens`](https://rerun.io/docs/concepts/query-and-transform/lenses#mutate-lenses), or to derive completely new components from them via [`DeriveLens`](https://rerun.io/docs/concepts/query-and-transform/lenses#derive-lenses).
|
||||
In both cases, we use [`Selector`](https://rerun.io/docs/concepts/query-and-transform/lenses#selectors)s to extract component fields we're interested in, and pipe them through custom transformation functions.
|
||||
|
||||
#### `MutateLens` example
|
||||
|
||||
A simple `MutateLens` used in this example is the one that fixes the swapped `Pinhole:resolution` component of the external camera streams:
|
||||
```python
|
||||
mcap_stream.lenses(
|
||||
MutateLens(
|
||||
"Pinhole:resolution",
|
||||
Selector(".").pipe(
|
||||
lambda resolution: pa.array(
|
||||
[(height, width) for width, height in resolution.to_pylist()], type=resolution.type
|
||||
)
|
||||
),
|
||||
),
|
||||
content=["/external/cam_low", "/external/cam_high"],
|
||||
output_mode="forward_unmatched",
|
||||
)
|
||||
```
|
||||
The `content` filter makes sure that this lens only gets applied to the external camera entities, while the `output_mode` makes sure we forward the other pinhole entities that don't match unchanged (here: the robot cameras that don't require the fix).
|
||||
|
||||
#### `DeriveLens` example
|
||||
|
||||
A more complex lens setup is required for the forward kinematics, i.e. to compute the 3D transforms from joint values (angles, distances).
|
||||
For this we need the recorded joint states from the MCAP, as well as the URDF for the kinematic structure.
|
||||
|
||||
Our MCAP file contains joint states encoded in a custom Protobuf schema:
|
||||
```proto
|
||||
message JointState {
|
||||
google.protobuf.Timestamp timestamp = 1;
|
||||
repeated string joint_names = 2;
|
||||
repeated double joint_positions = 3;
|
||||
repeated double joint_velocities = 4;
|
||||
repeated double joint_efforts = 5;
|
||||
}
|
||||
```
|
||||
This custom schema is not part of the [directly supported message types](https://rerun.io/docs/concepts/logging-and-ingestion/mcap/message-formats) of the MCAP importer (like e.g. the video streams). But thanks to [schema reflection](https://rerun.io/docs/concepts/logging-and-ingestion/mcap/message-formats#schema-reflection), we still get chunks with queryable Rerun components that we can process in our streams.
|
||||
|
||||
Each input row of joint states contains `N` joint values that map to `N` 3D transforms, for which we want to have a dedicated output row with [`Transform3D`](https://rerun.io/docs/reference/types/archetypes/transform3d) each.
|
||||
Due to this input-to-output row length mismatch, we use two sequential lenses:
|
||||
1. For each joint state message…
|
||||
* select the joint names and values
|
||||
* use [`UrdfTree.compute_joint_transform_batches`](https://ref.rerun.io/docs/python/stable/urdf/#rerun.urdf.UrdfTree.compute_joint_transform_batches)
|
||||
* output a single row with a list of `N` 3D transforms.
|
||||
2. Scatter each computed row into `N` rows with `Transform3D` component columns.
|
||||
|
||||
#### Others
|
||||
|
||||
Besides fixing camera data and computing forward kinematics, we also apply lenses for smaller things like URDF model colorization.
|
||||
|
||||
Finally, the streams are merged and written to two RRD files with the same recording ID to form layers of a single [logical recording](https://rerun.io/docs/concepts/logging-and-ingestion/recordings#logical-vs-physical-recordings).
|
||||
We use two RRDs for demonstration purposes, but merging into a single RRD would be also possible.
|
||||
|
||||
See the code for all implementation details.
|
||||
|
||||
## Run the code
|
||||
|
||||
```bash
|
||||
pip install -e examples/python/robot_data_preprocessing
|
||||
python -m robot_data_preprocessing
|
||||
```
|
||||
|
||||
The resulting RRDs can be opened in the viewer:
|
||||
```bash
|
||||
rerun examples/python/robot_data_preprocessing/output/*.rrd
|
||||
```
|
||||
Since we use consistent recording IDs, the two output RRD layers show up as a single recording.
|
||||
|
||||
<image src="https://static.rerun.io/e8b3975732ed5f42e125b0c80b487e97d8e99041_robot_postprocessing.gif" width=500/>
|
||||
<!-- MP4 version: https://static.rerun.io/5e0d6be2f9a1c21686ff8177a9085c6858ec3f74_robot_postprocessing.mp4 -->
|
||||
|
||||
## Summary
|
||||
|
||||
We showed how a non-trivial robotics problem can be solved through a structured data pipeline.
|
||||
The chunk processing API provides the tools to build such custom pipelines in a compact manner (here: < 200 lines of Python code) while having a powerful execution engine under the hood.
|
||||
|
||||
We also demonstrated how recording IDs can be used to structure RRDs into logical recordings, allowing also to potentially add more layers (e.g. for metadata or extra sensor data).
|
||||
|
||||
Documentation links for further reading:
|
||||
* [Chunk processing API](https://rerun.io/docs/concepts/logging-and-ingestion/chunk-processing-api)
|
||||
* [Lenses API](https://rerun.io/docs/concepts/query-and-transform/lenses)
|
||||
* [Recordings](https://rerun.io/docs/concepts/logging-and-ingestion/recordings)
|
||||
* [Working with MCAP](https://rerun.io/docs/howto/logging-and-ingestion/mcap)
|
||||
* [Loading URDF models](https://rerun.io/docs/howto/logging-and-ingestion/urdf)
|
||||
|
||||
## Going further
|
||||
|
||||
In a real-world setting, this kind of processing would be only the first step of data curation, to finalize multiple raw recordings before ingesting them to central storage.
|
||||
|
||||
With Rerun, this would mean registering a dataset to a [catalog server](https://rerun.io/docs/concepts/how-does-rerun-work#catalog-server) (either via Rerun Hub for enterprise scalability, or using the open-source `rerun server` for small-scale local development).
|
||||
This enables e.g. to perform [queries across recordings](https://rerun.io/docs/concepts/query-and-transform/dataframe-queries) for analytics or to export training data.
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:c27a3f89b57d662ba6a3146c9fd8e37fb1566320b46dbb35d5e2e28670a9707a
|
||||
size 33133865
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:10e6e10a1a5c1efcf743884f7f0dad14c6cc3e15790695bfe5b7038c43f8a4c4
|
||||
size 448184
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:f3540af6e3e1883f101417a8c8a0e3f6de829e2aaa965428cab7b3135c46c645
|
||||
size 34084
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:463d04b9b507e1c8c070e5545eb1b62a4ce9a6d9b7d207675a9596a084b0dc4b
|
||||
size 673284
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:9cb543c0258cb01fc10d408595585ba9171389061cb218c6d53fbcaf36783ce7
|
||||
size 140684
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3b22a0e420f0265c4fbed1238f8550f460fe6b6b70e931b1699c657626373de5
|
||||
size 140684
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:75e40b0b807f51d238638cd44e787413d8590c1fcd28dff7beaa549c126f1fd0
|
||||
size 234284
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:8a9a84dd9c9a67687e5ad0b37a7ed334dbfea6b0f995fef28524cc8da5f21ed2
|
||||
size 1242284
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:7a8b99ef84ffd7e91382af9aa733f815a7747c207f64e0499c5889a58d58c8de
|
||||
size 40084
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:5b5203c92dfe03b94ed5a679c6cec1f61e48d6d202b5fa03abb032d63313676e
|
||||
size 40084
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:39029c2f95faef3dab1340e4beec221d6866e694e4df0b22795eed6f00b5f893
|
||||
size 40084
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:acaff3be15f686064b6ec8b985758866b2bf51aed0a9b793d8c8242dde43ee0a
|
||||
size 84084
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3bb04a77ca20a14178aa59f8ef1ef50d26695e748639677b24f0be11efc016d8
|
||||
size 84084
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:25c87bd6a975f1a561e547e8c490ffe4b59d8a962634548242b1f4a5c6607268
|
||||
size 40084
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:63ac36c2f4ab9d75341b857212edfa71f45b4eff66068b2dac4c509155cda6c8
|
||||
size 285284
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:64540fa6c62d8358336bc3cddfe2aebb40c4323822737544edb67b1ad25800fc
|
||||
size 282984
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:bd3231c1656b55b871e777825897056368f9934230998c4d92931c2a7eca4150
|
||||
size 29884
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:633fb1643f22b9c1f9b6b8c5d16c4b1216a77a3c216efd0e70d5138492c7a44d
|
||||
size 64484
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:c2a9fe629c68a475409ecd536761e6cebcdc5a6c858fddbe01bc266017c079c0
|
||||
size 46284
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:413868d347309dff92d67d45144a7311ef833354c0ea1c26bea0f450d3909dc7
|
||||
size 78484
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a140433a98d95c9dfe2e3ced4c581cd9fc97a1b3f84a2d3a13f14bfa8b10fdae
|
||||
size 124984
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:d72ff84f6a405796834e2f297fbe453557fdd7a8e31691cd6f1110ef0aecca69
|
||||
size 51084
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:54d031aae8599f7ae69adaa834bd5435928484e5815ce4acce964522a115c2db
|
||||
size 1899284
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:2514d9d8cd8974e29d95116d9f5b50c1beba5e746e3955f6ccb5ccb594ac6796
|
||||
size 10174
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3d88abdc152ebe3dabec44100fcbcbb179c132b1b903c97220225a8ba0232687
|
||||
size 4702
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:0c939865a1adfc8a148817d6e12d8a04541a192ab9b6af37c94e3cae8330aa64
|
||||
size 1955684
|
||||
@@ -0,0 +1,16 @@
|
||||
{
|
||||
"transforms": [
|
||||
{
|
||||
"parent": "world",
|
||||
"child": "left_base_link",
|
||||
"translation": [-0.4575, -0.019, 0.02],
|
||||
"quaternion_xyzw": [0.0, 0.0, 0.0, 1.0]
|
||||
},
|
||||
{
|
||||
"parent": "world",
|
||||
"child": "right_base_link",
|
||||
"translation": [0.4575, -0.019, 0.02],
|
||||
"quaternion_xyzw": [0.0, 0.0, -1.0, 0.0]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:f5773d3a0b306bedebe2fa174fd234d331d81988dcf32042ad93065ce74508ce
|
||||
size 12818
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:673d769291613be4abcc286a7a2cd213b5e28de8d9238cd0d69bf6b7d0e49e8a
|
||||
size 18386
|
||||
@@ -0,0 +1,15 @@
|
||||
[project]
|
||||
name = "robot_data_preprocessing"
|
||||
version = "0.1.0"
|
||||
readme = "README.md"
|
||||
dependencies = ["pyarrow", "rerun-sdk"]
|
||||
|
||||
[project.scripts]
|
||||
robot_data_preprocessing = "robot_data_preprocessing:main"
|
||||
|
||||
[tool.rerun-example]
|
||||
skip = true
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
@@ -0,0 +1,196 @@
|
||||
"""
|
||||
Demonstrates how to use Rerun's chunk processing API to assemble a robot recording
|
||||
from multiple file sources (MCAP, custom data, URDF, …):
|
||||
|
||||
- fix recording errors
|
||||
- add external static data
|
||||
- compute joint transforms using URDF
|
||||
- insert URDF assets
|
||||
- …
|
||||
|
||||
The resulting merged stream is saved to an RRD file, which can be
|
||||
opened in the Rerun viewer or registered to a dataset catalog.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
import rerun as rr
|
||||
from rerun.experimental import Chunk, DeriveLens, LazyChunkStream, McapReader, MutateLens, OptimizationProfile, Selector
|
||||
from rerun.urdf import UrdfTree
|
||||
|
||||
PARENT_DIR = Path(__file__).parent
|
||||
DATA_DIR = PARENT_DIR / "input_data"
|
||||
OUTPUT_DIR = PARENT_DIR / "output"
|
||||
|
||||
|
||||
def json_transforms_stream(json_path: Path) -> LazyChunkStream:
|
||||
"""Loads transform data saved in JSON as a chunk stream of static Transform3D."""
|
||||
with json_path.open() as f:
|
||||
transforms = json.load(f)["transforms"]
|
||||
|
||||
chunk = Chunk.from_columns(
|
||||
"/tf_static/robot_offsets",
|
||||
indexes=[],
|
||||
columns=rr.Transform3D.columns(
|
||||
translation=[transform["translation"] for transform in transforms],
|
||||
quaternion=[transform["quaternion_xyzw"] for transform in transforms],
|
||||
parent_frame=[transform["parent"] for transform in transforms],
|
||||
child_frame=[transform["child"] for transform in transforms],
|
||||
),
|
||||
)
|
||||
return LazyChunkStream.from_iter([chunk])
|
||||
|
||||
|
||||
def change_albedo_factor_lens(new_albedo: rr.components.AlbedoFactor) -> MutateLens:
|
||||
"""Replaces Asset3D albedo factors with a fixed color."""
|
||||
|
||||
return MutateLens(
|
||||
"Asset3D:albedo_factor",
|
||||
Selector(".").pipe(lambda old_albedo: pa.array([new_albedo] * len(old_albedo), type=old_albedo.type)),
|
||||
)
|
||||
|
||||
|
||||
def joints_batch_lens(robot_urdf: UrdfTree, to_entity: str = "/tmp") -> DeriveLens:
|
||||
"""Computes intermediate transform batches from each joint state message using the URDF."""
|
||||
return DeriveLens("schemas.proto.JointState:message", output_entity=to_entity).to_component(
|
||||
"rerun.urdf.JointTransformBatch",
|
||||
Selector(".").pipe(
|
||||
lambda joint_state_messages: robot_urdf.compute_joint_transform_batches(
|
||||
names=Selector(".joint_names").execute(joint_state_messages),
|
||||
values=Selector(".joint_positions").execute(joint_state_messages),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def output_transforms_lens() -> DeriveLens:
|
||||
"""Scatters transform batches into final Transform3D rows per joint."""
|
||||
return (
|
||||
DeriveLens("rerun.urdf.JointTransformBatch", output_entity="/tf", scatter=True)
|
||||
.to_component(
|
||||
rr.Transform3D.descriptor_translation(),
|
||||
Selector(".[].translation"),
|
||||
)
|
||||
.to_component(
|
||||
rr.Transform3D.descriptor_quaternion(),
|
||||
Selector(".[].quaternion"),
|
||||
)
|
||||
.to_component(
|
||||
rr.Transform3D.descriptor_parent_frame(),
|
||||
Selector(".[].parent_frame"),
|
||||
)
|
||||
.to_component(
|
||||
rr.Transform3D.descriptor_child_frame(),
|
||||
Selector(".[].child_frame"),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Run the main chunk-processing pipeline for this example."""
|
||||
OUTPUT_DIR.mkdir(exist_ok=True)
|
||||
|
||||
# Create a chunk stream from the MCAP file.
|
||||
# The reader uses Rerun's MCAP importer (like the viewer or `rerun mcap convert` CLI),
|
||||
# so we get Rerun components that we can process in-stream.
|
||||
mcap_stream = McapReader(DATA_DIR / "episode.mcap").stream()
|
||||
|
||||
# The world-to-base transform offsets of the two robots are stored in a separate JSON file.
|
||||
robot_offsets_stream = json_transforms_stream(DATA_DIR / "offsets.json")
|
||||
|
||||
# Load the same robot URDF twice, with distinct entity path and frame name prefixes for each robot.
|
||||
robot_urdf_left = UrdfTree.from_file_path(
|
||||
DATA_DIR / "robot.urdf",
|
||||
entity_path_prefix="robot_left",
|
||||
frame_prefix="left_",
|
||||
static_transform_entity_path="/tf_static/left_robot",
|
||||
)
|
||||
robot_urdf_right = UrdfTree.from_file_path(
|
||||
DATA_DIR / "robot.urdf",
|
||||
entity_path_prefix="robot_right",
|
||||
frame_prefix="right_",
|
||||
static_transform_entity_path="/tf_static/right_robot",
|
||||
)
|
||||
# Load the scene URDF (table & external cameras).
|
||||
scene_urdf = UrdfTree.from_file_path(DATA_DIR / "scene.urdf", static_transform_entity_path="/tf_static/scene")
|
||||
|
||||
# The external camera calibration in our example MCAP has swapped width/height.
|
||||
# We can fix this with a MutateLens.
|
||||
mcap_stream = mcap_stream.lenses(
|
||||
MutateLens(
|
||||
"Pinhole:resolution",
|
||||
Selector(".").pipe(
|
||||
lambda resolution: pa.array(
|
||||
[(height, width) for width, height in resolution.to_pylist()], type=resolution.type
|
||||
)
|
||||
),
|
||||
),
|
||||
content=["/external/cam_low", "/external/cam_high"],
|
||||
output_mode="forward_unmatched",
|
||||
)
|
||||
|
||||
# For each robot, compute the joint transforms in batches and convert to the final Transform3D chunks.
|
||||
# We keep the original joint states in the stream ("forward_all") while dropping the temporary batch values.
|
||||
mcap_stream = (
|
||||
mcap_stream
|
||||
.lenses(joints_batch_lens(robot_urdf_left), content="/robot_left/joint_states", output_mode="forward_all")
|
||||
.lenses(output_transforms_lens(), content="/tmp", output_mode="drop_unmatched")
|
||||
.lenses(joints_batch_lens(robot_urdf_right), content="/robot_right/joint_states", output_mode="forward_all")
|
||||
.lenses(output_transforms_lens(), content="/tmp", output_mode="drop_unmatched")
|
||||
)
|
||||
|
||||
# We also modify each robot's visual meshes to have custom colors / transparency by mutating the albedo factor.
|
||||
robot_urdf_left_stream = robot_urdf_left.stream().lenses(
|
||||
change_albedo_factor_lens(rr.components.AlbedoFactor([80, 120, 175, 125])),
|
||||
content="/robot_left/wxai/visual_geometries/**",
|
||||
output_mode="forward_unmatched",
|
||||
)
|
||||
robot_urdf_right_stream = robot_urdf_right.stream().lenses(
|
||||
change_albedo_factor_lens(rr.components.AlbedoFactor([200, 120, 90, 125])),
|
||||
content="/robot_right/wxai/visual_geometries/**",
|
||||
output_mode="forward_unmatched",
|
||||
)
|
||||
|
||||
# Drop the collision meshes from each URDF.
|
||||
# (you can also disable them in the viewer, but here we demonstrate how to drop them entirely)
|
||||
robot_urdf_left_stream = robot_urdf_left_stream.drop(content="/robot_left/wxai/collision_geometries/**")
|
||||
robot_urdf_right_stream = robot_urdf_right_stream.drop(content="/robot_right/wxai/collision_geometries/**")
|
||||
|
||||
# Merge the streams in logical groups (base recording and URDF data).
|
||||
# (alternatively we could also merge everything in one stream here, if desired)
|
||||
data_stream = LazyChunkStream.merge(
|
||||
mcap_stream,
|
||||
robot_offsets_stream,
|
||||
)
|
||||
urdf_stream = LazyChunkStream.merge(
|
||||
robot_urdf_left_stream,
|
||||
robot_urdf_right_stream,
|
||||
scene_urdf.stream(),
|
||||
)
|
||||
|
||||
# Run the pipeline, materialize into a ChunkStore and optimize it before writing to an RRD.
|
||||
# Here we use an optimization profile suited for object-store (query & stream applications).
|
||||
data_stream.collect(optimize=OptimizationProfile.OBJECT_STORE).write_rrd(
|
||||
OUTPUT_DIR / "data.rrd",
|
||||
application_id="rerun_example_robot_data_preprocessing",
|
||||
recording_id="episode",
|
||||
)
|
||||
# Write also the URDF streams to an RRD.
|
||||
# Note how we use the same `recording_id` here to group the two RRD layers into the same logical recording.
|
||||
# https://rerun.io/docs/concepts/logging-and-ingestion/recordings#logical-vs-physical-recordings
|
||||
urdf_stream.collect(optimize=OptimizationProfile.OBJECT_STORE).write_rrd(
|
||||
OUTPUT_DIR / "urdf.rrd",
|
||||
application_id="rerun_example_robot_data_preprocessing",
|
||||
recording_id="episode",
|
||||
)
|
||||
|
||||
print(f"\nWrote output RRDs to: {OUTPUT_DIR}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user