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 index 08e9ec0..d73ebc8 100644 --- a/examples/krec_to_rrd.py +++ b/examples/krec_to_rrd.py @@ -1,6 +1,12 @@ """Usage: # Visualize synthetic KRec data: python examples/krec_to_rrd.py --synthetic --spawn-viewer -v +python examples/krec_to_rrd.py --input /home/kasm-user/ali_repos/kmodel/data/datasets/krec_data/dec_3__11_10am_og_krecs_edited/2024-12-03_17-47-30/recording_20241125_184919_04bf0bae-c5d7-46bb-b1ef-e021c1ad85f8.krec_edited.krec.mkv --spawn-viewer -v +python examples/krec_to_rrd.py --input /home/kasm-user/ali_repos/kmodel/data/datasets/krec_data/dec_3__11_10am_og_krecs_edited/2024-12-03_17-47-30/temp/temp_2024-12-04_18-25-38/recording_20241125_184841_aad3f390-2589-4fb7-a4b1-c9fab8dd8cc8.krec_edited_from_mkv.krec --spawn-viewer -v +python examples/krec_to_rrd.py --input /home/kasm-user/ali_repos/kmodel/data/datasets/krec_data/dec_3__11_10am_og_krecs/recording_20241125_184810_c249e9f6-4ebf-48c7-b8ea-4aaad721a4f8.krec.mkv --spawn-viewer -v + +python examples/krec_to_rrd.py --input /home/kasm-user/ali_repos/kmodel/data/datasets/krec_data/dec_3__11_10am_og_krecs/recording_20241125_184810_c249e9f6-4ebf-48c7-b8ea-4aaad721a4f8.krec.mkv --spawn-viewer -v + # Save to RRD file: python examples/krec_to_rrd.py --synthetic --output output.rrd -v @@ -10,124 +16,20 @@ import logging import time from pathlib import Path -import numpy as np -import rerun as rr -import krec -import uuid - -def create_sine_wave_krec(num_frames=100, fps=30): - # Create timestamps - timestamps = np.arange(num_frames) / fps - - # Create sine waves - position_waves = { - 0: np.sin(2 * np.pi * 0.5 * timestamps), - 1: np.sin(2 * np.pi * 0.5 * timestamps), - 2: np.sin(2 * np.pi * 0.5 * timestamps), - } - velocity_waves = { - 0: np.sin(2 * np.pi * 0.5 * timestamps), - 1: np.sin(2 * np.pi * 0.5 * timestamps), - 2: np.sin(2 * np.pi * 0.5 * timestamps), - } - torque_waves = { - 0: 0.5 * np.sin(2 * np.pi * 0.5 * timestamps), - 1: 0.5 * np.sin(2 * np.pi * 0.5 * timestamps), - 2: 0.5 * np.sin(2 * np.pi * 0.5 * timestamps), - } - - accel_waves = { - "x": 0.1 * np.sin(2 * np.pi * 0.5 * timestamps), - "y": 0.1 * np.sin(2 * np.pi * 0.5 * timestamps), - "z": 9.81 + 0.1 * np.sin(2 * np.pi * 0.5 * timestamps), - } - gyro_waves = { - "x": np.sin(2 * np.pi * 0.5 * timestamps), - "y": np.sin(2 * np.pi * 0.5 * timestamps), - "z": np.sin(2 * np.pi * 0.5 * timestamps), - } - - # Create KRec with header - start_time = int(time.time_ns()) - header = krec.KRecHeader( - uuid=str(uuid.uuid4()), - task="sine_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), - ) - - # Add actuator configs - for i in range(3): - 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) - # Add frames with sine wave data - for i in range(num_frames): - frame = krec.KRecFrame( - video_timestamp=start_time + int(timestamps[i] * 1e9), - video_frame_number=i, - inference_step=i, - ) +import krec +import rerun as rr +import rerun.blueprint as rrb - # Add actuator states and commands - for j in range(3): - state = krec.ActuatorState( - actuator_id=j, - online=True, - position=position_waves[j][i], - velocity=velocity_waves[j][i], - torque=torque_waves[j][i], - temperature=25.0 + np.sin(2 * np.pi * 0.1 * timestamps[i]), # Slowly varying temperature - voltage=12.0 + np.sin(2 * np.pi * 0.05 * timestamps[i]), # Slowly varying voltage - current=1.0 + 0.1 * np.sin(2 * np.pi * 0.2 * timestamps[i]), # Slowly varying current - ) - frame.add_actuator_state(state) - - command = krec.ActuatorCommand( - actuator_id=j, - position=position_waves[j][i], - velocity=velocity_waves[j][i], - torque=torque_waves[j][i], - ) - frame.add_actuator_command(command) - - # Add IMU values - accel = krec.Vec3(x=accel_waves["x"][i], y=accel_waves["y"][i], z=accel_waves["z"][i]) - gyro = krec.Vec3(x=gyro_waves["x"][i], y=gyro_waves["y"][i], z=gyro_waves["z"][i]) - # Simple rotation quaternion (could be made more complex if needed) - quaternion = krec.IMUQuaternion(x=0.0, y=0.0, z=np.sin(timestamps[i] / 2), w=np.cos(timestamps[i] / 2)) - imu_values = krec.IMUValues(accel=accel, gyro=gyro, quaternion=quaternion) - frame.set_imu_values(imu_values) - - krec_obj.add_frame(frame) - - return krec_obj, { - "position_waves": position_waves, - "velocity_waves": velocity_waves, - "torque_waves": torque_waves, - "accel_waves": accel_waves, - "gyro_waves": gyro_waves, - "timestamps": timestamps, - } +from generate_synthetic_krec import generate_synthetic_krec def log_frame_data(frame, frame_idx): """Log frame data to Rerun visualization.""" - + # 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)) @@ -135,7 +37,7 @@ def log_frame_data(frame, frame_idx): # Log actuator states for state in frame.get_actuator_states(): - prefix = f"actuators/state_{state.actuator_id}" + 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)) @@ -146,7 +48,7 @@ def log_frame_data(frame, frame_idx): # Log actuator commands for cmd in frame.get_actuator_commands(): - prefix = f"actuators/command_{cmd.actuator_id}" + 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)) @@ -159,19 +61,134 @@ def log_frame_data(frame, frame_idx): rr.log("imu/acceleration/x", rr.Scalar(imu_values.accel.x)) rr.log("imu/acceleration/y", rr.Scalar(imu_values.accel.y)) rr.log("imu/acceleration/z", rr.Scalar(imu_values.accel.z)) - - if imu_values.gyro: + if imu_values.gyro: rr.log("imu/angular_velocity/x", rr.Scalar(imu_values.gyro.x)) rr.log("imu/angular_velocity/y", rr.Scalar(imu_values.gyro.y)) rr.log("imu/angular_velocity/z", rr.Scalar(imu_values.gyro.z)) - + if imu_values.quaternion: rr.log("imu/orientation/x", rr.Scalar(imu_values.quaternion.x)) rr.log("imu/orientation/y", rr.Scalar(imu_values.quaternion.y)) rr.log("imu/orientation/z", rr.Scalar(imu_values.quaternion.z)) rr.log("imu/orientation/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") @@ -179,43 +196,43 @@ def main(args): # 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: - # Use synthetic data - # extracted_krec = load_krec("path to krec.krec") - extracted_krec, _ = create_sine_wave_krec() + extracted_krec = generate_synthetic_krec() logging.info("Created synthetic KRec data") else: - # Extract from file extracted_krec = krec.extract_from_video(args.input, verbose=args.verbose) - logging.info("Extraction from file successful") - logging.info(f"Processing {len(extracted_krec)} frames") + logging.info("Processing %d frames" % len(extracted_krec)) # Log header information header = extracted_krec.header - rr.log("metadata/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 - )) # 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(f"Saved RRD to: {output_path}") + logging.info("Saved RRD to: %s" % output_path) elif args.spawn_viewer: - # Keep the viewer running try: while True: time.sleep(1) @@ -223,9 +240,10 @@ def main(args): logging.info("Ctrl-C received. Exiting.") except Exception as e: - logging.error(f"Error: {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") @@ -237,9 +255,9 @@ def main(args): 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) \ No newline at end of file + main(args)