控制视觉语言模型进行通用图像修复
DA-CLIP的官方PyTorch实现
我们的后续工作使用控制视觉语言模型在实际环境中进行照片级图像修复(CVPRW 2024) 提出了用于更好图像生成的后验采样,并处理类似于Real-ESRGAN的真实世界混合退化图像。
更新
[2024.04.16] 我们的后续论文"使用控制视觉语言模型在实际环境中进行照片级图像修复"现在已经在ArXiv上发表了!
[2024.04.15] 更新了一个用于处理真实世界退化图像的野外-IR模型和用于更好图像生成的后验采样。还提供了预训练权重wild-ir.pth和wild-daclip_ViT-L-14.pt用于野外-IR。
[2024.01.20] 🎉🎉🎉 我们的DA-CLIP论文被ICLR 2024接收了 🎉🎉🎉 我们进一步在模型卡中提供了一个更强大的模型。
[2023.10.25] 增加了训练和测试数据集链接。
[2023.10.13] 添加了Replicate的演示和api🔥。感谢@chenxwh!!!我们更新了Hugging Face的演示🔥和在线Colab演示🔥。感谢@fffiloni和@camenduru!!!我们在Hugging Face中也制作了一个模型卡🤗并提供了更多测试示例。
[2023.10.09] 发布了DA-CLIP 和通用IR模型的预训练权重,分别是link1 和 link2。此外,我们还提供了一个Gradio应用文件,用于测试你自己的图像。
如何运行代码?
依赖项
- 操作系统: Ubuntu 20.04
- nvidia:
- cuda: 11.4
- python 3.8
安装
我们建议您先创建一个虚拟环境:
python3 -m venv .env
source .env/bin/activate
pip install -U pip
pip install -r requirements.txt
DA-CLIP使用
进入universal-image-restoration
目录并运行:
import torch
from PIL import Image
import open_clip
checkpoint = 'pretrained/daclip_ViT-B-32.pt'
model, preprocess = open_clip.create_model_from_pretrained('daclip_ViT-B-32', pretrained=checkpoint)
tokenizer = open_clip.get_tokenizer('ViT-B-32')
image = preprocess(Image.open("haze_01.png")).unsqueeze(0)
degradations = ['motion-blurry','hazy','jpeg-compressed','low-light','noisy','raindrop','rainy','shadowed','snowy','uncompleted']
text = tokenizer(degradations)
with torch.no_grad(), torch.cuda.amp.autocast():
text_features = model.encode_text(text)
image_features, degra_features = model.encode_image(image, control=True)
degra_features /= degra_features.norm(dim=-1, keepdim=True)
text_features /= text_features.norm(dim=-1, keepdim=True)
text_probs = (100.0 * degra_features @ text_features.T).softmax(dim=-1)
index = torch.argmax(text_probs[0])
print(f"任务: {task_name}: {degradations[index]} - {text_probs[0][index]}")
数据集准备
按照我们论文的数据集构建部分准备训练和测试数据集:
#### 训练数据集 ####
####(未完成表示补全)####
datasets/universal/train
|--motion-blurry
| |--LQ/*.png
| |--GT/*.png
|--hazy
|--jpeg-compressed
|--low-light
|--noisy
|--raindrop
|--rainy
|--shadowed
|--snowy
|--uncompleted
#### 测试数据集 ####
####(与训练相同的结构)####
datasets/universal/val
...
#### 清晰的标签 ####
datasets/universal/daclip_train.csv
datasets/universal/daclip_val.csv
然后进入universal-image-restoration/config/daclip-sde
目录并修改options/train.yml
和options/test.yml
中的数据集路径。
你可以向train
和val
目录中添加更多任务或数据集,并将退化词添加到distortion
中。
数据集链接
退化类型 | motion-blurry | hazy | jpeg-compressed* | low-light | noisy* (与jpeg相同) |
---|---|---|---|---|---|
数据集 | Gopro | RESIDE-6k | DIV2K+Flickr2K | LOL | DIV2K+Flickr2K |
你应该只提取训练数据集用于训练,所有验证数据集可以在Google Drive下载。对于jpeg和noisy数据集,你可以使用这个脚本生成低质量图像。
训练
DA-CLIP
详见DA-CLIP.md。
通用图像修复
训练的主要代码在universal-image-restoration/config/daclip-sde
,DA-CLIP的核心网络在universal-image-restoration/open_clip/daclip_model.py
。
-
将预训练的DA-CLIP权重放入
pretrained
目录并检查daclip
路径。 -
然后你可以按照以下bash脚本训练模型:
cd universal-image-restoration/config/daclip-sde
# 单GPU:
python3 train.py -opt=options/train.yml
# 分布式训练,需要更改选项文件中的gpu_ids
python3 -m torch.distributed.launch --nproc_per_node=2 --master_port=4321 train.py -opt=options/train.yml --launcher pytorch
模型和训练日志将保存在log/universal-ir
中。
你可以通过运行tail -f log/universal-ir/train_universal-ir_***.log -n 100
随时打印日志。
相同的训练步骤可用于在野外进行图像修复(wild-ir)。
预训练模型
模型名称 | 描述 | GoogleDrive | HuggingFace |
---|---|---|---|
DA-CLIP | 退化感知的CLIP模型 | 下载 | 下载 |
Universal-IR | 基于DA-CLIP的通用图像还原模型 | 下载 | 下载 |
DA-CLIP-mix | 退化感知的CLIP模型(添加高斯模糊+面部修复和高斯模糊+下雨) | 下载 | 下载 |
Universal-IR-mix | 基于DA-CLIP的通用图像还原模型(添加稳健训练和混合退化) | 下载 | 下载 |
Wild-DA-CLIP | 野外场景中的退化感知CLIP模型(ViT-L-14) | 下载 | 下载 |
Wild-IR | 野外场景中的基于DA-CLIP的图像还原模型 | 下载 | 下载 |
评价
要评估我们的方法在图像还原上的效果,请修改benchmark路径和模型路径然后运行
cd universal-image-restoration/config/universal-ir
python test.py -opt=options/test.yml
Gradio
这里我们提供一个app.py文件来测试你自己的图像。在此之前,你需要下载预训练的权重(DA-CLIP和UIR)并修改options/test.yml
中的模型路径。然后只需运行python app.py
,你可以打开http://localhost:7860
来测试模型(我们还在images
目录中提供了几个具有不同退化的图像)。我们还在谷歌云盘上提供了我们的测试数据集中更多的示例。
同样的步骤也适用于野外场景中的图像还原(wild-ir)。
结果
统一图像还原(点击展开)
特定退化还原(点击展开)
野外场景图像还原(点击展开)
注意!!
🙁 在测试中我们发现,目前的预训练模型仍然难以处理一些真实世界中的图像,这些图像可能与我们的训练数据集在分布上存在偏移(捕获于不同设备或不同分辨率或退化)。我们将其视为未来的工作,并将努力使我们的模型更实用!我们也鼓励对我们工作感兴趣的用户使用更大的数据集和更多样化的退化类型来训练自己的模型。
🙁 顺便提一句,我们还发现,直接调整输入图像的尺寸会导致大多数任务的性能较差。我们可以尝试将调整尺寸步骤添加到训练中,但由于插值,这总是会破坏图像质量。
🙁 对于修复任务,我们当前的模型仅支持面部修复,这是由于数据集的限制。我们提供了我们的掩码示例,你可以使用generate_masked_face脚本生成不完整的面部。
声明: 我们的DA-CLIP基于IR-SDE和open_clip。感谢他们的代码!
联系
如果你有任何问题,请联系:ziwei.luo@it.uu.se
引用
如果我们的代码对你的研究或工作有帮助,请考虑引用我们的论文。 以下是BibTeX引用:
@article{luo2023controlling,
title={Controlling Vision-Language Models for Universal Image Restoration},
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:2310.01018},
year={2023}
}
@article{luo2024photo,
title={Photo-Realistic Image Restoration in the Wild with Controlled Vision-Language Models},
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:2404.09732},
year={2024}
}