From 9012146511c52f4ee88dc6f41a659286a928897d Mon Sep 17 00:00:00 2001 From: Alex Butler Date: Sun, 14 Jan 2024 16:30:15 +0000 Subject: [PATCH] Use a single ffmpeg process to calculate VMAF (#177) * Use a single ffmpeg process to calculate VMAF * do format first & setpts last --- .github/workflows/ci.yml | 17 ++-- CHANGELOG.md | 3 + Cargo.lock | 49 +--------- Cargo.toml | 7 -- src/command/args/vmaf.rs | 103 ++++++++++++++------ src/command/sample_encode.rs | 12 ++- src/command/vmaf.rs | 10 +- src/vmaf.rs | 177 +---------------------------------- 8 files changed, 106 insertions(+), 272 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d0af7e9..13b6095 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,14 +20,15 @@ jobs: - run: cargo run --locked -- print-completions fish - run: cargo run --locked -- print-completions zsh - test-windows: - runs-on: windows-latest - env: - RUST_BACKTRACE: 1 - steps: - - run: rustup update stable - - uses: actions/checkout@v4 - - run: cargo check + # Disabled while we no longer need cfg(windows) code + # test-windows: + # runs-on: windows-latest + # env: + # RUST_BACKTRACE: 1 + # steps: + # - run: rustup update stable + # - uses: actions/checkout@v4 + # - run: cargo check rustfmt: runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index 8112466..6e40f4c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,6 @@ +# Unreleased +* Use a single ffmpeg process to calculate VMAF replacing multi process piping. + # v0.7.12 * Improve eta stability. diff --git a/Cargo.lock b/Cargo.lock index af22ea6..1c5839d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -26,7 +26,6 @@ dependencies = [ "tokio", "tokio-process-stream", "tokio-stream", - "unix-named-pipe", ] [[package]] @@ -318,17 +317,6 @@ version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f" -[[package]] -name = "errno" -version = "0.2.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f639046355ee4f37944e44f60642c6f3a7efa3cf6b78c78a0d989a8ce6c396a1" -dependencies = [ - "errno-dragonfly", - "libc", - "winapi", -] - [[package]] name = "errno" version = "0.3.8" @@ -339,16 +327,6 @@ dependencies = [ "windows-sys 0.52.0", ] -[[package]] -name = "errno-dragonfly" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf" -dependencies = [ - "cc", - "libc", -] - [[package]] name = "fastrand" version = "2.0.1" @@ -753,7 +731,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "322394588aaf33c24007e8bb3238ee3e4c5c09c084ab32bc73890b99ff326bca" dependencies = [ "bitflags 2.4.1", - "errno 0.3.8", + "errno", "libc", "linux-raw-sys", "windows-sys 0.52.0", @@ -844,19 +822,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4dccd0940a2dcdf68d092b8cbab7dc0ad8fa938bf95787e1b916b0e3d0e8e970" - -[[package]] -name = "socket2" -version = "0.5.5" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b5fac59a5cb5dd637972e5fca70daf0523c9067fcdc4842f053dae04a18f8e9" -dependencies = [ - "libc", - "windows-sys 0.48.0", -] +checksum = "2593d31f82ead8df961d8bd23a64c2ccf2eb5dd34b0a34bfb4dd54011c72009e" [[package]] name = "strsim" @@ -946,7 +914,6 @@ dependencies = [ "num_cpus", "pin-project-lite", "signal-hook-registry", - "socket2", "tokio-macros", "windows-sys 0.48.0", ] @@ -1013,16 +980,6 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" -[[package]] -name = "unix-named-pipe" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ad653da8f36ac5825ba06642b5a3cce14a4e52c6a5fab4a8928d53f4426dae2" -dependencies = [ - "errno 0.2.8", - "libc", -] - [[package]] name = "utf8parse" version = "0.2.1" diff --git a/Cargo.toml b/Cargo.toml index ed45151..42543f4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,13 +31,6 @@ tokio = { version = "1.15", features = ["rt", "macros", "process", "fs", "signal tokio-process-stream = "0.4" tokio-stream = "0.1" -[target.'cfg(unix)'.dependencies] -unix-named-pipe = "0.2" - -[target.'cfg(windows)'.dependencies.tokio] -version = "1.15" -features = ["net"] - [profile.release] lto = true opt-level = "s" diff --git a/src/command/args/vmaf.rs b/src/command/args/vmaf.rs index 63ffefd..b504a0d 100644 --- a/src/command/args/vmaf.rs +++ b/src/command/args/vmaf.rs @@ -1,6 +1,7 @@ +use crate::command::args::PixelFormat; use anyhow::Context; use clap::Parser; -use std::{fmt::Display, sync::Arc, thread}; +use std::{borrow::Cow, fmt::Display, sync::Arc, thread}; /// Common vmaf options. #[derive(Parser, Clone, Hash)] @@ -39,7 +40,12 @@ impl Vmaf { } /// Returns ffmpeg `filter_complex`/`lavfi` value for calculating vmaf. - pub fn ffmpeg_lavfi(&self, distorted_res: Option<(u32, u32)>) -> String { + pub fn ffmpeg_lavfi( + &self, + distorted_res: Option<(u32, u32)>, + pix_fmt: PixelFormat, + ref_vfilter: Option<&str>, + ) -> String { let mut args = self.vmaf_args.clone(); if !args.iter().any(|a| a.contains("n_threads")) { // default n_threads to all cores @@ -63,14 +69,30 @@ impl Vmaf { } } - if let Some((w, h)) = self.vf_scale(model.unwrap_or_default(), distorted_res) { - // scale both streams to the vmaf width - lavfi.insert_str( - 0, - &format!("[0:v]scale={w}:{h}:flags=bicubic[dis];[1:v]scale={w}:{h}:flags=bicubic[ref];[dis][ref]"), - ); - } + let ref_vf: Cow<_> = match ref_vfilter { + None => "".into(), + Some(vf) if vf.ends_with(',') => vf.into(), + Some(vf) => format!("{vf},").into(), + }; + + // prefix: + // * Add reference-vfilter if any + // * convert both streams to common pixel format + // * scale to vmaf width if necessary + // * sync presentation timestamp + let prefix = if let Some((w, h)) = self.vf_scale(model.unwrap_or_default(), distorted_res) { + format!( + "[0:v]format={pix_fmt},scale={w}:{h}:flags=bicubic,setpts=PTS-STARTPTS[dis];\ + [1:v]format={pix_fmt},{ref_vf}scale={w}:{h}:flags=bicubic,setpts=PTS-STARTPTS[ref];[dis][ref]" + ) + } else { + format!( + "[0:v]format={pix_fmt},setpts=PTS-STARTPTS[dis];\ + [1:v]format={pix_fmt},{ref_vf}setpts=PTS-STARTPTS[ref];[dis][ref]" + ) + }; + lavfi.insert_str(0, &prefix); lavfi } @@ -176,7 +198,12 @@ fn vmaf_lavfi() { vmaf_args: vec!["n_threads=5".into(), "n_subsample=4".into()], vmaf_scale: VmafScale::Auto, }; - assert_eq!(vmaf.ffmpeg_lavfi(None), "libvmaf=n_threads=5:n_subsample=4"); + assert_eq!( + vmaf.ffmpeg_lavfi(None, PixelFormat::Yuv420p, Some("scale=1280:-1,fps=24")), + "[0:v]format=yuv420p,setpts=PTS-STARTPTS[dis];\ + [1:v]format=yuv420p,scale=1280:-1,fps=24,setpts=PTS-STARTPTS[ref];\ + [dis][ref]libvmaf=n_threads=5:n_subsample=4" + ); } #[test] @@ -186,10 +213,15 @@ fn vmaf_lavfi_default() { vmaf_scale: VmafScale::Auto, }; let expected = format!( - "libvmaf=n_threads={}", + "[0:v]format=yuv420p10le,setpts=PTS-STARTPTS[dis];\ + [1:v]format=yuv420p10le,setpts=PTS-STARTPTS[ref];\ + [dis][ref]libvmaf=n_threads={}", thread::available_parallelism().map_or(1, |p| p.get()) ); - assert_eq!(vmaf.ffmpeg_lavfi(None), expected); + assert_eq!( + vmaf.ffmpeg_lavfi(None, PixelFormat::Yuv420p10le, None), + expected + ); } #[test] @@ -199,10 +231,15 @@ fn vmaf_lavfi_include_n_threads() { vmaf_scale: VmafScale::Auto, }; let expected = format!( - "libvmaf=log_path=output.xml:n_threads={}", + "[0:v]format=yuv420p,setpts=PTS-STARTPTS[dis];\ + [1:v]format=yuv420p,setpts=PTS-STARTPTS[ref];\ + [dis][ref]libvmaf=log_path=output.xml:n_threads={}", thread::available_parallelism().map_or(1, |p| p.get()) ); - assert_eq!(vmaf.ffmpeg_lavfi(None), expected); + assert_eq!( + vmaf.ffmpeg_lavfi(None, PixelFormat::Yuv420p, None), + expected + ); } /// Low resolution videos should be upscaled to 1080p @@ -213,9 +250,9 @@ fn vmaf_lavfi_small_width() { vmaf_scale: VmafScale::Auto, }; assert_eq!( - vmaf.ffmpeg_lavfi(Some((1280, 720))), - "[0:v]scale=1920:-1:flags=bicubic[dis];\ - [1:v]scale=1920:-1:flags=bicubic[ref];\ + vmaf.ffmpeg_lavfi(Some((1280, 720)), PixelFormat::Yuv420p, None), + "[0:v]format=yuv420p,scale=1920:-1:flags=bicubic,setpts=PTS-STARTPTS[dis];\ + [1:v]format=yuv420p,scale=1920:-1:flags=bicubic,setpts=PTS-STARTPTS[ref];\ [dis][ref]libvmaf=n_threads=5:n_subsample=4" ); } @@ -228,8 +265,10 @@ fn vmaf_lavfi_4k() { vmaf_scale: VmafScale::Auto, }; assert_eq!( - vmaf.ffmpeg_lavfi(Some((3840, 2160))), - "libvmaf=n_threads=5:n_subsample=4:model=version=vmaf_4k_v0.6.1" + vmaf.ffmpeg_lavfi(Some((3840, 2160)), PixelFormat::Yuv420p, None), + "[0:v]format=yuv420p,setpts=PTS-STARTPTS[dis];\ + [1:v]format=yuv420p,setpts=PTS-STARTPTS[ref];\ + [dis][ref]libvmaf=n_threads=5:n_subsample=4:model=version=vmaf_4k_v0.6.1" ); } @@ -241,9 +280,9 @@ fn vmaf_lavfi_3k_upscale_to_4k() { vmaf_scale: VmafScale::Auto, }; assert_eq!( - vmaf.ffmpeg_lavfi(Some((3008, 1692))), - "[0:v]scale=3840:-1:flags=bicubic[dis];\ - [1:v]scale=3840:-1:flags=bicubic[ref];\ + vmaf.ffmpeg_lavfi(Some((3008, 1692)), PixelFormat::Yuv420p, None), + "[0:v]format=yuv420p,scale=3840:-1:flags=bicubic,setpts=PTS-STARTPTS[dis];\ + [1:v]format=yuv420p,scale=3840:-1:flags=bicubic,setpts=PTS-STARTPTS[ref];\ [dis][ref]libvmaf=n_threads=5:model=version=vmaf_4k_v0.6.1" ); } @@ -260,8 +299,10 @@ fn vmaf_lavfi_small_width_custom_model() { vmaf_scale: VmafScale::Auto, }; assert_eq!( - vmaf.ffmpeg_lavfi(Some((1280, 720))), - "libvmaf=model=version=foo:n_threads=5:n_subsample=4" + vmaf.ffmpeg_lavfi(Some((1280, 720)), PixelFormat::Yuv420p, None), + "[0:v]format=yuv420p,setpts=PTS-STARTPTS[dis];\ + [1:v]format=yuv420p,setpts=PTS-STARTPTS[ref];\ + [dis][ref]libvmaf=model=version=foo:n_threads=5:n_subsample=4" ); } @@ -280,10 +321,10 @@ fn vmaf_lavfi_custom_model_and_width() { }, }; assert_eq!( - vmaf.ffmpeg_lavfi(Some((1280, 720))), - "[0:v]scale=123:-1:flags=bicubic[dis];\ - [1:v]scale=123:-1:flags=bicubic[ref];\ - [dis][ref]libvmaf=model=version=foo:n_threads=5:n_subsample=4" + vmaf.ffmpeg_lavfi(Some((1280, 720)), PixelFormat::Yuv420p, None), + "[0:v]format=yuv420p,scale=123:-1:flags=bicubic,setpts=PTS-STARTPTS[dis];\ + [1:v]format=yuv420p,scale=123:-1:flags=bicubic,setpts=PTS-STARTPTS[ref];\ + [dis][ref]libvmaf=model=version=foo:n_threads=5:n_subsample=4" ); } @@ -294,7 +335,9 @@ fn vmaf_lavfi_1080p() { vmaf_scale: VmafScale::Auto, }; assert_eq!( - vmaf.ffmpeg_lavfi(Some((1920, 1080))), - "libvmaf=n_threads=5:n_subsample=4" + vmaf.ffmpeg_lavfi(Some((1920, 1080)), PixelFormat::Yuv420p, None), + "[0:v]format=yuv420p,setpts=PTS-STARTPTS[dis];\ + [1:v]format=yuv420p,setpts=PTS-STARTPTS[ref];\ + [dis][ref]libvmaf=n_threads=5:n_subsample=4" ); } diff --git a/src/command/sample_encode.rs b/src/command/sample_encode.rs index 280d850..0f4add0 100644 --- a/src/command/sample_encode.rs +++ b/src/command/sample_encode.rs @@ -214,12 +214,14 @@ pub async fn run( bar.set_message("vmaf running,"); let mut vmaf = vmaf::run( &sample, - args.vfilter.as_deref(), &encoded_sample, - &vmaf.ffmpeg_lavfi(encoded_probe.resolution), - enc_args - .pix_fmt - .max(input_pixel_format.unwrap_or(PixelFormat::Yuv444p10le)), + &vmaf.ffmpeg_lavfi( + encoded_probe.resolution, + enc_args + .pix_fmt + .max(input_pixel_format.unwrap_or(PixelFormat::Yuv444p10le)), + args.vfilter.as_deref(), + ), )?; let mut vmaf_score = -1.0; while let Some(vmaf) = vmaf.next().await { diff --git a/src/command/vmaf.rs b/src/command/vmaf.rs index 3f02225..3eed6c8 100644 --- a/src/command/vmaf.rs +++ b/src/command/vmaf.rs @@ -29,7 +29,7 @@ pub struct Args { pub reference: PathBuf, /// Ffmpeg video filter applied to the reference before analysis. - /// E.g. --vfilter "scale=1280:-1,fps=24". + /// E.g. --reference-vfilter "scale=1280:-1,fps=24". #[arg(long)] pub reference_vfilter: Option, @@ -68,10 +68,12 @@ pub async fn vmaf( let mut vmaf = vmaf::run( &reference, - reference_vfilter.as_deref(), &distorted, - &vmaf.ffmpeg_lavfi(dprobe.resolution), - dpix_fmt.max(rpix_fmt), + &vmaf.ffmpeg_lavfi( + dprobe.resolution, + dpix_fmt.max(rpix_fmt), + reference_vfilter.as_deref(), + ), )?; let mut vmaf_score = -1.0; while let Some(vmaf) = vmaf.next().await { diff --git a/src/vmaf.rs b/src/vmaf.rs index 3dbfd88..904a98e 100644 --- a/src/vmaf.rs +++ b/src/vmaf.rs @@ -1,8 +1,5 @@ //! vmaf logic -use crate::{ - command::args::PixelFormat, - process::{exit_ok_stderr, Chunks, CommandExt, FfmpegOut}, -}; +use crate::process::{exit_ok_stderr, Chunks, CommandExt, FfmpegOut}; use anyhow::Context; use std::path::Path; use tokio::process::Command; @@ -13,39 +10,18 @@ use tokio_stream::{Stream, StreamExt}; /// This can produce more accurate results than testing directly from original source. pub fn run( reference: &Path, - reference_vfilter: Option<&str>, distorted: &Path, filter_complex: &str, - pix_fmt: PixelFormat, ) -> anyhow::Result> { - // convert reference & distorted to yuv streams of the same pixel format - // frame rate and presentation timestamp to improve vmaf accuracy - let (yuv_out, yuv_pipe) = yuv::pipe(reference, pix_fmt, reference_vfilter)?; - let yuv_pipe = yuv_pipe.filter_map(VmafOut::ignore_ok); - - #[cfg(unix)] - let (distorted_fifo, distorted_yuv_pipe) = yuv::unix::pipe_to_fifo(distorted, pix_fmt)?; - #[cfg(unix)] - let (distorted, yuv_pipe) = ( - &distorted_fifo, - yuv_pipe.merge(distorted_yuv_pipe.filter_map(VmafOut::ignore_ok)), - ); - #[cfg(windows)] - let (distorted_npipe, distorted_yuv_pipe) = yuv::windows::named_pipe(distorted, pix_fmt)?; - #[cfg(windows)] - let (distorted, yuv_pipe) = ( - &distorted_npipe, - yuv_pipe.merge(distorted_yuv_pipe.filter_map(VmafOut::ignore_ok)), - ); - let vmaf: ProcessChunkStream = Command::new("ffmpeg") .kill_on_drop(true) + .arg2("-r", "24") .arg2("-i", distorted) - .arg2("-i", "-") + .arg2("-r", "24") + .arg2("-i", reference) .arg2("-filter_complex", filter_complex) .arg2("-f", "null") .arg("-") - .stdin(yuv_out) .try_into() .context("ffmpeg vmaf")?; @@ -56,7 +32,7 @@ pub fn run( Item::Done(code) => VmafOut::ignore_ok(exit_ok_stderr("ffmpeg vmaf", code, &chunks)), }); - Ok(yuv_pipe.merge(vmaf)) + Ok(vmaf) } #[derive(Debug)] @@ -89,146 +65,3 @@ impl VmafOut { None } } - -mod yuv { - use super::*; - use std::process::Stdio; - - /// ffmpeg yuv4mpegpipe returning the stdout & [`FfmpegProgress`] stream. - pub fn pipe( - input: &Path, - pix_fmt: PixelFormat, - vfilter: Option<&str>, - ) -> anyhow::Result<(Stdio, impl Stream>)> { - // sync presentation timestamp - let vfilter: std::borrow::Cow<'_, str> = match vfilter { - None => "setpts=PTS-STARTPTS".into(), - Some(vf) if vf.contains("setpts=") => vf.into(), - Some(vf) => format!("{vf},setpts=PTS-STARTPTS").into(), - }; - - let mut yuv4mpegpipe = Command::new("ffmpeg") - .kill_on_drop(true) - // Use 24fps to match vmaf models - .arg2("-r", "24") - .arg2("-i", input) - .arg2("-pix_fmt", pix_fmt.as_str()) - .arg2("-vf", vfilter.as_ref()) - .arg2("-strict", "-1") - .arg2("-f", "yuv4mpegpipe") - .arg("-") - .stdin(Stdio::null()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .context("ffmpeg yuv4mpegpipe")?; - let stdout = yuv4mpegpipe.stdout.take().unwrap().try_into().unwrap(); - let stream = FfmpegOut::stream(yuv4mpegpipe, "ffmpeg yuv4mpegpipe"); - Ok((stdout, stream)) - } - - #[cfg(windows)] - pub mod windows { - use super::*; - - pub fn named_pipe( - input: &Path, - pix_fmt: PixelFormat, - ) -> anyhow::Result<(String, impl Stream>)> { - let in_name = { - let mut n = String::from(r"\\.\pipe\ab-av1-in-"); - n.extend(std::iter::repeat_with(fastrand::alphanumeric).take(12)); - n - }; - - let in_server = tokio::net::windows::named_pipe::ServerOptions::new() - .access_outbound(false) - .first_pipe_instance(true) - .max_instances(1) - .create(&in_name)?; - - let out_name = in_name.replacen("-in-", "-out-", 1); - let out_server = tokio::net::windows::named_pipe::ServerOptions::new() - .access_inbound(false) - .first_pipe_instance(true) - .max_instances(1) - .create(&out_name)?; - - async fn copy_in_pipe_to_out( - mut in_pipe: tokio::net::windows::named_pipe::NamedPipeServer, - mut out_pipe: tokio::net::windows::named_pipe::NamedPipeServer, - ) -> tokio::io::Result<()> { - in_pipe.connect().await?; - in_pipe.readable().await?; - out_pipe.connect().await?; - out_pipe.writable().await?; - tokio::io::copy(&mut in_pipe, &mut out_pipe).await?; - Ok(()) - } - tokio::spawn(async move { - if let Err(err) = copy_in_pipe_to_out(in_server, out_server).await { - eprintln!("Error copy_in_pipe_to_out: {err}"); - } - }); - - let yuv4mpegpipe = Command::new("ffmpeg") - .kill_on_drop(true) - .arg2("-r", "24") - .arg2("-i", input) - .arg2("-pix_fmt", pix_fmt.as_str()) - .arg2("-vf", "setpts=PTS-STARTPTS") - .arg2("-strict", "-1") - .arg2("-f", "yuv4mpegpipe") - .arg("-y") - .arg(&in_name) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::piped()) - .spawn() - .context("ffmpeg yuv4mpegpipe")?; - let stream = FfmpegOut::stream(yuv4mpegpipe, "ffmpeg yuv4mpegpipe"); - - Ok((out_name, stream)) - } - } - - #[cfg(unix)] - pub mod unix { - use super::*; - use crate::temporary::{self, TempKind}; - use std::path::PathBuf; - - /// ffmpeg yuv4mpegpipe returning the temporary fifo path & [`FfmpegProgress`] stream. - pub fn pipe_to_fifo( - input: &Path, - pix_fmt: PixelFormat, - ) -> anyhow::Result<(PathBuf, impl Stream>)> { - let fifo = PathBuf::from(format!( - "/tmp/ab-av1-{}.fifo", - std::iter::repeat_with(fastrand::alphanumeric) - .take(12) - .collect::() - )); - unix_named_pipe::create(&fifo, None)?; - temporary::add(&fifo, TempKind::NotKeepable); - - let yuv4mpegpipe = Command::new("ffmpeg") - .kill_on_drop(true) - .arg2("-r", "24") - .arg2("-i", input) - .arg2("-pix_fmt", pix_fmt.as_str()) - .arg2("-vf", "setpts=PTS-STARTPTS") - .arg2("-strict", "-1") - .arg2("-f", "yuv4mpegpipe") - .arg("-y") - .arg(&fifo) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::piped()) - .spawn() - .context("ffmpeg yuv4mpegpipe")?; - let stream = FfmpegOut::stream(yuv4mpegpipe, "ffmpeg yuv4mpegpipe"); - Ok((fifo, stream)) - } - } -}