Denoise noisy digits and make predictions using an autoencoder made in Flax

5 min readFeb 7, 2024

In my most previous blog post I discussed how to denoise an image, and that post can be read here:-

In this post I intend to denoise images to use them in an autoencoder written in Flax and Jax.

Autoencoders are trained on encoding input data such as images into a smaller feature vector, and afterward, reconstruct it by a second neural network, called a decoder. An block diagram of an autoencoder can be seen below:-

In order to create the autoencoder in flax, I chose a noisy digit dataset from Kaggle, and it can be accessed here:-

I have written the program in Kaggle’s free online Jupyter Notebook and stored it in my account for the data science company.

Once I created the Jupyter Notebook and defined the problem statement, I imported the libraries that I would need to execute the program, being:-

  1. Numpy to create numpy arrays and carry out numerical computations,
  2. Pandas to create dataframes and process data, os to go into the operating system of the computer,
  3. Pathlib to create a path that the images can be located in,
  4. Matplotlib to visualise the images,
  5. Seaborn to aid in the visualisation of the images,
  6. PIL to process images,
  7. Cv2 to carry out computer vision functionality,
  8. Skimage to aid in processing images,
  9. Jax to create jax arrays and carry out numerical computations,
  10. flax , which is written on top of jax and creates the neural network,
  11. Optax supplies the optimizer in the neural network, and
  12. Sklearn to provide machine learning functionality.




I have close to five decades experience in the world of work, being in fast food, the military, business, non-profits, and the healthcare sector.