denoising.py

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray
from scipy.signal import convolve2d
from curvelet import CurveletCoefficients, fdct_wrapping, ifdct_wrapping


def act_filter(
    noisy_img: NDArray[np.float64],
    noise_var: float | NDArray[np.float64] | None = None,
    threshold_setting: str = "s",
) -> NDArray[np.float64]:
    """
    Denoise an image using adaptive curvelet thresholding (ACT).

    Parameters
    ----------
    noisy_img:
        Noisy input image.
    noise_var:
        Either a scalar noise variance for AWGN or a 2-D FFT-PSD array for
        stationary colored noise. When omitted, the noise variance is
        estimated internally from the highest-frequency curvelet band.
    threshold_setting:
        One of ``"s"``, ``"h"``, or ``"ksigma"``.

    Returns
    -------
    numpy.ndarray
        Denoised image.
        
    Reference
    ----------
    N. Eslahi and A. Aghagolzadeh, "Compressive Sensing Image Restoration
    Using Adaptive Curvelet Thresholding and Nonlocal sparse Regularization,"
    IEEE Trans. Image Process., vol.25, no.7, pp. 3126-3140, Jul. 2016
    https://doi.org/10.1109/TIP.2016.2562563        
    """

    noisy_img = np.asarray(noisy_img, dtype=np.float64)
    threshold_setting = (threshold_setting or "s").lower()
    if threshold_setting not in {"s", "h", "ksigma"}:
        raise ValueError(
            f"Unknown threshold setting {threshold_setting!r}. Expected 's', 'h', or 'ksigma'."
        )

    estimate_noise_var = noise_var is None or (
        isinstance(noise_var, np.ndarray) and noise_var.size == 0
    )
    ## applying forward curvelet transform to the noisy image
    noisy_dcut_coeffs = fdct_wrapping(noisy_img, is_real=True)

    if estimate_noise_var: ## blind AWGN denoising
        ## Estimating noise std by applying the MAD over the high-passed
        ## noisy image to discard the influence of outliers.
        highest_freq_coeff = noisy_dcut_coeffs[-1][0]
        noise_std = np.median(np.abs(highest_freq_coeff - np.median(highest_freq_coeff))) / 0.6745
        noise_var_array = float(noise_std**2)
    else: ##non-blind denoising
        noise_var_array = np.asarray(noise_var, dtype=np.float64)
        if noise_var_array.size == 1:
            noise_var_array = float(noise_var_array.reshape(()))
    
    ## global FFT-PSD of the noise
    if np.isscalar(noise_var_array):
        ## flat FFT-PSD of white noise
        fft_psd = float(noise_var_array) * np.ones_like(noisy_img) * noisy_img.size
        noise_type = "white" ## uncorrelated (a.k.a. white) Gaussian noise
    else:
        fft_psd = np.asarray(noise_var_array, dtype=np.float64)
        if fft_psd.shape != noisy_img.shape:
            raise ValueError(
                f"noise_var must match the input image shape {noisy_img.shape}, got {fft_psd.shape}."
            )
        if float(np.max(fft_psd) - np.min(fft_psd)) < 0.015 * fft_psd.size:
            noise_type = "white" 
        else:
            noise_type = "colored" ## correlated (a.k.a. colored) Gaussian noise
    
    ## computing the noise root-PSD in DCuT domain (or the std of the noise ...
    ## within each subband) given the noise FFT-PSD (or given the noise ...
    ## variance if it is modeled as AWGN).
    noise_dcut_root_psd = compute_dcut_root_psd(fft_psd)
    ## Adaptive Curvelet Thresholding stage
    denoised_coeffs = adaptive_curvelet_thresholding(
        noisy_dcut_coeffs,
        noise_dcut_root_psd,
        threshold_setting,
        noise_type,
    )
    ## applying inverse curvelet transform to the denoised coefficients
    return np.asarray(ifdct_wrapping(denoised_coeffs, is_real=True), dtype=np.float64)

def adaptive_curvelet_thresholding(
    noisy_dcut_coeffs: CurveletCoefficients,
    noise_dcut_root_psd: list[NDArray[np.float64]],
    threshold_setting: str,
    noise_type: str,
) -> CurveletCoefficients:
    """
    Apply the ACT thresholding rule in the curvelet domain.

    Parameters
    ----------
    noisy_dcut_coeffs:
        Curvelet coefficients of the noisy image.
    noise_dcut_root_psd:
        Noise standard deviation estimate for each curvelet subband.
    threshold_setting:
        ``"s"``, ``"h"``, or ``"ksigma"``.
    noise_type:
        ``"white"`` or ``"colored"``.

    Returns
    -------
    CurveletCoefficients
        Thresholded coefficients.
        
    Reference
    ----------
    N. Eslahi and A. Aghagolzadeh, "Compressive Sensing Image Restoration
    Using Adaptive Curvelet Thresholding and Nonlocal sparse Regularization,"
    IEEE Trans. Image Process., vol.25, no.7, pp. 3126-3140, Jul. 2016
    https://doi.org/10.1109/TIP.2016.2562563
    
    """

    denoised_coeffs: CurveletCoefficients = [
        [np.array(subband, copy=True) for subband in scale] for scale in noisy_dcut_coeffs
    ]
    ## computing the noise DCuT-root-PSD (i.e. the std of noise in each DCuT subband)
    ## through scaling the norm2 of the DCuT basis/frame by the noise std.
    for scale_index in range(1, len(noisy_dcut_coeffs)):
        for angle_index in range(len(noisy_dcut_coeffs[scale_index])):
            subband = np.asarray(noisy_dcut_coeffs[scale_index][angle_index], dtype=np.float64)
            noise_root = float(noise_dcut_root_psd[scale_index][angle_index])

            if threshold_setting in {"s", "h"}:
            ## estimating the sample std of the clean (noise-free) DCuT coefficients
            ## via the maximum likelihood estimator given the noisy DCuT coefficients
            ## and the noise root-PSD in DCuT domain (see Eq.(12) of (Eslahi & Aghagolzadeh, 2016)).
                clean_subband_std = ml_estimator(subband, noise_root, noise_type)
                with np.errstate(divide="ignore", invalid="ignore"):
                    ## computing the adaptive threshold for the selected subband and then using it within the selected filter
                    threshold = np.sqrt(2.0) * (noise_root**2 / clean_subband_std)

                if threshold_setting == "s": ## soft-shrinkage
                    denoised_coeffs[scale_index][angle_index] = np.sign(subband) * np.maximum(
                        np.abs(subband) - threshold,
                        0,
                    )
                else: ## hard-shrinkage
                    threshold = (3.0 + float(scale_index == len(noisy_dcut_coeffs) - 1)) * threshold / np.sqrt(
                        2.0
                    )
                    denoised_coeffs[scale_index][angle_index] = subband * (
                        np.abs(subband) > threshold
                    )
            else: ## k-sigma thresholding (Starck, Candes, Donoho, 2002)
                threshold = (3.0 + float(scale_index == len(noisy_dcut_coeffs) - 1)) * noise_root
                denoised_coeffs[scale_index][angle_index] = subband * (np.abs(subband) > threshold)

    return denoised_coeffs

def ml_estimator(
    noisy_dcut_coeffs: NDArray[np.float64],
    noise_root_psd: float,
    noise_type: str,
) -> NDArray[np.float64]:
    """
    Estimate the local clean-coefficient standard deviation by maximum likelihood (ML).
    More precisely, ``ml_estimator`` estimates the sample std of the clean (noise-free) coefficients
    in the DCuT domain at a specific subband (scale and rotation angle) via the ML estimator given
    the noisy DCuT coefficients and the noise std.

    Parameters
    ----------
    noisy_dcut_coeffs:
        Noisy curvelet coefficients at a single subband.
    noise_root_psd:
        Noise standard deviation of that subband.
    noise_type:
        ``"white"`` or ``"colored"``.

    Returns
    -------
    numpy.ndarray
        Estimated local standard deviation map of the clean coefficients.
    """

    ## convolutional kernel for ML averaging (local averaging window)
    if noise_type.lower() == "white":
        kernel = np.ones((7, 7), dtype=np.float64) / 48.0
        kernel[3, 3] = 0.0
    else:
        kernel = np.ones((31, 31), dtype=np.float64) / 960.0
        kernel[15, 15] = 0.0
    ## estimating sample variance of noisy coefficients as the average of
    ## squared coefficients over a local moving window
    est_sample_var_noisy = convolve2d(
        np.asarray(noisy_dcut_coeffs, dtype=np.float64) ** 2,
        kernel,
        mode="same",
        boundary="wrap",
    )
    ## estimating the sample variance of clean coefficients by subtracting
    ## the noise variance from the estimated sample variance of noisy coefficients    
    est_sample_var_clean = est_sample_var_noisy - noise_root_psd**2
    ## imposing non-negativity (because it is PSD/variance!)
    est_sample_var_clean[est_sample_var_clean < 0] = 0
    ## taking the square-root for sample variance
    return np.sqrt(est_sample_var_clean)

def compute_dcut_root_psd(fft_psd: NDArray[np.float64]) -> list[NDArray[np.float64]]:
    """
    Compute the noise root-PSD in DCuT domain (or the std of noise within 
    each curvelet subband) given the noise FFT-PSD (or given the noise variance
    if noise is modeled as AWGN).


    Parameters
    ----------
    fft_psd:
        Noise FFT-PSD over the image grid.

    Returns
    -------
    list[numpy.ndarray]
        The noise root-PSD in DCuT domain or the std of the noise within each subband.
    """

    ## approximating the convolutional kernel of the noise
    kernel_noise = np.fft.fftshift(np.fft.ifft2(np.sqrt(np.asarray(fft_psd, dtype=np.float64))))
    ## DCuT frames for all subbands (overcomplete representation)
    dcut_frames = fdct_wrapping(kernel_noise, is_real=True)
    ## computing the root-PSD in DCuT domain
    root_psd: list[NDArray[np.float64]] = []
    for scale in dcut_frames:
        values = [
            float(np.sqrt(np.mean(np.square(np.real_if_close(np.asarray(subband)))))) for subband in scale
        ]
        root_psd.append(np.asarray(values, dtype=np.float64))
    return root_psd