論文閱讀：AutoDIR: Automatic All-in-One Image Restoration with Latent Diffusion
這是 ECCV 2024 的一篇文章，利用擴散模型實現(xiàn)圖像恢復(fù)的任務(wù)。

Abstract

這篇文章提出了一個創(chuàng)新的 all-in-one 的圖像恢復(fù)框架，融合了隱擴散技術(shù)，各種不同的圖像退化都可以用這一個模型搞定，簡稱 AutoDIR。AutoDIR 模型可以自動的識別以及恢復(fù)一系列未知的圖像退化。AutoDIR 提供了直觀的開放式詞匯圖像編輯功能，使用戶能夠根據(jù)自己的喜好定制和增強圖像。AutoDIR 由兩個關(guān)鍵階段組成：一個基于語義無關(guān)的視覺語言模型的盲圖像質(zhì)量評估（BIQA）階段，它會自動檢測輸入圖像中的未知圖像退化類型；一個一體化圖像恢復(fù)（AIR）階段，利用結(jié)構(gòu)校正的隱擴散來處理多種類型的圖像退化。大量的實驗評估表明，AutoDIR 在更廣泛的圖像恢復(fù)任務(wù)中優(yōu)于當(dāng)前最先進的方法。AutoDIR 的設(shè)計還能夠?qū)崿F(xiàn)靈活的用戶控制（通過文本提示），并作為圖像恢復(fù)的基礎(chǔ)模型推廣到新的任務(wù)中。

Introduction

文章作者探索了一種能夠處理單個圖像的多種未知退化的通用模型。為實現(xiàn)這一目標，相應(yīng)的模型應(yīng)具備以下能力：（1）分解和區(qū)分未知的退化類型，（2）一個與具體任務(wù)無關(guān)的框架，能夠恢復(fù)各種退化，（3）理想情況下，允許用戶根據(jù)自己的視覺偏好自由調(diào)整恢復(fù)結(jié)果。為解決類似問題已經(jīng)有了很多相關(guān)的工作，但沒有一個能夠同時滿足這三點。

為了能同時解決上述三個問題，文章作者提出了一個名為 AutoDIR 的流程，它滿足上述所有三種能力，并且能夠自動檢測和恢復(fù)具有多種未知退化的圖像。AutoDIR 由兩個階段組成：語義無關(guān)的盲圖像質(zhì)量評估（SA-BIQA）階段和由 SA-BIQA 中生成的文本提示引導(dǎo)的一體化圖像恢復(fù)（AIR）階段。

在 SA-BIQA 階段，我們能夠以開放詞匯的方式準確識別未知偽影情況下的每種退化。這是通過我們提出的語義無關(guān)的 CLIP（SA-CLIP）模型實現(xiàn)的，該模型采用了語義無關(guān)的正則化項，將原始的語義識別 CLIP 轉(zhuǎn)換為語義無關(guān)的形式，重點關(guān)注圖像的結(jié)構(gòu)質(zhì)量而非語義內(nèi)容。此外，我們可以利用 SA-BIQA 階段生成的文本嵌入作為指令來引導(dǎo)進一步的恢復(fù)模型。這種方法不僅能夠?qū)崿F(xiàn)有效的恢復(fù)，而且通過提供開放詞匯的指令，還允許在運行時進行靈活的用戶控制和編輯。

AIR 階段是使用在廣泛任務(wù)上聯(lián)合訓(xùn)練的多任務(wù)圖像恢復(fù)模型來處理退化。鑒于不同任務(wù)的多樣性（例如，像超分辨率這樣的一些任務(wù)需要生成紋理，但像低光增強這樣的其他任務(wù)需要保留除亮度之外的一切），我們提出了一種混合方法，在通過引入額外的結(jié)構(gòu)歸納偏差來保持圖像結(jié)構(gòu)一致性的同時，最大限度地發(fā)揮擴散模型的生成能力。

為了評估 AutoDIR 的有效性和泛化能力，我們進行了一組全面的實驗，涵蓋了七個圖像恢復(fù)任務(wù)，包括去噪、運動去模糊、低光增強、去霧、去雨、去雨滴和超分辨率。實驗結(jié)果表明，AutoDIR 始終優(yōu)于當(dāng)前最先進的方法。AutoDIR 還針對屏下攝像頭和水下攝像頭拍攝的圖像恢復(fù)進行了評估，這些是具有多種未知退化的成像系統(tǒng)的例子。

Method

圖 2 展示了所提出的 AutoDIR（具有潛在擴散的自動一體化圖像恢復(fù)）的總體流程圖，這是一個能夠自動檢測和處理圖像中多種未知退化的統(tǒng)一模型。AutoDIR 包括兩個主要階段：

• 語義無關(guān)的盲圖像質(zhì)量評估（SA-BIQA）：此階段會自動識別輸入圖像中主要存在的退化，如噪聲、模糊、霧霾，并生成相應(yīng)的文本提示，記為 $e_{auto}$，隨后在圖像恢復(fù)過程中使用。
• 一體化圖像恢復(fù) (AIR)：此階段利用結(jié)構(gòu)校正潛在擴散模型（SC-LDM），在來自 SA-BIQA 的文本嵌入 $e_{auto}$ 或用戶自定義的開放詞匯指令 $e_{user}$ 的引導(dǎo)下，生成恢復(fù)后的圖像 $I_{res}$。

Semantic-Agnostic Blind Image Quality Assessment (SA-BIQA)

正如之前的工作所表明的，評估圖像退化的一種常見方法是專門為此目的訓(xùn)練一個圖像分類器。然而，這種簡單的方法在處理包含廣泛偽影的大型數(shù)據(jù)集時面臨挑戰(zhàn)。即使使用重型的 ViT 編碼器，僅基于圖像信息對圖像退化進行分類的準確率也限制在 77.65%。

為了解決這一限制，我們提出利用人類語言知識來增強對圖像退化的檢測。我們引入了一個語義無關(guān)的 CLIP（SA-CLIP）模型作為我們盲圖像質(zhì)量評估（BIQA）的主干。SA-CLIP 基于 CLIP 模型構(gòu)建，該模型在人類語言知識和圖像質(zhì)量之間建立了聯(lián)系。然而，我們觀察到，直接應(yīng)用 CLIP 或為 BIQA 任務(wù)對其進行簡單微調(diào)并不能產(chǎn)生可靠的結(jié)果。我們深入研究了這個問題，并找出了這個問題背后的原因。預(yù)訓(xùn)練的 CLIP 模型主要是為視覺識別任務(wù)而訓(xùn)練的，這些任務(wù)優(yōu)先考慮語義信息而不是圖像質(zhì)量。因此，它在 BIQA 任務(wù)中的準確率較低。例如，該模型可能難以區(qū)分低光的狗圖像和有噪點的狗圖像，因為它更關(guān)注 “狗” 的方面，而不是噪聲或光照的存在。

為了克服這個問題，我們分兩步解決：（i）我們?yōu)槲⒄{(diào) CLIP 構(gòu)建了一個新的圖像質(zhì)量評估任務(wù)。（ii）我們提出了一個新的正則化項，用于語義無關(guān)和圖像質(zhì)量感知的訓(xùn)練，以導(dǎo)出 SA-CLIP 模型。

如圖 2 所示，假設(shè) $\mathcal{C}$ 表示文章中考慮的圖像退化類型的集合， C = { c 1 , c 2 , . . . , c K ? 1 , c K } \mathcal{C}=\{c_1, c_2, ..., c_{K-1}, c_K \} C={c1?,c2?,...,cK?1?,cK?}，其中 $c_i$ 表示某種退化類型，$K?1$ 表示總的退化類型的數(shù)量，我們還添加了一種特殊類型 $c_K="no"$ 表示多步圖像恢復(fù)的結(jié)束標識。文本提示描述集 T = { T ∣ T = "A?photo?needs?ci?artifact?reduction , c ∈ C } \mathcal{T}=\{T|T = \text{"A photo needs ci artifact reduction}, c \in \mathcal{C} \} T={T∣T="A?photo?needs?ci?artifact?reduction,c∈C}。給定一張包含了若干未知偽影的受損圖像 $I$，我們的語義無關(guān) CLIP 旨在識別 $I$ 的主要退化并提取相應(yīng)的文本嵌入。SA-CLIP 包含一個圖像編碼器 ${\epsilon }_I$ 和一個文本編碼器 ${\epsilon }_T$。首先獲得圖像嵌入 ${\epsilon }_I\in {\mathbb{R}}^d$ 和文本嵌入 ${\epsilon }_T\in {\mathbb{R}}^{K\times d}$，然后計算圖像嵌入與每個文本嵌入的余弦相似度。

$$\text{logit}(c_i|I)=\frac{{\epsilon }_I(I)?{\epsilon }_T(T_i)}{{\Vert {\epsilon }_I(I)\Vert }_2{\Vert {\epsilon }_T(T)\Vert }_2}\text{\hspace{1em}(1)}$$

其中，$T_i$ 表示第 $i$ 個文本嵌入，對計算得到的余弦相似度用 softmax 計算每個相似度量的概率 $\hat{p}(c_i|I)$：

$$\hat{p}(c_i|I)=\frac{\exp (\text{logit}(c_i|I))}{{\sum }_{i=1}^K\exp (\text{logit}(c_i|I))}\text{\hspace{1em}(2)}$$

$$e_{auto}=\sum _{i=1}^K\hat{p}(c_i|I){\epsilon }_T(T_i)\text{\hspace{1em}(3)}$$

圖像質(zhì)量評估的簡單微調(diào)，在 CLIP 模型的優(yōu)化期間，我們凍結(jié)文本編碼器 ${\epsilon }_T$ 的參數(shù)，并使用多類別保真度損失微調(diào)圖像編碼器 ${\epsilon }_I$。保真度損失可以表示為：

$$L_{FID}=1?\sum _{i=1}^K\sqrt{y(c_i|I)\hat{p}(c_i|I)}\text{\hspace{1em}(4)}$$

其中，$y(c_i|I)$ 表示一個二分類的變量，如果某個退化類型占主導(dǎo)，那么該變量值為 1，否則為 0 。

圖像質(zhì)量評估的語義無關(guān)約束微調(diào)，由于原始的 CLIP 模型是在諸如圖像分類等任務(wù)上進行預(yù)訓(xùn)練的，其相應(yīng)的 ${\epsilon }_I$ 編碼器傾向于根據(jù)圖像的語義信息（例如，貓或狗）而不是圖像質(zhì)量（例如，有噪點或清晰）對圖像進行編碼。當(dāng)我們根據(jù)圖像質(zhì)量微調(diào) CLIP 模型以生成用于 BIQA 的文本時，這成為一個顯著的限制。如圖 3 a）和 b）所示，由原始 CLIP 和在有霧圖像上微調(diào)的 CLIP 提取的圖像嵌入，以及它們相應(yīng)的真實干凈圖像，無法分開，這表明其重點在于語義信息而非圖像質(zhì)量差異。

為了解決這個問題，我們提出了一種稱為語義無關(guān)約束損失 $L_{SA}$ 的新方法來規(guī)范微調(diào)過程，并防止模型僅僅依賴語義信息而非圖像質(zhì)量。當(dāng) CLIP 模型表明在真實干凈圖像 $I_{gt}$（對應(yīng)于退化圖像 $I$）中存在偽影 $c_i$ 時，語義無關(guān)損失 $L_{SA}$ 會施加懲罰。這種懲罰迫使 CLIP 模型根據(jù)圖像質(zhì)量區(qū)分 $I_{gt}$ 和 $I$，鼓勵 CLIP 圖像編碼器 ${\epsilon }_I$ 專注于提取圖像質(zhì)量信息而非語義信息。這種約束損失可以通過以下等式推導(dǎo)得出：

$$L_{SA}=\sum _{i=1}^K\sqrt{y(c_i|I)\hat{p}(c_i|I_{gt})}\text{\hspace{1em}(5)}$$

將 $L_{SA}$ 與 $L_{FID}$ 結(jié)合，得到最終的微調(diào) loss：

$$L_{BIQA}=L_{FID}+\lambda L_{SA}\text{\hspace{1em}(6)}$$

All-in-one Image Restoration (AIR)

一體化圖像恢復(fù)（AIR）階段旨在在一個共享的框架中處理多種退化?；跀U散的生成模型的最新進展已經(jīng)展示了它們生成多樣化圖像的卓越能力，使其適用于多任務(wù)圖像恢復(fù)。先前的研究已經(jīng)表明，生成模型具有生成缺失或扭曲細節(jié)的卓越能力，特別是對于需要虛構(gòu)的任務(wù)，例如超分辨率?；谶@些見解，我們基于隱擴散模型（LDM）進行 AIR 階段。LDM 結(jié)合了文本和圖像嵌入條件，使用生成先驗來恢復(fù)圖像 $I_{sd}$。文本嵌入條件 e = { e a u t o , e u s e r } e = \{ e_{auto}, e_{user} \} e={eauto?,euser?} 旨在區(qū)分不同類型的圖像退化，而來自 LDM 的圖像編碼器 ${\epsilon }_{ldm}$ 的潛在圖像嵌入條件 $z_I={\epsilon }_{ldm}(I)$ 提供了結(jié)構(gòu)信息。

然而，盡管基于 LDM 的生成模型可以為多任務(wù)圖像恢復(fù)提供基礎(chǔ)，但由于具有變分自編碼器（VAE）的壓縮重建過程，它們在重建具有復(fù)雜和小結(jié)構(gòu)的圖像方面存在局限性，有工作試圖通過在特定類別的圖像（例如，人臉）上重新訓(xùn)練變分自編碼器（VAE）網(wǎng)絡(luò)以學(xué)習(xí)專門的概率分布來減少壓縮重建過程引起的失真。然而，由于圖像內(nèi)容的多樣性，這種方法不適用于圖像恢復(fù)任務(wù)。為了解決這些限制，我們向 LDM 引入了一個輕量級的插件式結(jié)構(gòu)校正模塊，增強了其在圖像恢復(fù)期間處理復(fù)雜和小結(jié)構(gòu)的能力。

結(jié)構(gòu)校正潛在擴散模型（SC-LDM）， 雖然基于 LDM 的生成模型可以為多任務(wù)圖像恢復(fù)提供基礎(chǔ)，但人們普遍注意到它們可能無法保持原始圖像結(jié)構(gòu)，例如人臉和文本，如圖 4 所示。為了解決結(jié)構(gòu)失真問題，我們采用了一個有效的結(jié)構(gòu)校正模塊（SCM），記為 $\mathcal{F}$。SCM 的目的是以殘差的方式從原始圖像中提取上下文信息 $\mathcal{R}$，并將其與中間圖像恢復(fù)結(jié)果 $I_{sd}$ 相結(jié)合。這通過以下等式實現(xiàn)

$$I_{res}=I_{sd}+w?\mathcal{F}([I_{sd},I])\text{\hspace{1em}(7)}$$

其中 [ ] 表示連接，并且 w 是一個可調(diào)節(jié)系數(shù)，其范圍在 0 到 1 之間。w 的值決定了利用上下文信息來恢復(fù)最終結(jié)果的程度。w 的值較大時強調(diào)上下文信息的使用，這對于需要結(jié)構(gòu)一致性的任務(wù)（例如低光增強）是有益的。相反，w 的值較小時通常用于保持潛在擴散模型對于像超分辨率這樣的任務(wù)的生成能力。通過集成 SCM，AutoDIR 有效地恢復(fù)了原始圖像的失真上下文，如 圖 4 所示，無縫地整合了在編輯階段所做的增強。

在訓(xùn)練階段，我們?yōu)閳D像恢復(fù)任務(wù)對潛在擴散模型（LDM）的 UNet 的 backbone ${?}_{\theta }(e,[z_t,z_I],t)$ 進行微調(diào)，其目標函數(shù)為：

$$L_{LD}={\mathbb{E}}_{{\epsilon }_{ldm}(x),c_I,e,?,t}[{\Vert ??{?}_{\theta }(e,[z_t,z_I],t)\Vert }_2^2]\text{\hspace{1em}(8)}$$

對于結(jié)構(gòu)校正潛在擴散模型（SC-LDM），我們不使用耗時的完整逆向采樣過程來生成編輯后的隱變量 ${\hat{z}}_t$，而是利用通過以下方式計算得到的估計編輯后的隱變量 $\tilde{z}$：

$$\tilde{z}=\frac{z_t}{\sqrt{\bar{\alpha }}}?\frac{\sqrt{1?\bar{\alpha }}({?}_{\theta }(e,[z_t,z_I],t))}{\sqrt{\bar{\alpha }}}\text{\hspace{1em}(9)}$$

其中 $\alpha$ 表示引入的噪聲調(diào)度器。結(jié)構(gòu)校正潛在擴散模型（SC-LDM）的損失函數(shù)進一步定義為：

$$L_{AIR}=\Vert I_{gt}?(\mathcal{F}(\mathcal{D}(\tilde{z}),I)+\mathcal{D}(\tilde{z}))\Vert \text{\hspace{1em}(10)}$$

處理多任務(wù)圖像恢復(fù)的機制，

圖 5 展示了我們在逆向擴散過程中探索文本條件解開不同圖像恢復(fù)任務(wù)機制的實驗。我們發(fā)現(xiàn)不同的文本條件會產(chǎn)生不同的交叉注意力圖。如圖 5 所示，改變文本提示會導(dǎo)致交叉注意力圖發(fā)生顯著變化。該圖與文本提示緊密對齊，對于 “去霧” 提示，注意力在整個圖像上均勻分布；對于 “低分辨率” 提示，注意力集中在具有邊緣或紋理的部分；對于 “去雨滴” 提示，注意力集中在諸如雨滴等特定區(qū)域。這表明 AutoDIR 可以將擴散注意力引導(dǎo)到更有可能存在圖像偽影的區(qū)域。