GridFormer: Residual Dense Transformer with Grid Structure for Image Restoration in Adverse Weather Conditions (IJCV 2024)

Tao Wang, Kaihao Zhang, Ziqin Shao, Wenhan Luo, Bjorn Stenger, Tong Lu, Tae-Kyun Kim, Wei Liu, Hongdong Li

News

• Jun 24, 2024: Pre-trained models are released!
• Jun 24, 2024: Codes is released!

This repository contains the dataset, code and pre-trained models for our paper.

---

Abstract: Image restoration in adverse weather conditions is a difficult task in computer vision. In this paper, we propose a novel transformer-based framework called GridFormer which serves as a backbone for image restoration under adverse weather conditions. GridFormer is designed in a grid structure using a residual dense transformer block, and it introduces two core designs. First, it uses an enhanced attention mechanism in the transformer layer. The mechanism includes stages of the sampler and compact self-attention to improve efficiency, and a local enhancement stage to strengthen local information. Second, we introduce a residual dense transformer block (RDTB) as the final GridFormer layer. This design further improves the network’s ability to learn effective features from both preceding and current local features. The GridFormer framework achieves state-of-the-art results on five diverse image restoration tasks in adverse weather conditions, including image deraining, dehazing, deraining & dehazing, desnowing, and multi-weather restoration.

---

Network Architecture

Datasets

Task	Dataset	Links
Image Dehazing	ITS	Baidu cloud plate
Image Dehazing	SOTS-indoor	Baidu cloud plate
GridFormer-Haze4K	Haze4K	Baidu cloud plate
GridFormer-Snow100K	Snow100k	Baidu cloud plate
GridFormer-RainDrop	RainDrop	Baidu cloud plate
GridFormer-Outdoor-Rain	Outdoor-Rain	Baidu cloud plate
GridFormer-Multi-weather-Restoration	Mix	Baidu cloud plate

Pre-trained Models

Name	Dataset	Pre-trained models	Configs
GridFormer-SOTS-Indoor	SOTS-Indoor	gdrive | Baidu cloud plate	train | test
GridFormer-Haze4K	Haze4K	gdrive | Baidu cloud plate	train | test
GridFormer-Snow100K	Snow100k	gdrive | Baidu cloud plate	train | test
GridFormer-RainDrop	RainDrop	gdrive | Baidu cloud plate	train | test
GridFormer-Outdoor-Rain	Outdoor-Rain	gdrive | Baidu cloud plate	train | test
GridFormer-Multi-Weather	Mix	gdrive | Baidu cloud plate	train | test

Results

Experiments are performed for image restoration in adverser weather cnditions tasks including, image dehazing, image desnowing, Raindrop removal, image deraining & image dehazing, and multi-weather restoration.

Image Dehazing (click to expand)

Image Desnowing (click to expand)

Raindrop Removal (click to expand)

Image Deraining & image Dehazing (click to expand)

Get Started

Dependencies and Installation

1. Create Conda Environment

conda create -n GridFormer python=3.7
conda activate GridFormer
conda install pytorch=1.8 torchvision cudatoolkit=10.2 -c pytorch
pip install matplotlib scikit-learn scikit-image opencv-python yacs joblib natsort h5py tqdm
pip install einops gdown addict future lmdb numpy pyyaml requests scipy tb-nightly yapf lpips

2. Clone Repo

git clone https://github.com/TaoWangzj/GridFormer.git

3. Install basicsr

cd GridFormer
python setup.py develop 

Train (we use the dehazing for an example)

1. To download training and testing data

2. To train GridFormer, run

cd GridFormer
./train.sh options/train/GridFormer/Dehazing/SOTS-Indoor/GridFormer_SOTS-Indoor.yml

Note: The above training script uses 8 GPUs by default. To use any other number of GPUs, modify GridFormer/train.sh and GridFormer_SOTS-Indoor.yml

Test

1. Download the pre-trained model and place it in ./checkpoints/

2. Testing

python basic/test.py -opt  options/test/GridFormer/Dehazing/SOTS-Indoor/GridFormer_SOTS-Indoor.yml

Citation

If you use GridFormer, please consider citing:

@article{wang2024gridformer,
    title={Gridformer: Residual dense transformer with grid structure for image restoration in adverse weather conditions}, 
    author={Gridformer: Residual dense transformer with grid structure for image restoration in adverse weather conditions},
    journal={International Journal of Computer Vision},
    pages={1--23},
    year={2024}
}

Contact

If you have any questions, please contact taowangzj@gamil.com

Acknowledgment: This code is based on the BasicSR toolbox and Restormer.

---

statistics