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

• 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.

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.

Raindrop Removal (click to expand)

Image Deraining & image Dehazing (click to expand)

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 

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

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

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}
}

If you have any questions, please contact [email protected]

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