# Image restoration of Anonymous NeurIPS submission:  Improving Diffusion-based Inverse Algorithms under Few Steps Constraints via Linear Extrapolation (Paper ID: 15962)

This project is based on:

\- https://github.com/bahjat-kawar/ddrm (DDRM),

\- https://github.com/wyhuai/DDNM (DDNM), and

\- https://github.com/DPS2022/diffusion-posterior-sampling (DPS)

## Environment

You can set up the environment using the `environment.yml` (the requirement is the same as 

[DDRM]: https://github.com/bahjat-kawar/ddrm

). Run

```bash
conda env create -f environment.yml
conda activate LLE
```

## Experiments in the paper

### Pre-trained models

We use pre-trained models on CelebA-HQ, and FFHQ. Please download the pre-trained models from https://github.com/ermongroup/SDEdit for CelebA-HQ (https://drive.google.com/file/d/1wSoA5fm_d6JBZk4RZ1SzWLMgev4WqH21/view?usp=share_link), and from https://github.com/DPS2022/diffusion-posterior-sampling  for FFHQ (https://drive.google.com/drive/folders/1jElnRoFv7b31fG0v6pTSQkelbSX3xGZh). Place them into exp/logs/[dataset_name] respectively.

### Test datasets

Download the CelebA-HQ test set from https://github.com/wyhuai/DDNM (https://drive.google.com/drive/folders/1cSCTaBtnL7OIKXT4SVME88Vtk4uDd_u4) and place it into exp/datasets/celeba. Download the FFHQ test set from https://github.com/NVlabs/ffhq-dataset and place it into exp/datasets/ffhq (Note: you only need to download the folder 00000 for testing). Thus the exp/ folder should look as follows:

```bash
exp
├── logs
├── datasets
│   ├── celeba # all CelebA files
│   ├── ffhq # out of distribution ImageNet images
```

### Reproduce the results

Please run the following code:

```
python main_inverse.py --ni --config {CONFIG}.yml --doc {DATASET} --algo {ALGO} --timesteps {STEPS} --deg {DEGRADATION} --sigma_0 {SIGMA_0} -i {IMAGE_FOLDER} --learned
```


```bash
CUDA_VISIBLE_DEVICES=0 python main_inverse.py --ni --config celeba_hq.yml --doc celeba --algo ddnm --timesteps 3 --deg sr4 --sigma_0 0.00 -i celeba/sr4_noiseless/ddnm/3steps --learned

CUDA_VISIBLE_DEVICES=0 python main_inverse.py --ni --config ffhq.yml --doc ffhq --algo ddnm --timesteps 3 --deg inpainting --sigma_0 0.00 -i ffhq/inpainting_noiseless/ddnm/3steps --learned
```

### Calculate the metrics

We provide the example code for calculating the metrics in `calculate_metrics/`. Please first copy all the images using `mv_files.py` (mainly for FID calculation), and then run `cal_metrics.py` to calculate the metrics. Note that you may need to adjust the paths in the files.

