deepshrink_ffmpeg/
vmaf.rs

1//! VMAF quality measurement via ffmpeg's `libvmaf` filter.
2//!
3//! `libvmaf` is not present in every ffmpeg build, so callers should gate on
4//! [`has_libvmaf`] and degrade gracefully (skip the measurement) when it is
5//! absent — see the CLI's `--vmaf` handling.
6
7use std::path::Path;
8use std::process::Command;
9
10use crate::FfmpegError;
11
12/// Whether this ffmpeg build exposes the `libvmaf` filter.
13pub fn has_libvmaf(ffmpeg: &Path) -> bool {
14    Command::new(ffmpeg)
15        .args(["-hide_banner", "-filters"])
16        .output()
17        .map(|o| {
18            let text = String::from_utf8_lossy(&o.stdout);
19            text.contains("libvmaf")
20        })
21        .unwrap_or(false)
22}
23
24/// Measure the VMAF of `distorted` against the `reference` (original) video.
25///
26/// The distorted stream is scaled up to the reference resolution and both
27/// streams are normalized to `fps` so their frame counts align (a downscale or
28/// fps cap would otherwise desynchronize the comparison). Returns the pooled
29/// mean VMAF score in 0..=100.
30pub fn measure_vmaf(
31    ffmpeg: &Path,
32    distorted: &Path,
33    reference: &Path,
34    ref_width: u32,
35    ref_height: u32,
36    fps: f64,
37    n_threads: usize,
38) -> Result<f64, FfmpegError> {
39    // Input 0 = distorted, input 1 = reference. libvmaf consumes `[main][ref]`.
40    let fps_term = if fps.is_finite() && fps > 0.0 {
41        format!(",fps={fps}")
42    } else {
43        String::new()
44    };
45    let threads = n_threads.max(1);
46    let filter = format!(
47        "[0:v]scale={ref_width}:{ref_height}:flags=bicubic,setsar=1{fps_term},\
48         settb=AVTB,setpts=PTS-STARTPTS[dist];\
49         [1:v]setsar=1{fps_term},settb=AVTB,setpts=PTS-STARTPTS[ref];\
50         [dist][ref]libvmaf=n_threads={threads}"
51    );
52
53    let output = Command::new(ffmpeg)
54        .args(["-hide_banner", "-nostdin", "-i"])
55        .arg(distorted)
56        .arg("-i")
57        .arg(reference)
58        .args(["-lavfi", &filter, "-f", "null", "-"])
59        .output()
60        .map_err(|source| FfmpegError::Spawn {
61            tool: "ffmpeg",
62            source,
63        })?;
64
65    if !output.status.success() {
66        return Err(FfmpegError::CommandFailed {
67            tool: "ffmpeg",
68            status: output.status.to_string(),
69            stderr: String::from_utf8_lossy(&output.stderr).trim().to_string(),
70        });
71    }
72
73    // libvmaf logs `VMAF score: NN.NN` to stderr.
74    let stderr = String::from_utf8_lossy(&output.stderr);
75    parse_vmaf_score(&stderr)
76        .ok_or_else(|| FfmpegError::Parse("could not find a VMAF score in ffmpeg output".into()))
77}
78
79/// Extract the `VMAF score: <n>` value from ffmpeg's stderr.
80fn parse_vmaf_score(stderr: &str) -> Option<f64> {
81    const MARKER: &str = "VMAF score:";
82    for line in stderr.lines() {
83        if let Some(idx) = line.find(MARKER) {
84            let tail = line[idx + MARKER.len()..].trim();
85            if let Some(score) = tail
86                .split_whitespace()
87                .next()
88                .and_then(|t| t.parse::<f64>().ok())
89            {
90                return Some(score);
91            }
92        }
93    }
94    None
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn parses_score_line() {
103        let s = "frame=  150 fps=...\n[libvmaf @ 0x7f] VMAF score: 91.234567\n";
104        assert_eq!(parse_vmaf_score(s), Some(91.234567));
105    }
106
107    #[test]
108    fn parses_bare_score_line() {
109        assert_eq!(parse_vmaf_score("VMAF score: 100.000000"), Some(100.0));
110    }
111
112    #[test]
113    fn none_when_absent() {
114        assert_eq!(parse_vmaf_score("frame=1 fps=2\nEncoding done\n"), None);
115    }
116}
deepshrink_ffmpeg/vmaf.rs

deepshrink_ffmpeg/
vmaf.rs
