Removing Noise from Images using a CNN model in Pytorch — Part 2

Photo by SCREEN POST on Unsplash

This is a part 2; as such, it is required that you have gone through part 1 before you proceed. If you haven’t, here’s a link

Removing Noise from Images using a CNN model in Pytorch — Part 1

Now that you have gone through part 1 and you understand the approach we’ll be taking, it’s time to get our hands dirty with some code.

Import Packages

We’ll start by importing the required packages

Model Definition

Now, let’s define our model’s class.

Helper Function

We’ll need a function to help us read through the image folder and ensure that we only load images into our dataset.

Create Dataset

We need to create a dataset, so we can load our data.

Preprocess Image

Putting all the functions we have created together, we can now preprocess our data (images)

Gaussian Noise

To train our model, we will be adding gaussian noise to our images

Train Function

As we prepare to train our model, let’s define the train function

Validate function

We’ll need to test our model at intervals, so let’s create a validation function

Save checkpoint

We’ll also save our model

Let’s train

Finally, let’s train our model

Test Function

Now that we have trained our model, let’s prepare to try it out on an image. Here are the functions we’ll be using:

It’s time to put to use the above functions

Here’s what I got after 10 epochs

Output image

For better result, one would have to train for a longer period of time.

Viola! That’s it.