eboigne / PyTV-4D

Python routines to compute the Total Variation (TV) of 2D, 3D and 4D images on CPU & GPU. Compatible with proximal algorithms (ADMM, Chambolle & Pock, ...)
GNU General Public License v3.0
31 stars 6 forks source link

Differences between PyTV-4D and skimage #2

Closed zxdawn closed 11 months ago

zxdawn commented 11 months ago

As mentioned in #1, PyTV-4D doesn't work for not square data. So, I would use skimage for 2D data. Before switching to that, I checked the output of these two packages:

import pytv
import numpy as np
import matplotlib.pyplot as plt

noise_level = 100
nb_it = 300
regularization = 25
step_size = 5e-3

np.random.seed(0)
cameraman_truth = pytv.utils.cameraman()
cameraman_truth = np.reshape(cameraman_truth, (1,1,)+cameraman_truth.shape)
cameraman_noisy = cameraman_truth + noise_level * np.random.rand(*cameraman_truth.shape) # Add noise
cameraman_estimate = np.copy(cameraman_noisy)

# ---- pytv-4d ---- #
loss_fct_GD = np.zeros([nb_it,])
for it in range(nb_it): # A simple sub-gradient descent algorithm for image denoising
    tv, G = pytv.tv_CPU.tv_hybrid(cameraman_estimate)
    cameraman_estimate += - step_size * ((cameraman_estimate - cameraman_noisy) + regularization * G)
    loss_fct_GD[it] = 0.5 * np.sum(np.square(cameraman_estimate - cameraman_noisy)) + regularization * tv

# ---- skimage ---- #
from skimage.restoration import denoise_tv_chambolle

cameraman_denoise = denoise_tv_chambolle(cameraman_noisy, weight=regularization, max_num_iter=nb_it)

# ---- plt ---- # 
fig, axs = plt.subplots(ncols=3, figsize=(9, 3))
axs = axs.flat

axs[0].imshow(cameraman_noisy.squeeze(), cmap = plt.get_cmap('gray'))
axs[1].imshow(cameraman_estimate.squeeze(), cmap = plt.get_cmap('gray'))
axs[2].imshow(cameraman_denoise.squeeze(), cmap = plt.get_cmap('gray'))

axs[0].set_title('Nosy image')
axs[1].set_title('Denosied image (pytv-4d)')
axs[2].set_title('Denosied image (skimage)')

image

It looks like the results are different. It would be nice to know how to reproduce Py4D-TV with skimage for a 2D image.

eboigne commented 11 months ago

I like the suggestion, thanks for sharing. This would be great to have, and if you have the time to help develop this feature, I can assist in pushing this to PyTV-4D.

To get the two implementations to retrieve the exact same result, note that you'd need to match both:

  1. The optimization algorithm: right now, you're comparing gradient descent from pytv with an algorithm by Chambolle (2004) from skimage. You'd want to use the Chambolle (2004) algorithm with pytv instead of gradient descent to get the same result.

Note that there's a snippet of code available here to use pytv with a proximal algorithm by Chambolle & Pock (2011), which may be helpful in writing this, although it is not the same algorithm. Alternatively, skimage also has a split-Bregman implementation which is another route to consider.

  1. The discretization: you're using the hybrid discretization from pytv. It seems the equivalent implementation from the Chambolle (2004) algorithm would be closest to the upwind discretization instead: pytv.tv_CPU.tv_upwind. There may be other differences to identify, I haven't gone in the paper in depth.
eboigne commented 11 months ago

A quick follow-up: the code below will get you 90% there.

I made a copy of the scikit-image implementation of the Chambolle (2004) algorithm, and modified it so it uses the PyTV-4D routines:

def denoise_tv_chambolle_2D_matching_scikit(image, weight=0.1, eps=2.e-4, max_num_iter=200, D = pytv.tv_operators_CPU.D_upwind, D_T = pytv.tv_operators_CPU.D_T_upwind):
    """A 2D implementation of the Chambolle (2004) algorithm using PyTV-4D routines, that matches the function _denoise_tv_chambolle_nd from scikit-image
    Refer to: https://github.com/scikit-image/scikit-image/blob/v0.22.0/skimage/restoration/_denoise.py#L367

    Parameters
    ----------
    image : ndarray
        n-D input data to be denoised.
    weight : float, optional
        Denoising weight. The greater `weight`, the more denoising (at
        the expense of fidelity to `input`).
    eps : float, optional
        Relative difference of the value of the cost function that determines
        the stop criterion. The algorithm stops when:

            (E_(n-1) - E_n) < eps * E_0

    max_num_iter : int, optional
        Maximal number of iterations used for the optimization.

    Returns
    -------
    out : ndarray
        Denoised array of floats.

    Notes
    -----
    Rudin, Osher and Fatemi algorithm.
    """

    image = image.squeeze()
    ndim = 2

    p = np.zeros((image.ndim, ) + image.shape, dtype=image.dtype)
    g = np.zeros_like(p)
    d = np.zeros_like(image)
    i = 0
    while i < max_num_iter:
        if i > 0:
            d = D_T(np.transpose(np.reshape(p, (1,1,)+p.shape), [0,2,1,3,4])).squeeze()
            out = image + d
        else:
            out = image
        E = (d ** 2).sum()

        g = D(np.reshape(out, (1,1,)+out.shape)).squeeze()
        # In 2D with upwind scheme, this line is equivalent to:
        # g[0,:-1,:] = np.diff(out, axis = 0)
        # g[1,:,:-1] = np.diff(out, axis = 1)

        norm = np.sqrt((g ** 2).sum(axis=0))[np.newaxis, ...]
        E += weight * norm.sum()
        tau = 1. / (2.*ndim)
        norm *= tau / weight
        norm += 1.
        p -= tau * g
        p /= norm
        E /= float(image.size)
        if i == 0:
            E_init = E
            E_previous = E
        else:
            if np.abs(E_previous - E) < eps * E_init:
                break
            else:
                E_previous = E
        i += 1
    return out

Then I used it to compare with the original implementation from Scikit-image:

import pytv
import numpy as np
import matplotlib.pyplot as plt
from skimage.restoration import denoise_tv_chambolle

noise_level = 100
max_num_iter = 300
regularization = 50

np.random.seed(0)
cameraman_truth = pytv.utils.cameraman()
cameraman_truth = np.reshape(cameraman_truth, (1,1,)+cameraman_truth.shape)
cameraman_noisy = cameraman_truth + noise_level * np.random.rand(*cameraman_truth.shape) # Add noise
cameraman_estimate = np.copy(cameraman_noisy)

cameraman_chambolle2004_pytv_upwind = denoise_tv_chambolle_2D_matching_scikit(cameraman_noisy.squeeze(), weight=regularization, eps=2.e-4, max_num_iter=max_num_iter)
cameraman_chambolle2004_scikit = denoise_tv_chambolle(cameraman_noisy.squeeze(), weight=regularization, eps=2.e-4, max_num_iter=max_num_iter)

max_diff = np.nanmax(np.abs(cameraman_chambolle2004_scikit-cameraman_chambolle2004_pytv_upwind))
print(max_diff)

# ---- plt ---- # 
fig, axs = plt.subplots(ncols=3, figsize=(9, 3))
axs = axs.flat

axs[0].imshow(cameraman_noisy.squeeze(), cmap = plt.get_cmap('gray'))
axs[1].imshow(cameraman_chambolle2004_pytv_upwind.squeeze(), cmap = plt.get_cmap('gray'))
axs[2].imshow(cameraman_chambolle2004_scikit.squeeze(), cmap = plt.get_cmap('gray'))

axs[0].set_title('Noisy image')
axs[1].set_title('Denoised image (pytv-4d)')
axs[2].set_title('Denoised image (skimage)')

output

This function computes the gradients exactly as is done in scikit-image. There's still very small differences in the divergence calculations, i.e. the line d = D_T(np.transpose(np.reshape(p, (1,1,)+p.shape), [0,2,1,3,4])).squeeze() is not exactly equivalent to what's in the scikit function, but within a very small error it's correct. I didn't bother fixing the indices to get this 100% right at this point. Visually, you won't notice any differences in the images at this level.

Final note: these are only equal with the upwind scheme, as that's what's used in the Chambolle (2004) algorithm.

zxdawn commented 11 months ago

Thanks for the feedback, Emeric. Because I'm not familiar with the TV filter, I don't know how to apply it quickly. Nice to see you have done that!

May I ask what's the benefit of using tv_hybrid compared to the Chambolle (2004) algorithm from scikit-image?

eboigne commented 11 months ago

I have found that you get less patchy and smoother images with the hybrid scheme. It's second-order accurate whereas upwind/downwind are only first-order accurate, so numerically it's a more precise discretization of the continuous TV function.

The image you put above illustrates that nicely. In my research, I had similar improvements with it which is why I added this feature.

There's a slight additional computational/storage cost to using the hybrid scheme, which may be relevant if you run large cases.