Krec to rrd conversion

examples/generate_synthetic_krec.py

@@ -0,0 +1,108 @@
+"""Functions for generating synthetic KRec data."""
+
+import logging
+import time
+import uuid
+
+import krec
+import numpy as np
+
+
+def generate_synthetic_krec(num_frames=100, fps=30, frequency=10.0, num_actuators=10):
+    """Generate a synthetic KRec object with simple wave data for 10 actuators.
+    States follow commands with a one-frame delay.
+
+    Args:
+        num_frames: Number of frames to generate
+        fps: Frames per second
+        frequency: Frequency multiplier for the sine waves (higher = faster oscillation)
+        num_actuators: Number of actuators to simulate (default: 10)
+    """
+    timestamps = np.arange(num_frames) / fps
+
+    # Generate command waves for each actuator
+    # Each actuator gets the same frequency (which decreases over time) but different phase offset
+    actuator_wave = {}
+    for i in range(num_actuators):
+        time_varying_frequency = frequency / (1 + timestamps / 3)
+        phase = time_varying_frequency * timestamps
+        phase_with_offset = phase + i
+        actuator_wave[i] = np.sin(phase_with_offset)
+
+    start_time = int(time.time_ns())
+    header = krec.KRecHeader(
+        uuid=str(uuid.uuid4()),
+        task="simple_wave_test",
+        robot_platform="test_platform",
+        robot_serial="test_serial_001",
+        start_timestamp=start_time,
+        end_timestamp=start_time + int(num_frames * (1 / fps) * 1e9),
+    )
+
+    for i in range(num_actuators):
+        actuator_config = krec.ActuatorConfig(
+            actuator_id=i,
+            kp=1.0,
+            kd=0.1,
+            ki=0.01,
+            max_torque=10.0,
+            name=f"Joint{i}",
+        )
+        header.add_actuator_config(actuator_config)
+
+    krec_obj = krec.KRec(header)
+
+    for i in range(num_frames):
+        frame = krec.KRecFrame(
+            video_timestamp=start_time + int(timestamps[i] * 1e9),
+            video_frame_number=i,
+            inference_step=i,
+        )
+
+        for j in range(num_actuators):
+            command_val = actuator_wave[j][i]
+            state_val = actuator_wave[j][i - 1] if i > 0 else 0.0
+
+            command = krec.ActuatorCommand(
+                actuator_id=j,
+                position=command_val,
+                velocity=0.0,
+                torque=0.0,
+            )
+            frame.add_actuator_command(command)
+
+            state = krec.ActuatorState(
+                actuator_id=j,
+                online=True,
+                position=state_val,
+                velocity=0.0,
+                torque=0.0,
+                temperature=25.0,
+                voltage=12.0,
+                current=1.0,
+            )
+            frame.add_actuator_state(state)
+
+        imu_values = krec.IMUValues(
+            accel=krec.Vec3(x=0.0, y=0.0, z=1.0),
+            gyro=krec.Vec3(x=0.0, y=0.0, z=0.0),
+            quaternion=krec.IMUQuaternion(x=0.0, y=0.0, z=0.0, w=1.0),
+        )
+        frame.set_imu_values(imu_values)
+
+        krec_obj.add_frame(frame)
+
+    return krec_obj
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+    logging.info("Generating synthetic KRec data...")
+    krec_obj = generate_synthetic_krec()
+
+    logging.info("Generated KRec with %d frames" % len(krec_obj))
+    logging.info("UUID: %s" % krec_obj.header.uuid)
+    logging.info("Task: %s" % krec_obj.header.task)
+    logging.info("Robot Platform: %s" % krec_obj.header.robot_platform)
+    logging.info("Start Timestamp: %d" % krec_obj.header.start_timestamp)

examples/krec_to_rrd.py

@@ -0,0 +1,272 @@
+"""
+Convert KRec data (.krec) to Rerun visualization format (.rrd) and optionally save the rrd file to disk.
+
+Usage:
+# Visualize synthetic KRec data:
+python examples/krec_to_rrd.py --synthetic --spawn-viewer -v
+
+# Visualize KRec file:
+python examples/krec_to_rrd.py --input path/to/recording.krec --spawn-viewer -v
+
+# Save to RRD file:
+python examples/krec_to_rrd.py --synthetic --output output.rrd -v
+"""
+
+import argparse
+import logging
+import time
+from pathlib import Path
+
+import krec
+import rerun as rr
+import rerun.blueprint as rrb
+
+from generate_synthetic_krec import generate_synthetic_krec
+
+
+def log_frame_data(frame, frame_idx):
+    """Log KRec frame data to Rerun visualization.
+
+    Logs metadata, actuator states/commands, and IMU data for each frame.
+    Data is organized hierarchically in the Rerun viewer under metadata/,
+    actuators/, and imu/ paths.
+
+    Args:
+        frame: KRecFrame object containing the frame data
+        frame_idx: Integer index of the current frame for time sequencing
+    """
+
+    # Set time sequence
+    rr.set_time_sequence("frame_idx", frame_idx)
+
+    # Log metadata
+    rr.log("metadata/video_frame_number", rr.Scalar(frame.video_frame_number))
+    rr.log("metadata/timestamp", rr.Scalar(frame.video_timestamp))
+    rr.log("metadata/inference_step", rr.Scalar(frame.inference_step))
+
+    # Log actuator states
+    for state in frame.get_actuator_states():
+        prefix = f"actuators/actuator_{state.actuator_id}/state"
+        rr.log(f"{prefix}/position", rr.Scalar(state.position))
+        rr.log(f"{prefix}/velocity", rr.Scalar(state.velocity))
+        rr.log(f"{prefix}/torque", rr.Scalar(state.torque))
+        rr.log(f"{prefix}/temperature", rr.Scalar(state.temperature))
+        rr.log(f"{prefix}/voltage", rr.Scalar(state.voltage))
+        rr.log(f"{prefix}/current", rr.Scalar(state.current))
+        rr.log(f"{prefix}/online", rr.Scalar(float(state.online)))
+
+    # Log actuator commands
+    for cmd in frame.get_actuator_commands():
+        prefix = f"actuators/actuator_{cmd.actuator_id}/command"
+        rr.log(f"{prefix}/position", rr.Scalar(cmd.position))
+        rr.log(f"{prefix}/velocity", rr.Scalar(cmd.velocity))
+        rr.log(f"{prefix}/torque", rr.Scalar(cmd.torque))
+
+    # Log IMU values if present
+    imu_values = frame.get_imu_values()
+    if imu_values:
+        if imu_values.accel:
+            # Only scalar logging for acceleration
+            rr.log("imu/accel/x", rr.Scalar(imu_values.accel.x))
+            rr.log("imu/accel/y", rr.Scalar(imu_values.accel.y))
+            rr.log("imu/accel/z", rr.Scalar(imu_values.accel.z))
+
+        if imu_values.gyro:
+            rr.log("imu/gyro/x", rr.Scalar(imu_values.gyro.x))
+            rr.log("imu/gyro/y", rr.Scalar(imu_values.gyro.y))
+            rr.log("imu/gyro/z", rr.Scalar(imu_values.gyro.z))
+
+        if imu_values.quaternion:
+            rr.log("imu/quaternion/x", rr.Scalar(imu_values.quaternion.x))
+            rr.log("imu/quaternion/y", rr.Scalar(imu_values.quaternion.y))
+            rr.log("imu/quaternion/z", rr.Scalar(imu_values.quaternion.z))
+            rr.log("imu/quaternion/w", rr.Scalar(imu_values.quaternion.w))
+
+
+def get_grid_blueprint():
+    """Returns a blueprint with a 2x5 grid layout of actuator data."""
+    return rrb.Vertical(
+        # Header info at top
+        rrb.TextDocumentView(name="Metadata", contents=["+ /krec/header/**"]),
+        # Two rows of 5 actuators each
+        rrb.Vertical(
+            rrb.Horizontal(  # First row of 5
+                contents=[
+                    rrb.TimeSeriesView(
+                        name=f"Actuator {i}",
+                        contents=[f"+ /actuators/actuator_{i}/**"],
+                        axis_y=rrb.ScalarAxis(range=(-2.0, 2.0)),  # Slightly larger than ±1 for sine waves
+                    )
+                    for i in range(5)  # Actuators 0-4
+                ],
+            ),
+            rrb.Horizontal(  # Second row of 5
+                contents=[
+                    rrb.TimeSeriesView(
+                        name=f"Actuator {i}",
+                        contents=[f"+ /actuators/actuator_{i}/**"],
+                        axis_y=rrb.ScalarAxis(range=(-2.0, 2.0)),
+                    )
+                    for i in range(5, 10)  # Actuators 5-9
+                ],
+            ),
+        ),
+        row_shares=[1, 8],
+        name="Grid View",
+    )
+
+
+def get_grouped_blueprint():
+    """Returns a blueprint with actuator data grouped by measurement type."""
+    return rrb.Vertical(
+        # Header info at top
+        rrb.TextDocumentView(name="Metadata", contents=["+ /krec/header/**"]),
+        # Main content area split into left and right
+        rrb.Horizontal(
+            # Left side: Commands and States
+            rrb.Vertical(
+                contents=[
+                    rrb.TimeSeriesView(
+                        name="Position Commands",
+                        contents=[f"+ /actuators/actuator_{i}/command/position" for i in range(10)],
+                        axis_y=rrb.ScalarAxis(range=(-2.0, 2.0)),
+                    ),
+                    rrb.TimeSeriesView(
+                        name="Position States",
+                        contents=[f"+ /actuators/actuator_{i}/state/position" for i in range(10)],
+                        axis_y=rrb.ScalarAxis(range=(-2.0, 2.0)),
+                    ),
+                    rrb.TimeSeriesView(
+                        name="Velocity Commands",
+                        contents=[f"+ /actuators/actuator_{i}/command/velocity" for i in range(10)],
+                        axis_y=rrb.ScalarAxis(range=(-2.0, 2.0)),
+                    ),
+                    rrb.TimeSeriesView(
+                        name="Velocity States",
+                        contents=[f"+ /actuators/actuator_{i}/state/velocity" for i in range(10)],
+                        axis_y=rrb.ScalarAxis(range=(-2.0, 2.0)),
+                    ),
+                ]
+            ),
+            # Right side: Other measurements
+            rrb.Vertical(
+                contents=[
+                    rrb.TimeSeriesView(
+                        name="Torque Commands",
+                        contents=[f"+ /actuators/actuator_{i}/command/torque" for i in range(10)],
+                        axis_y=rrb.ScalarAxis(range=(-2.0, 2.0)),
+                    ),
+                    rrb.TimeSeriesView(
+                        name="Torque States",
+                        contents=[f"+ /actuators/actuator_{i}/state/torque" for i in range(10)],
+                        axis_y=rrb.ScalarAxis(range=(-2.0, 2.0)),
+                    ),
+                    rrb.TimeSeriesView(
+                        name="Temperature",
+                        contents=[f"+ /actuators/actuator_{i}/state/temperature" for i in range(10)],
+                        axis_y=rrb.ScalarAxis(range=(0.0, 50.0)),  # Temperature typically 20-30°C
+                    ),
+                    rrb.TimeSeriesView(
+                        name="Voltage",
+                        contents=[f"+ /actuators/actuator_{i}/state/voltage" for i in range(10)],
+                        axis_y=rrb.ScalarAxis(range=(0.0, 15.0)),  # Voltage typically around 12V
+                    ),
+                    rrb.TimeSeriesView(
+                        name="Current",
+                        contents=[f"+ /actuators/actuator_{i}/state/current" for i in range(10)],
+                        axis_y=rrb.ScalarAxis(range=(0.0, 2.0)),  # Current typically around 1A
+                    ),
+                ]
+            ),
+        ),
+        row_shares=[1, 8],  # Header takes 1/9 of space, main content takes 8/9
+        name="Grouped View",
+    )
+
+
+def get_combined_blueprint():
+    """Returns a blueprint with both grouped and grid views in separate tabs."""
+    return rrb.Blueprint(
+        rrb.Tabs(
+            get_grouped_blueprint(),
+            get_grid_blueprint(),
+            active_tab=0,  # Start with the grouped view
+        ),
+        rrb.BlueprintPanel(state="expanded"),
+        rrb.SelectionPanel(state="collapsed"),
+        rrb.TimePanel(state="expanded"),
+    )
+
+
+def main(args):
+    logging.info("Starting visualization process")
+
+    try:
+        # Initialize Rerun
+        spawn_viewer = args.spawn_viewer and not args.output
+        rr.init("krec_visualization", spawn=spawn_viewer)
+        rr.send_blueprint(get_combined_blueprint())
+
+        if args.synthetic:
+            extracted_krec = generate_synthetic_krec()
+            logging.info("Created synthetic KRec data")
+        else:
+            extracted_krec = krec.extract_from_video(args.input, verbose=args.verbose)
+
+        logging.info("Processing %d frames" % len(extracted_krec))
+
+        # Log header information
+        header = extracted_krec.header
+
+        # Log each frame
+        for idx, frame in enumerate(extracted_krec):
+            log_frame_data(frame, idx)
+
+        rr.log(
+            "/krec/header",
+            rr.TextDocument(
+                f"UUID: {header.uuid}\n"
+                f"Task: {header.task}\n"
+                f"Robot Platform: {header.robot_platform}\n"
+                f"Robot Serial: {header.robot_serial}\n"
+                f"Start Timestamp: {header.start_timestamp}\n"
+                f"End Timestamp: {header.end_timestamp}",
+                media_type=rr.MediaType.MARKDOWN,
+            ),
+        )
+        logging.info("Logged header information")
+
+        # Save to RRD file if output path provided
+        if args.output:
+            output_path = Path(args.output)
+            rr.save(str(output_path))
+            logging.info("Saved RRD to: %s" % output_path)
+        elif args.spawn_viewer:
+            try:
+                while True:
+                    time.sleep(1)
+            except KeyboardInterrupt:
+                logging.info("Ctrl-C received. Exiting.")
+
+    except Exception as e:
+        logging.error("Error: %s" % e)
+        raise
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Visualize KRec data using Rerun")
+    parser.add_argument("--input", type=str, help="Input KREC or KREC.MKV file path")
+    parser.add_argument("-o", "--output", type=str, help="Output RRD file path (optional)")
+    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output")
+    parser.add_argument("--synthetic", action="store_true", help="Use synthetic data instead of input file")
+    parser.add_argument("--spawn-viewer", action="store_true", help="Spawn Rerun viewer (ignored if --output is set)")
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+    args = parser.parse_args()
+
+    # Validate args
+    if not args.synthetic and not args.input:
+        parser.error("Either --synthetic or --input must be provided")
+
+    main(args)