Add option to print scores of each frame pair

I wanted to do more stats breakdowns, and being able to print out the score for each frame would be really convenient. A friend made a patch for an earlier build, and I didn't see any reason the feature couldn't be added to the main repo in some form.

diff --git a/crates/turbo-metrics/src/main.rs b/crates/turbo-metrics/src/main.rs
index 569955d..e92eab7 100644
--- a/crates/turbo-metrics/src/main.rs
+++ b/crates/turbo-metrics/src/main.rs
@@ -39,6 +39,9 @@ struct CliArgs {
     /// Compute ssimulacra2 score
     #[arg(long)]
     ssimulacra2: bool,
+    /// frame score output only
+    #[arg(long)]
+    frame_score_only: bool,

     /// Only compute metrics every few frames, effectively down-sampling the measurements.
     /// Still, this tool will decode all frames, hence increasing overhead. Check Mpx/s to see what I mean.
@@ -52,15 +55,17 @@ struct CliArgs {
 }

 fn main() {
-    let args = dbg!(CliArgs::parse());
+    let args = CliArgs::parse();

     cudarse_driver::init_cuda().expect("Could not initialize the CUDA API");
     let dev = CuDevice::get(0).unwrap();
-    println!(
-        "Using device {} with CUDA version {}",
-        dev.name().unwrap(),
-        cudarse_driver::cuda_driver_version().unwrap()
-    );
+    if !args.frame_score_only {
+        println!(
+            "Using device {} with CUDA version {}",
+            dev.name().unwrap(),
+            cudarse_driver::cuda_driver_version().unwrap()
+        );
+    }
     // Bind to main thread
     let ctx = dev.retain_primary_ctx().unwrap();
     ctx.set_current().unwrap();
@@ -80,8 +85,9 @@ fn main() {
         }
         let format = cb_ref.format().unwrap();
         let colors_ref = color_characteristics_from_format(&format);
-        println!("ref: {}", video_format_line(&format));
-
+        if !args.frame_score_only {
+            println!("ref: {}", video_format_line(&format));
+        }
         let size = format.size();

         while cb_dis.format().is_none() {
@@ -89,7 +95,9 @@ fn main() {
         }
         let format_dis = cb_dis.format().unwrap();
         let colors_dis = color_characteristics_from_format(&format_dis);
-        println!("dis: {}", video_format_line(&format_dis));
+        if !args.frame_score_only {
+            println!("dis: {}", video_format_line(&format_dis));
+        }

         assert_eq!(size, format_dis.size());

@@ -146,7 +154,9 @@ fn main() {
         };

         let (mut ssimu, mut scores_ssimu) = if args.ssimulacra2 {
-            println!("Initializing SSIMULACRA2");
+            if !args.frame_score_only {
+                println!("Using SSIMULACRA2");
+            }
             (
                 Some(Ssimulacra2::new(&lrgb_ref, &lrgb_dis, &streams[0]).unwrap()),
                 Some(Vec::with_capacity(4096)),
@@ -157,7 +167,9 @@ fn main() {

         streams[0].sync().unwrap();

-        println!("Initialized, now processing ...");
+        if !args.frame_score_only {
+            println!("Initialized, now processing ...");
+        }
         let start = Instant::now();
         let mut decode_count = 0;
         let mut compute_count = 0;
@@ -266,6 +278,8 @@ fn main() {

                     streams[0].sync().unwrap();

+                    compute_count += 1;
+
                     if let Some(scores_psnr) = &mut scores_psnr {
                         scores_psnr.push(psnr as f64);
                     }
@@ -276,10 +290,13 @@ fn main() {
                         scores_msssim.push(msssim as f64);
                     }
                     if let Some(scores_ssimu) = &mut scores_ssimu {
+                        // print frame number and score to stdout
+                        if args.frame_score_only {
+                            println!("{} {}", decode_count, ssimu.as_mut().unwrap().get_score());
+                        }
                         scores_ssimu.push(ssimu.as_mut().unwrap().get_score());
                     }

-                    compute_count += 1;
                 } else {
                     break 'main;
                 }
@@ -292,25 +309,27 @@ fn main() {
             format.size().width as f64 * format.size().height as f64 * compute_count as f64
                 / total as f64
                 / 1000.0;
-        println!("Done !");
-        println!(
-            "Decoded: {}, processed: {} frame pairs in {total} ms ({} fps) (Mpx/s: {:.3})",
-            decode_count, compute_count, fps, perf_score
-        );
-        println!("Stats :");
-        if let Some(scores) = &scores_psnr {
-            println!("  psnr: {:#?}", Stats::compute(scores));
-        }
-        if let Some(scores) = &scores_ssim {
-            println!("  ssim: {:#?}", Stats::compute(scores));
-        }
-        if let Some(scores) = &scores_msssim {
-            println!("  msssim: {:#?}", Stats::compute(scores));
-        }
-        if let Some(scores) = &scores_ssimu {
-            println!("  ssimulacra2: {:#?}", Stats::compute(scores));
+        // if not args.frame_score_only print this stuff
+        if !args.frame_score_only {
+            println!("Done !");
+            println!(
+                "Decoded: {}, processed: {} frame pairs in {total} ms ({} fps) (Mpx/s: {:.3})",
+                decode_count, compute_count, fps, perf_score
+            );
+            println!("Stats :");
+            if let Some(scores) = &scores_psnr {
+                println!("  psnr: {:#?}", Stats::compute(scores));
+            }
+            if let Some(scores) = &scores_ssim {
+                println!("  ssim: {:#?}", Stats::compute(scores));
+            }
+            if let Some(scores) = &scores_msssim {
+                println!("  msssim: {:#?}", Stats::compute(scores));
+            }
+            if let Some(scores) = &scores_ssimu {
+                println!("  ssimulacra2: {:#?}", Stats::compute(scores));
+            }
         }
     }
-
     dev.release_primary_ctx().unwrap();
 }
Gui-Yom / turbo-metrics

Add option to print scores of each frame pair #4
