fast_ssim2/
input.rs

1//! Input image types and conversion to linear RGB.
2//!
3//! This module provides the [`ToLinearRgb`] trait for converting various image
4//! formats to the internal linear RGB representation used by SSIMULACRA2.
5//!
6//! ## Supported input formats (with `imgref` feature)
7//!
8//! | Type | Color Space | Conversion |
9//! |------|-------------|------------|
10//! | `ImgRef<[u8; 3]>` | sRGB (gamma) | `/255` + linearize |
11//! | `ImgRef<[u16; 3]>` | sRGB (gamma) | `/65535` + linearize |
12//! | `ImgRef<[f32; 3]>` | Linear RGB | none |
13//! | `ImgRef<u8>` | sRGB grayscale | `/255` + linearize + expand |
14//! | `ImgRef<f32>` | Linear grayscale | expand to RGB |
15//!
16//! ## Convention
17//!
18//! - Integer types (u8, u16) are assumed to be **sRGB** (gamma-encoded)
19//! - Float types (f32) are assumed to be **linear**
20
21/// Internal linear RGB image representation.
22///
23/// Stores pixels as `[f32; 3]` in linear RGB color space (0.0-1.0 range).
24#[derive(Clone)]
25pub struct LinearRgbImage {
26    pub(crate) data: Vec<[f32; 3]>,
27    pub(crate) width: usize,
28    pub(crate) height: usize,
29}
30
31impl LinearRgbImage {
32    /// Creates a new linear RGB image from raw data.
33    pub fn new(data: Vec<[f32; 3]>, width: usize, height: usize) -> Self {
34        debug_assert_eq!(data.len(), width * height);
35        Self {
36            data,
37            width,
38            height,
39        }
40    }
41
42    /// Returns the image width.
43    pub fn width(&self) -> usize {
44        self.width
45    }
46
47    /// Returns the image height.
48    pub fn height(&self) -> usize {
49        self.height
50    }
51
52    /// Returns the pixel data.
53    pub fn data(&self) -> &[[f32; 3]] {
54        &self.data
55    }
56
57    /// Returns mutable pixel data.
58    pub fn data_mut(&mut self) -> &mut [[f32; 3]] {
59        &mut self.data
60    }
61}
62
63/// Trait for converting image types to linear RGB.
64///
65/// Implement this trait to add support for custom image types.
66pub trait ToLinearRgb {
67    /// Convert to linear RGB image.
68    fn to_linear_rgb(&self) -> LinearRgbImage;
69}
70
71// =============================================================================
72// sRGB conversion functions
73// =============================================================================
74
75/// Convert sRGB (gamma-encoded) value to linear.
76///
77/// Uses the standard sRGB transfer function.
78#[inline]
79pub fn srgb_to_linear(s: f32) -> f32 {
80    if s <= 0.04045 {
81        s / 12.92
82    } else {
83        ((s + 0.055) / 1.055).powf(2.4)
84    }
85}
86
87/// Convert 8-bit sRGB value to linear f32.
88#[inline]
89pub fn srgb_u8_to_linear(v: u8) -> f32 {
90    // Use lookup table for performance
91    SRGB_TO_LINEAR_LUT[v as usize]
92}
93
94/// Convert 16-bit sRGB value to linear f32.
95#[inline]
96pub fn srgb_u16_to_linear(v: u16) -> f32 {
97    srgb_to_linear(v as f32 / 65535.0)
98}
99
100// Precomputed lookup table for sRGB u8 -> linear f32
101// Generated with: (0..256).map(|i| srgb_to_linear(i as f32 / 255.0))
102static SRGB_TO_LINEAR_LUT: std::sync::LazyLock<[f32; 256]> = std::sync::LazyLock::new(|| {
103    let mut lut = [0.0f32; 256];
104    for (i, entry) in lut.iter_mut().enumerate() {
105        *entry = srgb_to_linear(i as f32 / 255.0);
106    }
107    lut
108});
109
110// =============================================================================
111// imgref implementations
112// =============================================================================
113
114#[cfg(feature = "imgref")]
115mod imgref_impl {
116    use super::*;
117    use imgref::ImgRef;
118
119    /// RGB u8 (sRGB) -> Linear RGB
120    impl ToLinearRgb for ImgRef<'_, [u8; 3]> {
121        fn to_linear_rgb(&self) -> LinearRgbImage {
122            let data: Vec<[f32; 3]> = self
123                .pixels()
124                .map(|[r, g, b]| {
125                    [
126                        srgb_u8_to_linear(r),
127                        srgb_u8_to_linear(g),
128                        srgb_u8_to_linear(b),
129                    ]
130                })
131                .collect();
132            LinearRgbImage::new(data, self.width(), self.height())
133        }
134    }
135
136    /// RGB u16 (sRGB) -> Linear RGB
137    impl ToLinearRgb for ImgRef<'_, [u16; 3]> {
138        fn to_linear_rgb(&self) -> LinearRgbImage {
139            let data: Vec<[f32; 3]> = self
140                .pixels()
141                .map(|[r, g, b]| {
142                    [
143                        srgb_u16_to_linear(r),
144                        srgb_u16_to_linear(g),
145                        srgb_u16_to_linear(b),
146                    ]
147                })
148                .collect();
149            LinearRgbImage::new(data, self.width(), self.height())
150        }
151    }
152
153    /// RGB f32 (already linear) -> Linear RGB
154    impl ToLinearRgb for ImgRef<'_, [f32; 3]> {
155        fn to_linear_rgb(&self) -> LinearRgbImage {
156            let data: Vec<[f32; 3]> = self.pixels().collect();
157            LinearRgbImage::new(data, self.width(), self.height())
158        }
159    }
160
161    /// Grayscale u8 (sRGB) -> Linear RGB
162    impl ToLinearRgb for ImgRef<'_, u8> {
163        fn to_linear_rgb(&self) -> LinearRgbImage {
164            let data: Vec<[f32; 3]> = self
165                .pixels()
166                .map(|v| {
167                    let l = srgb_u8_to_linear(v);
168                    [l, l, l]
169                })
170                .collect();
171            LinearRgbImage::new(data, self.width(), self.height())
172        }
173    }
174
175    /// Grayscale f32 (linear) -> Linear RGB
176    impl ToLinearRgb for ImgRef<'_, f32> {
177        fn to_linear_rgb(&self) -> LinearRgbImage {
178            let data: Vec<[f32; 3]> = self.pixels().map(|v| [v, v, v]).collect();
179            LinearRgbImage::new(data, self.width(), self.height())
180        }
181    }
182}
183
184// =============================================================================
185// yuvxyb compatibility
186// =============================================================================
187
188impl ToLinearRgb for yuvxyb::LinearRgb {
189    fn to_linear_rgb(&self) -> LinearRgbImage {
190        LinearRgbImage::new(self.data().to_vec(), self.width(), self.height())
191    }
192}
193
194// =============================================================================
195// Conversion to yuvxyb::LinearRgb (for internal pipeline)
196// =============================================================================
197
198impl From<LinearRgbImage> for yuvxyb::LinearRgb {
199    fn from(img: LinearRgbImage) -> Self {
200        yuvxyb::LinearRgb::new(img.data, img.width, img.height)
201            .expect("LinearRgbImage dimensions are always valid")
202    }
203}
204
205impl ToLinearRgb for yuvxyb::Rgb {
206    fn to_linear_rgb(&self) -> LinearRgbImage {
207        // yuvxyb::Rgb handles the sRGB -> linear conversion internally via TryFrom
208        let linear: yuvxyb::LinearRgb = yuvxyb::LinearRgb::try_from(self.clone())
209            .expect("Rgb to LinearRgb conversion should not fail");
210        linear.to_linear_rgb()
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    #[test]
219    fn test_srgb_to_linear_bounds() {
220        assert!((srgb_to_linear(0.0) - 0.0).abs() < 1e-6);
221        assert!((srgb_to_linear(1.0) - 1.0).abs() < 1e-6);
222    }
223
224    #[test]
225    fn test_srgb_to_linear_midpoint() {
226        // sRGB 0.5 should be approximately 0.214 in linear
227        let linear = srgb_to_linear(0.5);
228        assert!((linear - 0.214).abs() < 0.01);
229    }
230
231    #[test]
232    fn test_srgb_u8_to_linear() {
233        assert!((srgb_u8_to_linear(0) - 0.0).abs() < 1e-6);
234        assert!((srgb_u8_to_linear(255) - 1.0).abs() < 1e-6);
235    }
236
237    #[test]
238    fn test_linear_rgb_image_accessors() {
239        let data = vec![[0.5, 0.3, 0.1], [0.2, 0.4, 0.6]];
240        let img = LinearRgbImage::new(data.clone(), 2, 1);
241
242        assert_eq!(img.width(), 2);
243        assert_eq!(img.height(), 1);
244        assert_eq!(img.data(), &data[..]);
245    }
246
247    #[test]
248    fn test_yuvxyb_linearrgb_roundtrip() {
249        let data = vec![[0.5, 0.3, 0.1]; 4];
250        let yuvxyb_img = yuvxyb::LinearRgb::new(data.clone(), 2, 2).expect("valid dimensions");
251
252        let our_img = yuvxyb_img.to_linear_rgb();
253        assert_eq!(our_img.width(), 2);
254        assert_eq!(our_img.height(), 2);
255        assert_eq!(our_img.data(), &data[..]);
256
257        // Convert back
258        let back: yuvxyb::LinearRgb = our_img.into();
259        assert_eq!(back.data(), &data[..]);
260    }
261}
262
263#[cfg(all(test, feature = "imgref"))]
264mod imgref_tests {
265    use super::*;
266    use imgref::{Img, ImgVec};
267
268    #[test]
269    fn test_imgref_u8_srgb_conversion() {
270        // Create a 2x2 sRGB image
271        let pixels: Vec<[u8; 3]> = vec![
272            [0, 0, 0],       // black
273            [255, 255, 255], // white
274            [128, 128, 128], // mid gray
275            [255, 0, 0],     // red
276        ];
277        let img: ImgVec<[u8; 3]> = Img::new(pixels, 2, 2);
278
279        let linear = img.as_ref().to_linear_rgb();
280        assert_eq!(linear.width(), 2);
281        assert_eq!(linear.height(), 2);
282
283        // Black should be [0, 0, 0]
284        assert!((linear.data()[0][0] - 0.0).abs() < 1e-6);
285        // White should be [1, 1, 1]
286        assert!((linear.data()[1][0] - 1.0).abs() < 1e-6);
287        assert!((linear.data()[1][1] - 1.0).abs() < 1e-6);
288        // Mid gray (sRGB 128) should be ~0.215 in linear
289        assert!((linear.data()[2][0] - 0.215).abs() < 0.01);
290        // Red should have R=1, G=B=0
291        assert!((linear.data()[3][0] - 1.0).abs() < 1e-6);
292        assert!((linear.data()[3][1] - 0.0).abs() < 1e-6);
293    }
294
295    #[test]
296    fn test_imgref_f32_passthrough() {
297        // f32 is assumed to already be linear - should pass through unchanged
298        let pixels: Vec<[f32; 3]> = vec![[0.5, 0.3, 0.1], [0.9, 0.8, 0.7]];
299        let img: ImgVec<[f32; 3]> = Img::new(pixels.clone(), 2, 1);
300
301        let linear = img.as_ref().to_linear_rgb();
302        assert_eq!(linear.data(), &pixels[..]);
303    }
304
305    #[test]
306    fn test_imgref_grayscale_u8_expansion() {
307        // Grayscale u8 should expand to R=G=B and apply sRGB conversion
308        let pixels: Vec<u8> = vec![0, 255, 128];
309        let img: ImgVec<u8> = Img::new(pixels, 3, 1);
310
311        let linear = img.as_ref().to_linear_rgb();
312
313        // Black
314        let black = linear.data()[0];
315        assert!((black[0] - 0.0).abs() < 1e-6);
316        assert_eq!(black[0], black[1]);
317        assert_eq!(black[1], black[2]);
318
319        // White
320        let white = linear.data()[1];
321        assert!((white[0] - 1.0).abs() < 1e-6);
322        assert_eq!(white[0], white[1]);
323
324        // Mid gray
325        let gray = linear.data()[2];
326        assert!((gray[0] - 0.215).abs() < 0.01);
327        assert_eq!(gray[0], gray[1]);
328    }
329
330    #[test]
331    fn test_imgref_grayscale_f32_expansion() {
332        // Grayscale f32 should expand to R=G=B (already linear)
333        let pixels: Vec<f32> = vec![0.0, 1.0, 0.5];
334        let img: ImgVec<f32> = Img::new(pixels, 3, 1);
335
336        let linear = img.as_ref().to_linear_rgb();
337
338        assert_eq!(linear.data()[0], [0.0, 0.0, 0.0]);
339        assert_eq!(linear.data()[1], [1.0, 1.0, 1.0]);
340        assert_eq!(linear.data()[2], [0.5, 0.5, 0.5]);
341    }
342
343    #[test]
344    fn test_compute_ssimulacra2_with_imgref_u8() {
345        use crate::compute_ssimulacra2;
346
347        // Create two 16x16 images (minimum viable for SSIMULACRA2)
348        let pixels1: Vec<[u8; 3]> = vec![[128, 128, 128]; 16 * 16];
349        let pixels2: Vec<[u8; 3]> = vec![[130, 130, 130]; 16 * 16]; // slightly different
350
351        let img1: ImgVec<[u8; 3]> = Img::new(pixels1, 16, 16);
352        let img2: ImgVec<[u8; 3]> = Img::new(pixels2, 16, 16);
353
354        // Should compute successfully
355        let score = compute_ssimulacra2(img1.as_ref(), img2.as_ref()).unwrap();
356        // Small difference should result in high score (close to 100)
357        assert!(
358            score > 90.0,
359            "Score {score} should be > 90 for very similar images"
360        );
361    }
362
363    #[test]
364    fn test_compute_ssimulacra2_identical_imgref() {
365        use crate::compute_ssimulacra2;
366
367        // Identical images should score 100
368        let pixels: Vec<[u8; 3]> = vec![[100, 150, 200]; 16 * 16];
369        let img: ImgVec<[u8; 3]> = Img::new(pixels, 16, 16);
370
371        let score = compute_ssimulacra2(img.as_ref(), img.as_ref()).unwrap();
372        assert!(
373            (score - 100.0).abs() < 0.01,
374            "Identical images should score 100, got {score}"
375        );
376    }
377
378    #[test]
379    fn test_precompute_with_imgref() {
380        use crate::Ssimulacra2Reference;
381
382        // Create source and distorted images
383        let source_pixels: Vec<[u8; 3]> = vec![[128, 128, 128]; 32 * 32];
384        let distorted_pixels: Vec<[u8; 3]> = vec![[130, 128, 126]; 32 * 32];
385
386        let source: ImgVec<[u8; 3]> = Img::new(source_pixels, 32, 32);
387        let distorted: ImgVec<[u8; 3]> = Img::new(distorted_pixels, 32, 32);
388
389        // Use precompute API with imgref
390        let reference = Ssimulacra2Reference::new(source.as_ref()).unwrap();
391        let score = reference.compare(distorted.as_ref()).unwrap();
392
393        // Should compute successfully with reasonable score
394        assert!(
395            score > 80.0,
396            "Score {score} should be > 80 for similar images"
397        );
398    }
399}