Uformer

Uformer

Python

1.0.0

下载

Uformer：用于图像恢复的一般U形变压器（CVPR 2022）

Zhendong Wang，Xiaodong Cun，Jianmin Bao，Wengang Zhou，Jianzhuang Liu，Houqiang li

纸链接：[arxiv] [CVPR]

更新：

2024.10.24修复过期的链接。
2022.07.06上传我们的Uformer的新代码和模型。
2022.04.09上传uformer在denoising（Sidd，dnd），运动脱脂（GoPro，hide，realblur-j/-r）和defocus deblurring（dpdd）上上传。
2022.03.02 uformer已被CVPR 2022所接受！
2021.11.30 arxiv链接中的更新uformer。新的代码，模型和结果将上传。
2021.10.28在sidd和dnd上释放uformer32的结果。
2021.09.30释放预先训练的Uformer16用于sidd denoising。
2021.08.19发布预训练的模型（Uformer32）！添加脚本以进行flop/gmac计算。
2021.07.29添加了一个用于测试任意分辨率图像的预训练模型的脚本。

在本文中，我们提出了一种有效，有效的基于变压器的体系结构，在该体系结构中，我们使用变压器块构建了一个层次结构编码器网络进行图像恢复。 Uformer具有两个核心设计，使其适合此任务。第一个关键要素是局部增强的窗口变压器块，我们在其中使用基于窗口的非重叠窗口自我注意来减少计算要求，并在馈送前进网络中采用深度卷积，以进一步提高其捕获本地环境的潜力。第二个关键要素是我们探索三个跳过连接方案，以有效地从编码器传递信息到解码器。在这两种设计的支持下，Uformer具有很高的能力，可以捕获有用的图像修复依赖项。在几个图像恢复任务上进行了广泛的实验，证明了Uformer的优越性，包括图像DeNoising，deraining，Deblurring和Demoinering。我们希望我们的工作将鼓励进一步的研究探索基于变形金刚的架构，以完成低级视觉任务。

Uformer

包装依赖性

该项目由Pytorch 1.9.0，Python3.7，Cuda11.1构建。对于包装依赖项，您可以通过：

pip install -r requirements.txt

预验证的模型

uformer_b：sidd | Gopro

验证模型的结果

uformer_b：sidd | DND | Gopro |隐藏| realblur-j | Realblur-r | DPDD

数据准备

Denoising

对于SIDD的培训数据，您可以从官方URL下载Sidd-Medium数据集。然后生成培训补丁以供培训：

 python3 generate_patches_SIDD . py - - src_dir .. / SIDD_Medium_Srgb / Data - - tar_dir .. / datasets / denoising / sidd / train

要评估SIDD和DND，您可以从此处下载数据。

deblurring

对于GoPro的培训，以及对GoPro，hide，realblur-j和realblur-r进行评估，您可以从此处下载数据。

然后将所有deno的数据放入../datasets/denoising ，然后将所有DeBlurring数据放入../datasets/deblurring deblurring中。

训练

Denoising

要在SIDD上培训Uformer，您可以通过以下方式开始培训

sh script/train_denoise.sh

deblurring

要在GoPro上培训Uformer，您可以通过以下方式开始培训：

sh script/train_motiondeblur.sh

评估

要评估Uformer，您可以运行：

sh script/test.sh

要在每个数据集上进行评估，您应该取消点击相应的行。

计算成本

我们提供了一个简单的脚本来自行计算拖鞋，在model.py中添加了一个简单的脚本。您可以更改配置并运行：

 python3 model . py

该回购中GMAC的手动计算与主论文略有不同，但不影响结论。我们稍后将纠正论文。

引用

如果您发现此项目在您的研究中有用，请考虑引用：

 @InProceedings{Wang_2022_CVPR,
    author    = {Wang, Zhendong and Cun, Xiaodong and Bao, Jianmin and Zhou, Wengang and Liu, Jianzhuang and Li, Houqiang},
    title     = {Uformer: A General U-Shaped Transformer for Image Restoration},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2022},
    pages     = {17683-17693}
}