Fading Coder

One Final Commit for the Last Sprint

Home > Tech > Content

Implementing the À Trous Wavelet Transform for Multi-Scale Image Processing in Rust

Tech May 13 3

Multi-Scale Decomposition

Multi-scale analysis decomposes input data into multiple signals, each representing information at a specific scale. In image processing, a scale corresponds to the pixel dimensions of various structures or details. For instance, an astronomical photograph contains structures ranging from tiny singular-pixel noise and sharp stars to expansive galactic formations.

By separating these scales, targeted operations can be applied without altering other structural details. Denoising is a prime example: random noise predominantly exists in the smallest scale layers. Applying a denoising filter exclusively to these initial layers and subsequently recombining them with the unmodified large-scale layers preserves the sharpness of broader structures that would otherwise blur.

The À Trous Algorithm

The À Trous (with holes) algorithm is a discrete wavelet transform method widely adopted in astronomical data processing. It isolates image structures based purely on their spatial scale without performing downsampling or upsampling, keeping the image dimensions constant throughout the process.

The decomposition process follows these steps:

  1. Begin with the source image as the initial input and define the total number of decomposition levels n.
  2. For each level i from 0 to n-1:
    • Convolve the current input using a scaling function, where adjacent pixels are spaced 2i units apart. This produces a smoothed approximation Ci.
    • Calculate the wavelet layer Wi by subtracting the smoothed result from the input: Wi = Input - Ci.
    • Update the Input to equal Ci for the next iteration.
  3. After generating n wavelet layers, the remaining Input serves as the final residual layer.

Reconstruction is achieved by summing all wavelet layers and the residual layer. Optional bias multipliers can be applied to individual layers during recombination to enhance or suppress specific features.

Scaling Functions

Scaling functions act as convolution kernels tailored to represent specific scale data effectively:

  • B3-spline: A highly smooth kernel suited for isolating large-scale structures.
  • Small-scale: A sharp kernel optimized for fine details.
  • Linear interpolation: A balanced kernel appropriate when both small and large structures require processing.

Convolution Pixel Spacing and Boundaries

At scale i, the distance between the target pixel and its neighbors is 2i. For a central pixel, the algorithm skips 2i - 1 pixels in every direction, creating the characteristic holes.

When the sampling coordinates fall outside the image boundaries, a mirroring technique is used to reflect the image logically at the edges. This avoids creating physical expanded copies of the buffer.

Maximum Scale Calculation

The maximum number of decomposition levels is bounded by the image dimensions. It is determined by the base-2 logarithm of the smaller dimension, discarding the decimal portion: floor(log2(min(width, height))). Exceeding this limit would require sampling beyond the mirrored boundaries, rendering convolution impossible.

Rust Implementation

To implement this transform, we define a struct to hold the image state and iteration metadata. The ndarray crate handles multidimensional array operations, while image manages file I/O and pixel conversions.

// wavelet.rs
use ndarray::Array2;
use image::GenericImageView;

pub struct WaveletDecomposer {
    source: Array2<f32>,
    max_levels: usize,
    active_level: usize,
    cols: usize,
    rows: usize,
}

impl WaveletDecomposer {
    pub fn build(img: &image::DynamicImage, levels: usize) -> Self {
        let (w, h) = img.dimensions();
        let (cols, rows) = (w as usize, h as usize);
        let mut pixel_grid = Array2::<f32>::zeros((rows, cols));

        for (x, y, px) in img.to_luma32f().enumerate_pixels() {
            pixel_grid[[y as usize, x as usize]] = px.0[0];
        }

        Self {
            source: pixel_grid,
            max_levels: levels,
            active_level: 0,
            cols,
            rows,
        }
    }
}

Scaling Kernel

// scaling.rs
#[derive(Clone, Copy)]
pub struct InterpolationScalingKernel {
    matrix: [[f32; 3]; 3]
}

impl Default for InterpolationScalingKernel {
    fn default() -> Self {
        Self {
            matrix: [
                [1.0/16.0, 2.0/16.0, 1.0/16.0],
                [2.0/16.0, 4.0/16.0, 2.0/16.0],
                [1.0/16.0, 2.0/16.0, 1.0/16.0],
            ]
        }
    }
}

Convolution Logic

A trait defines the convolution interface, allowing for future extension to different color models. The implementation handles the pixel spacing and mirroring boundary conditions mathematically.

// convolve.rs
pub trait SpatialConvolution {
    fn find_neighbor_coords(&self, gap: usize, kernel_offset: [isize; 2], center: [usize; 2]) -> [usize; 2];
    fn calculate_convolved_value(&self, gap: usize, pos: [usize; 2]) -> f32;
}

impl SpatialConvolution for WaveletDecomposer {
    fn find_neighbor_coords(&self, gap: usize, kernel_offset: [isize; 2], center: [usize; 2]) -> [usize; 2] {
        let [kx, ky] = kernel_offset;
        let offset_x = kx * gap as isize;
        let offset_y = ky * gap as isize;
        let [cx, cy] = center;

        let mut nx = cx as isize + offset_x;
        let mut ny = cy as isize + offset_y;

        // Mirroring boundary logic
        if nx < 0 { nx = -nx; }
        if ny < 0 { ny = -ny; }
        if nx >= self.cols as isize { nx = 2 * (self.cols as isize - 1) - nx; }
        if ny >= self.rows as isize { ny = 2 * (self.rows as isize - 1) - ny; }

        [ny as usize, nx as usize]
    }

    fn calculate_convolved_value(&self, gap: usize, pos: [usize; 2]) -> f32 {
        let mut accum = 0.0;
        let kernel = InterpolationScalingKernel::default();

        for dx in -1..=1 {
            for dy in -1..=1 {
                let coord = self.find_neighbor_coords(gap, [dx, dy], pos);
                let weight = kernel.matrix[(dy + 1) as usize][(dx + 1) as usize];
                accum += weight * self.source[coord];
            }
        }
        accum
    }
}

Iterator Implementation

Implementing the Iterator trait provides a natural interface for generating layers, enabling standard functional combinators like map and filter.

// wavelet.rs
impl Iterator for WaveletDecomposer {
    type Item = (Array2<f32>, Option<usize>);

    fn next(&mut self) -> Option<Self::Item> {
        let lvl = self.active_level;
        self.active_level += 1;

        if lvl > self.max_levels {
            return None;
        }

        if lvl == self.max_levels {
            return Some((self.source.clone(), None));
        }

        let spacing = 2_usize.pow(lvl as u32);
        let mut blurred = Array2::<f32>::zeros((self.rows, self.cols));

        for x in 0..self.cols {
            for y in 0..self.rows {
                blurred[[y, x]] = self.calculate_convolved_value(spacing, [x, y]);
            }
        }

        let detail_layer = self.source.clone() - &blurred;
        self.source = blurred;

        Some((detail_layer, Some(lvl)))
    }
}

Recomposition

A trait handles the summation of layers and normalization back into a viewable image format.

// recompose.rs
use image::{DynamicImage, ImageBuffer, Luma};
use ndarray::Array2;

pub trait WaveletRecomposer: Iterator<Item = (Array2<f32>, Option<usize>)> {
    fn merge_to_image(self, w: usize, h: usize) -> DynamicImage where Self: Sized {
        let mut composite = Array2::<f32>::zeros((h, w));

        for (layer, _) in self {
            composite += &layer;
        }

        let min_val = composite.iter().cloned().reduce(f32::min).unwrap_or(0.0);
        let max_val = composite.iter().cloned().reduce(f32::max).unwrap_or(1.0);
        let range = max_val - min_val;

        let mut img_buf: ImageBuffer<Luma<u16>, Vec<u16>> = ImageBuffer::new(w as u32, h as u32);

        for (x, y, px) in img_buf.enumerate_pixels_mut() {
            let val = composite[[y as usize, x as usize]];
            let normalized = (val - min_val) / range;
            *px = Luma([(normalized * u16::MAX as f32) as u16]);
        }

        DynamicImage::ImageLuma16(img_buf)
    }
}

impl<I: Iterator<Item = (Array2<f32>, Option<usize>)>> WaveletRecomposer for I {}

Practical Application

By chaining the iterator with a mapping function, small-scale layers (containing noise) can be selectively filtered before recomposition.

// main.rs
use image::{DynamicImage, ImageBuffer, Luma};
use atrous_rs::recompose::WaveletRecomposer;
use atrous_rs::wavelet::WaveletDecomposer;

fn main() {
    let source_img = image::open("noisy_galaxy.tif").expect("Failed to load image");
    let decomposer = WaveletDecomposer::build(&source_img, 9);

    let result = decomposer
        .map(|(mut layer_data, scale)| {
            if scale.map_or(false, |s| s < 2) {
                let mut dynamic = DynamicImage::ImageLuma16(
                    ImageBuffer::from_fn(
                        layer_data.ncols() as u32,
                        layer_data.nrows() as u32,
                        |x, y| {
                            let v = (layer_data[[y as usize, x as usize]] * u16::MAX as f32) as u16;
                            Luma([v])
                        }
                    )
                ).to_luma8();

                dynamic = imageproc::filter::bilateral_filter(&dynamic, 10, 10.0, 3.0);

                for (x, y, px) in dynamic.enumerate_pixels() {
                    layer_data[[y as usize, x as usize]] = px.0[0] as f32 / u8::MAX as f32;
                }
            }
            (layer_data, scale)
        })
        .merge_to_image(source_img.width() as usize, source_img.height() as usize)
        .to_luma8();

    result.save("denoised_output.png").unwrap();
}
Tags: Rust

Related Articles

Understanding Strong and Weak References in Java

Strong References Strong reference are the most prevalent type of object referencing in Java. When an object has a strong reference pointing to it, the garbage collector will not reclaim its memory. F...

Comprehensive Guide to SSTI Explained with Payload Bypass Techniques

Introduction Server-Side Template Injection (SSTI) is a vulnerability in web applications where user input is improper handled within the template engine and executed on the server. This exploit can r...

Implement Image Upload Functionality for Django Integrated TinyMCE Editor

Django’s Admin panel is highly user-friendly, and pairing it with TinyMCE, an effective rich text editor, simplifies content management significantly. Combining the two is particular useful for bloggi...

Leave a Comment

Anonymous

◎Feel free to join the discussion and share your thoughts.