Skip to main content

fast_ssim2/
lib.rs

1//! # fast-ssim2
2//!
3//! Fast SIMD-accelerated implementation of [SSIMULACRA2](https://github.com/cloudinary/ssimulacra2),
4//! a perceptual image quality metric.
5//!
6//! ## Quick Start
7//!
8//! The simplest way to compare two images:
9//!
10//! ```ignore
11//! use fast_ssim2::compute_ssimulacra2;
12//! use imgref::ImgVec;
13//!
14//! // Load your images (8-bit sRGB)
15//! let source: ImgVec<[u8; 3]> = load_image("source.png");
16//! let distorted: ImgVec<[u8; 3]> = load_image("distorted.png");
17//!
18//! let score = compute_ssimulacra2(source.as_ref(), distorted.as_ref())?;
19//! // score: 100 = identical, 90+ = imperceptible, <50 = significant degradation
20//! ```
21//!
22//! ## Score Interpretation
23//!
24//! | Score | Quality |
25//! |-------|---------|
26//! | **100** | Identical (no difference) |
27//! | **90+** | Imperceptible difference |
28//! | **70-90** | Minor, subtle difference |
29//! | **50-70** | Noticeable difference |
30//! | **<50** | Significant degradation |
31//!
32//! ## Supported Input Formats
33//!
34//! ### With `imgref` feature (recommended for most users)
35//!
36//! | Type | Color Space | Notes |
37//! |------|-------------|-------|
38//! | `ImgRef<[u8; 3]>` | sRGB | Standard 8-bit RGB images |
39//! | `ImgRef<[u16; 3]>` | sRGB | 16-bit RGB (HDR workflows) |
40//! | `ImgRef<[f32; 3]>` | **Linear RGB** | Already linearized data |
41//! | `ImgRef<u8>` | sRGB grayscale | Expanded to R=G=B |
42//! | `ImgRef<f32>` | Linear grayscale | Expanded to R=G=B |
43//!
44//! **Convention:** Integer types assume sRGB gamma encoding. Float types assume linear RGB.
45//!
46//! ### Without features (using `yuvxyb` types)
47//!
48//! ```
49//! use fast_ssim2::{compute_ssimulacra2, Rgb, TransferCharacteristic, ColorPrimaries};
50//! use std::num::NonZeroUsize;
51//!
52//! let data: Vec<[f32; 3]> = vec![[0.5, 0.5, 0.5]; 64 * 64];
53//! let w = NonZeroUsize::new(64).unwrap();
54//! let h = NonZeroUsize::new(64).unwrap();
55//! let source = Rgb::new(data.clone(), w, h,
56//!     TransferCharacteristic::SRGB, ColorPrimaries::BT709)?;
57//! let distorted = Rgb::new(data, w, h,
58//!     TransferCharacteristic::SRGB, ColorPrimaries::BT709)?;
59//!
60//! let score = compute_ssimulacra2(source, distorted)?;
61//! // compute_ssimulacra2 accepts yuvxyb::Rgb, yuvxyb::LinearRgb, and more
62//! # Ok::<(), Box<dyn std::error::Error>>(())
63//! ```
64//!
65//! ## Batch Comparisons (2x Faster)
66//!
67//! When comparing multiple images against the same reference (e.g., evaluating
68//! different compression levels), precompute the reference data once:
69//!
70//! ```
71//! use fast_ssim2::{Ssimulacra2Reference, Rgb, TransferCharacteristic, ColorPrimaries};
72//! use std::num::NonZeroUsize;
73//!
74//! // Create test data
75//! let data: Vec<[f32; 3]> = vec![[0.5, 0.5, 0.5]; 64 * 64];
76//! let w = NonZeroUsize::new(64).unwrap();
77//! let h = NonZeroUsize::new(64).unwrap();
78//! let source = Rgb::new(data.clone(), w, h,
79//!     TransferCharacteristic::SRGB, ColorPrimaries::BT709)?;
80//!
81//! // Precompute reference data (~50% of the work)
82//! let reference = Ssimulacra2Reference::new(source)?;
83//!
84//! // Compare multiple distorted versions efficiently
85//! let distorted = Rgb::new(data, w, h,
86//!     TransferCharacteristic::SRGB, ColorPrimaries::BT709)?;
87//! let score = reference.compare(distorted)?;
88//! # Ok::<(), Box<dyn std::error::Error>>(())
89//! ```
90//!
91//! ## Custom Input Types
92//!
93//! Implement [`ToLinearRgb`] to support your own image types:
94//!
95//! ```
96//! use fast_ssim2::{ToLinearRgb, LinearRgbImage, srgb_u8_to_linear};
97//!
98//! struct MyImage {
99//!     pixels: Vec<[u8; 3]>,
100//!     width: usize,
101//!     height: usize,
102//! }
103//!
104//! impl ToLinearRgb for MyImage {
105//!     fn to_linear_rgb(&self) -> LinearRgbImage {
106//!         let data: Vec<[f32; 3]> = self.pixels.iter()
107//!             .map(|[r, g, b]| [
108//!                 srgb_u8_to_linear(*r),
109//!                 srgb_u8_to_linear(*g),
110//!                 srgb_u8_to_linear(*b),
111//!             ])
112//!             .collect();
113//!         LinearRgbImage::new(data, self.width, self.height)
114//!     }
115//! }
116//! ```
117//!
118//! Helper functions for sRGB conversion:
119//! - [`srgb_u8_to_linear`] - 8-bit lookup table (fastest)
120//! - [`srgb_u16_to_linear`] - 16-bit conversion
121//! - [`srgb_to_linear`] - General f32 conversion
122//!
123//! ## SIMD Configuration
124//!
125//! SIMD is enabled by default via the `archmage` crate, providing cross-platform
126//! acceleration on x86_64 (AVX2, AVX-512), AArch64 (NEON), and WASM (SIMD128).
127//!
128//! | Backend | Speed | Platforms |
129//! |---------|-------|-----------|
130//! | `Scalar` | 1.0× (baseline) | All |
131//! | `Simd` (default) | 2-3× | x86_64, AArch64, WASM |
132//!
133//! To explicitly select a backend:
134//!
135//! ```
136//! use fast_ssim2::{compute_ssimulacra2_with_config, Ssimulacra2Config};
137//!
138//! # let source = fast_ssim2::LinearRgbImage::new(vec![[0.0; 3]; 64], 8, 8);
139//! # let distorted = fast_ssim2::LinearRgbImage::new(vec![[0.0; 3]; 64], 8, 8);
140//! let score = compute_ssimulacra2_with_config(
141//!     source,
142//!     distorted,
143//!     Ssimulacra2Config::scalar(), // or ::simd()
144//! )?;
145//! # Ok::<(), fast_ssim2::Ssimulacra2Error>(())
146//! ```
147//!
148//! ## Features
149//!
150//! | Feature | Default | Description |
151//! |---------|---------|-------------|
152//! | `imgref` | | Support for `imgref` image types |
153//! | `rayon` | | Parallel computation |
154//!
155//! ## Requirements
156//!
157//! - **Minimum image size:** 8×8 pixels
158//! - **MSRV:** 1.89.0
159
160#![forbid(unsafe_code)]
161
162mod blur;
163mod input;
164mod precompute;
165// Reference data for parity testing (hidden from docs but accessible for tests)
166#[doc(hidden)]
167pub mod reference_data;
168#[allow(clippy::too_many_arguments)] // arcane macro generates dispatchers inheriting param count
169mod simd_ops;
170mod xyb_simd;
171
172pub use blur::Blur;
173pub use input::{LinearRgbImage, ToLinearRgb};
174pub use precompute::Ssimulacra2Reference;
175// Re-export commonly used types from yuvxyb for convenience
176pub use yuvxyb::{
177    ColorPrimaries, Frame, LinearRgb, MatrixCoefficients, Pixel, Plane, Rgb,
178    TransferCharacteristic, Yuv, YuvConfig,
179};
180
181// Re-export sRGB conversion functions for users implementing custom input types
182pub use input::{srgb_to_linear, srgb_u8_to_linear, srgb_u16_to_linear};
183
184// Internal imports for XYB color space
185use yuvxyb::Xyb;
186
187// How often to downscale and score the input images.
188// Each scaling step will downscale by a factor of two.
189pub(crate) const NUM_SCALES: usize = 6;
190
191/// SIMD implementation backend for all operations (blur, XYB conversion, SSIM computation).
192#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
193pub enum SimdImpl {
194    /// Scalar implementation (baseline, most portable)
195    Scalar,
196    /// Cross-platform SIMD via archmage (default, AVX2/AVX-512/NEON/WASM128)
197    #[default]
198    Simd,
199}
200
201impl SimdImpl {
202    /// Returns the name of this implementation
203    pub fn name(&self) -> &'static str {
204        match self {
205            SimdImpl::Scalar => "scalar",
206            SimdImpl::Simd => "simd (archmage)",
207        }
208    }
209}
210
211/// Configuration for SSIMULACRA2 computation.
212#[derive(Debug, Clone, Copy, Default)]
213pub struct Ssimulacra2Config {
214    /// Implementation backend for all operations
215    pub impl_type: SimdImpl,
216}
217
218impl Ssimulacra2Config {
219    /// Create configuration with specified implementation
220    pub fn new(impl_type: SimdImpl) -> Self {
221        Self { impl_type }
222    }
223
224    /// Default configuration using SIMD for all operations
225    pub fn simd() -> Self {
226        Self::new(SimdImpl::Simd)
227    }
228
229    /// Scalar configuration (baseline, most compatible)
230    pub fn scalar() -> Self {
231        Self::new(SimdImpl::Scalar)
232    }
233}
234
235/// Errors which can occur when attempting to calculate a SSIMULACRA2 score from two input images.
236#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
237pub enum Ssimulacra2Error {
238    /// The conversion from input image to [LinearRgb] (via [TryFrom]) returned an [Err].
239    #[error("Failed to convert input image to linear RGB")]
240    LinearRgbConversionFailed,
241
242    /// The two input images do not have the same width and height.
243    #[error("Source and distorted image width and height must be equal")]
244    NonMatchingImageDimensions,
245
246    /// One of the input images has a width and/or height of less than 8 pixels.
247    #[error("Images must be at least 8x8 pixels")]
248    InvalidImageSize,
249
250    /// Gaussian blur operation failed.
251    #[error("Gaussian blur operation failed")]
252    GaussianBlurError,
253}
254
255/// Computes the SSIMULACRA2 score with default configuration (safe SIMD).
256pub fn compute_frame_ssimulacra2<T, U>(source: T, distorted: U) -> Result<f64, Ssimulacra2Error>
257where
258    LinearRgb: TryFrom<T> + TryFrom<U>,
259{
260    compute_frame_ssimulacra2_impl(source, distorted, Ssimulacra2Config::default())
261}
262
263/// Computes the SSIMULACRA2 score with custom implementation configuration.
264pub fn compute_frame_ssimulacra2_with_config<T, U>(
265    source: T,
266    distorted: U,
267    config: Ssimulacra2Config,
268) -> Result<f64, Ssimulacra2Error>
269where
270    LinearRgb: TryFrom<T> + TryFrom<U>,
271{
272    compute_frame_ssimulacra2_impl(source, distorted, config)
273}
274
275/// Computes the SSIMULACRA2 score from any input type implementing [`ToLinearRgb`].
276///
277/// This is the recommended API for new code. It supports:
278/// - `imgref` types (with the `imgref` feature): `ImgRef<[u8; 3]>`, `ImgRef<[f32; 3]>`, etc.
279/// - `yuvxyb` types: `Rgb`, `LinearRgb`
280/// - Custom types implementing [`ToLinearRgb`]
281///
282/// # Color space conventions
283/// - Integer types (`u8`, `u16`) are assumed to be sRGB (gamma-encoded)
284/// - Float types (`f32`) are assumed to be linear RGB
285/// - Grayscale types are expanded to RGB (R=G=B)
286///
287/// # Example
288/// ```ignore
289/// use imgref::ImgVec;
290/// use fast_ssim2::compute_ssimulacra2;
291///
292/// let source: ImgVec<[u8; 3]> = /* ... */;
293/// let distorted: ImgVec<[u8; 3]> = /* ... */;
294/// let score = compute_ssimulacra2(&source, &distorted)?;
295/// ```
296pub fn compute_ssimulacra2<S, D>(source: S, distorted: D) -> Result<f64, Ssimulacra2Error>
297where
298    S: ToLinearRgb,
299    D: ToLinearRgb,
300{
301    compute_ssimulacra2_with_config(source, distorted, Ssimulacra2Config::default())
302}
303
304/// Computes the SSIMULACRA2 score with custom configuration from [`ToLinearRgb`] inputs.
305pub fn compute_ssimulacra2_with_config<S, D>(
306    source: S,
307    distorted: D,
308    config: Ssimulacra2Config,
309) -> Result<f64, Ssimulacra2Error>
310where
311    S: ToLinearRgb,
312    D: ToLinearRgb,
313{
314    let img1: LinearRgb = source.into_linear_rgb().into();
315    let img2: LinearRgb = distorted.into_linear_rgb().into();
316    compute_frame_ssimulacra2_impl(img1, img2, config)
317}
318
319fn compute_frame_ssimulacra2_impl<T, U>(
320    source: T,
321    distorted: U,
322    config: Ssimulacra2Config,
323) -> Result<f64, Ssimulacra2Error>
324where
325    LinearRgb: TryFrom<T> + TryFrom<U>,
326{
327    let Ok(mut img1) = LinearRgb::try_from(source) else {
328        return Err(Ssimulacra2Error::LinearRgbConversionFailed);
329    };
330
331    let Ok(mut img2) = LinearRgb::try_from(distorted) else {
332        return Err(Ssimulacra2Error::LinearRgbConversionFailed);
333    };
334
335    if img1.width() != img2.width() || img1.height() != img2.height() {
336        return Err(Ssimulacra2Error::NonMatchingImageDimensions);
337    }
338
339    if img1.width().get() < 8 || img1.height().get() < 8 {
340        return Err(Ssimulacra2Error::InvalidImageSize);
341    }
342
343    let mut width = img1.width().get();
344    let mut height = img1.height().get();
345    let impl_type = config.impl_type;
346
347    // Pre-allocate reusable buffers (sized for initial dimensions, shrunk per scale)
348    let alloc_plane = || vec![0.0f32; width * height];
349    let alloc_3planes = || [alloc_plane(), alloc_plane(), alloc_plane()];
350
351    let mut mul = alloc_3planes();
352    let mut sigma1_sq = alloc_3planes();
353    let mut sigma2_sq = alloc_3planes();
354    let mut sigma12 = alloc_3planes();
355    let mut mu1 = alloc_3planes();
356    let mut mu2 = alloc_3planes();
357    let mut img1_planar = alloc_3planes();
358    let mut img2_planar = alloc_3planes();
359
360    let mut blur = Blur::with_simd_impl(width, height, impl_type);
361    let mut msssim = Msssim::default();
362
363    for scale in 0..NUM_SCALES {
364        if width < 8 || height < 8 {
365            break;
366        }
367
368        if scale > 0 {
369            img1 = downscale_by_2(&img1);
370            img2 = downscale_by_2(&img2);
371            width = img1.width().get();
372            height = img2.height().get();
373        }
374
375        // Shrink all buffers to current scale size
376        let size = width * height;
377        for buf in [
378            &mut mul,
379            &mut sigma1_sq,
380            &mut sigma2_sq,
381            &mut sigma12,
382            &mut mu1,
383            &mut mu2,
384            &mut img1_planar,
385            &mut img2_planar,
386        ] {
387            for c in buf.iter_mut() {
388                c.truncate(size);
389            }
390        }
391        blur.shrink_to(width, height);
392
393        let mut img1_xyb = linear_rgb_to_xyb(img1.clone(), impl_type);
394        let mut img2_xyb = linear_rgb_to_xyb(img2.clone(), impl_type);
395
396        make_positive_xyb(&mut img1_xyb);
397        make_positive_xyb(&mut img2_xyb);
398
399        xyb_to_planar_into(&img1_xyb, &mut img1_planar);
400        xyb_to_planar_into(&img2_xyb, &mut img2_planar);
401
402        image_multiply(&img1_planar, &img1_planar, &mut mul, impl_type);
403        blur.blur_into(&mul, &mut sigma1_sq);
404
405        image_multiply(&img2_planar, &img2_planar, &mut mul, impl_type);
406        blur.blur_into(&mul, &mut sigma2_sq);
407
408        image_multiply(&img1_planar, &img2_planar, &mut mul, impl_type);
409        blur.blur_into(&mul, &mut sigma12);
410
411        blur.blur_into(&img1_planar, &mut mu1);
412        blur.blur_into(&img2_planar, &mut mu2);
413
414        let avg_ssim = ssim_map(
415            width, height, &mu1, &mu2, &sigma1_sq, &sigma2_sq, &sigma12, impl_type,
416        );
417        let avg_edgediff = edge_diff_map(
418            width,
419            height,
420            &img1_planar,
421            &mu1,
422            &img2_planar,
423            &mu2,
424            impl_type,
425        );
426        msssim.scales.push(MsssimScale {
427            avg_ssim,
428            avg_edgediff,
429        });
430    }
431
432    Ok(msssim.score())
433}
434
435/// Convert LinearRgb to Xyb using the specified implementation
436fn linear_rgb_to_xyb(linear_rgb: LinearRgb, impl_type: SimdImpl) -> Xyb {
437    match impl_type {
438        SimdImpl::Scalar => Xyb::from(linear_rgb),
439        SimdImpl::Simd => {
440            let width = linear_rgb.width(); // NonZeroUsize
441            let height = linear_rgb.height(); // NonZeroUsize
442            let mut data = linear_rgb.into_data();
443            xyb_simd::linear_rgb_to_xyb_simd(&mut data);
444            Xyb::new(data, width, height).expect("XYB construction should not fail")
445        }
446    }
447}
448
449// For backwards compatibility
450pub(crate) fn linear_rgb_to_xyb_simd(linear_rgb: LinearRgb) -> Xyb {
451    linear_rgb_to_xyb(linear_rgb, SimdImpl::Simd)
452}
453
454pub(crate) fn make_positive_xyb(xyb: &mut Xyb) {
455    for pix in xyb.data_mut().iter_mut() {
456        pix[2] = (pix[2] - pix[1]) + 0.55;
457        pix[0] = (pix[0]).mul_add(14.0, 0.42);
458        pix[1] += 0.01;
459    }
460}
461
462pub(crate) fn xyb_to_planar(xyb: &Xyb) -> [Vec<f32>; 3] {
463    let size = xyb.width().get() * xyb.height().get();
464    let mut out = [vec![0.0f32; size], vec![0.0f32; size], vec![0.0f32; size]];
465    xyb_to_planar_into(xyb, &mut out);
466    out
467}
468
469/// Convert XYB to planar format into pre-allocated buffers (zero-allocation)
470pub(crate) fn xyb_to_planar_into(xyb: &Xyb, out: &mut [Vec<f32>; 3]) {
471    let [out0, out1, out2] = out;
472    for (((i, o0), o1), o2) in xyb
473        .data()
474        .iter()
475        .copied()
476        .zip(out0.iter_mut())
477        .zip(out1.iter_mut())
478        .zip(out2.iter_mut())
479    {
480        *o0 = i[0];
481        *o1 = i[1];
482        *o2 = i[2];
483    }
484}
485
486pub(crate) fn image_multiply(
487    img1: &[Vec<f32>; 3],
488    img2: &[Vec<f32>; 3],
489    out: &mut [Vec<f32>; 3],
490    impl_type: SimdImpl,
491) {
492    match impl_type {
493        SimdImpl::Scalar => image_multiply_scalar(img1, img2, out),
494        SimdImpl::Simd => simd_ops::image_multiply_simd(img1, img2, out),
495    }
496}
497
498fn image_multiply_scalar(img1: &[Vec<f32>; 3], img2: &[Vec<f32>; 3], out: &mut [Vec<f32>; 3]) {
499    for ((plane1, plane2), out_plane) in img1.iter().zip(img2.iter()).zip(out.iter_mut()) {
500        for ((&p1, &p2), o) in plane1.iter().zip(plane2.iter()).zip(out_plane.iter_mut()) {
501            *o = p1 * p2;
502        }
503    }
504}
505
506pub(crate) fn downscale_by_2(in_data: &LinearRgb) -> LinearRgb {
507    use std::num::NonZeroUsize;
508    const SCALE: usize = 2;
509    let in_w = in_data.width().get();
510    let in_h = in_data.height().get();
511    let out_w = in_w.div_ceil(SCALE);
512    let out_h = in_h.div_ceil(SCALE);
513    let mut out_data = vec![[0.0f32; 3]; out_w * out_h];
514    let normalize = 1.0f32 / (SCALE * SCALE) as f32;
515
516    let in_data = &in_data.data();
517    for oy in 0..out_h {
518        for ox in 0..out_w {
519            for c in 0..3 {
520                let mut sum = 0f32;
521                for iy in 0..SCALE {
522                    for ix in 0..SCALE {
523                        let x = (ox * SCALE + ix).min(in_w - 1);
524                        let y = (oy * SCALE + iy).min(in_h - 1);
525                        sum += in_data[y * in_w + x][c];
526                    }
527                }
528                out_data[oy * out_w + ox][c] = sum * normalize;
529            }
530        }
531    }
532
533    LinearRgb::new(
534        out_data,
535        NonZeroUsize::new(out_w).expect("out_w must be nonzero"),
536        NonZeroUsize::new(out_h).expect("out_h must be nonzero"),
537    )
538    .expect("Resolution and data size match")
539}
540
541#[allow(clippy::too_many_arguments)]
542pub(crate) fn ssim_map(
543    width: usize,
544    height: usize,
545    m1: &[Vec<f32>; 3],
546    m2: &[Vec<f32>; 3],
547    s11: &[Vec<f32>; 3],
548    s22: &[Vec<f32>; 3],
549    s12: &[Vec<f32>; 3],
550    impl_type: SimdImpl,
551) -> [f64; 3 * 2] {
552    match impl_type {
553        SimdImpl::Scalar => ssim_map_scalar(width, height, m1, m2, s11, s22, s12),
554        SimdImpl::Simd => simd_ops::ssim_map_simd(width, height, m1, m2, s11, s22, s12),
555    }
556}
557
558fn ssim_map_scalar(
559    width: usize,
560    height: usize,
561    m1: &[Vec<f32>; 3],
562    m2: &[Vec<f32>; 3],
563    s11: &[Vec<f32>; 3],
564    s22: &[Vec<f32>; 3],
565    s12: &[Vec<f32>; 3],
566) -> [f64; 3 * 2] {
567    const C2: f32 = 0.0009f32;
568
569    let one_per_pixels = 1.0f64 / (width * height) as f64;
570    let mut plane_averages = [0f64; 3 * 2];
571
572    for c in 0..3 {
573        let mut sum_d = 0.0f64;
574        let mut sum_d4 = 0.0f64;
575        for (row_m1, (row_m2, (row_s11, (row_s22, row_s12)))) in m1[c].chunks_exact(width).zip(
576            m2[c].chunks_exact(width).zip(
577                s11[c]
578                    .chunks_exact(width)
579                    .zip(s22[c].chunks_exact(width).zip(s12[c].chunks_exact(width))),
580            ),
581        ) {
582            for x in 0..width {
583                let mu1 = row_m1[x];
584                let mu2 = row_m2[x];
585                let mu11 = mu1 * mu1;
586                let mu22 = mu2 * mu2;
587                let mu12 = mu1 * mu2;
588                let mu_diff = mu1 - mu2;
589
590                let num_m = mu_diff.mul_add(-mu_diff, 1.0f32);
591                let num_s = 2.0f32.mul_add(row_s12[x] - mu12, C2);
592                let denom_s = (row_s11[x] - mu11) + (row_s22[x] - mu22) + C2;
593                let d = (1.0f32 - (num_m * num_s) / denom_s).max(0.0f32);
594                let d2 = d * d;
595                let d4 = d2 * d2;
596                sum_d += f64::from(d);
597                sum_d4 += f64::from(d4);
598            }
599        }
600        plane_averages[c * 2] = one_per_pixels * sum_d;
601        plane_averages[c * 2 + 1] = (one_per_pixels * sum_d4).sqrt().sqrt();
602    }
603
604    plane_averages
605}
606
607pub(crate) fn edge_diff_map(
608    width: usize,
609    height: usize,
610    img1: &[Vec<f32>; 3],
611    mu1: &[Vec<f32>; 3],
612    img2: &[Vec<f32>; 3],
613    mu2: &[Vec<f32>; 3],
614    impl_type: SimdImpl,
615) -> [f64; 3 * 4] {
616    match impl_type {
617        SimdImpl::Scalar => edge_diff_map_scalar(width, height, img1, mu1, img2, mu2),
618        SimdImpl::Simd => simd_ops::edge_diff_map_simd(width, height, img1, mu1, img2, mu2),
619    }
620}
621
622fn edge_diff_map_scalar(
623    width: usize,
624    height: usize,
625    img1: &[Vec<f32>; 3],
626    mu1: &[Vec<f32>; 3],
627    img2: &[Vec<f32>; 3],
628    mu2: &[Vec<f32>; 3],
629) -> [f64; 3 * 4] {
630    let one_per_pixels = 1.0f64 / (width * height) as f64;
631    let mut plane_averages = [0f64; 3 * 4];
632
633    for c in 0..3 {
634        let mut sum1 = [0.0f64; 4];
635        for (row1, (row2, (rowm1, rowm2))) in img1[c].chunks_exact(width).zip(
636            img2[c]
637                .chunks_exact(width)
638                .zip(mu1[c].chunks_exact(width).zip(mu2[c].chunks_exact(width))),
639        ) {
640            for x in 0..width {
641                let d1: f64 = (1.0 + f64::from((row2[x] - rowm2[x]).abs()))
642                    / (1.0 + f64::from((row1[x] - rowm1[x]).abs()))
643                    - 1.0;
644
645                let artifact = d1.max(0.0);
646                sum1[0] += artifact;
647                sum1[1] += artifact.powi(4);
648
649                let detail_lost = (-d1).max(0.0);
650                sum1[2] += detail_lost;
651                sum1[3] += detail_lost.powi(4);
652            }
653        }
654        plane_averages[c * 4] = one_per_pixels * sum1[0];
655        plane_averages[c * 4 + 1] = (one_per_pixels * sum1[1]).sqrt().sqrt();
656        plane_averages[c * 4 + 2] = one_per_pixels * sum1[2];
657        plane_averages[c * 4 + 3] = (one_per_pixels * sum1[3]).sqrt().sqrt();
658    }
659
660    plane_averages
661}
662
663#[derive(Debug, Clone, Default)]
664pub(crate) struct Msssim {
665    pub scales: Vec<MsssimScale>,
666}
667
668#[derive(Debug, Clone, Copy, Default)]
669pub(crate) struct MsssimScale {
670    pub avg_ssim: [f64; 3 * 2],
671    pub avg_edgediff: [f64; 3 * 4],
672}
673
674impl Msssim {
675    #[allow(clippy::too_many_lines)]
676    pub fn score(&self) -> f64 {
677        const WEIGHT: [f64; 108] = [
678            0.0,
679            0.000_737_660_670_740_658_6,
680            0.0,
681            0.0,
682            0.000_779_348_168_286_730_9,
683            0.0,
684            0.0,
685            0.000_437_115_573_010_737_9,
686            0.0,
687            1.104_172_642_665_734_6,
688            0.000_662_848_341_292_71,
689            0.000_152_316_327_837_187_52,
690            0.0,
691            0.001_640_643_745_659_975_4,
692            0.0,
693            1.842_245_552_053_929_8,
694            11.441_172_603_757_666,
695            0.0,
696            0.000_798_910_943_601_516_3,
697            0.000_176_816_438_078_653,
698            0.0,
699            1.878_759_497_954_638_7,
700            10.949_069_906_051_42,
701            0.0,
702            0.000_728_934_699_150_807_2,
703            0.967_793_708_062_683_3,
704            0.0,
705            0.000_140_034_242_854_358_84,
706            0.998_176_697_785_496_7,
707            0.000_319_497_559_344_350_53,
708            0.000_455_099_211_379_206_3,
709            0.0,
710            0.0,
711            0.001_364_876_616_324_339_8,
712            0.0,
713            0.0,
714            0.0,
715            0.0,
716            0.0,
717            7.466_890_328_078_848,
718            0.0,
719            17.445_833_984_131_262,
720            0.000_623_560_163_404_146_6,
721            0.0,
722            0.0,
723            6.683_678_146_179_332,
724            0.000_377_244_079_796_112_96,
725            1.027_889_937_768_264,
726            225.205_153_008_492_74,
727            0.0,
728            0.0,
729            19.213_238_186_143_016,
730            0.001_140_152_458_661_836_1,
731            0.001_237_755_635_509_985,
732            176.393_175_984_506_94,
733            0.0,
734            0.0,
735            24.433_009_998_704_76,
736            0.285_208_026_121_177_57,
737            0.000_448_543_692_383_340_8,
738            0.0,
739            0.0,
740            0.0,
741            34.779_063_444_837_72,
742            44.835_625_328_877_896,
743            0.0,
744            0.0,
745            0.0,
746            0.0,
747            0.0,
748            0.0,
749            0.0,
750            0.0,
751            0.000_868_055_657_329_169_8,
752            0.0,
753            0.0,
754            0.0,
755            0.0,
756            0.0,
757            0.000_531_319_187_435_874_7,
758            0.0,
759            0.000_165_338_141_613_791_12,
760            0.0,
761            0.0,
762            0.0,
763            0.0,
764            0.0,
765            0.000_417_917_180_325_133_6,
766            0.001_729_082_823_472_283_3,
767            0.0,
768            0.002_082_700_584_663_643_7,
769            0.0,
770            0.0,
771            8.826_982_764_996_862,
772            23.192_433_439_989_26,
773            0.0,
774            95.108_049_881_108_6,
775            0.986_397_803_440_068_2,
776            0.983_438_279_246_535_3,
777            0.001_228_640_504_827_849_3,
778            171.266_725_589_730_7,
779            0.980_785_887_243_537_9,
780            0.0,
781            0.0,
782            0.0,
783            0.000_513_006_458_899_067_9,
784            0.0,
785            0.000_108_540_578_584_115_37,
786        ];
787
788        let mut ssim = 0.0f64;
789
790        let mut i = 0usize;
791        for c in 0..3 {
792            for scale in &self.scales {
793                for n in 0..2 {
794                    ssim = WEIGHT[i].mul_add(scale.avg_ssim[c * 2 + n].abs(), ssim);
795                    i += 1;
796                    ssim = WEIGHT[i].mul_add(scale.avg_edgediff[c * 4 + n].abs(), ssim);
797                    i += 1;
798                    ssim = WEIGHT[i].mul_add(scale.avg_edgediff[c * 4 + n + 2].abs(), ssim);
799                    i += 1;
800                }
801            }
802        }
803
804        ssim *= 0.956_238_261_683_484_4_f64;
805        ssim = (6.248_496_625_763_138e-5 * ssim * ssim).mul_add(
806            ssim,
807            2.326_765_642_916_932f64.mul_add(ssim, -0.020_884_521_182_843_837 * ssim * ssim),
808        );
809
810        if ssim > 0.0f64 {
811            ssim = ssim
812                .powf(0.627_633_646_783_138_7)
813                .mul_add(-10.0f64, 100.0f64);
814        } else {
815            ssim = 100.0f64;
816        }
817
818        ssim
819    }
820}
821
822#[cfg(test)]
823mod tests {
824    use std::path::PathBuf;
825
826    use super::*;
827    use yuvxyb::{ColorPrimaries, Rgb, TransferCharacteristic};
828
829    #[test]
830    fn test_ssimulacra2() {
831        let source = image::open(
832            PathBuf::from(env!("CARGO_MANIFEST_DIR"))
833                .join("test_data")
834                .join("tank_source.png"),
835        )
836        .unwrap();
837        let distorted = image::open(
838            PathBuf::from(env!("CARGO_MANIFEST_DIR"))
839                .join("test_data")
840                .join("tank_distorted.png"),
841        )
842        .unwrap();
843        let source_data = source
844            .to_rgb32f()
845            .chunks_exact(3)
846            .map(|chunk| [chunk[0], chunk[1], chunk[2]])
847            .collect::<Vec<_>>();
848        let source_data = Xyb::try_from(
849            Rgb::new(
850                source_data,
851                std::num::NonZeroUsize::new(source.width() as usize).unwrap(),
852                std::num::NonZeroUsize::new(source.height() as usize).unwrap(),
853                TransferCharacteristic::SRGB,
854                ColorPrimaries::BT709,
855            )
856            .unwrap(),
857        )
858        .unwrap();
859        let distorted_data = distorted
860            .to_rgb32f()
861            .chunks_exact(3)
862            .map(|chunk| [chunk[0], chunk[1], chunk[2]])
863            .collect::<Vec<_>>();
864        let distorted_data = Xyb::try_from(
865            Rgb::new(
866                distorted_data,
867                std::num::NonZeroUsize::new(distorted.width() as usize).unwrap(),
868                std::num::NonZeroUsize::new(distorted.height() as usize).unwrap(),
869                TransferCharacteristic::SRGB,
870                ColorPrimaries::BT709,
871            )
872            .unwrap(),
873        )
874        .unwrap();
875        let result = compute_frame_ssimulacra2(source_data, distorted_data).unwrap();
876        let expected = 17.398_505_f64;
877        assert!(
878            (result - expected).abs() < 0.25f64,
879            "Result {result:.6} not equal to expected {expected:.6}",
880        );
881    }
882
883    #[test]
884    fn test_xyb_simd_vs_yuvxyb() {
885        use yuvxyb::{ColorPrimaries, TransferCharacteristic};
886
887        let source = image::open(
888            PathBuf::from(env!("CARGO_MANIFEST_DIR"))
889                .join("test_data")
890                .join("tank_source.png"),
891        )
892        .unwrap();
893
894        let source_data: Vec<[f32; 3]> = source
895            .to_rgb32f()
896            .chunks_exact(3)
897            .map(|chunk| [chunk[0], chunk[1], chunk[2]])
898            .collect();
899
900        let width = source.width() as usize;
901        let height = source.height() as usize;
902        let nz_width = std::num::NonZeroUsize::new(width).unwrap();
903        let nz_height = std::num::NonZeroUsize::new(height).unwrap();
904
905        let rgb_for_yuvxyb = Rgb::new(
906            source_data.clone(),
907            nz_width,
908            nz_height,
909            TransferCharacteristic::SRGB,
910            ColorPrimaries::BT709,
911        )
912        .unwrap();
913        let lrgb_for_yuvxyb = yuvxyb::LinearRgb::try_from(rgb_for_yuvxyb).unwrap();
914        let xyb_yuvxyb = yuvxyb::Xyb::from(lrgb_for_yuvxyb);
915
916        let rgb_for_simd = Rgb::new(
917            source_data,
918            nz_width,
919            nz_height,
920            TransferCharacteristic::SRGB,
921            ColorPrimaries::BT709,
922        )
923        .unwrap();
924        let lrgb_for_simd = LinearRgb::try_from(rgb_for_simd).unwrap();
925        let xyb_simd = linear_rgb_to_xyb_simd(lrgb_for_simd);
926
927        let mut max_diff = [0.0f32; 3];
928        for (yuvxyb_pix, simd_pix) in xyb_yuvxyb.data().iter().zip(xyb_simd.data().iter()) {
929            for c in 0..3 {
930                let diff = (yuvxyb_pix[c] - simd_pix[c]).abs();
931                max_diff[c] = max_diff[c].max(diff);
932            }
933        }
934
935        assert!(
936            max_diff[0] < 1e-5 && max_diff[1] < 1e-5 && max_diff[2] < 1e-5,
937            "SIMD XYB differs from yuvxyb: max_diff={:?}",
938            max_diff
939        );
940    }
941}