1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

use image::{DynamicImage, ImageBuffer, Luma, Rgb};
use ndarray::{Array2, Array3};

use crate::aggregate::Aggregate;
use crate::layer::{WaveletLayer, WaveletLayerBuffer};

#[derive(Copy, Clone)]
pub enum OutputLayer {
    Grayscale,
    Rgb,
}

impl OutputLayer {
    fn to_num_channels(self) -> usize {
        match self {
            OutputLayer::Grayscale => 1,
            OutputLayer::Rgb => 3,
        }
    }
}

pub trait RecomposableWaveletLayers: Iterator<Item = WaveletLayer> {
    fn recompose_into_image(
        self,
        width: usize,
        height: usize,
        output_layer: OutputLayer,
    ) -> DynamicImage
    where
        Self: Sized,
    {
        match output_layer {
            OutputLayer::Grayscale => {
                let mut result = Array2::<f32>::zeros((height, width));

                for layer in self {
                    match layer.buffer {
                        WaveletLayerBuffer::Grayscale { data } => {
                            result += &data;
                        }
                        _ => unimplemented!(),
                    }
                }

                let min_pixel = result.min();
                let max_pixel = result.max();

                let mut result_img: ImageBuffer<Luma<u16>, Vec<u16>> =
                    ImageBuffer::new(width as u32, height as u32);

                let rescale_ratio = max_pixel - min_pixel;

                for (x, y, pixel) in result_img.enumerate_pixels_mut() {
                    let intensity = result[(y as usize, x as usize)];

                    *pixel =
                        Luma([((intensity - min_pixel) / rescale_ratio * u16::MAX as f32) as u16]);
                }

                DynamicImage::ImageLuma16(result_img)
            }
            OutputLayer::Rgb => {
                let mut result =
                    Array3::<f32>::zeros((height, width, output_layer.to_num_channels()));

                for layer in self {
                    match layer.buffer {
                        WaveletLayerBuffer::Rgb { data } => {
                            result += &data;
                        }
                        _ => unimplemented!(),
                    }
                }

                let min_pixel = result.min();
                let max_pixel = result.max();

                let mut result_img: ImageBuffer<Rgb<f32>, Vec<f32>> =
                    ImageBuffer::new(width as u32, height as u32);

                let rescale_ratio = max_pixel - min_pixel;

                for (x, y, pixel) in result_img.enumerate_pixels_mut() {
                    let red = result[(y as usize, x as usize, 0)];
                    let green = result[(y as usize, x as usize, 1)];
                    let blue = result[(y as usize, x as usize, 2)];

                    let scaled_red = (red - min_pixel) / rescale_ratio;
                    let scaled_green = (green - min_pixel) / rescale_ratio;
                    let scaled_blue = (blue - min_pixel) / rescale_ratio;

                    *pixel = Rgb([scaled_red, scaled_green, scaled_blue]);
                }

                DynamicImage::ImageRgb32F(result_img)
            }
        }
    }
}

impl<T> RecomposableWaveletLayers for T where T: Iterator<Item = WaveletLayer> {}