Uformer - 在源代碼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}
}