From 2239e6c2624fe9febb8797ff1eaa9b4e4bf6a565 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damian=20H=C3=B6ster?= <damian.hoester@posteo.de>
Date: Sat, 19 Oct 2024 16:42:23 +0200
Subject: [PATCH] Make --vmaf parameter work independently again

Solves issue #824
---
 av1an-core/src/context.rs  | 51 +++++++++++++++++++++++++-------------
 av1an-core/src/scenes.rs   |  4 +++
 av1an-core/src/settings.rs |  4 +++
 av1an/src/main.rs          |  6 ++++-
 site/src/Cli/vmaf.md       |  2 +-
 5 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/av1an-core/src/context.rs b/av1an-core/src/context.rs
index c554cab7..b57aae80 100644
--- 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
index e9aa39fd..b2983e6c 100644
--- 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
index 9400409e..2eeaf712 100644
--- 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
index a61d95a7..6690058c 100644
--- 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
index 1453316e..ddd2c212 100644
--- 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