首頁>編程相關>其他源碼
下載
達克利普烏爾

控制視覺語言模型以實現通用影像恢復
DA-CLIP 的官方 PyTorch 實作。

項目頁面|紙|型號卡 ?

達克利普

我們的後續工作《使用受控視覺語言模型進行野外照片逼真影像復原》(CVPRW 2024) 提出了後取樣，以實現更好的影像生成，並處理類似於 Real-ESRGAN 的真實世界混合降解影像。

更新

[ 2024.04.16 ] 我們的後續論文「Photo-Realistic Image Restoration in the Wild with Controlled Vision-Language Models」現已在 ArXiv 上發布！
[ 2024.04.15 ] 更新了用於現實世界退化的野生紅外線模型和後驗採樣以更好地產生影像。也為wild-ir提供了預訓練權重wild-ir.pth和wild-daclip_ViT-L-14.pt。
[ 2024.01.20 ] ???我們的DA-CLIP論文被ICLR 2024接受了？我們進一步在模型卡中提供了更強大的模型。
[ 2023.10.25 ] 新增了用於訓練和測試的資料集連結。
[ 2023.10.13 ] 新增複製示範和 api。感謝@chenxwh！ ！我們更新了 Hugging Face 演示和線上 Colab 演示。感謝@fffiloni和@camenduru！我們還在抱臉裡做了一張模特兒卡？並提供了更多測試範例。
[ 2023.10.09 ] DA-CLIP和Universal IR模型的預訓練權重分別在link1和link2發布。此外，我們也為您想要測試自己的圖像的情況提供了一個Gradio應用程式檔案。

如何運行程式碼？

依賴關係

  • 作業系統：Ubuntu 20.04
  • 英偉達：
    • CUDA：11.4
  • 蟒蛇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: { task_name } : { degradations [ index ] } - { text_probs [ 0 ][ index ] } " )

資料集準備

按照我們的論文資料集建立部分準備訓練和測試資料集，如下所示： 

 # ### for training dataset ####
# ### (uncompleted means inpainting) ####
datasets/universal/train
| --motion-blurry
|  | --LQ/ * .png
|  | --GT/ * .png
| --hazy
| --jpeg-compressed
| --low-light
| --noisy
| --raindrop
| --rainy
| --shadowed
| --snowy
| --uncompleted

# ### for testing dataset ####
# ### (the same structure as train) ####
datasets/universal/val
...

# ### for clean captions ####
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 。

資料集連結

降解運動模糊朦朧jpeg 壓縮*弱光嘈雜*（與 jpeg 相同）
數據集戈普羅居住-6k DIV2K+Flickr2K哈哈DIV2K+Flickr2K
降解雨滴下雨的陰影下的下雪的未完成的
數據集雨滴Rain100H：訓練、測試SRD雪100K西萊巴HQ-256

您應該只提取用於訓練的訓練資料集，所有驗證資料集都可以在Google Drive中下載。對於 jpeg 和雜訊資料集，您可以使用此腳本產生 LQ 影像。

訓練

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

# For single GPU:
python3 train.py -opt=options/train.yml

# For distributed training, need to change the gpu_ids in option file
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）影像恢復。

預訓練模型

型號名稱描述Google雲端硬碟抱臉
DA-CLIP退化感知 CLIP 模型下載下載
通用紅外線基於DA-CLIP的通用影像修復模型下載下載
DA-CLIP-混合退化感知 CLIP 模型（加入高斯模糊 + 臉部修復和高斯模糊 + Rainy）下載下載
通用紅外線混合基於 DA-CLIP 的通用影像恢復模型（新增穩健訓練和混合降級）下載下載
野生DA-CLIP野外環境中的降解感知 CLIP 模型 (ViT-L-14)下載下載
野生紅外線基於DA-CLIP的野外影像復原模型下載下載

評估

若要評估我們的影像復原方法，請修改基準路徑和模型路徑並運行

 cd universal-image-restoration/config/universal-ir
python test.py -opt=options/test.yml

格拉迪奧

這裡我們提供了一個 app.py 檔案來測試您自己的圖像。在此之前，您需要下載預先訓練的權重（DA-CLIP和UIR）並修改options/test.yml中的模型路徑。然後只需執行python app.py ，您就可以開啟http://localhost:7860來測試模型。 （我們還在images目錄中提供了幾個具有不同降級的圖像）。我們還提供了來自Google驅動器中測試資料集的更多範例。

相同的步驟可用於野外 (wild-ir) 影像復原。

結果

達克利普

統一圖像修復（點擊展開）

達克利普

針對退化的恢復（點擊展開）

達克利普

野外影像復原（點擊展開）

達克利普

注意！

？在測試中，我們發現目前的預訓練模型仍然難以處理一些現實世界的影像，這些影像可能與我們的訓練資料集（從不同的裝置擷取或具有不同的解析度或退化）發生分佈變化。我們將其視為未來的工作，並將努力使我們的模型更加實用！我們也鼓勵對我們的工作感興趣的用戶使用更大的資料集和更多的退化類型來訓練自己的模型。

？順便說一句，我們還發現直接調整輸入影像的大小會導致大多數任務的表現不佳。我們可以嘗試在訓練中加入調整大小步驟，但由於插值，它總是會破壞影像品質。

？對於修復任務，由於資料集限制，我們目前的模型僅支援臉部修復。我們提供了蒙版範例，您可以使用generate_masked_face腳本來產生未完成的臉部。

致謝：我們的 DA-CLIP 是基於 IR-SDE 和 open_clip。感謝他們的代碼！

接觸

如有任何疑問，請聯絡：[email protected]

引文

如果我們的程式碼對您的研究或工作有幫助，請考慮引用我們的論文。以下是 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}
}

---感謝您的關注！ ---

統計數據
展開
附加信息
相關應用
爲您推薦
相關資訊 全部
網友評論