Make --vmaf parameter work independently again

Solves issue #824
master-of-zen · Oct 20, 2024 · b6ee383 · b6ee383
1 parent 96ee8b8
commit b6ee383
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 19 deletions.
diff --git a/av1an-core/src/context.rs b/av1an-core/src/context.rs
@@ -9,6 +9,7 @@ use std::path::{Path, PathBuf};
 use std::process::{exit, Command, Stdio};
 use std::sync::atomic::{self, AtomicBool, AtomicUsize};
 use std::sync::{mpsc, Arc};
+use std::thread::available_parallelism;
 use std::{cmp, fs, iter, thread};
 
 use ansi_term::{Color, Style};
@@ -349,30 +350,46 @@ impl Av1anContext {
         }
       }
 
-      if let Some(ref tq) = self.args.target_quality {
-        let mut temp_res = tq.vmaf_res.to_string();
-        if tq.vmaf_res == "inputres" {
-          let inputres = self.args.input.resolution()?;
-          temp_res.push_str(&format!(
-            "{}x{}",
-            &inputres.0.to_string(),
-            &inputres.1.to_string()
-          ));
-          temp_res.to_string();
+      if self.args.vmaf || self.args.target_quality.is_some() {
+        let vmaf_res = if let Some(ref tq) = self.args.target_quality {
+          if tq.vmaf_res == "inputres" {
+            let inputres = self.args.input.resolution()?;
+            format!("{}x{}", inputres.0, inputres.1)
+          } else {
+            tq.vmaf_res.clone()
+          }
         } else {
-          temp_res = tq.vmaf_res.to_string();
-        }
+          self.args.vmaf_res.clone()
+        };
+
+        let vmaf_model = self.args.vmaf_path.as_deref().or_else(|| {
+          self
+            .args
+            .target_quality
+            .as_ref()
+            .and_then(|tq| tq.model.as_deref())
+        });
+        let vmaf_scaler = "bicubic";
+        let vmaf_filter = self.args.vmaf_filter.as_deref().or_else(|| {
+          self
+            .args
+            .target_quality
+            .as_ref()
+            .and_then(|tq| tq.vmaf_filter.as_deref())
+        });
 
         if self.args.vmaf {
+          let vmaf_threads = available_parallelism().map_or(1, std::num::NonZero::get);
+
           if let Err(e) = vmaf::plot(
             self.args.output_file.as_ref(),
             &self.args.input,
-            tq.model.as_deref(),
-            temp_res.as_str(),
-            tq.vmaf_scaler.as_str(),
+            vmaf_model,
+            &vmaf_res,
+            vmaf_scaler,
             1,
-            tq.vmaf_filter.as_deref(),
-            tq.vmaf_threads,
+            vmaf_filter,
+            vmaf_threads,
           ) {
             error!("VMAF calculation failed with error: {}", e);
           }

diff --git a/av1an-core/src/scenes.rs b/av1an-core/src/scenes.rs
@@ -300,6 +300,10 @@ fn get_test_args() -> Av1anContext {
     zones: None,
     scaler: String::new(),
     ignore_frame_mismatch: false,
+    vmaf_path: None,
+    vmaf_res: "1920x1080".to_string(),
+    vmaf_threads: None,
+    vmaf_filter: None,
   };
   Av1anContext {
     vs_script: None,

diff --git a/av1an-core/src/settings.rs b/av1an-core/src/settings.rs
@@ -78,6 +78,10 @@ pub struct EncodeArgs {
   pub concat: ConcatMethod,
   pub target_quality: Option<TargetQuality>,
   pub vmaf: bool,
+  pub vmaf_path: Option<PathBuf>,
+  pub vmaf_res: String,
+  pub vmaf_threads: Option<usize>,
+  pub vmaf_filter: Option<String>,
 }
 
 impl EncodeArgs {

diff --git a/av1an/src/main.rs b/av1an/src/main.rs
@@ -492,7 +492,7 @@ pub struct CliOpts {
   #[clap(long, default_value = "1920x1080", help_heading = "VMAF")]
   pub vmaf_res: String,
 
-  /// Number of threads to use for VMAF calculation
+  /// Number of threads to use for target quality VMAF calculation
   #[clap(long, help_heading = "VMAF")]
   pub vmaf_threads: Option<usize>,
 
@@ -750,6 +750,10 @@ pub fn parse_cli(args: CliOpts) -> anyhow::Result<Vec<EncodeArgs>> {
       )?,
       target_quality: args.target_quality_params(temp, video_params, output_pix_format.format),
       vmaf: args.vmaf,
+      vmaf_path: args.vmaf_path.clone(),
+      vmaf_res: args.vmaf_res.clone(),
+      vmaf_threads: args.vmaf_threads,
+      vmaf_filter: args.vmaf_filter.clone(),
       verbosity: if args.quiet {
         Verbosity::Quiet
       } else if args.verbose {

diff --git a/site/src/Cli/vmaf.md b/site/src/Cli/vmaf.md
@@ -21,7 +21,7 @@
 		[default: 1920x1080]
 
 	--vmaf-threads <VMAF_THREADS>
-		Number of threads to use for VMAF calculation
+		Number of threads to use for target quality VMAF calculation
 
 	--vmaf-filter <VMAF_FILTER>
 		Filter applied to source at VMAF calcualation