Awesome

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

Pre-trained Models

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.

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.