Denoise image using total variationΒΆ

../_images/sphx_glr_total_variation_denoising_001.png

Out:

Loss in iteration 0 of 500: 3.074
/home/docs/checkouts/readthedocs.org/user_builds/kornia/envs/latest/lib/python3.6/site-packages/torch/autograd/__init__.py:132: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  allow_unreachable=True)  # allow_unreachable flag
Loss in iteration 25 of 500: 2.923
Loss in iteration 50 of 500: 2.717
Loss in iteration 75 of 500: 2.526
Loss in iteration 100 of 500: 2.352
Loss in iteration 125 of 500: 2.196
Loss in iteration 150 of 500: 2.056
Loss in iteration 175 of 500: 1.931
Loss in iteration 200 of 500: 1.820
Loss in iteration 225 of 500: 1.721
Loss in iteration 250 of 500: 1.633
Loss in iteration 275 of 500: 1.556
Loss in iteration 300 of 500: 1.488
Loss in iteration 325 of 500: 1.428
Loss in iteration 350 of 500: 1.376
Loss in iteration 375 of 500: 1.329
Loss in iteration 400 of 500: 1.288
Loss in iteration 425 of 500: 1.252
Loss in iteration 450 of 500: 1.219
Loss in iteration 475 of 500: 1.191

import torch
import kornia
import cv2
import numpy as np

import matplotlib.pyplot as plt

# read the image with OpenCV
img: np.ndarray = cv2.imread('./data/doraemon.png')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
img = img + np.random.normal(loc=0.0, scale=0.1, size=img.shape)
img = np.clip(img, 0.0, 1.0)

# convert to torch tensor
noisy_image: torch.tensor = kornia.image_to_tensor(img).squeeze()  # CxHxW


# define the total variation denoising network
class TVDenoise(torch.nn.Module):
    def __init__(self, noisy_image):
        super(TVDenoise, self).__init__()
        self.l2_term = torch.nn.MSELoss(reduction='mean')
        self.regularization_term = kornia.losses.TotalVariation()
        # create the variable which will be optimized to produce the noise free image
        self.clean_image = torch.nn.Parameter(data=noisy_image.clone(), requires_grad=True)
        self.noisy_image = noisy_image

    def forward(self):
        return self.l2_term(self.clean_image, self.noisy_image) + 0.0001 * self.regularization_term(self.clean_image)

    def get_clean_image(self):
        return self.clean_image

tv_denoiser = TVDenoise(noisy_image)

# define the optimizer to optimize the 1 parameter of tv_denoiser
optimizer = torch.optim.SGD(tv_denoiser.parameters(), lr=0.1, momentum=0.9)

# run the optimization loop
num_iters = 500
for i in range(num_iters):
    optimizer.zero_grad()
    loss = tv_denoiser()
    if i % 25 == 0:
        print("Loss in iteration {} of {}: {:.3f}".format(i, num_iters, loss.item()))
    loss.backward()
    optimizer.step()

# convert back to numpy
img_clean: np.ndarray = kornia.tensor_to_image(tv_denoiser.get_clean_image())

# Create the plot
fig, axs = plt.subplots(1, 2, figsize=(16, 10))
axs = axs.ravel()

axs[0].axis('off')
axs[0].set_title('Noisy image')
axs[0].imshow(img)

axs[1].axis('off')
axs[1].set_title('Cleaned image')
axs[1].imshow(img_clean)

plt.show()

Total running time of the script: ( 0 minutes 3.837 seconds)

Gallery generated by Sphinx-Gallery