Algolzw/image-restoration-sde

Image Restoration SDE
_{Official PyTorch Implementations of [IR-SDE](ICML 2023) and [Refusion](CVPRW 2023).}

Ziwei Luo, Fredrik K. Gustafsson, Zheng Zhao, Jens Sjölund, Thomas B. Schön
Department of Information Technology, Uppsala University

Image Restoration with Mean-Reverting Stochastic Differential Equations [Weights and Results]

Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models [Weights and Results]

Both methods can be used for synthetic and real-world datasets. You can find more details and results from this Project page.

Updates

[2023.07.06] All results and pre-trained models for IR-SDE and Refusion are provided in link1 and link2, respectively.
[2023.04.25] 🎉🎉🎉 Our IR-SDE paper was accepted by ICML 2023 🎉🎉🎉
[2023.04.23] We released the codes for various tasks (including IR-SDE, Refusion, and Latent-Refusion).
[2023.04.17] 🎉🎉🎉 Our method Refusion is the Winning Solution in NTIRE 2023 Image Shadow Removal Challenge [Paper].

Dependenices

• OS: Ubuntu 20.04
• nvidia :
  • cuda: 11.7
  • cudnn: 8.5.0
• python3
• pytorch >= 1.13.0
• Python packages: pip install -r requirements.txt

How to use our Code?

Here we provide an example for image deraining task, but can be changed to any problem with replacing the dataset. And the code for Denoising-SDE/ODE will be released later!

We retrained the deraining model from scratch on 4 Nvidia TITAN XP GPUs and find that it achieves a new SOTA even in terms of PSNR on Rain100H:

Note that we didn't tune any parameter, the last saved checkpoint was used to evaluation.

The pretrained models for all tasks are provided here, and the performances of other SOTAs can be find in here.

For Refusion Model

Note that if you want to train/test the Refusion model, the only change is to replace the "ir-sde.yml" to "refusion.yml". More details please refer to the "demo.sh" bash script in each task's directory.

Dataset Preparation

We employ Rain100H datasets for training (totally 1,800 images) and testing (100 images).

Download training and testing datasets and process it in a way such that rain images and no-rain images are in separately directories, as

#### for training dataset ####
datasets/rain/trainH/GT
datasets/rain/trainH/LQ


#### for testing dataset ####
datasets/rain/testH/GT
datasets/rain/testH/LQ

Then get into the codes/config/deraining directory and modify the dataset paths in option files in options/train/ir-sde.yml and options/test/ir-sde.yml.

Train

The main code for training is in codes/config/deraining and the core algorithms for IR-SDE is in codes/utils/sde_utils.py.

You can train the model following below bash scripts:

cd codes/config/deraining

# For single GPU:
python3 train.py -opt=options/train/ir-sde.yml

# For distributed training, need to change the gpu_ids in option file
python3 -m torch.distributed.launch --nproc_per_node=2 --master_poer=4321 train.py -opt=options/train/ir-sde.yml --launcher pytorch

Then the models and training logs will save in log/derain_sde/. You can print your log at time by running tail -f log/derain_sde/train_derain_sde_***.log -n 100.

Evaluation

To evaluate our method, please modify the benchmark path and model path and run

cd codes/config/deraining
python test.py -opt=options/test/ir-sde.yml

We also provide the results for different tasks here.

Gradio

Here we provide an app.py file for testing your own images. Before that, you need to download the pretrained weights and modify the model path in test_sde_derain.yml. Then by simply running python app.py, you can open http://localhost:7860 to test the model. (We also provide 3 raining images from Rain100H in images dir).

Some Results

Interpolation

We also provide a interpolation demo to perform interpolation between two images in codes/demos/interpolation.py, the usage is:

cd codes/demos
python interpolation.py -s source_image_path -t target_image_path --save save_dir

Example of interpolation:

Citations

If our code helps your research or work, please consider citing our paper. The following are BibTeX references:

@article{luo2023image,
  title={Image Restoration with Mean-Reverting Stochastic Differential Equations},
  author={Luo, Ziwei and Gustafsson, Fredrik K and Zhao, Zheng and Sj{\"o}lund, Jens and Sch{\"o}n, Thomas B},
  journal={International Conference on Machine Learning},
  year={2023},
  organization={PMLR}
}

@inproceedings{luo2023refusion,
  title={Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models},
  author={Luo, Ziwei and Gustafsson, Fredrik K and Zhao, Zheng and Sj{\"o}lund, Jens and Sch{\"o}n, Thomas B},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops},
  pages={1680--1691},
  year={2023}
}

Contact

E-mail: [email protected]

--- Thanks for your interest! ---