This repository contains the official implementation of the following paper:
Image Dehazing Transformer with Transmission-Aware 3D Position Embedding
Chun-Le Guo, Qixin Yan, Saeed Anwar, Runmin Cong, Wenqi Ren, Chongyi Li*
IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022
Paper Link: [official link]
Overview structure of our method. Our method consists of five key modules: a transmission-aware 3D position embedding module, a Transformer module, a CNN encoder module, a feature modulation module, and a CNN decoder module.
Clone Repo
git clone <code_link>
cd Dehamer-main/
Create Conda Environment and Install Dependencies
conda env create -f environment.yaml
conda activate dehamer
Before performing the following steps, please download our pretrained model first.
Download Links: [Google Drive] [Baidu Disk (password: 1tap)]
Then, unzip the file and place the models to ckpts/<dataset_name>
directory, separately.
The directory structure will be arranged as:
ckpts
|- dense
|- PSNR1662_SSIM05602.pt
|- NH
|- PSNR2066_SSIM06844.pt
|- indoor
|- PSNR3663_ssim09881.pt
|- outdoor
|- PSNR3518_SSIM09860.pt
We provide some classic test images in the classic_test_image
directory.
Run the following command to process them:
CUDA_VISIBLE_DEVICES=X python src/test_PSNR.py --dataset-name our_test
The dehazed images will be saved in the results/
directory.
You can also run the following command to process your own images:
CUDA_VISIBLE_DEVICES=X python src/test_PSNR.py \
--dataset-name our_test -t path/to/your/test/folder
Dataset | OTS-Train | ITS-Train | Dense-Haze | NH-HAZE | SOTS-Test |
---|---|---|---|---|---|
Details | about 45G | about 4.5G | about 250M | about 315M | about 415M |
Google drive | Google drive | Google drive | Google drive | Google drive | Google drive |
Baidu Cloud | Not available | Baidu Cloud | Baidu Cloud | Baidu Cloud | Baidu Cloud |
The train and valid split files are provided in data/<dataset_name>
.
The data
directory structure will be arranged as: (Note: please check it carefully)
data
|- classic_test_image
|- input
|- canyon.png
|- canyon1.png
|- val_list.txt
|-Dense-Haze
|- train_dense
|- haze
|- 01_hazy.png
|- 02_hazy.png
|- clear_images
|- 01_GT.png
|- 02_GT.png
|- trainlist.txt
|- valid_dense
|- input
|- 51_hazy.png
|- 52_hazy.png
|- gt
|- 51_GT.png
|- 52_GT.png
|- val_list.txt
|-NH-Haze
|- train_NH
|- haze
|- 01_hazy.png
|- 02_hazy.png
|- clear_images
|- 01_GT.png
|- 02_GT.png
|- trainlist.txt
|- valid_NH
|- input
|- 51_hazy.png
|- 52_hazy.png
|- gt
|- 51_GT.png
|- 52_GT.png
|- val_list.txt
|-ITS
|- train_indoor
|- haze
|- 1_1_0.90179.png
|- 2_1_0.99082.png
|- clear_images
|- 1.png
|- 2.png
|- trainlist.txt
|- valid_indoor
|- input
|- 1400_1.png
|- 1401_1.png
|- gt
|- 1400.png
|- 1401.png
|- val_list.txt
|-OTS
|- train_outdoor
|- haze
|- 0005_0.8_0.04.jpg
|- 0008_0.8_0.04.jpg
|- clear_images
|- 0005.jpg
|- 0008.jpg
|- trainlist.txt
|- valid_outdoor
|- input
|- 0001_0.8_0.2.jpg
|- 0002_0.8_0.08.jpg
|- gt
|- 0001.png
|- 0002.png
|- val_list.txt
See python3 src/train.py --h
for list of optional arguments, or train.sh
for examples.
An example of NH-HAZE dataset.
CUDA_VISIBLE_DEVICES=0,1 python src/train.py \
--dataset-name NH \
--train-dir ./data/train_NH/ \
--valid-dir ./data/valid_NH/ \
--ckpt-save-path ../ckpts \
--ckpt-overwrite \
--nb-epochs 5000 \
--batch-size 2\
--train-size 800 1200 \
--valid-size 800 1200 \
--loss l1 \
--plot-stats \
--cuda
See python3 test_PSNR.py --h
for list of optional arguments, or test.sh
for an example.
An example of NH-HAZE dataset.
CUDA_VISIBLE_DEVICES=1 python src/test_PSNR.py \
--dataset-name NH
We also provide the mindspore code at https://github.com/Dragoniss/mindspore-phase2-Dehamer.
The results of our dehamer on benchmarks can also be download from [Google drive] and [Baidu Disk (password: 13e2)].
If you find our repo useful for your research, please consider citing our paper:
@inproceedings{guo2022dehamer,
author = {Chun-Le Guo, Qixin Yan, Saeed Anwar, Runmin Cong, Wenqi Ren, Chongyi Li},
title = {Image Dehazing Transformer with Transmission-Aware 3D Position Embedding},
booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
year={2022}
}
If you have any question, please feel free to contact us via guochunle@nankai.edu.com
or lichongyi25@gmail.com
.
This code is based on Swin-Transformer , Noise2noise, and GridDehazeNet.