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

Abstract

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

Introduction

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

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

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

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

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

Method

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

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

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

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

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

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

如圖 2 所示，假設(shè) $\mathcal{C}$ 表示文章中考慮的圖像退化類(lèi)型的集合， 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$ 表示某種退化類(lèi)型，$K?1$ 表示總的退化類(lèi)型的數(shù)量，我們還添加了一種特殊類(lèi)型 $c_K="no"$ 表示多步圖像恢復(fù)的結(jié)束標(biāo)識(shí)。文本提示描述集 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$，我們的語(yǔ)義無(wú)關(guān) CLIP 旨在識(shí)別 $I$ 的主要退化并提取相應(yīng)的文本嵌入。SA-CLIP 包含一個(gè)圖像編碼器 ${\epsilon }_I$ 和一個(gè)文本編碼器 ${\epsilon }_T$。首先獲得圖像嵌入 ${\epsilon }_I\in {\mathbb{R}}^d$ 和文本嵌入 ${\epsilon }_T\in {\mathbb{R}}^{K\times d}$，然后計(jì)算圖像嵌入與每個(gè)文本嵌入的余弦相似度。

$$\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$ 個(gè)文本嵌入，對(duì)計(jì)算得到的余弦相似度用 softmax 計(jì)算每個(gè)相似度量的概率 $\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ì)量評(píng)估的簡(jiǎn)單微調(diào)，在 CLIP 模型的優(yōu)化期間，我們凍結(jié)文本編碼器 ${\epsilon }_T$ 的參數(shù)，并使用多類(lèi)別保真度損失微調(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)$ 表示一個(gè)二分類(lèi)的變量，如果某個(gè)退化類(lèi)型占主導(dǎo)，那么該變量值為 1，否則為 0 。

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

為了解決這個(gè)問(wèn)題，我們提出了一種稱(chēng)為語(yǔ)義無(wú)關(guān)約束損失 $L_{SA}$ 的新方法來(lái)規(guī)范微調(diào)過(guò)程，并防止模型僅僅依賴(lài)語(yǔ)義信息而非圖像質(zhì)量。當(dāng) CLIP 模型表明在真實(shí)干凈圖像 $I_{gt}$（對(duì)應(yīng)于退化圖像 $I$）中存在偽影 $c_i$ 時(shí)，語(yǔ)義無(wú)關(guān)損失 $L_{SA}$ 會(huì)施加懲罰。這種懲罰迫使 CLIP 模型根據(jù)圖像質(zhì)量區(qū)分 $I_{gt}$ 和 $I$，鼓勵(lì) CLIP 圖像編碼器 ${\epsilon }_I$ 專(zhuān)注于提取圖像質(zhì)量信息而非語(yǔ)義信息。這種約束損失可以通過(guò)以下等式推導(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）階段旨在在一個(gè)共享的框架中處理多種退化。基于擴(kuò)散的生成模型的最新進(jìn)展已經(jīng)展示了它們生成多樣化圖像的卓越能力，使其適用于多任務(wù)圖像恢復(fù)。先前的研究已經(jīng)表明，生成模型具有生成缺失或扭曲細(xì)節(jié)的卓越能力，特別是對(duì)于需要虛構(gòu)的任務(wù)，例如超分辨率?；谶@些見(jiàn)解，我們基于隱擴(kuò)散模型（LDM）進(jìn)行 AIR 階段。LDM 結(jié)合了文本和圖像嵌入條件，使用生成先驗(yàn)來(lái)恢復(fù)圖像 $I_{sd}$。文本嵌入條件 e = { e a u t o , e u s e r } e = \{ e_{auto}, e_{user} \} e={eauto?,euser?} 旨在區(qū)分不同類(lèi)型的圖像退化，而來(lái)自 LDM 的圖像編碼器 ${\epsilon }_{ldm}$ 的潛在圖像嵌入條件 $z_I={\epsilon }_{ldm}(I)$ 提供了結(jié)構(gòu)信息。

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

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

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

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

在訓(xùn)練階段，我們?yōu)閳D像恢復(fù)任務(wù)對(duì)潛在擴(kuò)散模型（LDM）的 UNet 的 backbone ${?}_{\theta }(e,[z_t,z_I],t)$ 進(jìn)行微調(diào)，其目標(biā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)}$$

對(duì)于結(jié)構(gòu)校正潛在擴(kuò)散模型（SC-LDM），我們不使用耗時(shí)的完整逆向采樣過(guò)程來(lái)生成編輯后的隱變量 ${\hat{z}}_t$，而是利用通過(guò)以下方式計(jì)算得到的估計(jì)編輯后的隱變量 $\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)校正潛在擴(kuò)散模型（SC-LDM）的損失函數(shù)進(jìn)一步定義為：

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

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

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