dssim_core/
dssim.rs

1#![allow(non_upper_case_globals)]
2#![allow(non_snake_case)]
3/*
4 * © 2011-2017 Kornel Lesiński. All rights reserved.
5 *
6 * This file is part of DSSIM.
7 *
8 * DSSIM is free software: you can redistribute it and/or modify
9 * it under the terms of the GNU Affero General Public License
10 * as published by the Free Software Foundation, either version 3
11 * of the License, or (at your option) any later version.
12 *
13 * DSSIM is distributed in the hope that it will be useful,
14 * but WITHOUT ANY WARRANTY; without even the implied warranty of
15 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
16 * GNU Affero General Public License for more details.
17 *
18 * You should have received a copy of the license along with DSSIM.
19 * If not, see <http://www.gnu.org/licenses/agpl.txt>.
20 */
21
22use crate::blur;
23use crate::image::*;
24use crate::linear::ToRGBAPLU;
25pub use crate::tolab::ToLABBitmap;
26pub use crate::val::Dssim as Val;
27use imgref::*;
28use itertools::multizip;
29#[cfg(not(feature = "threads"))]
30use crate::lieon as rayon;
31use rayon::prelude::*;
32use rgb::{RGB, RGBA};
33use std::borrow::Borrow;
34use std::mem::MaybeUninit;
35use std::ops;
36use std::ops::Deref;
37use std::sync::Arc;
38
39trait Channable<T, I> {
40    fn img1_img2_blur(&self, modified: &Self, tmp: &mut [MaybeUninit<I>]) -> Vec<T>;
41}
42
43#[derive(Clone)]
44struct DssimChan<T> {
45    pub width: usize,
46    pub height: usize,
47    pub img: Option<ImgVec<T>>,
48    pub mu: Vec<T>,
49    pub img_sq_blur: Vec<T>,
50    pub is_chroma: bool,
51}
52
53/// Configuration for the comparison
54#[derive(Clone, Debug)]
55pub struct Dssim {
56    scale_weights: Vec<f64>,
57    save_maps_scales: u8,
58}
59
60#[derive(Clone)]
61struct DssimChanScale<T> {
62    chan: Vec<DssimChan<T>>,
63}
64
65/// Abstract wrapper for images. See [`Dssim::create_image()`]
66#[derive(Clone)]
67pub struct DssimImage<T> {
68    scale: Vec<DssimChanScale<T>>,
69}
70
71impl<T> DssimImage<T> {
72    #[inline]
73    #[must_use]
74    pub fn width(&self) -> usize {
75        self.scale[0].chan[0].width
76    }
77
78    #[inline]
79    #[must_use]
80    pub fn height(&self) -> usize {
81        self.scale[0].chan[0].height
82    }
83}
84
85// Weighed scales are inspired by the IW-SSIM, but details of the algorithm and weights are different
86const DEFAULT_WEIGHTS: [f64; 5] = [0.028, 0.197, 0.322, 0.298, 0.155];
87
88/// Detailed comparison result
89#[derive(Clone)]
90pub struct SsimMap {
91    /// SSIM scores
92    pub map: ImgVec<f32>,
93    /// Average SSIM (not DSSIM)
94    pub ssim: f64,
95}
96
97/// Create new context for a comparison
98#[must_use]
99pub fn new() -> Dssim {
100    Dssim::new()
101}
102
103impl DssimChan<f32> {
104    pub fn new(bitmap: ImgVec<f32>, is_chroma: bool) -> Self {
105        debug_assert!(bitmap.pixels().all(|i| i.is_finite() && i >= 0.0 && i <= 1.0));
106
107        Self {
108            width: bitmap.width(),
109            height: bitmap.height(),
110            mu: Vec::new(),
111            img: Some(bitmap),
112            img_sq_blur: Vec::new(),
113            is_chroma,
114        }
115    }
116}
117
118impl DssimChan<f32> {
119    fn preprocess(&mut self, tmp: &mut [MaybeUninit<f32>]) {
120        let width = self.width;
121        let height = self.height;
122        assert!(width > 0);
123        assert!(height > 0);
124
125        let img = self.img.as_mut().unwrap();
126        debug_assert_eq!(width * height, img.pixels().count());
127        debug_assert!(img.pixels().all(f32::is_finite));
128
129        if self.is_chroma {
130            blur::blur_in_place(img.as_mut(), tmp);
131        }
132        let (mu, _, _) = blur::blur(img.as_ref(), tmp).into_contiguous_buf();
133        self.mu = mu;
134
135        self.img_sq_blur = img.pixels().map(|i| {
136            debug_assert!(i <= 1.0 && i >= 0.0);
137            i * i
138        }).collect();
139        blur::blur_in_place(ImgRefMut::new(&mut self.img_sq_blur[..], width, height), tmp);
140    }
141}
142
143impl Channable<LAB, f32> for [DssimChan<f32>] {
144    fn img1_img2_blur(&self, modified: &Self, tmp32: &mut [MaybeUninit<f32>]) -> Vec<LAB> {
145
146        let blurred:Vec<_> = self.iter().zip(modified.iter()).map(|(o,m)|{
147            o.img1_img2_blur(m, tmp32)
148        }).collect();
149
150        multizip((blurred[0].iter().copied(), blurred[1].iter().copied(), blurred[2].iter().copied())).map(|(l,a,b)| {
151            LAB {l,a,b}
152        }).collect()
153    }
154}
155
156impl Channable<f32, f32> for DssimChan<f32> {
157    fn img1_img2_blur(&self, modified: &Self, tmp32: &mut [MaybeUninit<f32>]) -> Vec<f32> {
158        let modified_img = modified.img.as_ref().unwrap();
159        let width = modified_img.width();
160        let height = modified_img.height();
161
162        let mut out = Vec::with_capacity(width * height);
163
164        for (row1, row2) in self.img.as_ref().unwrap().rows().zip(modified_img.rows()) {
165            debug_assert_eq!(width, row1.len());
166            debug_assert_eq!(width, row2.len());
167            let row1 = &row1[0..width];
168            let row2 = &row2[0..width];
169            for (px1, px2) in row1.iter().copied().zip(row2.iter().copied()) {
170                debug_assert!(px1 <= 1.0 && px1 >= 0.0);
171                debug_assert!(px2 <= 1.0 && px2 >= 0.0);
172                out.push(px1 * px2);
173            }
174        }
175
176        debug_assert_eq!(out.len(), width * height);
177        blur::blur_in_place(ImgRefMut::new(&mut out, width, height), tmp32);
178        out
179    }
180}
181
182impl Dssim {
183    /// Create new context for comparisons
184    #[must_use]
185    pub fn new() -> Self {
186        Self {
187            scale_weights: DEFAULT_WEIGHTS[..].to_owned(),
188            save_maps_scales: 0,
189        }
190    }
191
192    /// Set how many scales will be used, and weights of each scale
193    pub fn set_scales(&mut self, scales: &[f64]) {
194        self.scale_weights = scales.to_vec();
195    }
196
197    /// Set how many scales will be kept for saving
198    pub fn set_save_ssim_maps(&mut self, num_scales: u8) {
199        self.save_maps_scales = num_scales;
200    }
201
202    /// Create image from an array of RGBA pixels (sRGB, non-premultiplied, alpha last).
203    ///
204    /// If you have a slice of `u8`, then see `rgb` crate's `as_rgba()`.
205    #[must_use]
206    pub fn create_image_rgba(&self, bitmap: &[RGBA<u8>], width: usize, height: usize) -> Option<DssimImage<f32>> {
207        if width * height < bitmap.len() {
208            return None;
209        }
210        let img = ImgVec::new(bitmap.to_rgbaplu(), width, height);
211        self.create_image(&img)
212    }
213
214    /// Create image from an array of packed RGB pixels (sRGB).
215    ///
216    /// If you have a slice of `u8`, then see `rgb` crate's `as_rgb()`.
217    #[must_use]
218    pub fn create_image_rgb(&self, bitmap: &[RGB<u8>], width: usize, height: usize) -> Option<DssimImage<f32>> {
219        if width * height < bitmap.len() {
220            return None;
221        }
222        let img = ImgVec::new(bitmap.to_rgblu(), width, height);
223        self.create_image(&img)
224    }
225
226    /// The input image is defined using the `imgref` crate, and the pixel type can be:
227    ///
228    /// * `ImgVec<RGBAPLU>` — RGBA premultiplied alpha, linear, float scaled to 0..1
229    /// * `ImgVec<RGBLU>` — RGBA linear, float scaled to 0..1
230    /// * `ImgVec<f32>` — linear light grayscale, float scaled to 0..1
231    ///
232    /// And there's [`ToRGBAPLU::to_rgbaplu()`][crate::ToRGBAPLU::to_rgbaplu()] trait to convert the input pixels from
233    /// `[RGBA<u8>]`, `[RGBA<u16>]`, `[RGB<u8>]`, or `RGB<u16>`. See `lib.rs` for example how it's done.
234    ///
235    /// You can implement `ToLABBitmap` and `Downsample` traits on your own image type.
236    pub fn create_image<InBitmap, OutBitmap>(&self, src_img: &InBitmap) -> Option<DssimImage<f32>>
237    where
238        InBitmap: ToLABBitmap + Send + Sync + Downsample<Output = OutBitmap>,
239        OutBitmap: ToLABBitmap + Send + Sync + Downsample<Output = OutBitmap>,
240    {
241        let num_scales = self.scale_weights.len();
242        let mut scale = Vec::with_capacity(num_scales);
243        Self::make_scales_recursive(num_scales, MaybeArc::Borrowed(src_img), &mut scale);
244        scale.reverse(); // depth-first made smallest scales first
245
246        Some(DssimImage { scale })
247    }
248
249    #[inline(never)]
250    fn make_scales_recursive<InBitmap, OutBitmap>(scales_left: usize, image: MaybeArc<'_, InBitmap>, scales: &mut Vec<DssimChanScale<f32>>)
251    where
252        InBitmap: ToLABBitmap + Send + Sync + Downsample<Output = OutBitmap>,
253        OutBitmap: ToLABBitmap + Send + Sync + Downsample<Output = OutBitmap>,
254    {
255        // Run to_lab and next downsampling in parallel
256        let (chan, _) = rayon::join({
257            let image = image.clone();
258            move || {
259                let lab = image.to_lab();
260                drop(image); // Free larger RGB image ASAP
261                DssimChanScale {
262                    chan: lab.into_par_iter().enumerate().map(|(n,l)| {
263                        let w = l.width();
264                        let h = l.height();
265                        let mut ch = DssimChan::new(l, n > 0);
266
267                        let pixels = w * h;
268                        let mut tmp = Vec::with_capacity(pixels);
269                        ch.preprocess(&mut tmp.spare_capacity_mut()[..pixels]);
270                        ch
271                    }).collect(),
272                }
273            }
274        }, {
275            let scales = &mut *scales;
276            move || {
277                if scales_left > 0 {
278                    let down = image.downsample();
279                    drop(image);
280                    if let Some(downsampled) = down {
281                        Self::make_scales_recursive(scales_left - 1, MaybeArc::Owned(Arc::new(downsampled)), scales);
282                    }
283                }
284            }
285        });
286        scales.push(chan);
287    }
288
289    /// Compare original with another image. See `create_image`
290    ///
291    /// The `SsimMap`s are returned only if you've enabled them first.
292    ///
293    /// `Val` is a fancy wrapper for `f64`
294    #[inline(never)]
295    pub fn compare<M: Borrow<DssimImage<f32>>>(&self, original_image: &DssimImage<f32>, modified_image: M) -> (Val, Vec<SsimMap>) {
296        let modified_image = modified_image.borrow();
297        let res: Vec<_> = self.scale_weights.as_slice().par_iter().with_min_len(1).cloned().zip(
298                        modified_image.scale.as_slice().par_iter().with_min_len(1).zip(
299                        original_image.scale.as_slice().par_iter().with_min_len(1))
300        ).enumerate().map(|(n, (weight, (modified_image_scale, original_image_scale)))| {
301            let scale_width = original_image_scale.chan[0].width;
302            let scale_height = original_image_scale.chan[0].height;
303            let mut tmp = Vec::with_capacity(scale_width * scale_height);
304            let tmp = &mut tmp.spare_capacity_mut()[0 .. scale_width*scale_height];
305
306            let ssim_map = match original_image_scale.chan.len() {
307                3 => {
308                    let (original_lab, (img1_img2_blur, modified_lab)) = rayon::join(
309                    || Self::lab_chan(original_image_scale),
310                    || {
311                        let img1_img2_blur = original_image_scale.chan.img1_img2_blur(&modified_image_scale.chan, tmp);
312                        (img1_img2_blur, Self::lab_chan(modified_image_scale))
313                    });
314
315                    Self::compare_scale(&original_lab, &modified_lab, &img1_img2_blur)
316                },
317                1 => {
318                    let img1_img2_blur = original_image_scale.chan[0].img1_img2_blur(&modified_image_scale.chan[0], tmp);
319                    Self::compare_scale(&original_image_scale.chan[0], &modified_image_scale.chan[0], &img1_img2_blur)
320                },
321                _ => panic!(),
322            };
323
324            let sum = ssim_map.pixels().fold(0., |sum, i| sum + f64::from(i));
325            let len = (ssim_map.width()*ssim_map.height()) as f64;
326            let avg = (sum / len).max(0.0).powf((0.5_f64).powf(n as f64));
327            let score = 1.0 - (ssim_map.pixels().fold(0., |sum, i| sum + (avg - f64::from(i)).abs()) / len);
328
329            let map = if self.save_maps_scales as usize > n {
330                Some(SsimMap {
331                    map: ssim_map,
332                    ssim: score,
333                })
334            } else {
335                None
336            };
337            (score, weight, map)
338        }).collect();
339
340        let mut ssim_sum = 0.0;
341        let mut weight_sum = 0.0;
342        let mut ssim_maps = Vec::new();
343        for (score, weight, map) in res {
344            ssim_sum += score * weight;
345            weight_sum += weight;
346            if let Some(m) = map {
347                ssim_maps.push(m);
348            }
349        }
350
351        (to_dssim(ssim_sum / weight_sum).into(), ssim_maps)
352    }
353
354    fn lab_chan(scale: &DssimChanScale<f32>) -> DssimChan<LAB> {
355        let l = &scale.chan[0];
356        let a = &scale.chan[1];
357        let b = &scale.chan[2];
358        assert_eq!(l.width, a.width);
359        assert_eq!(b.width, a.width);
360        DssimChan {
361            img_sq_blur: multizip((l.img_sq_blur.iter().copied(), a.img_sq_blur.iter().copied(), b.img_sq_blur.iter().copied()))
362                .map(|(l,a,b)|LAB {l,a,b}).collect(),
363            img: if let (Some(l),Some(a),Some(b)) = (&l.img, &a.img, &b.img) {
364                let buf = multizip((l.pixels(), a.pixels(), b.pixels())).map(|(l,a,b)|{
365                    debug_assert!(l.is_finite() && a.is_finite() && b.is_finite());
366                    LAB {l,a,b}
367                }).collect();
368                Some(ImgVec::new(buf, l.width(), l.height()))
369            } else {None},
370            mu: multizip((l.mu.iter().copied(), a.mu.iter().copied(), b.mu.iter().copied())).map(|(l,a,b)|LAB {l,a,b}).collect(),
371            is_chroma: false,
372            width: l.width,
373            height: l.height,
374        }
375    }
376
377    #[inline(never)]
378    fn compare_scale<L>(original: &DssimChan<L>, modified: &DssimChan<L>, img1_img2_blur: &[L]) -> ImgVec<f32>
379    where
380        L: Send + Sync + Clone + Copy + ops::Mul<Output = L> + ops::Sub<Output = L> + 'static,
381        f32: From<L>,
382    {
383        assert_eq!(original.width, modified.width);
384        assert_eq!(original.height, modified.height);
385
386        let width = original.width;
387        let height = original.height;
388
389        let c1 = 0.01 * 0.01;
390        let c2 = 0.03 * 0.03;
391
392        debug_assert_eq!(original.mu.len(), modified.mu.len());
393        debug_assert_eq!(original.img_sq_blur.len(), modified.img_sq_blur.len());
394        debug_assert_eq!(img1_img2_blur.len(), original.mu.len());
395        debug_assert_eq!(img1_img2_blur.len(), original.img_sq_blur.len());
396
397        let mu_iter = original.mu.as_slice().par_iter().with_min_len(1<<10).cloned().zip_eq(modified.mu.as_slice().par_iter().with_min_len(1<<10).cloned());
398        let sq_iter = original.img_sq_blur.as_slice().par_iter().with_min_len(1<<10).cloned().zip_eq(modified.img_sq_blur.as_slice().par_iter().with_min_len(1<<10).cloned());
399        let map_out = img1_img2_blur.par_iter().with_min_len(1<<10).cloned().zip_eq(mu_iter).zip_eq(sq_iter)
400        .map(|((img1_img2_blur, (mu1, mu2)), (img1_sq_blur, img2_sq_blur))| {
401            let mu1mu1 = mu1 * mu1;
402            let mu1mu2 = mu1 * mu2;
403            let mu2mu2 = mu2 * mu2;
404            let mu1_sq: f32 = mu1mu1.into();
405            let mu2_sq: f32 = mu2mu2.into();
406            let mu1_mu2: f32 = mu1mu2.into();
407            let sigma1_sq: f32 = (img1_sq_blur - mu1mu1).into();
408            let sigma2_sq: f32 = (img2_sq_blur - mu2mu2).into();
409            let sigma12: f32 = (img1_img2_blur - mu1mu2).into();
410
411            2.0f32.mul_add(mu1_mu2, c1) * 2.0f32.mul_add(sigma12, c2) /
412                       ((mu1_sq + mu2_sq + c1) * (sigma1_sq + sigma2_sq + c2))
413        }).collect();
414
415        ImgVec::new(map_out, width, height)
416    }
417}
418
419fn to_dssim(ssim: f64) -> f64 {
420    1.0 / ssim.max(std::f64::EPSILON) - 1.0
421}
422
423#[test]
424fn png_compare() {
425    use crate::linear::*;
426    use imgref::*;
427
428    let d = new();
429    let file1 = lodepng::decode32_file("../tests/test1-sm.png").unwrap();
430    let file2 = lodepng::decode32_file("../tests/test2-sm.png").unwrap();
431
432    let buf1 = &file1.buffer.to_rgbaplu()[..];
433    let buf2 = &file2.buffer.to_rgbaplu()[..];
434    let img1 = d.create_image(&Img::new(buf1, file1.width, file1.height)).unwrap();
435    let img2 = d.create_image(&Img::new(buf2, file2.width, file2.height)).unwrap();
436
437    let (res, _) = d.compare(&img1, img2);
438    assert!((0.001 - res).abs() < 0.0005, "res is {res}");
439
440    let img1b = d.create_image(&Img::new(buf1, file1.width, file1.height)).unwrap();
441    let (res, _) = d.compare(&img1, img1b);
442
443    assert!(0.000000000000001 > res);
444    assert!(res < 0.000000000000001);
445    assert_eq!(res, res);
446
447    let sub_img1 = d.create_image(&Img::new(buf1, file1.width, file1.height).sub_image(2,3,44,33)).unwrap();
448    let sub_img2 = d.create_image(&Img::new(buf2, file2.width, file2.height).sub_image(17,9,44,33)).unwrap();
449    // Test passing second image directly
450    let (res, _) = d.compare(&sub_img1, sub_img2);
451    assert!(res > 0.1);
452
453    let sub_img1 = d.create_image(&Img::new(buf1, file1.width, file1.height).sub_image(22,8,61,40)).unwrap();
454    let sub_img2 = d.create_image(&Img::new(buf2, file2.width, file2.height).sub_image(22,8,61,40)).unwrap();
455    // Test passing second image as reference
456    let (res, _) = d.compare(&sub_img1, sub_img2);
457    assert!(res < 0.01);
458}
459
460enum MaybeArc<'a, T> {
461    Owned(Arc<T>),
462    Borrowed(&'a T),
463}
464
465impl<T> Clone for MaybeArc<'_, T> {
466    fn clone(&self) -> Self {
467        match self {
468            Self::Owned(t) => Self::Owned(t.clone()),
469            Self::Borrowed(t) => Self::Borrowed(t),
470        }
471    }
472}
473
474impl<T> Deref for MaybeArc<'_, T> {
475    type Target = T;
476
477    fn deref(&self) -> &Self::Target {
478        match self {
479            Self::Owned(t) => t,
480            Self::Borrowed(t) => t,
481        }
482    }
483}
484
485#[test]
486fn poison() {
487    let a = RGBAPLU::new(1.,1.,1.,1.);
488    let b = RGBAPLU::new(0.,0.,0.,0.);
489    let n = 1./0.;
490    let n = RGBAPLU::new(n,n,n,n);
491    let buf = vec![
492      b,a,a,b,n,n,
493      a,b,b,a,n,n,
494      b,a,a,b,n,
495    ];
496    let img = ImgVec::new_stride(buf, 4, 3, 6);
497    assert!(img.pixels().all(|p| p.r.is_finite() && p.a.is_finite()));
498    assert!(img.as_ref().pixels().all(|p| p.g.is_finite() && p.b.is_finite()));
499
500    let d = new();
501    let sub_img1 = d.create_image(&img.as_ref()).unwrap();
502    let sub_img2 = d.create_image(&img.as_ref()).unwrap();
503    let (res, _) = d.compare(&sub_img1, sub_img2);
504    assert!(res < 0.000001);
505}
dssim_core/dssim.rs

dssim_core/
dssim.rs
