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

weatherdiffusion's Issues

May I ask what this string of code means? I noticed that it will generate 64x64 images

in ddm.py 202 row
May I ask what this string of code means? I noticed that it will generate 64x64 images

About the date that the code will be released

Thanks for your impressive work which is inspiring. I would like to ask when the code will be released.

Many thanks again!

About the weather removal dataset?

In the Outdoor-Rain dataset, There seems to be an overlap between the training dataset and the test dataset in a data?

Are you split the dataset as default?

Best,

About training

I am in the process of training my own dataset and found that after 200000 iterations, the loss does not show a decreasing trend. May I ask if this is normal?

Running Derain on my own images?

Is there a way to run the derain operation on my own images? The examples show how to evaluate your model, but I don't see how to just run derain on other images.

I'm also not sure where to get the appropriate weights. I found the Allweather 64|128 files, but not the allweather.txt file, despite what it says in #10 .

Thanks!

evaluation metrics on diffrent tasks

Hi,
Thanks for sharing your great work! we wonder about some details of the evaluation metrics, i.e., which task is evaluated on Y channel.
Best,
Josh

about model training

Hello author! Thank you very much for your contribution! I came across your paper and now I would like to dive into the following experiments, but I am a novice, what are the instructions needed if I train your model? Or is there anything I need to prepare in advance? Thanks!

The question about validation

Hi, I don't see a step in your code about calculating the metrics for validation, so how did you choose the best model?
Also it takes about 2 minutes to recover a single snow image on 3090 and there are 16000+ images in the transweather snow test set which takes a lot of time, I was wondering if you tested on all the data or just selected some of it?

Could you provide the other model weights?

I would like to follow your work, and need the checkpoints to test on new datasets. But there are only checkpoints of $WeatherDiffusion_{64/128}$ while they may perform worse on specific tasks. Could you provide the other model weights. Thanks for your work.

Training time

Thanks for your interesting work! I would like to know, how long does it take to train WeatherDiff64？

The training crashes

Dear authors, thank you for sharing your codes for the community, I have a question about the model training, is it nomarl that the model crashes during the training, as I tried to reproduct the model on derain and desnow datasets, but during certain iters, as 800000, the visual results shows like pure noise or pure black. Should I wait more iters or is there any error in my reproduction? Thank you very much

The image results on the three test datasets

Hello, can you please provide the image results on the test datasets so that I can do some visual comparsion with other methods?

Could you provide train bash?

Thank you for your wonderful work.Could you provide train bash or correct my bash? I find some color differences when I use your code to train other images.Thus, I want to check if this is something wrong with train bash.

my train bash:
CUDA_VISIBLE_DEVICES=1,2 python train_diffusion.py --config "allweather.yml" --image_folder='results/all_weather'

What's the difference between 'mean estimated noise' and 'mean reconstructed image' during sampling?

Thank you for your patch-based DDM. But I have a question. What's the difference between 'mean estimated noise' and 'mean reconstructed image' during sampling? I think they should be the same. Because Eq. (19) in your paper is linear with respect to $\epsilon_{\theta}$.

Test datasets

Could you upload the test datasets? Because I can't find them in TransWeather

How to prepare the test dataset

I was trying to test the model but there lacks information about how to set up the test dataset. (especially the allweather.txt and raindroptesta.txt file).

Could you please add some instructions on how to set it up. Thanks!

About training

Dear author
I'd like to ask you three questions

Most importantly, the training code seems to be incomplete (for example, the training process does not match the description in the paper, especially regarding patch). Is there a more sound training code? Thank you very much.

Second, are data_transform and inverse_data_transform necessary?

Finally, when I was training 128*128 pictures on Nvidia 3090, a single card could actually run the batch of 24 data. Is this normal?

Could you provide the RaindropDiff128 model weight?

I'm interest to raindrop removal problem for my bachelor thesis. Could you provide the RainDropDiff128 model weight? Thank you in advance and thanks for this useful work. If you have more time, if you don't mind, could you provide RainDropDiff64 too?

Furthermore, the urgency is, unlike RainHazeDiff and DeSnowDiff that are still enough to use WeatherDiff, in RainDrop case, RainDropDiff is more powerful.

Loss and Convergence

Thanks for your interesting work! I find loss shocks violently,so i would like to know how we can prove that the network has converged and stop training.(maybe some steps perform so poor that some loss is huge?)

Where can I find the supplementary materials?

Thank you for your wonderful work. I would like to study this paper in more detail. Where can I find the supplementary materials?

Could you provide some checkpoints?

Thank you for your amazing work. I noticed that you trained all models for about 2, 000, 000 iterations. It may takes too much time to train the model from scratch. So can you release some checkpoints or provide some training details such as hardware and training time?

IndexError: index 32 is out of bounds for dimension 0 with size 32

Hi. first,Thank you for releasing the code.

When executing the code in sampling, there is an error, so I ask you a question.

When you enter the model, the size of the tensor is halved. Therefore, the index value does not match.
ex) model(x) ::: tensor[64,6,64,64] ==> tensor[32,3,64,64]
I don't know why 0 dimensions are halved.

Using device: cuda
Note: Currently supports evaluations (restoration) when run only on a single GPU!
=> using dataset 'AllWeather'
=> evaluating snowtest100K-L...
=> creating denoising-diffusion model with wrapper...
scratch/ozan/ckpts/AllWeather_ddpm.pth.tar
=> loaded checkpoint 'scratch/ozan/ckpts/AllWeather_ddpm.pth.tar' (epoch 209, step 470000)
starting processing from image ['beautiful_smile_00003']
Traceback (most recent call last):
File "eval_diffusion.py", line 83, in
main()
File "eval_diffusion.py", line 79, in main
model.restore(val_loader, validation=args.test_set, r=args.grid_r)
File "C:\Users\PIAI\Desktop\WeatherDiffusion\models\restoration.py", line 36, in restore
x_output = self.diffusive_restoration(x_cond, r=r)
File "C:\Users\PIAI\Desktop\WeatherDiffusion\models\restoration.py", line 45, in diffusive_restoration
x_output = self.diffusion.sample_image(x_cond, x, patch_locs=corners, patch_size=p_size)
File "C:\Users\PIAI\Desktop\WeatherDiffusion\models\ddm.py", line 194, in sample_image
xs = utils.sampling.generalized_steps_overlapping(x, x_cond, seq, self.model, self.betas, eta=0.,
File "C:\Users\PIAI\Desktop\WeatherDiffusion\utils\sampling.py", line 74, in generalized_steps_overlapping
et_output[0, :, hi:hi + p_size, wi:wi + p_size] += outputs[idx]
IndexError: index 32 is out of bounds for dimension 0 with size 32

Can you provide the WeatherDiff128 model weights?

Thanks for your interesting and amazing work. I wanna to compare the results of two different imagesize network and use them as the pre-trained weights. But I only find the WeatherDiff64 model weights in the readme. So can you provide the WeatherDiff128 model weights? Thank you!

Is this an unsupervised method？

About Dataset（Test1 of Outdoor-Rain)

Hi @oozdenizci , thanks for your amazing work!

I would like to ask where I can find the test set Test1 corresponding to Outdoor-Rain, because I did not find the test set in the official github of Outdoor-Rain.

Looking forward to your reply!

about the source code

Thank you so much for your awesome work! I would like to ask when the code will be released.
Best wishes.

Can you tell me the installed version of python/pytorch library functions for the code?

thanks