From 5ec173cc83d7d560501c11bde1a7033a17ca8173 Mon Sep 17 00:00:00 2001
From: Luis Garcia <git@luigi311.com>
Date: Tue, 23 Apr 2024 07:50:09 +0000
Subject: [PATCH] Add support for vmaf-subsample Add vmaf-subsample as an
 alternative to probing-rate as it is more accurate while still providing a
 speed benefit compared to without it

---
 av1an-core/src/context.rs        |  1 +
 av1an-core/src/target_quality.rs |  2 ++
 av1an-core/src/vmaf.rs           | 13 +++++++++----
 av1an/src/main.rs                |  5 +++++
 4 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/av1an-core/src/context.rs b/av1an-core/src/context.rs
index 85ffe58a..6e877274 100644
--- a/av1an-core/src/context.rs
+++ b/av1an-core/src/context.rs
@@ -368,6 +368,7 @@ impl Av1anContext {
             temp_res.as_str(),
             tq.vmaf_scaler.as_str(),
             1,
+            1,
             tq.vmaf_filter.as_deref(),
             tq.vmaf_threads,
           ) {
diff --git a/av1an-core/src/target_quality.rs b/av1an-core/src/target_quality.rs
index 9f6e4388..d1c4db3b 100644
--- a/av1an-core/src/target_quality.rs
+++ b/av1an-core/src/target_quality.rs
@@ -25,6 +25,7 @@ pub struct TargetQuality {
   pub vmaf_threads: usize,
   pub model: Option<PathBuf>,
   pub probing_rate: usize,
+  pub vmaf_subsample: usize,
   pub probes: u32,
   pub target: f64,
   pub min_q: u32,
@@ -252,6 +253,7 @@ impl TargetQuality {
       &self.vmaf_res,
       &self.vmaf_scaler,
       self.probing_rate,
+      self.vmaf_subsample,
       self.vmaf_filter.as_deref(),
       self.vmaf_threads,
     )?;
diff --git a/av1an-core/src/vmaf.rs b/av1an-core/src/vmaf.rs
index 1f63344b..408da1ef 100644
--- a/av1an-core/src/vmaf.rs
+++ b/av1an-core/src/vmaf.rs
@@ -122,6 +122,7 @@ pub fn plot(
   res: &str,
   scaler: &str,
   sample_rate: usize,
+  subsample: usize,
   filter: Option<&str>,
   threads: usize,
 ) -> Result<(), Box<EncoderCrash>> {
@@ -156,6 +157,7 @@ pub fn plot(
     res,
     scaler,
     sample_rate,
+    subsample,
     filter,
     threads,
   )?;
@@ -172,6 +174,7 @@ pub fn run_vmaf(
   res: &str,
   scaler: &str,
   sample_rate: usize,
+  subsample: usize,
   vmaf_filter: Option<&str>,
   threads: usize,
 ) -> Result<(), Box<EncoderCrash>> {
@@ -193,16 +196,18 @@ pub fn run_vmaf(
 
   let vmaf = if let Some(model) = model {
     format!(
-      "[distorted][ref]libvmaf=log_fmt='json':eof_action=endall:log_path={}:model_path={}:n_threads={}",
+      "[distorted][ref]libvmaf=log_fmt='json':eof_action=endall:log_path={}:model_path={}:n_threads={}:n_subsample={}",
       ffmpeg::escape_path_in_filter(stat_file),
       ffmpeg::escape_path_in_filter(&model),
-      threads
+      threads,
+      subsample,
     )
   } else {
     format!(
-      "[distorted][ref]libvmaf=log_fmt='json':eof_action=endall:log_path={}:n_threads={}",
+      "[distorted][ref]libvmaf=log_fmt='json':eof_action=endall:log_path={}:n_threads={}:n_subsample={}",
       ffmpeg::escape_path_in_filter(stat_file),
-      threads
+      threads,
+      subsample,
     )
   };
 
diff --git a/av1an/src/main.rs b/av1an/src/main.rs
index 702d569c..bc82bba0 100644
--- a/av1an/src/main.rs
+++ b/av1an/src/main.rs
@@ -514,6 +514,10 @@ pub struct CliOpts {
   #[clap(long, default_value_t = 1, help_heading = "Target Quality")]
   pub probing_rate: u32,
 
+  /// Vmaf subsample to speedup VMAF calculation
+  #[clap(long, default_value_t = 1, help_heading = "Target Quality")]
+  pub vmaf_subsample: usize,
+
   /// Use encoding settings for probes specified by --video-params rather than faster, less accurate settings
   ///
   /// Note that this always performs encoding in one-pass mode, regardless of --passes.
@@ -572,6 +576,7 @@ impl CliOpts {
         video_params: video_params.clone(),
         probe_slow: self.probe_slow,
         probing_rate: adapt_probing_rate(self.probing_rate as usize),
+        vmaf_subsample: self.vmaf_subsample,
       }
     })
   }