Image Wavelet Denoising Tool

function processImage(originalImg, threshold = 20, waveletType = 'haar') {

    // --- Start Helper Functions ---

    /**
     * Performs a 2D forward Haar wavelet transform on a matrix.
     * The dimensions of the matrix must be powers of 2.
     * This is an in-place transform.
     * @param {number[][]} matrix - The 2D array of data (e.g., a color channel).
     * @param {number} width - The width of the matrix.
     * @param {number} height - The height of the matrix.
     * @returns {number[][]} The transformed matrix.
     */
    const fwt2d = (matrix, width, height) => {
        let w = width, h = height;
        const temp = new Float64Array(Math.max(w, h));

        // The transform is applied level by level to the approximation sub-band (top-left).
        while (w > 1 || h > 1) {
            let halfW = Math.max(1, Math.floor(w / 2));
            let halfH = Math.max(1, Math.floor(h / 2));

            // Transform rows
            if (w > 1) {
                for (let i = 0; i < h; i++) {
                    for (let j = 0; j < halfW; j++) {
                        const a = matrix[i][2 * j];
                        const b = matrix[i][2 * j + 1];
                        temp[j] = (a + b) / 2;
                        temp[j + halfW] = (a - b) / 2;
                    }
                    for (let j = 0; j < w; j++) matrix[i][j] = temp[j];
                }
            }

            // Transform columns
            if (h > 1) {
                for (let i = 0; i < w; i++) {
                    for (let j = 0; j < halfH; j++) {
                        const a = matrix[2 * j][i];
                        const b = matrix[2 * j + 1][i];
                        temp[j] = (a + b) / 2;
                        temp[j + halfH] = (a - b) / 2;
                    }
                    for (let j = 0; j < h; j++) matrix[j][i] = temp[j];
                }
            }
            w = halfW;
            h = halfH;
        }
        return matrix;
    };

    /**
     * Performs a 2D inverse Haar wavelet transform.
     * Reverses the process of fwt2d. This is an in-place transform.
     * @param {number[][]} matrix - Transformed 2D wavelet coefficients.
     * @param {number} width - The original width of the matrix.
     * @param {number} height - The original height of the matrix.
     * @returns {number[][]} The reconstructed matrix.
     */
    const iwt2d = (matrix, width, height) => {
        // Determine the size of the smallest approximation sub-band
        let maxLevels = Math.floor(Math.log2(Math.min(width, height)));
        let w = width / Math.pow(2, maxLevels);
        let h = height / Math.pow(2, maxLevels);

        const temp = new Float64Array(Math.max(width, height));

        // Iteratively apply the inverse transform, doubling the size at each level
        while (w <= width && h <= height) {
            let nextW = w, nextH = h;
            if (w < width) nextW = w * 2;
            if (h < height) nextH = h * 2;
            if (nextW === w && nextH === h) break;

            // Inverse transform columns first
            if (nextH > h) {
                for (let i = 0; i < w; i++) {
                    for (let j = 0; j < h; j++) {
                        const avg = matrix[j][i];
                        const diff = matrix[j + h][i];
                        temp[2 * j] = avg + diff;
                        temp[2 * j + 1] = avg - diff;
                    }
                    for (let j = 0; j < nextH; j++) matrix[j][i] = temp[j];
                }
            }

            // Inverse transform rows
            if (nextW > w) {
                for (let i = 0; i < nextH; i++) { // Use nextH here
                    for (let j = 0; j < w; j++) {
                        const avg = matrix[i][j];
                        const diff = matrix[i][j + w];
                        temp[2 * j] = avg + diff;
                        temp[2 * j + 1] = avg - diff;
                    }
                    for (let j = 0; j < nextW; j++) matrix[i][j] = temp[j];
                }
            }
            w = nextW;
            h = nextH;
        }
        return matrix;
    };

    /**
     * Applies soft thresholding to a value.
     * @param {number} value - The input coefficient.
     * @param {number} thresh - The threshold.
     * @returns {number} The thresholded value.
     */
    const softThreshold = (value, thresh) => {
        return Math.sign(value) * Math.max(0, Math.abs(value) - thresh);
    };

    /**
     * Applies thresholding to the detail coefficients of a transformed matrix.
     * It avoids thresholding the final approximation sub-band (top-left corner).
     * @param {number[][]} coeffs - The matrix of wavelet coefficients.
     * @param {number} thresh - The threshold value.
     * @param {number} width - The width of the coefficient matrix.
     * @param {number} height - The height of the coefficient matrix.
     */
    const thresholdCoefficients = (coeffs, thresh, width, height) => {
        let maxLevels = Math.floor(Math.log2(Math.min(width, height)));
        let approxW = width / Math.pow(2, maxLevels);
        let approxH = height / Math.pow(2, maxLevels);

        for (let y = 0; y < height; y++) {
            for (let x = 0; x < width; x++) {
                // If the coefficient is NOT in the top-leftmost approximation band, threshold it
                if (x >= approxW || y >= approxH) {
                    coeffs[y][x] = softThreshold(coeffs[y][x], thresh);
                }
            }
        }
    };
    
    /**
     * Denoises a single color channel using the wavelet transform.
     * @param {Float64Array[]} channelData - 2D array for the color channel.
     * @param {number} threshold - The denoising threshold.
     * @param {string} waveletType - The type of wavelet (only 'haar' supported).
     * @param {number} width - The padded width.
     * @param {number} height - The padded height.
     * @returns {number[][]} The denoised 2D channel data.
     */
    const denoiseChannel = (channelData, threshold, waveletType, width, height) => {
        if (waveletType !== 'haar') {
            console.warn(`Wavelet type '${waveletType}' is not supported. Using 'haar'.`);
        }
        
        // Deep copy of the channel data to prevent modifying the original during in-place transforms
        const dataCopy = channelData.map(row => new Float64Array(row));

        const transformedData = fwt2d(dataCopy, width, height);
        thresholdCoefficients(transformedData, threshold, width, height);
        const denoisedData = iwt2d(transformedData, width, height);

        return denoisedData;
    };


    // --- End Helper Functions ---


    // 1. SETUP: PAD IMAGE TO POWER-OF-2 DIMENSIONS
    const originalWidth = originalImg.width;
    const originalHeight = originalImg.height;

    const paddedWidth = 1 << Math.ceil(Math.log2(originalWidth));
    const paddedHeight = 1 << Math.ceil(Math.log2(originalHeight));
    
    const padCanvas = document.createElement('canvas');
    padCanvas.width = paddedWidth;
    padCanvas.height = paddedHeight;
    const padCtx = padCanvas.getContext('2d', { willReadFrequently: true });
    padCtx.drawImage(originalImg, 0, 0, originalWidth, originalHeight);
    
    const imageData = padCtx.getImageData(0, 0, paddedWidth, paddedHeight);
    const data = imageData.data;

    // 2. SEPARATE CHANNELS
    // Initialize 2D arrays for R, G, B channels
    const R = [], G = [], B = [];
    for (let i = 0; i < paddedHeight; i++) {
        R[i] = new Float64Array(paddedWidth);
        G[i] = new Float64Array(paddedWidth);
        B[i] = new Float64Array(paddedWidth);
    }
    // Populate the channel arrays
    for (let y = 0; y < paddedHeight; y++) {
        for (let x = 0; x < paddedWidth; x++) {
            const i = (y * paddedWidth + x) * 4;
            R[y][x] = data[i];
            G[y][x] = data[i + 1];
            B[y][x] = data[i + 2];
        }
    }
    
    // 3. PROCESS EACH CHANNEL
    const denoisedR = denoiseChannel(R, threshold, waveletType, paddedWidth, paddedHeight);
    const denoisedG = denoiseChannel(G, threshold, waveletType, paddedWidth, paddedHeight);
    const denoisedB = denoiseChannel(B, threshold, waveletType, paddedWidth, paddedHeight);
    
    // 4. RECONSTRUCT IMAGE FROM DENOISED CHANNELS
    const newImageData = padCtx.createImageData(paddedWidth, paddedHeight);
    const newData = newImageData.data;

    for (let y = 0; y < paddedHeight; y++) {
        for (let x = 0; x < paddedWidth; x++) {
            const i = (y * paddedWidth + x) * 4;
            // Clamp values to the valid 0-255 range and reconstruct pixel
            newData[i] = Math.max(0, Math.min(255, denoisedR[y][x]));
            newData[i + 1] = Math.max(0, Math.min(255, denoisedG[y][x]));
            newData[i + 2] = Math.max(0, Math.min(255, denoisedB[y][x]));
            newData[i + 3] = data[i + 3]; // Preserve original alpha channel
        }
    }
    
    // Put the denoised (but still padded) data back onto the padded canvas
    padCtx.putImageData(newImageData, 0, 0);

    // 5. CROP TO ORIGINAL DIMENSIONS AND RETURN
    const finalCanvas = document.createElement('canvas');
    finalCanvas.width = originalWidth;
    finalCanvas.height = originalHeight;
    const finalCtx = finalCanvas.getContext('2d');
    finalCtx.drawImage(padCanvas, 0, 0, originalWidth, originalHeight, 0, 0, originalWidth, originalHeight);

    return finalCanvas;
}