1use crate::kernels::{B3SplineKernel, Kernel, LinearInterpolationKernel};
2use convolve_image::Convolution;
3use ndarray::{Array2, Array3};
4
5use crate::layer::WaveletLayerBuffer;
6
7pub trait WaveletDecompose {
8 fn wavelet_decompose(&mut self, kernel: Kernel, pixel_scale: usize) -> WaveletLayerBuffer;
9}
10
11impl WaveletDecompose for Array2<f32> {
12 fn wavelet_decompose(&mut self, kernel: Kernel, pixel_scale: usize) -> WaveletLayerBuffer {
13 let stride = 2_usize.pow(pixel_scale as u32);
14 let mut current_data = self.clone();
15 match kernel {
16 Kernel::LinearInterpolationKernel => {
17 current_data.convolve(LinearInterpolationKernel::new().into(), stride);
18 }
19 Kernel::LowScaleKernel => {
20 unimplemented!("Low scale is not a separable kernel");
21 }
22 Kernel::B3SplineKernel => current_data.convolve(B3SplineKernel::new().into(), stride),
23 }
24
25 let final_data = self.clone() - ¤t_data;
26 *self = current_data;
27
28 WaveletLayerBuffer::Grayscale { data: final_data }
29 }
30}
31
32impl WaveletDecompose for Array3<f32> {
33 fn wavelet_decompose(&mut self, kernel: Kernel, pixel_scale: usize) -> WaveletLayerBuffer {
34 let stride = 2_usize.pow(pixel_scale as u32);
35 let mut current_data = self.clone();
36 match kernel {
37 Kernel::LinearInterpolationKernel => {
38 current_data.convolve(LinearInterpolationKernel::new().into(), stride);
39 }
40 Kernel::LowScaleKernel => {
41 unimplemented!("Low scale is not a separable kernel");
42 }
43 Kernel::B3SplineKernel => current_data.convolve(B3SplineKernel::new().into(), stride),
44 }
45
46 let final_data = self.clone() - ¤t_data;
47 *self = current_data;
48
49 WaveletLayerBuffer::Rgb { data: final_data }
50 }
51}