Diffusion models for denoising astronomical images

This repository contains the code of the main project of my Honor Program in Data Science (a.y. 2023/24) I collaborated with Professor Pierpaolo Brutti and PhD student Riccardo Ceccaroni to explore the application of Diffusion Models for Denoising Astronomical Images. This research project was in collaboration with Astronomical Observatory of Rome (National Institute of Astrophysics) and involved the use advanced generative models to address the problem of noise in astronomical data, which is critical for a number of tasks such as image analysis and reconstruction.

Table of Contents

1. Project Overview
2. Brief Description
3. Disclaimer
4. Prerequisites
   • Setting Up the Environment
   • Clone the Repository
   • Install Anaconda Project
   • Create the Environment
5. Running the Project
   • Training Script
   • Denoising Script
   • Plotting Script
6. Final Results
7. Used Technologies

Brief Description

The focus of this project was the use of Denoising Diffusion Probabilistic Models (DDPMs) to restore noisy images by employing stochastic processes and Evidence Lower-Bound (ELBO) method. This approach has gained prominence in image restoration due to its ability to sample from the posterior distribution of images, producing reconstructions with high perceptual quality. Our work drew heavily on the methodology described in recent literature, such as Kawar et al., 2021, who introduced Denoising Diffusion Restoration Models (DDRM). This method efficiently addresses linear inverse problems and o!ers significant improvements in reconstruction quality and runtime over other unsupervised methods. We also assumed stochastic denoising techniques inspired by Kawar et al., 2020, which employ a combination of KL-diveregence based loss and Langevin dynamics, to produce high quality outputs from a so-called Matching Score Probabilistic Model receiving noisy input images.

Disclaimer

The reported code repository is built on the implementations of diffusion models from the official pytorch repository and the code from the work "Radio-astronomical Image Reconstruction with Conditional Denoising Diffusion Model"

Prerequisites

• Conda installed on your machine.
• Recommended: Miniconda or Anaconda distribution.

Setting Up the Environment

To set up the project environment and install all necessary dependencies, follow these steps:

1. Clone the repository

First, clone the repository to your local machine:

git clone https://github.com/Engrima18/DiffusionDenoising
cd DiffusionDenoising

2. Install the required dependencies and create the honor conda environment.

The file env.yml is used to manage the environment setup.

conda conda env create -f env.yml

3. Create the environment

Activate the environment.

conda activate honor

Running the Project

There are several commands available in this project, including running different scripts like training, denoising, and plotting.

1. Training Script

To run the train.py script and train a Denoising Diffusion model with all the specifications from the configuration file (you choose the dataset).

python train.py

2. Denoising Script

To run the denoising.py script and starting the denoising process starting from certain input noisy images. In practice with this command you will use some diffusion model checkpoint to generate new images starting from noisy test images.

python denoising.py

3. Plotting script

To run the visual.py script and plot some results from the above denoising/generation process.

python visual.py

Final Results

To validate our approach, we trained a diffusion model on several datasets, including the ILLUSTRIS dataset. This allowed us to simulate various noise scenarios and test the robustness of our model across different astronomical environments.

Here we report a sequence of images demonstrating the denoising performances of our best model trained for 20 epochs on some images cropped from ILLUSTRIS: left the noisy input image, centre the target ground-truth image, right the model output image.

Used Technologies