diff --git a/examples/generate_synthetic_krec.py b/examples/generate_synthetic_krec.py new file mode 100644 index 0000000..3eaf616 --- /dev/null +++ b/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) diff --git a/examples/krec_to_rrd.py b/examples/krec_to_rrd.py new file mode 100644 index 0000000..dae4ae9 --- /dev/null +++ b/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)