Wavelet U-net Dehazing

WAVELET U-NET AND THE CHROMATIC ADAPTATION TRANSFORM FOR SINGLE IMAGE DEHAZING - ICIP 2019

This repository shows implementation of Wavelet U-net for image dehazing. This work establishes the new network combining wavelet transrom for single image dehazing. We use RESIDEdataset for evaluation, and it outperforms the state-of-art algorithms.

Paper

WAVELET U-NET AND THE CHROMATIC ADAPTATION TRANSFORM FOR SINGLE IMAGE DEHAZING
Hao-Hsiang Yang ^1,2, Yanwei Fu ^2
1 Graduate Institute of Electrical Engineering, National Taiwan University, Taipei, Taiwan, ² School of Data Science, Fudan University, Shanghai, China
IEEE International Conference on Image Processing (ICIP), 2019

Dependencies

• Python 3.6+
• PyTorch 0.4.0+
• PyWt

Usage

1. Cloning the repository

$ git clone https://github.com/dectrfov/Wavelet-U-net-Dehazing.git
$ cd Wavelet-U-net-Dehazing

2. Downloading the RESIDE dataset

We only use images in clear and haze folders All clear images are divided as training images (train_clear), testing images (train_hazy). The hazy images are placed to corresponding folders (val_clear and val_hazy).

3. Training

$ CUDA_VISIBLE_DEVICES=0 python train.py --epochs 100 \
                --lr 1e-4 \
                --use_gpu true \
                --gpu 0 \
                --ori_data_path /train_clear/ \
                --haze_data_path /train_hazy \
                --val_ori_data_path /val_clear/ \
                --val_haze_data_path /val_hazy/ \
                --num_workers 4 \
                --batch_size 40 \
                --val_batch_size 4 \
                --print_gap 500 \
                --model_dir /model/ \
                --log_dir /model/ \
                --sample_output_folder /samples/ \
                --net_name /dehaze_chromatic_

4. Testing

To test dehazing on RESIDE:

$ python demo.py  --sample_output_folder samples/ \
                --use_gpu true \
                --gpu 0 \
                --model_dir model/ \
                --ckpt dehaze_chromatic_100.pkl

Using pre-trained model for evaluation

1. Download model

Download from googledrive and put it in the model folder

2. Place haze images

Place hazy images in the samples folder

3. Run following command

$ python demo.py  --sample_output_folder samples/ \
                --use_gpu true \
                --gpu 0 \
                --model_dir model/ \
                --ckpt dehaze_chromatic_100.pkl