oximedia_codec/
stats.rs

1//! Bitstream quality statistics.
2//!
3//! [`BitstreamStats`] accumulates per-frame quality metrics such as PSNR and
4//! SSIM that can be computed between the original (uncompressed) and
5//! reconstructed (compressed/decompressed) pixel buffers.
6//!
7//! These metrics are commonly used to evaluate codec quality and to drive
8//! rate-distortion optimization loops.
9
10#![allow(dead_code)]
11#![allow(clippy::cast_precision_loss)]
12#![allow(clippy::cast_possible_truncation)]
13
14// ---------------------------------------------------------------------------
15// BitstreamStats
16// ---------------------------------------------------------------------------
17
18/// Accumulated bitstream quality statistics.
19///
20/// Call [`BitstreamStats::update_psnr`] and/or [`BitstreamStats::update_ssim`]
21/// after each encoded frame to accumulate per-frame metrics.  Aggregated
22/// averages can then be retrieved via [`BitstreamStats::mean_psnr`] and
23/// [`BitstreamStats::mean_ssim`].
24#[derive(Debug, Clone, Default)]
25pub struct BitstreamStats {
26    /// Running sum of per-frame PSNR values (dB).
27    psnr_sum: f64,
28    /// Running sum of per-frame SSIM values [0, 1].
29    ssim_sum: f64,
30    /// Number of frames contributing to PSNR statistics.
31    psnr_count: u64,
32    /// Number of frames contributing to SSIM statistics.
33    ssim_count: u64,
34    /// Minimum PSNR observed (dB).
35    psnr_min: f64,
36    /// Maximum PSNR observed (dB).
37    psnr_max: f64,
38    /// Minimum SSIM observed.
39    ssim_min: f64,
40    /// Maximum SSIM observed.
41    ssim_max: f64,
42}
43
44impl BitstreamStats {
45    /// Create a new, empty statistics accumulator.
46    #[must_use]
47    pub fn new() -> Self {
48        Self {
49            psnr_min: f64::INFINITY,
50            psnr_max: f64::NEG_INFINITY,
51            ssim_min: f64::INFINITY,
52            ssim_max: f64::NEG_INFINITY,
53            ..Self::default()
54        }
55    }
56
57    // -----------------------------------------------------------------------
58    // PSNR
59    // -----------------------------------------------------------------------
60
61    /// Compute and record the PSNR between the original and reconstructed
62    /// pixel buffers for one frame.
63    ///
64    /// PSNR is computed on the luma (Y) plane:
65    /// ```text
66    /// MSE = sum((orig[i] - recon[i])^2) / (w * h)
67    /// PSNR = 10 * log10(255^2 / MSE)
68    /// ```
69    ///
70    /// An MSE of 0.0 (identical buffers) yields +∞ dB; in that case `100.0`
71    /// dB is recorded as a practical maximum.
72    ///
73    /// # Parameters
74    /// - `orig`  – original uncompressed pixel buffer (u8 luma, length `w * h`).
75    /// - `recon` – reconstructed pixel buffer (same layout as `orig`).
76    /// - `w`     – frame width in pixels.
77    /// - `h`     – frame height in pixels.
78    ///
79    /// # Panics
80    /// Panics if `orig.len() != w * h` or `recon.len() != w * h`.
81    pub fn update_psnr(&mut self, orig: &[u8], recon: &[u8], w: u32, h: u32) {
82        let n = (w * h) as usize;
83        assert_eq!(orig.len(), n, "update_psnr: orig length mismatch");
84        assert_eq!(recon.len(), n, "update_psnr: recon length mismatch");
85
86        let mse: f64 = orig
87            .iter()
88            .zip(recon.iter())
89            .map(|(&a, &b)| {
90                let d = a as f64 - b as f64;
91                d * d
92            })
93            .sum::<f64>()
94            / n as f64;
95
96        let psnr = if mse < f64::EPSILON {
97            100.0_f64
98        } else {
99            10.0 * (255.0_f64 * 255.0 / mse).log10()
100        };
101
102        self.psnr_sum += psnr;
103        self.psnr_count += 1;
104        if psnr < self.psnr_min {
105            self.psnr_min = psnr;
106        }
107        if psnr > self.psnr_max {
108            self.psnr_max = psnr;
109        }
110    }
111
112    /// Returns the mean PSNR across all recorded frames, or `None` if no
113    /// frames have been processed.
114    #[must_use]
115    pub fn mean_psnr(&self) -> Option<f64> {
116        if self.psnr_count == 0 {
117            None
118        } else {
119            Some(self.psnr_sum / self.psnr_count as f64)
120        }
121    }
122
123    /// Returns the minimum per-frame PSNR, or `None` if no frames recorded.
124    #[must_use]
125    pub fn min_psnr(&self) -> Option<f64> {
126        if self.psnr_count == 0 {
127            None
128        } else {
129            Some(self.psnr_min)
130        }
131    }
132
133    /// Returns the maximum per-frame PSNR, or `None` if no frames recorded.
134    #[must_use]
135    pub fn max_psnr(&self) -> Option<f64> {
136        if self.psnr_count == 0 {
137            None
138        } else {
139            Some(self.psnr_max)
140        }
141    }
142
143    // -----------------------------------------------------------------------
144    // SSIM
145    // -----------------------------------------------------------------------
146
147    /// Compute and record the mean SSIM between original and reconstructed
148    /// buffers for one frame.
149    ///
150    /// Uses the simplified single-window SSIM formula (per-image, not
151    /// per-patch) as a fast approximation:
152    /// ```text
153    /// μ_x, μ_y  — mean pixel values
154    /// σ_x, σ_y  — standard deviations
155    /// σ_xy      — cross-covariance
156    /// SSIM = (2μ_xμ_y + C1)(2σ_xy + C2) / ((μ_x²+μ_y²+C1)(σ_x²+σ_y²+C2))
157    /// ```
158    /// with `C1 = (0.01 * 255)^2` and `C2 = (0.03 * 255)^2`.
159    ///
160    /// # Parameters
161    /// - `orig`  – original pixel buffer (u8 luma, length `w * h`).
162    /// - `recon` – reconstructed pixel buffer (same layout).
163    /// - `w`     – frame width in pixels.
164    /// - `h`     – frame height in pixels.
165    ///
166    /// # Panics
167    /// Panics if `orig.len() != w * h` or `recon.len() != w * h`.
168    pub fn update_ssim(&mut self, orig: &[u8], recon: &[u8], w: u32, h: u32) {
169        let n = (w * h) as usize;
170        assert_eq!(orig.len(), n, "update_ssim: orig length mismatch");
171        assert_eq!(recon.len(), n, "update_ssim: recon length mismatch");
172
173        let n_f = n as f64;
174
175        let mu_x: f64 = orig.iter().map(|&v| v as f64).sum::<f64>() / n_f;
176        let mu_y: f64 = recon.iter().map(|&v| v as f64).sum::<f64>() / n_f;
177
178        let var_x: f64 = orig
179            .iter()
180            .map(|&v| {
181                let d = v as f64 - mu_x;
182                d * d
183            })
184            .sum::<f64>()
185            / n_f;
186        let var_y: f64 = recon
187            .iter()
188            .map(|&v| {
189                let d = v as f64 - mu_y;
190                d * d
191            })
192            .sum::<f64>()
193            / n_f;
194
195        let cov_xy: f64 = orig
196            .iter()
197            .zip(recon.iter())
198            .map(|(&a, &b)| (a as f64 - mu_x) * (b as f64 - mu_y))
199            .sum::<f64>()
200            / n_f;
201
202        const C1: f64 = (0.01 * 255.0) * (0.01 * 255.0);
203        const C2: f64 = (0.03 * 255.0) * (0.03 * 255.0);
204
205        let numerator = (2.0 * mu_x * mu_y + C1) * (2.0 * cov_xy + C2);
206        let denominator = (mu_x * mu_x + mu_y * mu_y + C1) * (var_x + var_y + C2);
207
208        let ssim = if denominator.abs() < f64::EPSILON {
209            1.0
210        } else {
211            (numerator / denominator).clamp(-1.0, 1.0)
212        };
213
214        self.ssim_sum += ssim;
215        self.ssim_count += 1;
216        if ssim < self.ssim_min {
217            self.ssim_min = ssim;
218        }
219        if ssim > self.ssim_max {
220            self.ssim_max = ssim;
221        }
222    }
223
224    /// Returns the mean SSIM across all recorded frames, or `None` if no
225    /// frames have been processed.
226    #[must_use]
227    pub fn mean_ssim(&self) -> Option<f64> {
228        if self.ssim_count == 0 {
229            None
230        } else {
231            Some(self.ssim_sum / self.ssim_count as f64)
232        }
233    }
234
235    /// Returns the minimum per-frame SSIM, or `None` if no frames recorded.
236    #[must_use]
237    pub fn min_ssim(&self) -> Option<f64> {
238        if self.ssim_count == 0 {
239            None
240        } else {
241            Some(self.ssim_min)
242        }
243    }
244
245    /// Returns the maximum per-frame SSIM, or `None` if no frames recorded.
246    #[must_use]
247    pub fn max_ssim(&self) -> Option<f64> {
248        if self.ssim_count == 0 {
249            None
250        } else {
251            Some(self.ssim_max)
252        }
253    }
254
255    // -----------------------------------------------------------------------
256    // Frame counts
257    // -----------------------------------------------------------------------
258
259    /// Returns the number of frames that contributed to PSNR statistics.
260    #[must_use]
261    pub fn psnr_frame_count(&self) -> u64 {
262        self.psnr_count
263    }
264
265    /// Returns the number of frames that contributed to SSIM statistics.
266    #[must_use]
267    pub fn ssim_frame_count(&self) -> u64 {
268        self.ssim_count
269    }
270
271    /// Reset all statistics.
272    pub fn reset(&mut self) {
273        *self = Self::new();
274    }
275}
276
277// ---------------------------------------------------------------------------
278// Tests
279// ---------------------------------------------------------------------------
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284
285    fn identical_frame(size: usize, val: u8) -> (Vec<u8>, Vec<u8>) {
286        (vec![val; size], vec![val; size])
287    }
288
289    fn noisy_frame(size: usize) -> (Vec<u8>, Vec<u8>) {
290        let orig: Vec<u8> = (0..size).map(|i| (i % 256) as u8).collect();
291        let recon: Vec<u8> = orig.iter().map(|&v| v.saturating_add(10)).collect();
292        (orig, recon)
293    }
294
295    #[test]
296    fn new_stats_no_data() {
297        let s = BitstreamStats::new();
298        assert!(s.mean_psnr().is_none());
299        assert!(s.mean_ssim().is_none());
300    }
301
302    #[test]
303    fn psnr_identical_frames_near_100() {
304        let mut s = BitstreamStats::new();
305        let (orig, recon) = identical_frame(64 * 64, 128);
306        s.update_psnr(&orig, &recon, 64, 64);
307        let psnr = s.mean_psnr().expect("should have PSNR");
308        assert!(
309            (psnr - 100.0).abs() < 1e-6,
310            "identical PSNR should be 100 dB"
311        );
312    }
313
314    #[test]
315    fn psnr_positive_for_different_frames() {
316        let mut s = BitstreamStats::new();
317        let (orig, recon) = noisy_frame(16 * 16);
318        s.update_psnr(&orig, &recon, 16, 16);
319        let psnr = s.mean_psnr().expect("should have PSNR");
320        assert!(
321            psnr > 0.0 && psnr < 100.0,
322            "PSNR {psnr} should be in (0, 100)"
323        );
324    }
325
326    #[test]
327    fn psnr_accumulated_over_multiple_frames() {
328        let mut s = BitstreamStats::new();
329        let (o1, r1) = identical_frame(16 * 16, 100);
330        let (o2, r2) = noisy_frame(16 * 16);
331        s.update_psnr(&o1, &r1, 16, 16);
332        s.update_psnr(&o2, &r2, 16, 16);
333        assert_eq!(s.psnr_frame_count(), 2);
334    }
335
336    #[test]
337    fn ssim_identical_frames_near_one() {
338        let mut s = BitstreamStats::new();
339        let (orig, recon) = identical_frame(32 * 32, 128);
340        s.update_ssim(&orig, &recon, 32, 32);
341        let ssim = s.mean_ssim().expect("should have SSIM");
342        assert!(ssim > 0.99, "identical SSIM should be ~1.0, got {ssim}");
343    }
344
345    #[test]
346    fn ssim_drops_for_noisy_frames() {
347        let mut s = BitstreamStats::new();
348        let (orig, recon) = noisy_frame(32 * 32);
349        s.update_ssim(&orig, &recon, 32, 32);
350        let ssim = s.mean_ssim().expect("should have SSIM");
351        // SSIM should drop below 1.0 for noisy frames
352        assert!(ssim < 1.0, "noisy SSIM should be < 1.0");
353    }
354
355    #[test]
356    fn stats_min_max_psnr() {
357        let mut s = BitstreamStats::new();
358        let (o1, r1) = identical_frame(16 * 16, 100); // PSNR = 100
359        let (o2, r2) = noisy_frame(16 * 16); // PSNR < 100
360        s.update_psnr(&o1, &r1, 16, 16);
361        s.update_psnr(&o2, &r2, 16, 16);
362        assert!(s.min_psnr().expect("min psnr") <= s.max_psnr().expect("max psnr"));
363    }
364
365    #[test]
366    fn reset_clears_all() {
367        let mut s = BitstreamStats::new();
368        let (orig, recon) = identical_frame(4 * 4, 200);
369        s.update_psnr(&orig, &recon, 4, 4);
370        s.reset();
371        assert!(s.mean_psnr().is_none());
372        assert_eq!(s.psnr_frame_count(), 0);
373    }
374
375    #[test]
376    #[should_panic(expected = "orig length mismatch")]
377    fn update_psnr_panics_on_wrong_length() {
378        let mut s = BitstreamStats::new();
379        s.update_psnr(&[0u8; 10], &[0u8; 16], 4, 4);
380    }
381}
oximedia_codec/stats.rs

oximedia_codec/
stats.rs
