utsav-akhaury / PnP-ADMM

Image Deconvolution with Plug-and-Play ADMM using an X-Dense-U-Net denoiser
10 stars 2 forks source link

Forward-backward test #5

Closed utsav-akhaury closed 3 years ago

utsav-akhaury commented 3 years ago

To investigate the problems arising from the ADMM + U-Net Plug&Play (see issue #4), I implemented a simple forward-backward splitting algorithm with a single gradient descent step instead of ADMM.

The update equations and derivations are as follows -

Screenshot 2021-03-29 at 14 48 47 (2)

For calculating the standard deviation (STD) for all cases, a patch of size 20 X 20 was extracted from the images such that they exclude the region containing the central blob (and only contain the noise map). The following results were obtained after 1 gradient descent step (without denoising) -

image

Subsequently, the input & output to the U-Net must be scaled according to the theoretical value to achieve proper results.

utsav-akhaury commented 3 years ago

Instead of the U-Net, I tried using the Starlet Transform as a proximal operator for denoising (code). The results are shown in this notebook.

By deconvolving by the telescope PSF, we obtain an image with an infinite resolution. However we don't have ground truth images with infinite resolution. The only ground truths we have are images with HST resolution. In practice, we need to divide the CFHT PSF by the HST PSF in the Fourier space to go from CFHT to HST. As a result, the deconvolution with the partial PSF (which is smaller than the CFHT one) makes the problem less ill-posed and gives better results.

The Starlet Transform code was modified to include the PSF contribution (by convolving the starlet scales with the HT operator - the PSF rotated by 180o).

When coupled with the forward-backward algorithm, the results show a negative dip around the central blob, which keeps on increasing in magnitude with the no. of iterations. image

The negative blob was eliminated by projecting all negative values in the final reconstruction to zero. The galaxy seems to be better reconstructed, although some artefacts/residuals still appear. However, with a simple forward-backward algorithm involving a single gradient-descent step, the reconstruction is not expected to be very accurate. Our aim was just to check whether it performs well when coupled with starlet denoising instead of a U-Net. image (1)

utsav-akhaury commented 3 years ago

After making the following modifications, the Forward Backward + Starlet code seems to converge -

  1. Removing all complex intermediate variables and only dealing with real numbers
  2. Computing the Spectral Norm of the gradient Hᵀ.H (instead of H)
  3. Dividing the gradient step size computed using the Spectral Norm by a factor of 5
  4. Using Soft thresholding instead of Hard thresholding in Starlet transform
  5. Adding the contribution of the Starlet term to the cost function

Without dividing the gradient step size by any factor, the cost function plot is less smooth (which is expected), although there is convergence. image

utsav-akhaury commented 3 years ago

Proper convergence is achieved for both real and parametric data when the U-Net is trained on a range of sigma values (instead of a constant value).