IGITUGraz / WeatherDiffusion

Code for "Restoring Vision in Adverse Weather Conditions with Patch-Based Denoising Diffusion Models" [TPAMI 2023]

Restoring Vision in Adverse Weather Conditions with Patch-Based Denoising Diffusion Models

This is the code repository of the following paper to train and perform inference with patch-based diffusion models for image restoration under adverse weather conditions.

"Restoring Vision in Adverse Weather Conditions with Patch-Based Denoising Diffusion Models"
Ozan Özdenizci, Robert Legenstein
IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI), 2023.
https://doi.org/10.1109/TPAMI.2023.3238179

Datasets

We perform experiments for image desnowing on Snow100K, combined image deraining and dehazing on Outdoor-Rain, and raindrop removal on the RainDrop datasets. To train multi-weather restoration, we used the AllWeather training set from TransWeather, which is composed of subsets of training images from these three benchmarks.

Saved Model Weights

We share a pre-trained diffusive multi-weather restoration model WeatherDiff₆₄ with the network configuration in configs/allweather.yml. To evaluate WeatherDiff₆₄ using the pre-trained model checkpoint with the current version of the repository:

python eval_diffusion.py --config "allweather.yml" --resume 'WeatherDiff64.pth.tar' --test_set 'raindrop' --sampling_timesteps 25 --grid_r 16
python eval_diffusion.py --config "allweather.yml" --resume 'WeatherDiff64.pth.tar' --test_set 'rainfog' --sampling_timesteps 25 --grid_r 16
python eval_diffusion.py --config "allweather.yml" --resume 'WeatherDiff64.pth.tar' --test_set 'snow' --sampling_timesteps 25 --grid_r 16

A smaller value for grid_r will yield slightly better results and higher image quality:

python eval_diffusion.py --config "allweather.yml" --resume 'WeatherDiff64.pth.tar' --test_set 'raindrop' --sampling_timesteps 25 --grid_r 4
python eval_diffusion.py --config "allweather.yml" --resume 'WeatherDiff64.pth.tar' --test_set 'rainfog' --sampling_timesteps 25 --grid_r 4
python eval_diffusion.py --config "allweather.yml" --resume 'WeatherDiff64.pth.tar' --test_set 'snow' --sampling_timesteps 25 --grid_r 4

We also share our pre-trained diffusive multi-weather restoration model WeatherDiff₁₂₈ with the network configuration in configs/allweather128.yml.

Check out below for some visualizations of our patch-based diffusive image restoration approach.

Image Desnowing

Input Condition	Restoration Process	Output

Image Deraining & Dehazing

Input Condition	Restoration Process	Output

Raindrop Removal

Input Condition	Restoration Process	Output

Reference

If you use this code or models in your research and find it helpful, please cite the following paper:

@article{ozdenizci2023,
  title={Restoring vision in adverse weather conditions with patch-based denoising diffusion models},
  author={Ozan \"{O}zdenizci and Robert Legenstein},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 
  pages={1-12},
  year={2023},
  doi={10.1109/TPAMI.2023.3238179}
}

Acknowledgments

Authors of this work are affiliated with Graz University of Technology, Institute of Theoretical Computer Science, and Silicon Austria Labs, TU Graz - SAL Dependable Embedded Systems Lab, Graz, Austria. This work has been supported by the "University SAL Labs" initiative of Silicon Austria Labs (SAL) and its Austrian partner universities for applied fundamental research for electronic based systems.

Parts of this code repository is based on the following works:

• https://github.com/ermongroup/ddim
• https://github.com/bahjat-kawar/ddrm
• https://github.com/JingyunLiang/SwinIR