Skip to main content

dssim_core/
image.rs

1#![allow(dead_code)]
2
3use imgref::*;
4use rgb::*;
5
6/// RGBA, but: premultiplied alpha, linear (using sRGB primaries, but not its gamma curve), f32 unit scale 0..1
7pub type RGBAPLU = RGBA<f32>;
8/// RGB, but: linear (using sRGB primaries, but not its gamma curve), f32 unit scale 0..1
9pub type RGBLU = RGB<f32>;
10
11/// L\*a\*b\*b, but using float units (values are 100× smaller than in usual integer representation)
12#[derive(Debug, Copy, Clone)]
13pub struct LAB {
14    pub l: f32,
15    pub a: f32,
16    pub b: f32,
17}
18
19impl std::ops::Mul<Self> for LAB {
20    type Output = Self;
21
22    fn mul(self, other: Self) -> Self::Output {
23        Self {
24            l: self.l * other.l,
25            a: self.a * other.a,
26            b: self.b * other.b,
27        }
28    }
29}
30
31impl std::ops::Mul<LAB> for f32 {
32    type Output = LAB;
33
34    fn mul(self, other: LAB) -> Self::Output {
35        LAB {
36            l: self * other.l,
37            a: self * other.a,
38            b: self * other.b,
39        }
40    }
41}
42
43impl std::ops::Mul<f32> for LAB {
44    type Output = Self;
45
46    fn mul(self, other: f32) -> Self::Output {
47        Self {
48            l: self.l * other,
49            a: self.a * other,
50            b: self.b * other,
51        }
52    }
53}
54
55impl std::ops::Add<Self> for LAB {
56    type Output = Self;
57
58    fn add(self, other: Self::Output) -> Self::Output {
59        Self {
60            l: self.l + other.l,
61            a: self.a + other.a,
62            b: self.b + other.b,
63        }
64    }
65}
66
67impl std::ops::Add<f32> for LAB {
68    type Output = Self;
69
70    fn add(self, other: f32) -> Self::Output {
71        Self {
72            l: self.l + other,
73            a: self.a + other,
74            b: self.b + other,
75        }
76    }
77}
78
79impl std::ops::Sub<Self> for LAB {
80    type Output = Self;
81
82    fn sub(self, other: Self) -> Self::Output {
83        Self {
84            l: self.l - other.l,
85            a: self.a - other.a,
86            b: self.b - other.b,
87        }
88    }
89}
90
91impl LAB {
92    pub(crate) fn avg(self) -> f32 {
93        (self.l + self.a + self.b) * (1. / 3.)
94    }
95}
96
97impl From<LAB> for f64 {
98    fn from(other: LAB) -> Self {
99        (Self::from(other.l) + Self::from(other.a) + Self::from(other.b)) * (1. / 3.)
100    }
101}
102
103impl From<LAB> for f32 {
104    fn from(other: LAB) -> Self {
105        other.avg()
106    }
107}
108
109impl std::ops::Div<Self> for LAB {
110    type Output = Self;
111
112    fn div(self, other: Self::Output) -> Self::Output {
113        Self {
114            l: self.l / other.l,
115            a: self.a / other.a,
116            b: self.b / other.b,
117        }
118    }
119}
120
121/// Component-wise averaging of pixel values used by `Downsample` to support arbitrary pixel types
122///
123/// Used to naively resample 4 high-res pixels into one low-res pixel
124#[doc(hidden)]
125pub trait Average4 {
126    fn average4(a: Self, b: Self, c: Self, d: Self) -> Self;
127}
128
129impl Average4 for f32 {
130    fn average4(a: Self, b: Self, c: Self, d: Self) -> Self {
131        (a + b + c + d) * 0.25
132    }
133}
134
135impl Average4 for RGBAPLU {
136    fn average4(a: Self, b: Self, c: Self, d: Self) -> Self {
137        RGBAPLU {
138            r: Average4::average4(a.r, b.r, c.r, d.r),
139            g: Average4::average4(a.g, b.g, c.g, d.g),
140            b: Average4::average4(a.b, b.b, c.b, d.b),
141            a: Average4::average4(a.a, b.a, c.a, d.a),
142        }
143    }
144}
145
146impl Average4 for RGBLU {
147    fn average4(a: Self, b: Self, c: Self, d: Self) -> Self {
148        RGBLU {
149            r: Average4::average4(a.r, b.r, c.r, d.r),
150            g: Average4::average4(a.g, b.g, c.g, d.g),
151            b: Average4::average4(a.b, b.b, c.b, d.b),
152        }
153    }
154}
155
156pub(crate) trait ToRGB {
157    fn to_rgb(self, n: usize) -> RGBLU;
158}
159
160impl ToRGB for RGBAPLU {
161    fn to_rgb(self, n: usize) -> RGBLU {
162        let mut r = self.r;
163        let mut g = self.g;
164        let mut b = self.b;
165        let a = self.a;
166        if a < 255.0 {
167            if (n & 16) != 0 {
168                r += 1.0 - a;
169            }
170            if (n & 8) != 0 {
171                g += 1.0 - a; // assumes premultiplied alpha
172            }
173            if (n & 32) != 0 {
174                b += 1.0 - a;
175            }
176        }
177
178        RGBLU { r, g, b }
179    }
180}
181
182/// You can customize how images are downsampled
183///
184/// Multi-scale DSSIM needs to scale images down. This is it. It's supposed to return the same type of image, but half the size.
185///
186/// There is a default implementation that just averages 4 neighboring pixels.
187#[doc(hidden)]
188pub trait Downsample {
189    type Output;
190    fn downsample(&self) -> Option<Self::Output>;
191}
192
193impl<T> Downsample for ImgVec<T> where T: Average4 + Copy + Sync + Send {
194    type Output = Self;
195
196    fn downsample(&self) -> Option<Self::Output> {
197        self.as_ref().downsample()
198    }
199}
200
201impl<T> Downsample for ImgRef<'_, T> where T: Average4 + Copy + Sync + Send {
202    type Output = ImgVec<T>;
203
204    fn downsample(&self) -> Option<Self::Output> {
205        let stride = self.stride();
206        let width = self.width();
207        let height = self.height();
208
209        if width < 8 || height < 8 {
210            return None;
211        }
212
213        let half_height = height / 2;
214        let half_width = width / 2;
215
216        let mut scaled = Vec::with_capacity(half_width * half_height);
217        scaled.extend(self.buf().chunks(stride * 2).take(half_height).flat_map(|pair| {
218            let (top, bot) = pair.split_at(stride);
219            let top = &top[0..half_width * 2];
220            let bot = &bot[0..half_width * 2];
221
222            top.as_chunks::<2>().0.iter()
223                .zip(bot.chunks_exact(2))
224                .map(|(a, b)| Average4::average4(a[0], a[1], b[0], b[1]))
225        }));
226
227        assert_eq!(half_width * half_height, scaled.len());
228        Some(Img::new(scaled, half_width, half_height))
229    }
230}
231
232#[allow(dead_code)]
233pub(crate) fn worst(input: ImgRef<'_, f32>) -> ImgVec<f32> {
234    let stride = input.stride();
235    let half_height = input.height() / 2;
236    let half_width = input.width() / 2;
237
238    if half_height < 4 || half_width < 4 {
239        return input.new_buf(input.buf().to_vec());
240    }
241
242    let mut scaled = Vec::with_capacity(half_width * half_height);
243    scaled.extend(input.buf().chunks(stride * 2).take(half_height).flat_map(|pair| {
244        let (top, bot) = pair.split_at(stride);
245        let top = &top[0..half_width * 2];
246        let bot = &bot[0..half_width * 2];
247
248        top.as_chunks::<2>().0.iter().zip(bot.chunks_exact(2)).map(|(a,b)| {
249            a[0].min(a[1]).min(b[0].min(b[1]))
250        })
251    }));
252
253    assert_eq!(half_width * half_height, scaled.len());
254    Img::new(scaled, half_width, half_height)
255}
256
257#[allow(dead_code)]
258pub(crate) fn avgworst(input: ImgRef<'_, f32>) -> ImgVec<f32> {
259    let stride = input.stride();
260    let half_height = input.height() / 2;
261    let half_width = input.width() / 2;
262
263    if half_height < 4 || half_width < 4 {
264        return input.new_buf(input.buf().to_vec());
265    }
266
267    let mut scaled = Vec::with_capacity(half_width * half_height);
268    scaled.extend(input.buf().chunks(stride * 2).take(half_height).flat_map(|pair| {
269        let (top, bot) = pair.split_at(stride);
270        let top = &top[0..half_width * 2];
271        let bot = &bot[0..half_width * 2];
272
273        top.as_chunks::<2>().0.iter()
274            .zip(bot.chunks_exact(2))
275            .map(|(a, b)| (a[0] + a[1] + b[0] + b[1]).mul_add(0.25, a[0].min(a[1]).min(b[0].min(b[1]))) * 0.5)
276    }));
277
278    assert_eq!(half_width * half_height, scaled.len());
279    Img::new(scaled, half_width, half_height)
280}
281
282#[allow(dead_code)]
283pub(crate) fn avg(input: ImgRef<'_, f32>) -> ImgVec<f32> {
284    let stride = input.stride();
285    let half_height = input.height() / 2;
286    let half_width = input.width() / 2;
287
288    if half_height < 4 || half_width < 4 {
289        return input.new_buf(input.buf().to_vec());
290    }
291
292    let mut scaled = Vec::with_capacity(half_width * half_height);
293    scaled.extend(input.buf().chunks(stride * 2).take(half_height).flat_map(|pair| {
294        let (top, bot) = pair.split_at(stride);
295        let top = &top[0..half_width * 2];
296        let bot = &bot[0..half_width * 2];
297
298        top.as_chunks::<2>().0.iter().zip(bot.chunks_exact(2)).map(|(a,b)| {
299            (a[0] + a[1] + b[0] + b[1]) * 0.25
300        })
301    }));
302
303    assert_eq!(half_width * half_height, scaled.len());
304    Img::new(scaled, half_width, half_height)
305}