image-restoration-sde

图像复原 SDE
_{[IR-SDE](ICML 2023)和[Refusion](CVPRW 2023)的官方 PyTorch 实现。}

Ziwei Luo, Fredrik K. Gustafsson, Zheng Zhao, Jens Sjölund, Thomas B. Schön
乌普萨拉大学信息技术系

Refusion: 使用潜空间扩散模型实现大尺寸真实图像复原 [权重和结果]

Refusion 这两种方法都可用于合成和真实世界数据集。你可以在这个**项目页面**找到更多详情和结果。

更新

[2024.04.15] 更新了一种后验采样方法以获得更好的图像生成效果。详见论文使用受控视觉语言模型进行真实世界中的照片级图像复原。
[2023.10.02] 我们提出了一种新的通用图像复原框架，通过控制视觉语言模型实现 [点击这里]。
[2023.07.06] IR-SDE 和 Refusion 的所有结果和预训练模型分别在链接1 和链接2 提供。
[2023.04.25] 🎉🎉🎉 我们的 IR-SDE 论文被 ICML 2023 接收 🎉🎉🎉
[2023.04.23] 我们发布了各种任务的代码（包括 IR-SDE、Refusion 和 Latent-Refusion）。
[2023.04.17] 🎉🎉🎉 我们的方法 Refusion 是 NTIRE 2023 图像阴影去除挑战赛的获胜方案 [论文]。

依赖

操作系统: Ubuntu 20.04
nvidia :
- cuda: 11.7
- cudnn: 8.5.0
python3
pytorch >= 1.13.0
Python 包: pip install -r requirements.txt

如何使用我们的代码？

这里我们以图像去雨任务为例，但可以通过替换数据集来改变为任何问题。去噪-SDE/ODE 的代码将在稍后发布！

我们在 4 块 Nvidia TITAN XP GPU 上从头开始重新训练了去雨模型，发现它在 Rain100H 上甚至在 PSNR 方面也达到了新的 SOTA:

方法	PSNR	SSIM	LPIPS	FID
IR-SDE	31.65	0.9041	0.047	18.64
Restormer	31.46	0.904	-	-
MPRNet	30.41	0.8906	0.158	61.59
PReNet	29.46	0.8990	0.128	52.67

注意我们没有调整任何参数，评估时使用的是最后保存的检查点。

所有任务的预训练模型在这里提供，其他 SOTA 方法的性能可以在这里找到。

Refusion 模型

注意，如果你想训练/测试 Refusion 模型，唯一的变化是将 "ir-sde.yml" 替换为 "refusion.yml"。更多详情请参考每个任务目录中的 "demo.sh" bash 脚本。

数据集准备

我们使用 Rain100H 数据集进行训练（总共 1,800 张图像）和测试（100 张图像）。

下载训练和测试数据集，并以雨天图像和无雨图像分别存放在不同目录的方式处理，如下所示：

#### 训练数据集 ####
datasets/rain/trainH/GT
datasets/rain/trainH/LQ


#### 测试数据集 ####
datasets/rain/testH/GT
datasets/rain/testH/LQ

然后进入 codes/config/deraining 目录，修改 options/train/ir-sde.yml 和 options/test/ir-sde.yml 中的数据集路径。

训练

训练的主要代码在 codes/config/deraining 中，IR-SDE 的核心算法在 codes/utils/sde_utils.py 中。

你可以按照以下 bash 脚本训练模型：

cd codes/config/deraining
# 单GPU:
python3 train.py -opt=options/train/ir-sde.yml

# 分布式训练，需要修改选项文件中的gpu_ids
python3 -m torch.distributed.launch --nproc_per_node=2 --master_poer=4321 train.py -opt=options/train/ir-sde.yml --launcher pytorch

模型和训练日志将保存在log/derain_sde/目录下。你可以通过运行tail -f log/derain_sde/train_derain_sde_***.log -n 100实时查看日志。

评估

要评估我们的方法，请修改基准路径和模型路径，然后运行

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

我们还在这里提供了不同任务的结果。

Gradio

我们在app.py文件中提供了一个用于测试您自己图像的应用。在此之前，您需要下载预训练权重并修改test_sde_derain.yml中的模型路径。然后只需运行python app.py，就可以打开http://localhost:7860来测试模型。（我们还在images目录中提供了3张来自Rain100H的雨天图像）。

部分结果

IR-SDE

高分辨率非均匀去雾

Refusion

图像阴影去除

Refusion

插值

我们还在codes/demos/interpolation.py中提供了一个插值演示，用于在两张图像之间进行插值，使用方法如下：

cd codes/demos
python interpolation.py -s 源图像路径 -t 目标图像路径 --save 保存目录

插值示例：

IR-SDE

引用

如果我们的代码对您的研究或工作有帮助，请考虑引用我们的论文。以下是BibTeX引用：

@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}
}