1use image::{DynamicImage, ImageBuffer, Luma, Rgb};
4use ndarray::Array3;
5
6pub trait NDArrayBuffer {
8 fn to_nd_array_buffer(&self) -> Array3<f32>;
10
11 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}