av_metrics/video/
psnr.rs

1//! Peak Signal-to-Noise Ratio metric.
2//!
3//! PSNR is most easily defined via the mean squared error between two images.
4//!
5//! See https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio for more details.
6
7use crate::video::decode::Decoder;
8use crate::video::pixel::CastFromPrimitive;
9use crate::video::pixel::Pixel;
10use crate::video::{PlanarMetrics, VideoMetric};
11use crate::MetricsError;
12use std::error::Error;
13use std::mem::size_of;
14use v_frame::frame::Frame;
15use v_frame::plane::Plane;
16use v_frame::prelude::ChromaSampling;
17
18use super::FrameCompare;
19
20/// Calculates the PSNR for two videos. Higher is better.
21///
22/// PSNR is capped at 100 in order to avoid skewed statistics
23/// from e.g. all black frames, which would
24/// otherwise show a PSNR of infinity.
25#[inline]
26pub fn calculate_video_psnr<D: Decoder, F: Fn(usize) + Send>(
27    decoder1: &mut D,
28    decoder2: &mut D,
29    frame_limit: Option<usize>,
30    progress_callback: F,
31) -> Result<PlanarMetrics, Box<dyn Error>> {
32    let metrics = Psnr.process_video(decoder1, decoder2, frame_limit, progress_callback)?;
33    Ok(metrics.psnr)
34}
35
36/// Calculates the APSNR for two videos. Higher is better.
37///
38/// APSNR is capped at 100 in order to avoid skewed statistics
39/// from e.g. all black frames, which would
40/// otherwise show a APSNR of infinity.
41#[inline]
42pub fn calculate_video_apsnr<D: Decoder, F: Fn(usize) + Send>(
43    decoder1: &mut D,
44    decoder2: &mut D,
45    frame_limit: Option<usize>,
46    progress_callback: F,
47) -> Result<PlanarMetrics, Box<dyn Error>> {
48    let metrics = Psnr.process_video(decoder1, decoder2, frame_limit, progress_callback)?;
49    Ok(metrics.apsnr)
50}
51
52/// Calculates the PSNR for two video frames. Higher is better.
53///
54/// PSNR is capped at 100 in order to avoid skewed statistics
55/// from e.g. all black frames, which would
56/// otherwise show a PSNR of infinity.
57#[inline]
58pub fn calculate_frame_psnr<T: Pixel>(
59    frame1: &Frame<T>,
60    frame2: &Frame<T>,
61    bit_depth: usize,
62    chroma_sampling: ChromaSampling,
63) -> Result<PlanarMetrics, Box<dyn Error>> {
64    let metrics = Psnr.process_frame(frame1, frame2, bit_depth, chroma_sampling)?;
65    Ok(PlanarMetrics {
66        y: calculate_psnr(metrics[0]),
67        u: calculate_psnr(metrics[1]),
68        v: calculate_psnr(metrics[2]),
69        avg: calculate_summed_psnr(&metrics),
70    })
71}
72
73#[derive(Debug, Clone, Copy)]
74struct PsnrResults {
75    psnr: PlanarMetrics,
76    apsnr: PlanarMetrics,
77}
78
79struct Psnr;
80
81impl VideoMetric for Psnr {
82    type FrameResult = [PsnrMetrics; 3];
83    type VideoResult = PsnrResults;
84
85    fn process_frame<T: Pixel>(
86        &self,
87        frame1: &Frame<T>,
88        frame2: &Frame<T>,
89        bit_depth: usize,
90        _chroma_sampling: ChromaSampling,
91    ) -> Result<Self::FrameResult, Box<dyn Error>> {
92        if (size_of::<T>() == 1 && bit_depth > 8) || (size_of::<T>() == 2 && bit_depth <= 8) {
93            return Err(Box::new(MetricsError::InputMismatch {
94                reason: "Bit depths does not match pixel width",
95            }));
96        }
97
98        frame1.can_compare(frame2)?;
99
100        let bit_depth = bit_depth;
101        let mut y = Default::default();
102        let mut u = Default::default();
103        let mut v = Default::default();
104
105        rayon::scope(|s| {
106            s.spawn(|_| {
107                y = calculate_plane_psnr_metrics(&frame1.planes[0], &frame2.planes[0], bit_depth)
108            });
109            s.spawn(|_| {
110                u = calculate_plane_psnr_metrics(&frame1.planes[1], &frame2.planes[1], bit_depth)
111            });
112            s.spawn(|_| {
113                v = calculate_plane_psnr_metrics(&frame1.planes[2], &frame2.planes[2], bit_depth)
114            });
115        });
116
117        Ok([y, u, v])
118    }
119
120    fn aggregate_frame_results(
121        &self,
122        metrics: &[Self::FrameResult],
123    ) -> Result<Self::VideoResult, Box<dyn Error>> {
124        let psnr = PlanarMetrics {
125            y: calculate_summed_psnr(&metrics.iter().map(|m| m[0]).collect::<Vec<_>>()),
126            u: calculate_summed_psnr(&metrics.iter().map(|m| m[1]).collect::<Vec<_>>()),
127            v: calculate_summed_psnr(&metrics.iter().map(|m| m[2]).collect::<Vec<_>>()),
128            avg: calculate_summed_psnr(&metrics.iter().flatten().copied().collect::<Vec<_>>()),
129        };
130        let apsnr = PlanarMetrics {
131            y: metrics.iter().map(|m| calculate_psnr(m[0])).sum::<f64>() / metrics.len() as f64,
132            u: metrics.iter().map(|m| calculate_psnr(m[1])).sum::<f64>() / metrics.len() as f64,
133            v: metrics.iter().map(|m| calculate_psnr(m[2])).sum::<f64>() / metrics.len() as f64,
134            avg: metrics
135                .iter()
136                .map(|m| calculate_summed_psnr(m))
137                .sum::<f64>()
138                / metrics.len() as f64,
139        };
140        Ok(PsnrResults { psnr, apsnr })
141    }
142}
143
144#[derive(Debug, Clone, Copy, Default)]
145struct PsnrMetrics {
146    sq_err: f64,
147    n_pixels: usize,
148    sample_max: usize,
149}
150
151fn calculate_summed_psnr(metrics: &[PsnrMetrics]) -> f64 {
152    calculate_psnr(
153        metrics
154            .iter()
155            .fold(PsnrMetrics::default(), |acc, plane| PsnrMetrics {
156                sq_err: acc.sq_err + plane.sq_err,
157                sample_max: plane.sample_max,
158                n_pixels: acc.n_pixels + plane.n_pixels,
159            }),
160    )
161}
162
163/// Calculate the PSNR metrics for a `Plane` by comparing the original (uncompressed) to
164/// the compressed version.
165fn calculate_plane_psnr_metrics<T: Pixel>(
166    plane1: &Plane<T>,
167    plane2: &Plane<T>,
168    bit_depth: usize,
169) -> PsnrMetrics {
170    let sq_err = calculate_plane_total_squared_error(plane1, plane2);
171    let max = (1 << bit_depth) - 1;
172    PsnrMetrics {
173        sq_err,
174        n_pixels: plane1.cfg.width * plane1.cfg.height,
175        sample_max: max,
176    }
177}
178
179fn calculate_psnr(metrics: PsnrMetrics) -> f64 {
180    if metrics.sq_err <= std::f64::EPSILON {
181        return 100.0;
182    }
183    10.0 * ((metrics.sample_max.pow(2) as f64).log10() + (metrics.n_pixels as f64).log10()
184        - metrics.sq_err.log10())
185}
186
187/// Calculate the squared error for a `Plane` by comparing the original (uncompressed)
188/// to the compressed version.
189fn calculate_plane_total_squared_error<T: Pixel>(plane1: &Plane<T>, plane2: &Plane<T>) -> f64 {
190    plane1
191        .data
192        .iter()
193        .zip(plane2.data.iter())
194        .map(|(a, b)| (i32::cast_from(*a) - i32::cast_from(*b)).unsigned_abs() as u64)
195        .map(|err| err * err)
196        .sum::<u64>() as f64
197}
av_metrics/video/psnr.rs

av_metrics/video/
psnr.rs
