MB-TaylorFormer: Multi-branch Efficient Transformer Expanded by Taylor Formula for Image Dehazing （ICCV 2023）

This repository contains the official implementation of the following paper:

MB-TaylorFormer: Multi-branch Efficient Transformer Expanded by Taylor Formula for Image Dehazing
Yuwei Qiu, Kaihao Zhang, Chenxi Wang, Wenhan Luo, Hongdong Li, Zhi Jin^*
International Conference on Computer Vision (ICCV), 2023
Paper Link: [official link]

Overview

Architecture of MB-TaylorFormer. (a) MB-TaylorFormer consists of a multi-branch hierarchical design based on multi-scale patch embedding. (b) Multi-scale patch embedding embeds coarse-to-fine patches. (c) TaylorFormer with linear computational complexity. (d) MSAR module compensates for errors in Taylor expansion.

Installation

See INSTALL.md for the installation of dependencies required to run MB-TaylorFormer.

Prepare pretrained models

Before performing the following steps, please download our pretrained model first.

Download Links: [Google Drive]

Then, place the models to dehazing/pretrained_models directory.

Run the following command to process them:

CUDA_VISIBLE_DEVICES=X python dehazing/test.py --size ['B' or 'L'] --input_dir [Input path] --target_dir [GT path] --result_dir [Result path] --weights [Model weighting path]

For example:

CUDA_VISIBLE_DEVICES=0 python dehazing/test.py --size 'B'  --input_dir '/data/QYW/ITS_SOTS/test/hazy/' --target_dir '/data/QYW/ITS_SOTS/test/GT/' --result_dir '/data/qiuyuwei/ITS_result' --weights '/home/qiuyuwei/MB-TaylorFormer-main/Dehazing/pretrained_models/ITS-MB-TaylorFormer-B.pth'

Prepare dataset for training and evaluation

The data directory structure will be arranged as:

data
   |- ITS
      |- Train
      	|- Haze
        	|- 1_1_0.90179.png 
        	|- 2_1_0.99082.png
        |- GT
        	|- 1.png 
        	|- 2.png
      |- Test
      	|- Haze
        	|- 00001.png
        	|- 00002.png
        |- GT
        	|- 00001.png
        	|- 00002.png
   |- OTS
      |- Train
      	|- Haze
        	|- 0001_0.85_0.04.jpg
        	|- 0002_0.85_0.04.jpg
        |- GT
        	|- 0001.jpg
        	|- 0002.jpg
      |- Test
      	|- Haze
        	|- 00501.png
        	|- 00502.png
        |- GT
        	|- 00501.png
        	|- 00502.png
   |- Dense-Haze
      |- Train
         |- Haze
            |- 01_hazy.png 
            |- 02_hazy.png
         |- GT
            |- 01_GT.png 
            |- 02_GT.png
      |- Test
         |- Haze
            |- 51_hazy.png 
            |- 52_hazy.png
         |- GT
            |- 51_GT.png 
            |- 52_GT.png

   |- O-HAZE
      |- Train
         |- Haze
            |- 01_outdoor_hazy.jpg
            |- 02_outdoor_hazy.jpg
         |- GT
            |- 01_outdoor_GT.jpg 
            |- 02_outdoor_GT.jpg
      |- Test
         |- Haze
            |- 41_outdoor_haze.jpg
            |- 42_outdoor_haze.jpg
         |- GT
            |- 41_outdoor_GT.jpg
            |- 42_outdoor_GT.jpg

Training

To train MB-TaylorFormer with default settings, run

sh /train.sh Dehazing/Options/MB-TaylorFormer-B.yml 

or

sh /train.sh Dehazing/Options/MB-TaylorFormer-L.yml

Testing

Run the following command to quick test:

CUDA_VISIBLE_DEVICES=X python dehazing/test.py --size ['B' or 'L'] --input_dir [Input path] --target_dir [GT path] --result_dir [Result path] --weights [Model weighting path]

For example:

CUDA_VISIBLE_DEVICES=0 python dehazing/test.py --size 'B'  --input_dir '/data/QYW/ITS_SOTS/test/hazy/' --target_dir '/data/QYW/ITS_SOTS/test/GT/' --result_dir '/data/qiuyuwei/ITS_result' --weights '/home/qiuyuwei/MB-TaylorFormer-main/Dehazing/pretrained_models/ITS-MB-TaylorFormer-B.pth'

Results

Comming soon

Citation

If you find our repo useful for your research, please consider citing our paper:

 @misc{2308.14036,
 Author = {Yuwei Qiu and Kaihao Zhang and Chenxi Wang and Wenhan Luo and Hongdong Li and Zhi Jin},
 Title = {MB-TaylorFormer: Multi-branch Efficient Transformer Expanded by Taylor Formula for Image Dehazing},
 Year = {2023},
 Eprint = {arXiv:2308.14036},
 }	

Contact

If you have any question, please feel free to contact us via [email protected] or [email protected].

Acknowledgments

This code is based on Restormer.