From 98330e46fe87d8a08477f8435c828b67dfa3454b Mon Sep 17 00:00:00 2001 From: Ali Kuwajerwala Date: Tue, 17 Dec 2024 15:09:48 -0800 Subject: [PATCH] update saving to overwrite and add verbose arg (#17) * make combine w mkv overwrite by default * move docs to top and formatting * format krec read as well * remove duplicate docstring * remove empty line * cleanup prints use logging instead of prints * remove return, not needed * add explainer at the top * colorlogging and fix fstrings * isort imports * add colorlogging as dependency * formatting and lint * move colorlogging to requirements.txt * move argparse logic into main * lint * update setup.py * remove unused import * update types * update test krec read with logging and verbose * simplify load from mkv * remove list import, not needed --- examples/test_extract.py | 50 ++++++++ examples/test_krec_read.py | 86 +++++++++++++ examples/test_krec_save.py | 248 +++++++++++++++++++++++++++++++++++++ krec/bindings/src/lib.rs | 14 ++- requirements.txt | 1 + setup.py | 14 +-- src/ffmpeg.rs | 51 +++++--- 7 files changed, 433 insertions(+), 31 deletions(-) create mode 100644 examples/test_extract.py create mode 100644 examples/test_krec_read.py create mode 100644 examples/test_krec_save.py create mode 100644 requirements.txt diff --git a/examples/test_extract.py b/examples/test_extract.py new file mode 100644 index 0000000..f3c8305 --- /dev/null +++ b/examples/test_extract.py @@ -0,0 +1,50 @@ +"""Extract KRec data from an MKV file and optionally save it to a new file. + +Usage: +# Extract and view +python examples/test_extract.py input.mkv -v + +# Extract and save +python examples/test_extract.py input.mkv -o output.krec -v +""" + +import argparse +import logging +import sys +from pathlib import Path + +import colorlogging + +import krec + + +def main() -> int: + parser = argparse.ArgumentParser(description="Extract KRec data from MKV file") + parser.add_argument("input_file", type=str, help="Input MKV file path") + parser.add_argument("-o", "--output-file", type=str, help="Output KRec file path (optional)") + parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output") + + colorlogging.configure() + args = parser.parse_args() + + logging.info("Extracting from: %s", args.input_file) + + try: + extracted_krec = krec.extract_from_video(args.input_file, verbose=args.verbose) + logging.info("Extraction successful") + logging.info("Extracted KRec has %d frames", len(extracted_krec)) + + if args.output_file: + output_path = Path(args.output_file) + extracted_krec.save(str(output_path)) + logging.info("Saved to: %s", output_path) + + except Exception as e: + logging.error("Error: %s", e) + return 1 + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/test_krec_read.py b/examples/test_krec_read.py new file mode 100644 index 0000000..e41f41c --- /dev/null +++ b/examples/test_krec_read.py @@ -0,0 +1,86 @@ +"""Script to read and display KRec file contents. + +Usage: +python examples/test_krec_read.py --krec_file /path/to/krec/file.krec + +python examples/test_krec_read.py --krec_file /path/to/krec/file.krec.mkv + +python examples/test_krec_read.py --krec_file /path/to/krec/file.krec.mkv -v +""" + +import argparse +import logging +import os +import sys +import traceback + +import colorlogging + +import krec + + +def get_krec_file_type(file_path: str) -> str: + """Determine if the file is a direct KREC file or MKV-embedded KREC. + + Returns: + 'krec' for .krec files + 'mkv' for .krec.mkv files + raises RuntimeError for invalid extensions + """ + if file_path.endswith(".krec"): + return "krec" + elif file_path.endswith(".krec.mkv"): + return "mkv" + else: + error_msg = f"Invalid file extension. Expected '.krec' or '.krec.mkv', got: {file_path}" + raise RuntimeError(error_msg) + + +def load_krec_direct(krec_file_path: str) -> krec.KRec: + """Load a KREC file directly.""" + return krec.KRec.load(krec_file_path) + + +def load_krec_from_mkv(mkv_file_path: str, verbose: bool) -> krec.KRec: + """Load a KREC file from an MKV file into a manually created temp directory.""" + if not os.path.exists(mkv_file_path): + raise FileNotFoundError(f"File not found: {mkv_file_path}") + + return krec.extract_from_video(mkv_file_path, verbose=verbose) + + +def load_krec(file_path: str, verbose: bool) -> krec.KRec: + """Smart loader that handles both direct KREC and MKV-embedded KREC files.""" + file_type = get_krec_file_type(file_path) + return load_krec_direct(file_path) if file_type == "krec" else load_krec_from_mkv(file_path, verbose) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Read and display KRec file contents") + parser.add_argument( + "--krec_file", + type=str, + default="kdatagen/sim/resources/stompypro/krec_out/test_krec_write_out.krec", + help="Path to KRec file to read", + ) + parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output for ffmpeg") + + colorlogging.configure() + args = parser.parse_args() + + try: + logging.info("Reading KRec file: %s", args.krec_file) + krec_obj = load_krec(args.krec_file, verbose=args.verbose) + logging.info("KRec object: %s", krec_obj) + logging.info("KRec header: %s", krec_obj.header) + logging.info("Number of frames: %d", len(krec_obj)) + logging.info("Successfully loaded KRec file!") + return 0 + except Exception as e: + logging.error("Error: %s", e) + traceback.print_exc() + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/test_krec_save.py b/examples/test_krec_save.py new file mode 100644 index 0000000..c518e4a --- /dev/null +++ b/examples/test_krec_save.py @@ -0,0 +1,248 @@ +"""Script to test saving and loading KRec data with optional video combination. + +Usage: +# Save directly to KRec file: +python examples/test_krec_save.py --output_krec output_save_test.krec -v + +# Combine with video: +python examples/test_krec_save.py \ + --output_krec output_save_test.krec \ + --input_video /path/to/video.mkv \ + --output_video output_combined.mkv \ + -v +""" + +import argparse +import logging +import sys +import time +import uuid + +import numpy as np + +import krec + + +def create_sine_wave_krec(num_frames: int = 100, fps: int = 30) -> tuple[krec.KRec, dict]: + """Create a synthetic KRec with sine wave data for testing.""" + # Create timestamps + timestamps = np.arange(num_frames) / fps + + 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, + ) + + # 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, + } + + +def verify_krec_data(original_data: dict, loaded_krec: krec.KRec) -> bool: + """Verify that the loaded KRec matches the original data.""" + logging.info("Verifying loaded KRec data...") + + num_frames = len(loaded_krec) + + for i in range(num_frames): + frame = loaded_krec[i] + + # Verify actuator states + for j, state in enumerate(frame.get_actuator_states()): + expected_pos = original_data["position_waves"][j][i] + expected_vel = original_data["velocity_waves"][j][i] + expected_torque = original_data["torque_waves"][j][i] + + if not np.allclose([state.position], [expected_pos], rtol=1e-5): + logging.error(f"Position mismatch at frame {i}, joint {j}: {state.position} != {expected_pos}") + return False + + if not np.allclose([state.velocity], [expected_vel], rtol=1e-5): + logging.error(f"Velocity mismatch at frame {i}, joint {j}: {state.velocity} != {expected_vel}") + return False + + if not np.allclose([state.torque], [expected_torque], rtol=1e-5): + logging.error(f"Torque mismatch at frame {i}, joint {j}: {state.torque} != {expected_torque}") + return False + + # Verify IMU data + imu = frame.get_imu_values() + if imu: + expected_accel = [ + original_data["accel_waves"]["x"][i], + original_data["accel_waves"]["y"][i], + original_data["accel_waves"]["z"][i], + ] + expected_gyro = [ + original_data["gyro_waves"]["x"][i], + original_data["gyro_waves"]["y"][i], + original_data["gyro_waves"]["z"][i], + ] + + actual_accel = [imu.accel.x, imu.accel.y, imu.accel.z] + actual_gyro = [imu.gyro.x, imu.gyro.y, imu.gyro.z] + + if not np.allclose(actual_accel, expected_accel, rtol=1e-5): + logging.error(f"Acceleration mismatch at frame {i}: {actual_accel} != {expected_accel}") + return False + + if not np.allclose(actual_gyro, expected_gyro, rtol=1e-5): + logging.error(f"Gyro mismatch at frame {i}: {actual_gyro} != {expected_gyro}") + return False + + logging.info("All data verified successfully!") + return True + + +def main(args: argparse.Namespace) -> int: + logging.info("Creating synthetic KRec with sine waves...") + synthetic_krec, original_data = create_sine_wave_krec() + + # Save KRec file + logging.info(f"Saving KRec to: {args.output_krec}") + synthetic_krec.save(args.output_krec) + + # If video paths are provided, combine with video + if args.input_video and args.output_video: + logging.info(f"Combining video from {args.input_video} with KRec...") + try: + krec.combine_with_video( + args.input_video, + args.output_krec, + args.output_video, + ) + logging.info(f"Successfully combined video and KRec to: {args.output_video}") + + # Verify the output by trying to extract the KRec + logging.info("Verifying output by extracting KRec...") + extracted_krec = krec.extract_from_video(args.output_video, verbose=args.verbose) + logging.info(f"Successfully extracted KRec with {len(extracted_krec)} frames") + + # Verify the data matches + if not verify_krec_data(original_data, extracted_krec): + logging.error("Data verification failed!") + return 1 + + except Exception as e: + logging.error(f"Error during video processing: {e}") + return 1 + else: + # Load and verify the direct KRec save + logging.info(f"Loading KRec from: {args.output_krec}") + loaded_krec = krec.KRec.load(args.output_krec) + logging.info(f"Loaded KRec with {len(loaded_krec)} frames") + + # Verify the data matches + if not verify_krec_data(original_data, loaded_krec): + logging.error("Data verification failed!") + return 1 + + return 0 + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Save and verify KRec data") + parser.add_argument("--output_krec", type=str, required=True, help="Output KRec file path") + parser.add_argument("--input_video", type=str, help="Input video file path (optional)") + parser.add_argument("--output_video", type=str, help="Output video file path (optional)") + parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output") + + args = parser.parse_args() + + # Validate args + if bool(args.input_video) != bool(args.output_video): + parser.error("--input_video and --output_video must be provided together") + + sys.exit(main(args)) diff --git a/krec/bindings/src/lib.rs b/krec/bindings/src/lib.rs index 03369f4..a7611a6 100644 --- a/krec/bindings/src/lib.rs +++ b/krec/bindings/src/lib.rs @@ -986,7 +986,7 @@ impl PyKRec { self.save(&temp_path)?; // Combine with video - ::krec::combine_with_video(video_path, &temp_path, output_path) + ::krec::combine_with_video(video_path, &temp_path, output_path, None) .map_err(|e| PyErr::new::(e.to_string()))?; // Clean up temporary file @@ -1168,7 +1168,7 @@ impl PyKRecFrame { } if items.len() != 4 { return Err(PyValueError::new_err( - "Iterable must contain exactly 4 values: [video_timestamp, frame_number, inference_step, real_timestamp]" + "Iterable must contain exactly 4 values: [video_timestamp, video_frame_number, inference_step, real_timestamp]" )); } let mut inner = KRecFrame::default(); @@ -1344,8 +1344,14 @@ impl FrameIterator { #[gen_stub_pyfunction] #[pyfunction] -fn combine_with_video(video_path: &str, krec_path: &str, output_path: &str) -> PyResult<()> { - ::krec::combine_with_video(video_path, krec_path, output_path) +#[pyo3(signature = (video_path, krec_path, output_path, verbose=None))] +fn combine_with_video( + video_path: &str, + krec_path: &str, + output_path: &str, + verbose: Option, +) -> PyResult<()> { + ::krec::combine_with_video(video_path, krec_path, output_path, verbose) .map_err(|e| PyErr::new::(e.to_string())) } diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..27e2c9f --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +colorlogging \ No newline at end of file diff --git a/setup.py b/setup.py index 27836f8..7cf351f 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,6 @@ import re import shutil import subprocess -from typing import List from setuptools import find_packages, setup from setuptools.command.build_ext import build_ext @@ -21,8 +20,11 @@ assert version_re is not None, "Could not find version in Cargo.toml" version: str = version_re.group(1) +with open("requirements.txt", "r", encoding="utf-8") as f: + requirements: list[str] = f.read().splitlines() + with open("krec/requirements-dev.txt", "r", encoding="utf-8") as f: - requirements_dev: List[str] = f.read().splitlines() + requirements_dev: list[str] = f.read().splitlines() class RustBuildExt(build_ext): @@ -52,7 +54,7 @@ def run(self) -> None: setup( name="krec", version=version, - description="Python bindings for K-Scale recordingss", + description="Python bindings for K-Scale recordings", author="K-Scale Labs", url="https://github.com/kscalelabs/krec", rust_extensions=[ @@ -63,16 +65,14 @@ def run(self) -> None: ), ], setup_requires=["setuptools-rust"], + install_requires=requirements, + extras_require={"dev": requirements_dev}, zip_safe=False, long_description=long_description, long_description_content_type="text/markdown", python_requires=">=3.8", include_package_data=True, packages=find_packages(include=["krec"]), - install_requires=requirements_dev, - extras_require={ - "dev": requirements_dev, - }, cmdclass={ "build_ext": RustBuildExt, "build_py": CustomBuild, diff --git a/src/ffmpeg.rs b/src/ffmpeg.rs index a0a159c..01b7fcf 100644 --- a/src/ffmpeg.rs +++ b/src/ffmpeg.rs @@ -18,6 +18,7 @@ pub fn combine_with_video( video_path: impl AsRef, krec_path: impl AsRef, output_path: impl AsRef, + verbose: Option, ) -> Result<()> { info!("Combining video with KRec data"); debug!( @@ -51,26 +52,36 @@ pub fn combine_with_video( return Err(eyre!("KRec file missing robot serial")); } - let status = std::process::Command::new("ffmpeg") - .args([ - "-i", - &video_path.as_ref().to_string_lossy(), - "-attach", - &krec_path.as_ref().to_string_lossy(), - "-metadata:s:t", - "mimetype=application/octet-stream", - "-metadata:s:t", - &format!("uuid={}", krec.header.uuid), - "-metadata:s:t", - &format!("task={}", krec.header.task), - "-metadata:s:t", - &format!("robot_platform={}", krec.header.robot_platform), - "-metadata:s:t", - &format!("robot_serial={}", krec.header.robot_serial), - "-c", - "copy", - &output_path.as_ref().to_string_lossy(), - ]) + let mut command = std::process::Command::new("ffmpeg"); + command.args([ + "-y", // Add -y flag to automatically overwrite files + "-i", + &video_path.as_ref().to_string_lossy(), + "-attach", + &krec_path.as_ref().to_string_lossy(), + "-metadata:s:t", + "mimetype=application/octet-stream", + "-metadata:s:t", + &format!("uuid={}", krec.header.uuid), + "-metadata:s:t", + &format!("task={}", krec.header.task), + "-metadata:s:t", + &format!("robot_platform={}", krec.header.robot_platform), + "-metadata:s:t", + &format!("robot_serial={}", krec.header.robot_serial), + "-c", + "copy", + &output_path.as_ref().to_string_lossy(), + ]); + + // Control ffmpeg output based on verbose flag + if !verbose.unwrap_or(false) { + command + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()); + } + + let status = command .status() .map_err(|e| eyre!("Failed to execute ffmpeg: {}", e))?;