image_hdr/
extensions.rs

1//! Extensions on top of dependencies to facilitate the implementations of this library
2
3use image::{DynamicImage, ImageBuffer, Luma, Rgb};
4use ndarray::Array3;
5
6/// Trait to add the ability to get a nd-array buffer from the target type
7pub trait NDArrayBuffer {
8    /// Get the buffer as `Array3<f32>`
9    fn to_nd_array_buffer(&self) -> Array3<f32>;
10
11    /// Generate a new instance of the target from a nd-array buffer.
12    fn from_nd_array_buffer(buffer: Array3<f32>) -> Self;
13}
14
15impl NDArrayBuffer for DynamicImage {
16    fn to_nd_array_buffer(&self) -> Array3<f32> {
17        match self {
18            DynamicImage::ImageLuma8(_)
19            | DynamicImage::ImageLumaA8(_)
20            | DynamicImage::ImageLuma16(_)
21            | DynamicImage::ImageLumaA16(_) => {
22                let mut buffer =
23                    Array3::<f32>::zeros((self.height() as usize, self.width() as usize, 1));
24
25                for (x, y, pixel) in self.to_luma32f().enumerate_pixels() {
26                    buffer[[y as usize, x as usize, 0]] = pixel.0[0];
27                }
28
29                buffer
30            }
31            _ => {
32                let mut buffer =
33                    Array3::<f32>::zeros((self.height() as usize, self.width() as usize, 3));
34
35                for (x, y, pixel) in self.to_rgb32f().enumerate_pixels() {
36                    let [red, green, blue] = pixel.0;
37
38                    buffer[[y as usize, x as usize, 0]] = red;
39                    buffer[[y as usize, x as usize, 1]] = green;
40                    buffer[[y as usize, x as usize, 2]] = blue;
41                }
42
43                buffer
44            }
45        }
46    }
47
48    #[allow(clippy::cast_possible_truncation)]
49    #[allow(clippy::cast_sign_loss)]
50    fn from_nd_array_buffer(buffer: Array3<f32>) -> Self {
51        if let (height, width, 1) = buffer.dim() {
52            let mut result = ImageBuffer::<Luma<u16>, Vec<u16>>::new(width as u32, height as u32);
53            for (x, y, pixel) in result.enumerate_pixels_mut() {
54                let intensity = buffer[[y as usize, x as usize, 0]] * f32::from(u16::MAX);
55                *pixel = Luma([intensity as u16]);
56            }
57
58            DynamicImage::ImageLuma16(result)
59        } else if let (height, width, 3) = buffer.dim() {
60            let mut result = ImageBuffer::<Rgb<f32>, Vec<f32>>::new(width as u32, height as u32);
61            for (x, y, pixel) in result.enumerate_pixels_mut() {
62                let red = buffer[[y as usize, x as usize, 0]];
63                let green = buffer[[y as usize, x as usize, 1]];
64                let blue = buffer[[y as usize, x as usize, 2]];
65
66                *pixel = Rgb([red, green, blue]);
67            }
68
69            DynamicImage::ImageRgb32F(result)
70        } else {
71            panic!("Unexpected dimensions encountered");
72        }
73    }
74}