azlanqazi2012 / image-restoration-sde Goto Github PK

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 dataset. You can find more details and results from this Project page.

[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), and the pre-trained weights for Refusion is available at here.
[2023.04.17] 🎉🎉🎉 Our method Refusion won the 2nd place in NTIRE 2023 Image Shadow Removal Challenge [Paper].

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

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.

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.

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.

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.

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.

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).

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

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={arXiv preprint arXiv:2301.11699},
  year={2023}
}

@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},
  year={2023}
}

E-mail: [email protected]