1use image::{DynamicImage, ImageBuffer, Luma, Rgb};
2use ndarray::{Array2, Array3};
3
4use crate::aggregate::Aggregate;
5use crate::layer::{WaveletLayer, WaveletLayerBuffer};
6
7#[derive(Copy, Clone)]
8pub enum OutputLayer {
9 Grayscale,
10 Rgb,
11}
12
13impl OutputLayer {
14 fn to_num_channels(self) -> usize {
15 match self {
16 OutputLayer::Grayscale => 1,
17 OutputLayer::Rgb => 3,
18 }
19 }
20}
21
22pub trait RecomposableWaveletLayers: Iterator<Item = WaveletLayer> {
23 fn recompose_into_image(
24 self,
25 width: usize,
26 height: usize,
27 output_layer: OutputLayer,
28 ) -> DynamicImage
29 where
30 Self: Sized,
31 {
32 match output_layer {
33 OutputLayer::Grayscale => {
34 let mut result = Array2::<f32>::zeros((height, width));
35
36 for layer in self {
37 match layer.buffer {
38 WaveletLayerBuffer::Grayscale { data } => {
39 result += &data;
40 }
41 _ => unimplemented!(),
42 }
43 }
44
45 let min_pixel = result.min();
46 let max_pixel = result.max();
47
48 let mut result_img: ImageBuffer<Luma<u16>, Vec<u16>> =
49 ImageBuffer::new(width as u32, height as u32);
50
51 let rescale_ratio = max_pixel - min_pixel;
52
53 for (x, y, pixel) in result_img.enumerate_pixels_mut() {
54 let intensity = result[(y as usize, x as usize)];
55
56 *pixel =
57 Luma([((intensity - min_pixel) / rescale_ratio * u16::MAX as f32) as u16]);
58 }
59
60 DynamicImage::ImageLuma16(result_img)
61 }
62 OutputLayer::Rgb => {
63 let mut result =
64 Array3::<f32>::zeros((height, width, output_layer.to_num_channels()));
65
66 for layer in self {
67 match layer.buffer {
68 WaveletLayerBuffer::Rgb { data } => {
69 result += &data;
70 }
71 _ => unimplemented!(),
72 }
73 }
74
75 let min_pixel = result.min();
76 let max_pixel = result.max();
77
78 let mut result_img: ImageBuffer<Rgb<f32>, Vec<f32>> =
79 ImageBuffer::new(width as u32, height as u32);
80
81 let rescale_ratio = max_pixel - min_pixel;
82
83 for (x, y, pixel) in result_img.enumerate_pixels_mut() {
84 let red = result[(y as usize, x as usize, 0)];
85 let green = result[(y as usize, x as usize, 1)];
86 let blue = result[(y as usize, x as usize, 2)];
87
88 let scaled_red = (red - min_pixel) / rescale_ratio;
89 let scaled_green = (green - min_pixel) / rescale_ratio;
90 let scaled_blue = (blue - min_pixel) / rescale_ratio;
91
92 *pixel = Rgb([scaled_red, scaled_green, scaled_blue]);
93 }
94
95 DynamicImage::ImageRgb32F(result_img)
96 }
97 }
98 }
99}
100
101impl<T> RecomposableWaveletLayers for T where T: Iterator<Item = WaveletLayer> {}