Denoising images with deep learning

Recently I’ve started looking into computer vision problems, and image denoising/image reconstruction is one of pretty funny problems: it has more than one viable solution. So, I’ve decided to take a stab on it, and build basic image denoiser with neural network.

Luckily, I know that U-net architectures typically serve good for this kind of tasks. U-net idea is pretty close to autoencoders: you have feature extraction part of neural network, followed by reconstruction part. And since convolution/deconvolution layers are used all the way through the network — output pixels are not independent of each others.

This is how neural network is going to look like:

Principal schema for the neural network

The project itself will consist of 3 major parts that will be implemented:

  • Build data pipeline for training and validation
  • Design and train neural network
  • Test results visually

So, let’s start.

Data pipeline:

Neural network I’m going to create will accept arbitrary sized noisy image as input, and will return denoised image of the same size as output. This means that I’ll need bunch of source images and I’ll need to apply some noise to these images, and that’s about it, since for this kind of tasks original (unmodified) image serves as target image.

For my own convenience, I’ve organized on-the-fly data pipeline: one big image is loaded, partitioned into bunch of 64x64 (or 128x128, or whatever) images and fed into neural network. This process will be repeated for all images I have. On-the-fly pipeline simplifies experiments: different image sizes, different strides for partitioning, different noise levels, and so on.

Here’s going this pipeline will look in Python:

Train/Validation sets generators

Noise is added on the fly as well: certain percentage of pixels will be filled with random “dark pixels” with size varying between 1x1 up to 2x2. Definitely not a 100% natural noise, but good enough noise for 1-evening toy problem.

Adding some noise to images

Now, when I have 2-in-1 generator compatible with TensorFlow Datasets I can easily set it up, and use it for actual training:

Setting up generator and TensorFlow Datasets

With all data pipeline bits uncovered, it’s time to look into actual neural network.

Neural network architecture:

As I’ve mentioned above, neural network for this kind of tasks is typically U-net architecture. Something one could call “Convolutional autoencoder” — I have 2-branch structure here: X branch for bigger feature maps, Y branch for smaller feature maps. Eventually they’re combined and used for final image reconstruction.

As you can see, some arguments reference params dictionary, I’ve used this for hyperparams tuning. Here’s the example of params:

model = build_model({'lr': 0.001,
'conv_activation': 'swish',
'last_activation': 'tanh',
'out_activation': 'sigmoid',
'loss': 'mean_squared_error'})

Now, with all bits set — let's hit Run button, and approximately 4714 seconds later (using RTX 2080 Ti) one will see that epoch is over.

76200/76200 — 4714s 62ms/step — loss: 5.0034e-04 — val_loss: 2.3591e-04 

Trying the model:

For quick experiments with trained model I’ve created minimalistic Jupyter notebook (available in the GitHub repo, link at the bottom). Jupyter is pretty good for such things: one can just edit 1 line in 1 cell, and re-run this particular cell alone — and that’s exactly what I need here for “visual validation”.

Since the model is fully-convolutional, we can feed different shapes into it. However, for me personally, it’s interesting to see how network handles different amounts of applied noise.

64x64 RGB, 10%, 20% and 30% noise respectively

I think It’s safe to say that results are not ideal (since we can see artifacts even on 10% noise level), but on the other hand — hey, it’s 100k params network, capable of denoising 256 image patches in 10–15 milliseconds! So, with some additional tuning applied, we could probably get it down to 30-50k params and 3–7 milliseconds.

Thanks for reading. Feel free to reach me out, if you have any questions :)

Link to the repo with project source code:

Deep Learning Developer.