Images denoising

Hi all,

I have a question about pytorch and image denoising.

I’d like to code an image denoising. The problem is that we usually take the original images and add the noise.

Here’s an example:


.....
for i in range(50):

    total_loss = 0.0
    total_iter = 0

    myModel.train()
    for image, Nolabel in train_loader:
    
        noise = torch.randn(image.shape[0], 3, 128, 128) * noise_level
        image_n = torch.add(image, noise)
    
        image = image.cuda()
        image_n = image_n.cuda()

        optimizer.zero_grad()
        output = myModel(image_n)
    loss = loss_fn(output, image)
    .....

I don’t want to use that method. I have two folders, the first folder contains the images without noise (jpg format), the second folder contains the same images but with noise (jpg format). The two folders has the same number of images. How can I change the previous example to use both folders?

Hi,

You need to create your own custom Dataset class and override __getitem__ method in the way that it returns two set of images.

Here is a snippet:

import torch
from torch.utils import data

class Dataset(data.Dataset):
    """
    Return Dataset class representing our data set
    """
    def __int__(self, path_inputs, path_labels):
        """
        Initialize data set as a list of paths corresponding to each item of data set and labels of each data

        Args:
            path_inputs: a list of paths for each data point in data set
            path_labels: label of an item in data set 
        """

        self.path_labels = path_labels  # list of paths to original images
        self.path_inputs= path_inputs  # list of paths to images with noise

    def __len__(self):
        """
        Return the length of data set using list of paths

        :return: number of samples in data set
        """
        return len(self.path_inputs)  # or path_labels

    def __getitem__(self, index):
        """
        Generate one item of data set. 

        :param item: index of item in paths list

        :return: a sample of data
        """

        # Code to load data using Image.open or cv2.imread
        X = cv2.imread(self.path_inputs[index])  # images without noise
        y = cv2.imread(self.path_labels[index])  # images with noise

        return X, y

Bests

Hi,

@Nikronic thank you very much for your answer.

Can you please give me an example on how to call your code in the example code to better understand?

It is just like the code you have provided. You use this custom dataset instead of previously constructed one and pass it to the Dataloader. Then in a loop, get batches of data which in this case would be two batch of input and noisy images.

train_dataset = Dataset(path_inputs='path_to_original_images',
                              img_dir='path_to_noisy_images/')

train_loader = DataLoader(dataset=train_dataset,
                          batch_size=128,
                          shuffle=True,
                          num_workers=4)

for origin, noisy in data_loader:
    x = origin  # [batch, channel, height, width]
    y = noisy  # [batch, channel, height, width]
    ## do your stuff

Note that I did not provided the code of loading images from paths or even how to get the path, that is actually more python rather than pytorch. You can use glob and PIL or ‘Open CV’ for loading images.

1 Like

@Nikronic thanks a lot for your precious help.