Image Super-Resolution via Iterative Refinement

Paper | Project

Brief

Code Submission for UCLA CS245: Big Data Analytics course offered in Fall 2023

This is an unofficial implementation of Image Super-Resolution via Iterative Refinement(SR3) by PyTorch.

The other branches contain updated code for each specific experiments (zero_snr, splitImage, lazcos, bilinear_unet, and area-bicubic-upsampling-unet). This master branch only contains the unedited version of the unofficial code for the SR3 Paper (https://github.com/Janspiry/Image-Super-Resolution-via-Iterative-Refinement)

Status

★★★ NEW: The follow-up Palette-Image-to-Image-Diffusion-Models is now available; See the details here ★★★

Conditional Generation (with Super Resolution)

• 16×16 -> 128×128 on FFHQ-CelebaHQ
• 64×64 -> 512×512 on FFHQ-CelebaHQ

Unconditional Generation

• 128×128 face generation on FFHQ
• 1024×1024 face generation by a cascade of 3 models

Training Step

• log / logger
• metrics evaluation
• multi-gpu support
• resume training / pretrained model
• validate alone script
• Weights and Biases Logging 🌟 NEW

Results

Note: We set the maximum reverse steps budget to $2000$. We limited the model parameters in image noise and hue deviation occasionally appear in high-resolution images, resulting in low scores. There is a lot of room for optimization.

Tasks/Metrics	SSIM(+)	PSNR(+)	FID(-)	IS(+)
16×16 -> 128×128	0.675	23.26	-	-
64×64 -> 512×512	0.445	19.87	-	-
128×128	-	-
1024×1024	-	-

• 16×16 -> 128×128 on FFHQ-CelebaHQ [More Results]

• 64×64 -> 512×512 on FFHQ-CelebaHQ [More Results]

• 128×128 face generation on FFHQ [More Results]

Usage

Environment

pip install -r requirement.txt

Pretrained Model

This paper is based on "Denoising Diffusion Probabilistic Models", and we build both DDPM/SR3 network structures, which use timesteps/gamma as model embedding inputs, respectively. In our experiments, the SR3 model can achieve better visual results with the same reverse steps and learning rate. You can select the JSON files with annotated suffix names to train the different models.

Tasks	Platform（Code：qwer)
16×16 -> 128×128 on FFHQ-CelebaHQ	Google Drive|Baidu Yun
64×64 -> 512×512 on FFHQ-CelebaHQ	Google Drive|Baidu Yun
128×128 face generation on FFHQ	Google Drive|Baidu Yun

# Download the pretrained model and edit [sr|sample]_[ddpm|sr3]_[resolution option].json about "resume_state":
"resume_state": [your pretrained model's path]

Data Prepare

New Start

If you didn't have the data, you can prepare it by following steps:

• FFHQ 128×128 | FFHQ 512×512
• CelebaHQ 256×256 | CelebaMask-HQ 1024×1024

Download the dataset and prepare it in LMDB or PNG format using script.

# Resize to get 16×16 LR_IMGS and 128×128 HR_IMGS, then prepare 128×128 Fake SR_IMGS by bicubic interpolation
python data/prepare_data.py  --path [dataset root]  --out [output root] --size 16,128 -l

then you need to change the datasets config to your data path and image resolution:

"datasets": {
    "train": {
        "dataroot": "dataset/ffhq_16_128", // [output root] in prepare.py script
        "l_resolution": 16, // low resolution need to super_resolution
        "r_resolution": 128, // high resolution
        "datatype": "lmdb", //lmdb or img, path of img files
    },
    "val": {
        "dataroot": "dataset/celebahq_16_128", // [output root] in prepare.py script
    }
},

Own Data

You also can use your image data by following steps, and we have some examples in dataset folder.

At first, you should organize the images layout like this, this step can be finished by data/prepare_data.py automatically:

# set the high/low resolution images, bicubic interpolation images path 
dataset/celebahq_16_128/
├── hr_128 # it's same with sr_16_128 directory if you don't have ground-truth images.
├── lr_16 # vinilla low resolution images
└── sr_16_128 # images ready to super resolution

# super resolution from 16 to 128
python data/prepare_data.py  --path [dataset root]  --out celebahq --size 16,128 -l

Note: Above script can be used whether you have the vanilla high-resolution images or not.

then you need to change the dataset config to your data path and image resolution:

"datasets": {
    "train|val": { // train and validation part
        "dataroot": "dataset/celebahq_16_128",
        "l_resolution": 16, // low resolution need to super_resolution
        "r_resolution": 128, // high resolution
        "datatype": "img", //lmdb or img, path of img files
    }
},

Training/Resume Training

# Use sr.py and sample.py to train the super resolution task and unconditional generation task, respectively.
# Edit json files to adjust network structure and hyperparameters
python sr.py -p train -c config/sr_sr3.json

Test/Evaluation

# Edit json to add pretrain model path and run the evaluation 
python sr.py -p val -c config/sr_sr3.json

# Quantitative evaluation alone using SSIM/PSNR metrics on given result root
python eval.py -p [result root]

Inference Alone

Set the image path like steps in Own Data, then run the script:

# run the script
python infer.py -c [config file]

Acknowledgements

Our work is based on the following theoretical works:

• Denoising Diffusion Probabilistic Models
• Image Super-Resolution via Iterative Refinement
• WaveGrad: Estimating Gradients for Waveform Generation
• Large Scale GAN Training for High Fidelity Natural Image Synthesis

Furthermore, we are benefitting a lot from the following projects:

• https://github.com/bhushan23/BIG-GAN
• https://github.com/lmnt-com/wavegrad
• https://github.com/rosinality/denoising-diffusion-pytorch
• https://github.com/lucidrains/denoising-diffusion-pytorch
• https://github.com/hejingwenhejingwen/AdaFM