文章目录
- 前言
- 代码及原文链接
- 主要的点
- 如何进行图像恢复
- README
前言
关于扩散模型以及条件扩散模型的介绍,大家可以前往我的上一篇博客:扩散模型diffusion model用于图像恢复任务详细原理 (去雨,去雾等皆可),附实现代码。
SR3是利用扩散模型进行图像超分辨率研究的,它在使用低分辨率图像作为条件来进行反向采样时先将低分辨率图像直接上采样到高分辨率图像。因此,它通过一些改进可以作为图像恢复这种输入与输出图像大小相等的任务。
代码及原文链接
我用它来做图像恢复的代码:图像去雨、去雾等恢复任务:代码简化的SR3扩散模型,有注释及实验流程
主要的点
SR3模型将DDPM中的残差块替换为BigGAN的残差块,并将跳跃连接重新缩放 。
增加了剩余块的数量,以及不同分辨率下的通道乘法器。
还有一个拼接操作,就是将低分辨图像直接上采样后的图像与噪声图像进行拼接,从而实现条件噪声估计。
如何进行图像恢复
- 源代码中包括大量的小文件,不适合初学者调试和使用。
- 我根据其核心代码进行简化,将一些不必要的代码文件和代码块删除,并在Rain13K (就是MPRNet、Restormer等方法使用的去雨数据集) 上进行了实验,目前取得不错的表现。可以参考下边的对比。
- 其他的任务应该也是可以直接使用的,只需要修改一下配置文件中的数据集路径即可。
- 需要使用SR3扩散模型做图像去雨、去雾等恢复任务的可以前往:图像去雨、去雾等恢复任务:代码简化的SR3扩散模型,有注释及实验流程。
- 附上一组在Rain13K训练集上训练,在Rain100H上测试的结果。源代码中的迭代次数为1, 000, 000,我这里是只迭代了 180, 000次后进行测试的结果。
README
由于之前上传的时候忘记加这个文件,后来有些朋友私聊我,其实也没写多少,我就直接复制在这里吧,有需要的可以参考一下:
- 更新 config.json 中的数据集路径,一般就是直接把测试集当成验证集,其实这样是不合适的
- config.json 中的 phase 参数表示阶段,可以选择 train 或者 val,即可以通过改变这个参数来利用 main.py 分别实验训练和测试
- 运行 main.py ,会在当前目录下建一个 experiments 文件夹,其中对应日期的文件夹名,下边包括 checkpoint 和 logs
- 输出结果的路径在最开始的 path 那里,同时在测试的时候需要将 resume_state 改为训练好的模型文件的路径
- 由于这个 json 文件的特殊性,里边的注释是英文,其他地方是中文
- 整个流程很简单,注意阅读代码
- 附带的这个数据集图片太小,控制台打印不频繁,建议把数据集的图片数量增加