使用去噪扩散零空间模型的零样本图像修复

📖论文|🖼️项目页面|

Yinhuai Wang*、Jiwen Yu*、Jian Zhang
北京大学与PCL
*表示共同贡献

本仓库包含“使用Denoising Diffusion Null-Space Model进行零样本图像修复”的代码发布。DDNM可以在无需任何优化或训练的情况下解决各种图像修复任务！是的，以零样本的方式即可实现。

支持的应用：

• 任意尺寸🆕
• 老照片修复🆕
• 超分辨率
• 去噪
• 上色
• 修复缺失区域
• 去模糊
• 压缩感知

🧩新闻

• 现已提供高质量结果的Colab演示！

安装

代码

git clone https://github.com/wyhuai/DDNM.git

环境

pip install numpy torch blobfile tqdm pyYaml pillow    # 例如torch 1.7.1+cu110。

预训练模型

要修复人脸图像，请下载此模型(来自SDEdit)并将其放入DDNM/exp/logs/celeba/。

https://drive.google.com/file/d/1wSoA5fm_d6JBZk4RZ1SzWLMgev4WqH21/view?usp=share_link

要修复一般图像，请下载此模型(来自guided-diffusion)并将其放入DDNM/exp/logs/imagenet/。

wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt

快速入门

运行以下命令可立即获得4倍超分辨率结果。结果应位于DDNM/exp/image_samples/demo。

python main.py --ni --simplified --config celeba_hq.yml --path_y celeba_hq --eta 0.85 --deg "sr_averagepooling" --deg_scale 4.0 --sigma_y 0 -i demo

设置

详细的采样命令如下：

python main.py --ni --simplified --config {CONFIG}.yml --path_y {PATH_Y} --eta {ETA} --deg {DEGRADATION} --deg_scale {DEGRADATION_SCALE} --sigma_y {SIGMA_Y} -i {IMAGE_FOLDER}

选项说明如下：

• 本仓库实现了两种版本的DDNM。一种基于SVD的版本，在处理噪声任务时更为精确。另一种是简化版，不涉及SVD，用户可以灵活定义自己的退化类型。使用--simplified可激活简化版DDNM。若不加--simplified则为基于SVD的DDNM。
• PATH_Y是测试数据集的文件夹名称，位于DDNM/exp/datasets。
• ETA是DDIM的超参数。（默认值：0.85）
• DEGREDATION是支持的任务，包括cs_walshhadamard、cs_blockbased、inpainting、denoising、deblur_uni、deblur_gauss、deblur_aniso、sr_averagepooling、sr_bicubic、colorization、mask_color_sr以及用户自定义的diy。
• DEGRADATION_SCALE是退化的尺度。例如，--deg sr_bicubic --deg_scale 4将产生4倍超分辨率效果。
• SIGMA_Y是观测到的噪声水平。
• CONFIG是配置文件的名称（详见configs/目录），其中包含批量大小和采样步数等超参数。
• IMAGE_FOLDER是保存结果的文件夹名称。

对于配置文件，例如celeba_hq.yml，您可以修改以下属性：

sampling:
    batch_size: 1
    
time_travel:
    T_sampling: 100     # 采样步数
    travel_length: 1    # 时间旅行参数l和s，参见论文第3.3节。
    travel_repeat: 1    # 时间旅行参数r，参见论文第3.3节。

复现论文中的结果

定量评估

数据集下载链接：[Google云端硬盘] [百度网盘]

下载CelebA测试集并将其放入DDNM/exp/datasets/celeba/。

下载ImageNet测试集并将其放入DDNM/exp/datasets/imagenet/，同时替换文件DDNM/exp/imagenet_val_1k.txt。

运行以下命令。您可以通过增加batch_size来加快评估速度。

sh evaluation.sh

高质量结果

您可以尝试这个Colab演示来获得高质量结果。请注意，首页图中展示的高质量结果大多是由将DDNM应用于RePaint中的模型生成的。

🔥真实世界应用

演示：真实世界的超分辨率。

运行以下命令

python main.py --ni --simplified --config celeba_hq.yml --path_y solvay --eta 0.85 --deg "sr_averagepooling" --deg_scale 4.0 --sigma_y 0.1 -i demo

演示：老照片修复。

运行以下命令

python main.py --ni --simplified --config oldphoto.yml --path_y oldphoto --eta 0.85 --deg "mask_color_sr" --deg_scale 2.0 --sigma_y 0.02 -i demo

尝试你自己的照片。

你可以使用 DDNM 来恢复你自己退化的图像。DDNM 为你提供了完全的灵活性，可以自定义退化算子和噪声水平。请注意，这些定义对获得良好的结果至关重要。你可以参考以下指导：

1. 如果你使用的是 CelebA 预训练模型，可以尝试使用 GFPGAN 中的工具来裁剪和对齐你的照片。
2. 如果你的照片上有局部伪影，可以尝试使用 LaMa 中的工具绘制遮罩来覆盖这些伪影，并将此遮罩保存到 DDNM/exp/inp_masks/mask.png。然后运行 DDNM/exp/inp_masks/get_mask.py 生成 mask.npy。
3. 如果你的照片褪色了，你需要一个灰度算子作为退化的一部分。
4. 如果你的照片模糊，你需要一个下采样算子作为退化的一部分。同时，还需要设置合适的超分辨率缩放比例 --deg_scale。
5. 如果你的照片存在全局性的伪影，例如类似 JPEG 的伪影或未知噪声，你需要设置合适的 sigma_y 来去除这些伪影。
6. 在 DDNM/guided_diffusion/diffusion.py 中搜索 args.deg =='diy'，并相应地修改 $\mathbf{A}$ 的定义。 然后运行：

python main.py --ni --simplified --config celeba_hq.yml --path_y {YOUR_OWN_PATH} --eta 0.85 --deg "diy" --deg_scale {YOUR_OWN_SCALE} --sigma_y {YOUR_OWN_LEVEL} -i diy

🆕DDNM 适用于任意尺寸

以上展示了使用 DDNM 将 64x256 输入图像超分辨率至 256x1024 结果的示例。相关理论细节可在该 论文 的第 3.3 节中找到。

我们在文件夹 hq_demo 中实现了基于 RePaint 的 掩码移位修复 功能。你可以尝试这个 Colab 演示。

或者，你也可以在自己的设备上尝试该功能，需要下载预训练模型：

wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_classifier.pt
wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion.pt

并将它们放入 hq_demo/data/pretrained 目录。然后运行：

cd hq_demo
sh evaluation.sh

该脚本包含了最高可达 2K 分辨率的超分辨率结果。某些演示可能需要数小时才能完成。通过在 hq_demo/confs/inet256.yml 中设置较小的采样步长或时间旅行参数，可以加快速度，但可能会降低生成质量。

😊将 DDNM 应用于你自己的扩散模型

在你自己的扩散模型上实现一个基础的 DDNM 非常简单！你可以参考以下步骤：

1. 将这些算子实现复制到核心扩散采样文件中，然后定义你的任务类型，例如设置 IR_mode="super resolution"。

def color2gray(x):
    coef=1/3
    x = x[:,0,:,:] * coef + x[:,1,:,:]*coef +  x[:,2,:,:]*coef
    return x.repeat(1,3,1,1)

def gray2color(x):
    x = x[:,0,:,:]
    coef=1/3
    base = coef**2 + coef**2 + coef**2
    return th.stack((x*coef/base, x*coef/base, x*coef/base), 1)    
    
def PatchUpsample(x, scale):
    n, c, h, w = x.shape
    x = torch.zeros(n,c,h,scale,w,scale) + x.view(n,c,h,1,w,1)
    return x.view(n,c,scale*h,scale*w)

# A 及其伪逆 Ap 的实现
    
if IR_mode=="colorization":
    A = color2gray
    Ap = gray2color
    
elif IR_mode=="inpainting":
    A = lambda z: z*mask
    Ap = A
      
elif IR_mode=="super resolution":
    A = torch.nn.AdaptiveAvgPool2d((256//scale,256//scale))
    Ap = lambda z: PatchUpsample(z, scale)

elif IR_mode=="old photo restoration":
    A1 = lambda z: z*mask
    A1p = A1
    
    A2 = color2gray
    A2p = gray2color
    
    A3 = torch.nn.AdaptiveAvgPool2d((256//scale,256//scale))
    A3p = lambda z: PatchUpsample(z, scale)
    
    A = lambda z: A3(A2(A1(z)))
    Ap = lambda z: A1p(A2p(A3p(z)))

2. 在目标代码中找到变量 $\mathbf{x}_{0|t}$，并使用该函数的结果来修改 $\mathbf{x}_{t-1}$ 的采样过程。你可能需要提供输入的退化图像 $\mathbf{y}$ 和对应的噪声水平 $\sigma_\mathbf{y}$。

# DDNM+ 的核心实现，简化的去噪方案（第 3.3 节）。
# 如需更精确的去噪，请参阅论文（附录 I）和源代码。

def ddnm_plus_core(x0t, y, sigma_y=0, sigma_t, a_t):

    #式 19
    if sigma_t >= a_t*sigma_y: 
        lambda_t = 1
        gamma_t = sigma_t**2 - (a_t*lambda_t*sigma_y)**2
    else:
        lambda_t = sigma_t/(a_t*sigma_y)
        gamma_t = 0
        
    #式 17    
    x0t= x0t + lambda_t*Ap(y - A(x0t))
    
    return x0t, gamma_t

3. 实际上，这个仓库就包含了上述简化实现：可以在 DDNM/guided_diffusion/diffusion.py 中搜索 arg.simplified 查看相关代码。

参考文献

如果你觉得这个仓库对你的研究有帮助，请引用以下工作。

@article{wang2022zero,
  title={Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model},
  author={Wang, Yinhuai and Yu, Jiwen and Zhang, Jian},
  journal={The Eleventh International Conference on Learning Representations},
  year={2023}
}

本实现基于或受以下项目启发：

• https://github.com/wyhuai/RND（零空间学习）
• https://github.com/andreas128/RePaint（时间旅行技巧）
• https://github.com/bahjat-kawar/ddrm（代码结构）

DDNM 快速上手指南

DDNM (Denoising Diffusion Null-Space Model) 是一个基于去噪扩散模型的零样本图像修复工具。它无需任何微调或训练，即可解决超分辨率、去噪、上色、修复、去模糊等多种图像恢复任务。

环境准备

• 系统要求: Linux / macOS / Windows (推荐 Linux)
• Python: 3.6+
• PyTorch: 1.7.1+ (需包含 CUDA 支持，例如 cu110)
• 其他依赖: numpy, blobfile, tqdm, pyYaml, pillow

安装步骤

1. 克隆代码库

git clone https://github.com/wyhuai/DDNM.git
cd DDNM

2. 安装依赖

建议使用国内镜像源加速安装：

pip install numpy torch blobfile tqdm pyYaml pillow -i https://pypi.tuna.tsinghua.edu.cn/simple
# 注意：如果上述命令未安装指定版本的 torch，请根据官方文档单独安装对应 CUDA 版本的 torch
# 例如：pip install torch==1.7.1+cu110 -f https://download.pytorch.org/whl/torch_stable.html

3. 下载预训练模型

根据需求下载对应的模型并放入指定目录：

• 人脸图像恢复 (CelebA): 下载 模型文件 并存放至 DDNM/exp/logs/celeba/ 目录下。

• 通用图像恢复 (ImageNet):

  wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt -P exp/logs/imagenet/
  

  (注：如果 wget 速度慢，可手动下载后放入 DDNM/exp/logs/imagenet/)

基本使用

以下命令演示如何立即运行一个 4 倍超分辨率 (4x SR) 任务。结果将保存在 DDNM/exp/image_samples/demo 目录中。

python main.py --ni --simplified --config celeba_hq.yml --path_y celeba_hq --eta 0.85 --deg "sr_averagepooling" --deg_scale 4.0 --sigma_y 0 -i demo

命令参数简析

• --simplified: 启用简化版 DDNM（无需 SVD 分解，更灵活）。
• --config: 配置文件名 (如 celeba_hq.yml, oldphoto.yml)。
• --deg: 退化类型，支持 sr_averagepooling (超分), inpainting (修复), denoising (去噪), colorization (上色), deblur_uni (去模糊) 等。
• --deg_scale: 退化尺度，例如超分倍数设为 4.0。
• --sigma_y: 观测噪声水平，无噪声设为 0。
• -i: 输出结果文件夹名称。

更多应用场景示例

老照片修复:

python main.py --ni --simplified --config oldphoto.yml --path_y oldphoto --eta 0.85 --deg "mask_color_sr" --deg_scale 2.0 --sigma_y 0.02 -i demo

自定义图像修复: 将待处理图片放入 exp/datasets/{YOUR_OWN_PATH} 目录，然后运行：

python main.py --ni --simplified --config celeba_hq.yml --path_y {YOUR_OWN_PATH} --eta 0.85 --deg "diy" --deg_scale {YOUR_OWN_SCALE} --sigma_y {YOUR_OWN_LEVEL} -i diy

(注：使用 diy 模式需在代码中自定义退化算子，详见原文档 "Try your own photos" 部分)