Skip to content

Commit

Permalink
Make --vmaf parameter work independently again
Browse files Browse the repository at this point in the history
Solves issue master-of-zen#824

rustfmt
  • Loading branch information
damster101 committed Sep 26, 2024
1 parent 46cd62c commit 4d79448
Show file tree
Hide file tree
Showing 4 changed files with 45 additions and 19 deletions.
52 changes: 35 additions & 17 deletions av1an-core/src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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);
}
Expand Down
4 changes: 4 additions & 0 deletions av1an-core/src/settings.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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: Option<String>,
pub vmaf_threads: Option<usize>,
pub vmaf_filter: Option<String>,
}

impl EncodeArgs {
Expand Down
6 changes: 5 additions & 1 deletion av1an/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,

Expand Down Expand Up @@ -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: 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 {
Expand Down
2 changes: 1 addition & 1 deletion site/src/Cli/vmaf.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 4d79448

Please sign in to comment.