wehub-resource-sync
105 lines
2.1 KiB
Python
105 lines
2.1 KiB
Python
# region: setup
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.compute as pc
|
|
|
|
import rerun as rr
|
|
from rerun.experimental import (
|
|
Chunk,
|
|
DeriveLens,
|
|
LazyChunkStream,
|
|
McapReader,
|
|
Selector,
|
|
)
|
|
|
|
MCAP = (
|
|
Path(__file__).resolve().parents[4]
|
|
/ "tests"
|
|
/ "assets"
|
|
/ "mcap"
|
|
/ "trossen_transfer_cube.mcap"
|
|
)
|
|
OUT = Path("chunk_processing.rrd")
|
|
# endregion: setup
|
|
|
|
|
|
# region: reading
|
|
stream = McapReader(MCAP).stream()
|
|
# endregion: reading
|
|
|
|
|
|
# region: processing
|
|
JOINTS = [
|
|
"waist",
|
|
"shoulder",
|
|
"elbow",
|
|
"forearm_roll",
|
|
"wrist_angle",
|
|
"wrist_rotate",
|
|
]
|
|
|
|
|
|
def pick_joint(i: int) -> Callable[[pa.Array], pa.Array]:
|
|
"""Extract joint `i` from a list<float64> column and convert rad → deg."""
|
|
return lambda arr: pc.multiply(pc.list_element(arr, i), 180.0 / math.pi)
|
|
|
|
|
|
def fan(side: str) -> list[DeriveLens]:
|
|
return [
|
|
DeriveLens(
|
|
"schemas.proto.JointState:message",
|
|
output_entity=f"/joints_deg/{side}/{name}",
|
|
).to_component(
|
|
rr.Scalars.descriptor_scalars(),
|
|
Selector(".joint_positions").pipe(pick_joint(i)),
|
|
)
|
|
for i, name in enumerate(JOINTS)
|
|
]
|
|
|
|
|
|
processed = (
|
|
stream
|
|
.drop(content="/video_raw/**")
|
|
.lenses(
|
|
fan("left"),
|
|
content="/robot_left/**",
|
|
output_mode="forward_unmatched",
|
|
)
|
|
.lenses(
|
|
fan("right"),
|
|
content="/robot_right/**",
|
|
output_mode="forward_unmatched",
|
|
)
|
|
)
|
|
# endregion: processing
|
|
|
|
# TODO(ab): change this to merge properties instead, when we have proper
|
|
# interop between logging SDK and py-chunk
|
|
|
|
# region: merging
|
|
metadata = Chunk.from_columns(
|
|
"/metadata",
|
|
indexes=[],
|
|
columns=rr.AnyValues.columns(
|
|
processing_type="ingestion",
|
|
processing_version="v1",
|
|
),
|
|
)
|
|
merged = LazyChunkStream.merge(processed, LazyChunkStream.from_iter([metadata]))
|
|
# endregion: merging
|
|
|
|
|
|
# region: write
|
|
merged.write_rrd(
|
|
OUT,
|
|
application_id="rerun_example_chunk_processing",
|
|
recording_id=str(uuid.uuid4()),
|
|
)
|
|
# endregion: write
|