diff --git a/av1an-core/src/context.rs b/av1an-core/src/context.rs index c554cab7..4f472202 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,47 @@ 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() + .unwrap_or_else(|| "1920x1080".to_string()) + }; + + 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, |p| p.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/settings.rs b/av1an-core/src/settings.rs index 9400409e..6782f54b 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, pub vmaf: bool, + pub vmaf_path: Option, + pub vmaf_res: Option, + pub vmaf_threads: Option, + pub vmaf_filter: Option, } impl EncodeArgs { diff --git a/av1an/src/main.rs b/av1an/src/main.rs index a61d95a7..f02ce1f5 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, @@ -750,6 +750,10 @@ pub fn parse_cli(args: CliOpts) -> anyhow::Result> { )?, target_quality: args.target_quality_params(temp, video_params, output_pix_format.format), vmaf: args.vmaf, + vmaf_path: args.vmaf_path.clone(), + vmaf_res: Some(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 - Number of threads to use for VMAF calculation + Number of threads to use for target quality VMAF calculation --vmaf-filter Filter applied to source at VMAF calcualation