image_dwt/
recompose.rs

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> {}