Wavelets decomposition for destriping

[patquem]

Following a discussion with @sklumpe during the JOSS review of the pystack3d paper (openjournals/joss-reviews#7079 (comment)), wavelet decomposition could be a fast and promising alternative to the pyvsnr algorithm for destriping.

Below is the code to compare the two approaches.
Feel free to share your feedback.

import numpy as np
import matplotlib.pyplot as plt
import pywt as pywt
from skimage.morphology import erosion
from skimage.filters import gaussian
from skimage.data import camera
from pyvsnr.utils import stripes_addition
from pyvsnr import vsnr2d
import time


def remove_stripes_from_wavelets(img, wavelet='coif3', level=8, sigma=6,
                                 sigma_blur=35, offset_blur=100, nite=3):
    """

    Parameters
    ----------
    img: numpy.ndarray
        Raw image to handle
    wavelet: Wavelet object or name string, or 2-tuple of wavelets
        Wavelet to use.  This can also be a tuple containing a wavelet to
        apply along each axis in ``axes``.
    level: int, optional
        Decomposition level (must be >= 0). If level is None (default) then it
        will be calculated using the ``dwt_max_level`` function.
    sigma: float, optional
        Standard deviation to be considered in the gaussian filter applied
        during the wavelets decomposition

    Returns
    -------
    img_r: numpy.ndarray
        Output image with the removed stripes
    """
    coeffs2 = pywt.wavedec2(img, wavelet, level=level)
    new_coeffs = [coeffs2[0]]

    for i in range(1, len(coeffs2)):
        Ch, Cv, Cd = coeffs2[i]
        fCv = np.fft.fft(Cv, axis=0)
        fCv2 = np.fft.fftshift(fCv, axes=[0])
        my = fCv2.shape[0]
        x = np.arange(-my // 2, my // 2)
        gauss1d = 1 - np.exp(-x ** 2 / (2 * sigma ** 2))
        fCv3 = fCv2 * gauss1d[:, np.newaxis]
        Cv2 = np.fft.ifftshift(fCv3, axes=[0])
        Cv3 = np.fft.ifft(Cv2, axis=0)
        new_coeffs.append((Ch, Cv3.real, Cd))

    img_r = pywt.waverec2(new_coeffs, wavelet)

    # img0 = img_r.copy()
    # img_r = gaussian(img_r, sigma_blur) - offset_blur
    # for _ in range(1, nite):
    #     img_r = erosion(img_r)
    #
    # img_r = img0 - img_r.astype(float)

    return img_r


def run():
    amplitude = 0.2
    filters = [{'name': 'Gabor', 'noise_level': 100, 'sigma': (1000, 0.1),
                'theta': 90}]

    # Original image
    img0 = camera() / 255
    # Original image + Stripes
    img = stripes_addition(img0.T, amplitude, seed=0, norm=False).T

    nite = 20
    t0 = time.time()
    for _ in range(nite):
        img_wavelets = remove_stripes_from_wavelets(img, sigma=2)
    t1 = time.time()
    for _ in range(nite):
        img_vsnr = vsnr2d(img, filters=filters, algo='numpy')
    t2 = time.time()
    for _ in range(nite):
        img_vsnr = vsnr2d(img, filters=filters, algo='cupy')
    t3 = time.time()

    _, ax = plt.subplots(1, 4, figsize=(16, 6))
    ax[0].imshow(img0)
    ax[0].set_title("Original image\n")
    ax[1].imshow(img)
    ax[1].set_title(f"Original image + Stripes\n")
    ax[2].imshow(img_wavelets)
    ax[2].set_title(f"Wavelets destriping\n"
                    f"CPU={t1 - t0:.2f}s")
    ax[3].imshow(img_vsnr)
    ax[3].set_title(f"VSNR destriping\n"
                    f"CPU={t2 - t1:.1f}/{t3 - t2:.1f}s (numpy/cupy)")

    plt.show()

run()

[Obi-Wan]

I used to have an implementation of this algorithm as well, but with a couple of interesting changes: stationary wavelets instead of decimated wavelets, padding of the different levels to avoid Fourier aliasing, and a parameter to select the stripe direction.
It is meant for sinograms (thus the vertical axis is called angle_axis), but maybe you want to have look at it:
https://github.com/cicwi/PyCorrectedEmissionCT/blob/53ffcc5f4947f9a3eb21551e073bdd5afb431377/corrct/processing/pre.py#L290

I would normally suggest sticking to 1 or 2 (max) decomposition levels because you would otherwise compromise features, bias, and scaling of the image. However, this might depend on the application.