From 65e974e87895fa86189bf3e4515e6f8c9020f971 Mon Sep 17 00:00:00 2001 From: Ali Kuwajerwala Date: Wed, 4 Dec 2024 14:11:37 -0800 Subject: [PATCH] Suppress ffmpeg output by default and add verbose argument to display it (#12) * suppress ffmpeg ouput with verbose argument * add example script to test change --- examples/test_krec_read.py | 79 ++++++++++++++++++++++++++++++++++++++ krec/bindings/src/lib.rs | 23 +++++++++-- src/ffmpeg.rs | 52 ++++++++++++++++++------- 3 files changed, 136 insertions(+), 18 deletions(-) create mode 100644 examples/test_krec_read.py diff --git a/examples/test_krec_read.py b/examples/test_krec_read.py new file mode 100644 index 0000000..7b2bf2a --- /dev/null +++ b/examples/test_krec_read.py @@ -0,0 +1,79 @@ +import argparse +import krec +import os +from datetime import datetime + +"""Usage: +python examples/test_krec_read.py --krec_file /path/to/krec/file + +python examples/test_krec_read.py --krec_file /home/kasm-user/ali_repos/kmodel/data/datasets/krec_data/nov_29__8_37_pm_krec_w_mkv_w_states/recording_20241125_184810_c249e9f6-4ebf-48c7-b8ea-4aaad721a4f8_edited.krec.mkv +""" + +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) -> 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}") + + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + temp_dir = os.path.join(os.path.dirname(mkv_file_path), f"temp_{timestamp}") + os.makedirs(temp_dir, exist_ok=True) + + base_name = os.path.basename(mkv_file_path).split('.krec.mkv')[0] + krec_file_path = os.path.join(temp_dir, f"{base_name}_from_mkv.krec") + + # Extract and load from temp directory + krec.extract_from_video(mkv_file_path, krec_file_path) + # krec.extract_from_video(mkv_file_path, krec_file_path, verbose=True) # for getting the ffmpeg output + return krec.KRec.load(krec_file_path) + +def load_krec(file_path: str) -> krec.KRec: + """Smart loader that handles both direct KREC and MKV-embedded KREC files.""" + file_type = get_krec_file_type(file_path) + + if file_type == 'krec': + return load_krec_direct(file_path) + else: # file_type == 'mkv' + return load_krec_from_mkv(file_path) + + +def main(args: argparse.Namespace) -> None: + + print(f"Reading KRec file: {args.krec_file}") + krec_obj = load_krec(args.krec_file) + print(f"krec_obj: {krec_obj}") + print(f"krec_obj.header: {krec_obj.header}") + print(f"num frames: {len(krec_obj)}") + print(f"succesfully loaded KRec file!") + + +if __name__ == "__main__": + 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", + ) + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/krec/bindings/src/lib.rs b/krec/bindings/src/lib.rs index 7c4f9c3..a782256 100644 --- a/krec/bindings/src/lib.rs +++ b/krec/bindings/src/lib.rs @@ -7,7 +7,7 @@ use pyo3::prelude::*; use pyo3::types::PyIterator; use pyo3_stub_gen::define_stub_info_gatherer; use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyfunction, gen_stub_pymethods}; -use tracing::{info, instrument}; +use tracing::{debug, info, instrument, warn}; /// A 3D vector with x, y, z components #[gen_stub_pyclass] @@ -1332,9 +1332,24 @@ fn combine_with_video(video_path: &str, krec_path: &str, output_path: &str) -> P #[gen_stub_pyfunction] #[pyfunction] -fn extract_from_video(video_path: &str, output_path: &str) -> PyResult<()> { - ::krec::extract_from_video(video_path, output_path) - .map_err(|e| PyErr::new::(e.to_string())) +#[pyo3(signature = (video_path, output_path, verbose=None))] +fn extract_from_video(video_path: &str, output_path: &str, verbose: Option) -> PyResult<()> { + info!("Python binding: extract_from_video called"); + debug!( + "Python binding: video_path={}, output_path={}, verbose={:?}", + video_path, output_path, verbose + ); + + let result = ::krec::extract_from_video(video_path, output_path, verbose).map_err(|e| { + warn!("Python binding: extract_from_video failed: {}", e); + PyErr::new::(e.to_string()) + }); + + if result.is_ok() { + info!("Python binding: extract_from_video completed successfully"); + } + + result } #[pymodule] diff --git a/src/ffmpeg.rs b/src/ffmpeg.rs index 7ebfafd..71966ba 100644 --- a/src/ffmpeg.rs +++ b/src/ffmpeg.rs @@ -81,23 +81,47 @@ pub fn combine_with_video( } } -pub fn extract_from_video(video_path: &str, output_path: &str) -> Result<(), FFmpegError> { - let status = std::process::Command::new("ffmpeg") - .args([ - "-dump_attachment:t:0", - output_path, - "-i", - video_path, - "-f", - "null", - "/dev/null", - ]) - .status() - .map_err(|e| FFmpegError::FFmpeg(e.to_string()))?; +pub fn extract_from_video( + video_path: &str, + output_path: &str, + verbose: Option, +) -> Result<(), FFmpegError> { + info!("Starting extract_from_video"); + debug!("Input video path: {}", video_path); + debug!("Output path: {}", output_path); + debug!("Verbose mode: {}", verbose.unwrap_or(false)); + + let mut command = std::process::Command::new("ffmpeg"); + command.args([ + "-dump_attachment:t:0", + output_path, + "-i", + video_path, + "-f", + "null", + "/dev/null", + ]); + + // Control ffmpeg output based on verbose flag + if !verbose.unwrap_or(false) { + command + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()); + } + + debug!("Constructed ffmpeg command: {:?}", command); + + let status = command.status().map_err(|e| { + warn!("Failed to execute ffmpeg: {}", e); + FFmpegError::FFmpeg(e.to_string()) + })?; if status.success() { + info!("Successfully extracted KRec from video"); Ok(()) } else { - Err(FFmpegError::FFmpeg("FFmpeg command failed".to_string())) + let error_msg = format!("FFmpeg command failed with status: {}", status); + warn!("{}", error_msg); + Err(FFmpegError::FFmpeg(error_msg)) } }