ProRes 🌆

Exploring Degradation-aware Visual Prompt for Universal Image Restoration

Jiaqi Ma^1,✢, Tianheng Cheng^2,✢ Guoli Wang³ Xinggang Wang², Qian Zhang³, Lefei Zhang^1,📧

¹School of Computer Science, Wuhan University
² School of EIC, Huazhong University of Science & Technology
³ Horizon Robotics

(✢) Equal contribution. (📧) corresponding author.

Updates

This project is under active development, please stay tuned! ☕

June 26, 2023: We've released the arXiv paper of ProRes! Code & models are coming soon!

Introduction

Image restoration aims to reconstruct degraded images, e.g., denoising or deblurring. Existing works focus on designing task-specific methods and there are inadequate attempts at universal methods. However, simply unifying multiple tasks into one universal architecture suffers from uncontrollable and undesired predictions. To address those issues, we explore prompt learning in universal architectures for image restoration tasks.

In this paper, we present Degradation-aware Visual Prompts, which encode various types of image degradation, e.g., noise and blur, into unified visual prompts. These degradation-aware prompts provide control over image processing and allow weighted combinations for customized image restoration. We then leverage degradation-aware visual Prompts to establish a controllable and universal model for image Restoration, called ProRes, which is applicable to an extensive range of image restoration tasks. ProRes leverages the vanilla Vision Transformer (ViT) without any task-specific designs. Furthermore, the pre-trained ProRes can easily adapt to new tasks through efficient prompt tuning with only a few images. Without bells and whistles, ProRes achieves competitive performance compared to task-specific methods and experiments can demonstrate its ability for controllable restoration and adaptation for new tasks.

Highlights

• ProRes addresses universal image restoration with degradation-aware prompts, which is the first prompt-based versatile framework for image restoration.

• ProRes demonstrate two superior capabilities: (1) control ability for desired outputs and (2) transferability based on prompt tuning.

• ProRes can be easily adapted for new tasks or new datasets through effective and efficient prompt tuning.

• Specific prompts can control the output of ProRes. Moreover, combining different prompts can tackle the images with multiple corruptions.

Experimental Results

Performance on Image Restorations Tasks

	denoising		deraining		enhance		deblurring
	SIDD		5 datasets		LoL		4 datasets
	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
Task-specific models
Uformer	39.89	0.960	-	-	-	-	32.31	0.941
MPRNet	39.71	0.958	32.73	0.921	-	-	33.67	0.948
MIRNet-v2	39.84	0.959	-	-	24.74	0.851	-	-
Restormer	40.02	0.960	33.96	0.935	-	-	32.32	0.935
MAXIM	39.96	0.960	33.24	0.933	23.43	0.863	34.50	0.954
Universal models
Painter	38.88	0.954	29.49	0.868	22.40	0.872	-	-
ViT-Large	39.28	0.967	30.75	0.893	21.69	0.850	20.57	0.680
ProRes	39.28	0.967	30.67	0.891	22.73	0.877	28.03	0.897

Notes:

• The works we has use for reference including Uformer(paper,code), MPRNet(paper,code), MIRNet-v2(paper,code), Restormer(paper,code), MAXIM(paper,code) and Painter(paper,code).
• For both Painter and ProRes, we adopt ViT-Large with MAE pre-trained weights.
• More experimental results are listed in the paper!

---

Visualizations

Control Ability

1.Independent Control

Visualization results processed from images of different corruptions. Compared with the original inputs, the outputs are consistent with the given visual prompts.

2. Sensitive to Irrelevant Task-specific Prompts

Visualization results processed by different prompts. Compared with the original inputs, the outputs remain unchanged with irrelevant visual prompts.

3. Tackle Complicated Corruptions

Visualization results processed by ProRes from images of mixed types of degradation, i.e., low-light and rainy. ProRes adopts two visual prompts for low-light enhancement (E) and deraining (D) and combines the two visual prompts by linear weighted sum, i.e., αD + (1 − α)E, to control the restoration process.

Adaptation on New Datasets & Task

1. Low-light Enhancement Results

Visualization results of ProRes on the FiveK dataset. We adopt two settings, i.e., direct inference and prompt tuning, to evaluate ProRes on the FiveK dataset (a new dataset for low-light enhancement).

2. Dehazing Results

Visualization results of ProRes on the RESIDE-6K dataset via prompt tuning for image dehazing (a new task).

Getting Started

Project Structure

coming soon!

---

Requirement

• Linux, CUDA>=9.2, GCC>=5.4
• PyTorch >= 1.8.1
• MATLAB for evaluation
• Other requirements

  pip install -r requirements.txt

Data Preparation

Download Path

• Download the denoising dataset from SIDD.
• Download the low-light enhancement dataset from LoL.
• Download the deraining dataset from Synthetic Rain Datasets.
• Download the deblurring dataset from Synthetic Blur Datasets.

Preprocess Data

Run the following commands to generate corresponding JSON files for each dataset.

#denoising
python data/sidd/gen_json_sidd.py --split train
python data/sidd/gen_json_sidd.py --split val

# low-light enhancement
python data/lol/gen_json_lol.py --split train
python data/lol/gen_json_lol.py --split val

# deraining
python data/derain/gen_json_rain.py --split train
python data/derain/gen_json_rain.py --split val

# derblurring
python data/derain/gen_json_blur.py --split train
python data/derain/gen_json_blur.py --split val

Dataset Structure

We recommend the dataset directory structure to be the following:

$ProRes/datasets/
    denoise/
        train/
        val/
    enhance/
        our485/
            low/
            high/
        eval15/
            low/
            high/
    derain/
        train/
            input/
            target/
        test/
            Rain100H/
            Rain100L/
            Test100/
            Test1200/
            Test2800/
    deblur/
        train/
            input/
            target/
        test/
            GoPro/
            HIDE/
            RealBlur_J/
            RealBlur_R/
    json/
        train_denoise.json
        val_denoise.json
        train_enhance.json
        val_enhance.json
        train_derain.json
        val_derain.json
        train_deblur.json
        val_deblur.json

Training

coming soon!

Evaluation

coming soon!

Citation

If you find our paper and code useful for your research, please consider giving a star ⭐ and citation 📝 :

@article{
      title={ProRes: Exploring Degradation-aware Visual Prompt for Universal Image Restoration}, 
      author={Jiaqi Ma and Tianheng Cheng and Guoli Wang and Xinggang Wang and Qian Zhang and Lefei Zhang},
      journal={arXiv preprint arXiv:2306.13653},
      year={2023}
}

Acknowledgement

This project is based on MAE, BEiT, MIRNet, MPRNet, Uformer and Painter. Thanks for their wonderful work!