diff --git a/src/main.rs b/src/main.rs index 69ed81a..ec8a5ad 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,6 +18,7 @@ struct Cli { } #[derive(Subcommand, Debug)] +#[allow(clippy::large_enum_variant)] enum Commands { /// Compare two still images. Resolutions must be identical. Image { @@ -42,9 +43,17 @@ enum Commands { /// How many worker threads to use for decoding & calculating scores. /// Note: Memory usage increases linearly with the number of workers. - #[arg(long, short)] + #[arg(long, short, verbatim_doc_comment)] frame_threads: Option, + /// The amount of frames to skip. + #[arg(long, default_value_t = 0)] + skip_frames: usize, + + /// Limit the amount of frames to compare. + #[arg(long)] + frames: Option, + /// How to increment current frame count; e.g. 10 will read every 10th frame. #[arg(long, short)] increment: Option, @@ -99,6 +108,8 @@ fn main() { source, distorted, frame_threads, + skip_frames, + frames, increment, graph, verbose, @@ -135,6 +146,8 @@ fn main() { &source, &distorted, frame_threads, + skip_frames, + frames, inc, graph, verbose, diff --git a/src/video.rs b/src/video.rs index 0c77a0b..a5ffd0c 100644 --- a/src/video.rs +++ b/src/video.rs @@ -115,34 +115,53 @@ fn pretty_spinner_style() -> ProgressStyle { .progress_chars(PROGRESS_CHARS) } +type VideoCompareMutex = Arc>>; + +struct VideoCompare { + current_frame: usize, + next_frame: usize, + source: E, + distorted: F, +} + fn calc_score( - mtx: &Mutex<(usize, (E, F))>, + mtx: &VideoCompareMutex, src_yuvcfg: &YuvConfig, dst_yuvcfg: &YuvConfig, inc: usize, + end_frame: Option, verbose: bool, ) -> Option<(usize, f64)> { let (frame_idx, (src_frame, dst_frame)) = { let mut guard = mtx.lock().unwrap(); - let curr_frame = guard.0; - let src_frame = guard.1 .0.read_video_frame::(); - let dst_frame = guard.1 .1.read_video_frame::(); + let distance_to_next = guard.next_frame - guard.current_frame; - if let (Some(sf), Some(df)) = (src_frame, dst_frame) { - // skip remaining frames in increment size - for ii in 1..inc { - let _src_frame = guard.1 .0.read_video_frame::(); - let _dst_frame = guard.1 .1.read_video_frame::(); - if _src_frame.is_none() || _dst_frame.is_none() { - break; - } - if verbose { - println!("Frame {}: skip", curr_frame + ii); - } + if let Some(end_frame) = end_frame { + if guard.next_frame >= end_frame { + return None; } + } - guard.0 += inc; + for ii in 1..distance_to_next { + let _src_frame = guard.source.read_video_frame::(); + let _dst_frame = guard.distorted.read_video_frame::(); + if _src_frame.is_none() || _dst_frame.is_none() { + break; + } + if verbose { + println!("Frame {}: skip", guard.current_frame + ii); + } + } + + let src_frame = guard.source.read_video_frame::(); + let dst_frame = guard.distorted.read_video_frame::(); + + let curr_frame = guard.next_frame; + guard.current_frame = guard.next_frame; + guard.next_frame += inc; + + if let (Some(sf), Some(df)) = (src_frame, dst_frame) { (curr_frame, (sf, df)) } else { return None; @@ -163,6 +182,8 @@ pub fn compare_videos( source: &str, distorted: &str, frame_threads: usize, + skip_frames: usize, + frames_to_compare: Option, inc: usize, graph: bool, verbose: bool, @@ -196,6 +217,8 @@ pub fn compare_videos( None, distorted_frame_count, frame_threads, + skip_frames, + frames_to_compare, inc, graph, verbose, @@ -227,6 +250,8 @@ pub fn compare_videos( source_frame_count, None, frame_threads, + skip_frames, + frames_to_compare, inc, graph, verbose, @@ -267,6 +292,8 @@ pub fn compare_videos( source_frame_count, distorted_frame_count, frame_threads, + skip_frames, + frames_to_compare, inc, graph, verbose, @@ -288,6 +315,8 @@ fn compare_videos_inner( source_frame_count: Option, distorted_frame_count: Option, frame_threads: usize, + skip_frames: usize, + frames_to_compare: Option, inc: usize, graph: bool, verbose: bool, @@ -362,40 +391,57 @@ fn compare_videos_inner( let dst_bd = dst_config.bit_depth; let current_frame = 0usize; - let decoders = Arc::new(Mutex::new((current_frame, (source, distorted)))); + let end_frame = frames_to_compare + .map(|frames_to_compare| skip_frames + (frames_to_compare * inc)); + + let video_compare = Arc::new( + Mutex::new( + VideoCompare { + current_frame, + next_frame: skip_frames, + source, + distorted, + } + ) + ); + for _ in 0..frame_threads { - let decoders = Arc::clone(&decoders); + let video_compare = Arc::clone(&video_compare); let result_tx = result_tx.clone(); std::thread::spawn(move || { loop { let score = match (src_bd, dst_bd) { (8, 8) => calc_score::( - &decoders, + &video_compare, &src_config, &dst_config, inc, + end_frame, verbose, ), (8, _) => calc_score::( - &decoders, + &video_compare, &src_config, &dst_config, inc, + end_frame, verbose, ), (_, 8) => calc_score::( - &decoders, + &video_compare, &src_config, &dst_config, inc, + end_frame, verbose, ), (_, _) => calc_score::( - &decoders, + &video_compare, &src_config, &dst_config, inc, + end_frame, verbose, ), }; @@ -416,7 +462,10 @@ fn compare_videos_inner( let progress = if stderr().is_tty() && !verbose { let frame_count = source_frame_count.or(distorted_frame_count); let pb = if let Some(frame_count) = frame_count { - ProgressBar::new(frame_count as u64) + let fc = frames_to_compare.unwrap_or(frame_count - skip_frames) + .min(((frame_count - skip_frames) as f64 / inc as f64).ceil() as usize); + + ProgressBar::new(fc as u64) .with_style(pretty_progress_style()) .with_message(", avg: N/A") } else { @@ -443,7 +492,7 @@ fn compare_videos_inner( results.insert(score.0, score.1); avg = avg + (score.1 - avg) / (min(results.len(), 10) as f64); progress.set_message(format!(", avg: {:.1$}", avg, 2)); - progress.inc(inc.try_into().unwrap()); + progress.inc(1); } progress.finish();