2D TV Denoising

Import libraries

import numpy as np
import cv2
from matplotlib import pyplot as plt
import scipy.signal as signal
from scipy.signal import convolve2d
import scipy.fft as fft
import urllib.request
from skimage.metrics import peak_signal_noise_ratio as PSNR
from IPython.display import Image, HTML
from matplotlib.animation import FuncAnimation
import time

Import image

# Reading image (grayscale)
url = "https://i.stack.imgur.com/kP0u2.png"
with urllib.request.urlopen(url) as url_response:
    img_array = np.asarray(bytearray(url_response.read()), dtype=np.uint8)
    img = cv2.imdecode(img_array, cv2.IMREAD_GRAYSCALE)
# img is a 3-dimensional numpy array (third number indicates channel)
# Converting to (0,1)
x = img.astype(float) / 255.0
print(type(img))
print(img.shape)
plt.imshow(x, cmap="gray")

<class 'numpy.ndarray'>
(512, 512)

<matplotlib.image.AxesImage at 0x7f38c8a9a730>

Define conv and fft functions

# Define some of the operators that we need...
def conv2d_fft(x, h):
    p0 = x.shape[0] - h.shape[0]
    p1 = x.shape[1] - h.shape[1]
    h_pad = np.pad(h, ((0, p0), (0, p1)))
    Fh = fft.fft2(h_pad)
    Fx = fft.fft2(x)
    return np.real(fft.ifft2(Fx * Fh))


def conv2dT_fft(x, h):
    p0 = x.shape[0] - h.shape[0]
    p1 = x.shape[1] - h.shape[1]
    h_pad = np.pad(h, ((0, p0), (0, p1)))
    Fh = fft.fft2(h_pad)
    Fx = fft.fft2(x)
    return np.real(fft.ifft2(Fx * np.conj(Fh)))

Noise function

def awgn(img, n):
    """Generating Gaussian Noise
    with 0 mean and standard deviation n
    choose n between 0,1 for normalized image"""

    noise = np.random.randn(*img.shape) * n
    # Add the noise to the input image
    noisy_image = img + noise

    return noisy_image

Add noise to the image

"""
y1 = img + n
"""
# Add noise to the image
y1 = awgn(x, 0.2)

fig = plt.figure(figsize=(15, 15))
plt.subplot(121)
plt.imshow(x, cmap="gray", clim=[0, 1])
plt.title("image x")
plt.subplot(122)
plt.imshow(y1, cmap="gray", clim=[0, 1])
plt.title("Noisy image y = img + n")
plt.tight_layout()
plt.show()

Gradiant operator

# define gradient operators
"""
I am testing it with y1 for now
"""
dh = np.array([[1, -1], [0, 0]])  # horizontal gradient filter
dv = np.array([[1, 0], [-1, 0]])  # vertical gradient filter

Dh = lambda x: conv2d_fft(x, dh)
Dv = lambda x: conv2d_fft(x, dv)

DhT = lambda x: conv2dT_fft(x, dh)
DvT = lambda x: conv2dT_fft(x, dv)

# plot the image x and the gradient images Dh x and Dv x
fig = plt.figure(figsize=(15, 15))
plt.subplot(131)
plt.imshow(x, cmap="gray", clim=[0, 1])
plt.title("image x")
plt.subplot(132)
plt.imshow(np.abs(Dh(x)), cmap="gray", clim=[0, 1])
plt.title(r"$|D_hx|$")
plt.subplot(133)
plt.imshow(np.abs(Dv(x)), cmap="gray", clim=[0, 1])
plt.title(r"$|D_vx|$")
plt.tight_layout()
plt.show()

TV-denoising Solver

Write TV-denoising formulation with explanation here later

def TV_denoising(y, lamb=2, rho=1e1, maxiter=200, return_history=False):
    """
    TV-denoising solver to solve
    minimize 0.5 |x-y|_2^2 + lambda|Dx|_1
    """
    # define the soft-thresholding function
    """
    In the TVD case we have : 
    Vector (v): Dx_(k+1) + u_(k)
    Threshold (t): lamb/rho
    """
    start = time.time()
    soft_thresh = lambda v, t: np.maximum(np.abs(v) - t, 0.0) * np.sign(v)

    # DDT
    """
        DDT = please check notes for fourier transform format
    """
    # Calculating the difference between sizes x and d for padding purpose
    p0 = x.shape[0] - dh.shape[0]
    p1 = x.shape[1] - dh.shape[1]
    dh_pad = np.pad(dh, ((0, p0), (0, p1)))

    p0 = x.shape[0] - dv.shape[0]
    p1 = x.shape[1] - dv.shape[1]
    dv_pad = np.pad(dv, ((0, p0), (0, p1)))

    # Refer to Parisima's notes for computing DDT using FFT
    DDT = np.abs(fft.fft2(dh_pad)) ** 2 + np.abs(fft.fft2(dv_pad)) ** 2

    # -----------------------------
    # initilize iteration variables
    zh = np.zeros_like(y)
    zv = np.zeros_like(y)
    uh = np.zeros_like(zh)
    uv = np.zeros_like(zv)
    x_hat = np.zeros_like(y)
    # For computing error
    J = np.zeros(maxiter)
    history = []

    for k in range(maxiter):
        # solve the L2-L2 problem (update x)
        rhs = y + rho * ((DhT(zh) + DvT(zv)) - DhT(uh) - DvT(uv))
        F_rhs = fft.fft2(rhs)
        x_hat = np.real(fft.ifft2(F_rhs / (rho * DDT + 1)))

        # solve the TV problem (update z)
        zh = soft_thresh(Dh(x_hat) + uh, lamb / rho)
        zv = soft_thresh(Dv(x_hat) + uv, lamb / rho)

        # update u
        dual_h = Dh(x_hat) - zh
        dual_v = Dv(x_hat) - zv
        uh = uh + dual_h
        uv = uv + dual_v

        # compute the error
        J[k] = (dual_h**2).sum() + (dual_v**2).sum()
        history.append(x_hat.copy())
    end = time.time()
    if return_history:
        return x_hat, J, end - start, history
    return x_hat, J, end - start

x_hat, J, duration, history = TV_denoising(
    y1, lamb=0.2, rho=2, maxiter=50, return_history=True
)

fig = plt.figure()
fig.set_size_inches(5, 7.5)
ax = fig.add_subplot(321)
ax.imshow(x, cmap="gray", clim=[0, 1])
plt.title("original image")
ax2 = fig.add_subplot(322)
ax2.imshow(y1, cmap="gray", clim=[0, 1])
plt.title("noisy image")
ax3 = fig.add_subplot(323)
ax3.imshow(x_hat, cmap="gray", clim=[0, 1])
plt.title("recovered image")
ax4 = fig.add_subplot(324)
ax4.semilogy(range(len(J)), J, "b-", lw=2)
plt.title("Convergence")
plt.xlabel("iteration (k)")
plt.tight_layout()
plt.show()
print(f"Time taken = {duration}")

Time taken = 4.432166576385498

TVD using L2 regularizer animated

fig, axs = plt.subplots(1, 2, figsize=(10, 4), dpi=80)
axis_img = axs[0].imshow(y1, cmap="gray", clim=[0.0, 1.0])

(line,) = axs[1].semilogy(range(len(J)), J)
axs[1].set_xlabel("Iteration")
axs[1].set_ylabel("Convergence")
axs[1].set_title("Convergence vs Iteration")
plt.close()


def animate(i):
    axs[0].set_title(f"Recovered image at iteration {i}")
    axis_img.set(cmap="gray", clim=[0.0, 1.0])
    axis_img.set_data(history[i])
    line.set_data(np.array(range(i + 1)), J[: (i + 1)])
    return axis_img, line


animation = FuncAnimation(fig, animate, frames=50, interval=120, blit=True)
# let animation load
time.sleep(1)
plt.show();

display(HTML(f'<div style="text-align:center;">{animation.to_html5_video()}</div>'))