這次這篇文章介紹一篇很有意思的工作：SegRefiner，來自 NeurIPS, 2023，目前代碼已開源。

SegRefiner 提出一種新的的任務(wù)解釋，將分割細(xì)化視為一個數(shù)據(jù)生成過程。因此，細(xì)化可以通過一系列去噪擴(kuò)散步驟來實現(xiàn)，其中 Coarse Mask 是 Ground Truth 的噪聲版本。此外，為了處理二值掩模，進(jìn)一步設(shè)計了一種新穎的離散擴(kuò)散過程，在該過程中，每個像素執(zhí)行單向隨機(jī)狀態(tài)轉(zhuǎn)換。所提出的過程可以在訓(xùn)練期間逐漸將 Ground Truth 轉(zhuǎn)換為 Coarse Mask，并在推理期間使用 Coarse Mask 作為采樣起點。換句話說，SegRefiner 將掩模細(xì)化任務(wù)表述為一個條件生成問題，其中輸入圖像（指原圖）作為條件，用于迭代更新/細(xì)化 Coarse Mask 中的錯誤預(yù)測。

SegRefiner 是模型不可知的，因此適用于不同的分割模型和任務(wù)。SegRefiner 驗證的任務(wù)包括：語義分割，實例分割和二分圖像分割。

相關(guān)工作

分割細(xì)化的目標(biāo)是提高現(xiàn)有分割模型中掩模的質(zhì)量。一些工作專注于增強(qiáng)特定的分割模型。還有一些模型不可知方法的細(xì)化方法，這些策略強(qiáng)調(diào)使用多種形式的輸入，包括整個圖像、邊界補(bǔ)丁和邊緣條帶等。盡管這些技術(shù)可以細(xì)化來自不同模型的粗糙掩模，但它們的適用性仍然局限于特定的分割任務(wù)。

擴(kuò)散模型在檢測和分割任務(wù)中的應(yīng)用也成為越來越多研究的焦點，這些研究主要遵循 DDPM 的高斯擴(kuò)散過程，并利用額外的圖像編碼器來提取圖像特征作為生成掩模的條件。SegRefiner 是第一個將擴(kuò)散模型應(yīng)用于圖像分割細(xì)化任務(wù)的工作，它還在基于擴(kuò)散的分割任務(wù)中首次放棄連續(xù)的高斯假設(shè)，轉(zhuǎn)而采用新設(shè)計的離散擴(kuò)散過程。

Forward diffusion process

在介紹前向過程之前，先看一下整體框架的實現(xiàn)：

請?zhí)砑訄D片描述

右側(cè)是提出的 transition sample 模塊，它基于輸入狀態(tài)轉(zhuǎn)換概率從當(dāng)前掩模中隨機(jī)采樣像素，并改變它們的值以匹配目標(biāo)掩模中的值（也就是完成我們在引言中提到的“每個像素執(zhí)行單向隨機(jī)狀態(tài)轉(zhuǎn)換”）?？催@張框架圖的時候，注意區(qū)分不同 pipeline 的顏色區(qū)分，這里吐槽一下，我第一看 transition sample 沒有清楚的看到三種輸入。

在訓(xùn)練期間，轉(zhuǎn)換樣本模塊將 Ground Truth 轉(zhuǎn)換為 Coarse Mask，因此 Coarse Mask 是目標(biāo)掩模。在推理過程中，目標(biāo)掩模指的是預(yù)測的細(xì)致掩模，這個模塊根據(jù)預(yù)測的細(xì)致掩模和轉(zhuǎn)換概率在每個時間步中更新 Coarse Mask 中的值。

接下來仔細(xì)介紹下前向過程：

這種方法涉及將一個精細(xì)的 mask 逐漸轉(zhuǎn)變?yōu)橐粋€更粗糙的 mask，包含以下關(guān)鍵要素：

起始狀態(tài) (mo)：與 ground truth 的精細(xì) mask 相對應(yīng)。

目標(biāo)狀態(tài) (mr)：一個粗糙的 mask。

中間狀態(tài) (mt)：隨著時間步 t 的增加，逐漸從 mo 向 mr 演變的中間狀態(tài)。

轉(zhuǎn)移采樣模塊

引入了“轉(zhuǎn)移采樣”模塊，該模塊負(fù)責(zé)根據(jù)當(dāng)前 mask mt、粗 mask mr 以及狀態(tài)轉(zhuǎn)移概率來進(jìn)行狀態(tài)的轉(zhuǎn)移。狀態(tài)的轉(zhuǎn)移是單向的，保證了最終會收斂到粗糙 mask mr。

重參數(shù)技巧

此外，SegRefiner 還提到了重參數(shù)技巧（reparameterization trick），通過引入二元隨機(jī)變量 x 來描述這一過程。該過程允許直接獲取任何中間時間步的 mask mt，而無需逐步采樣。具體表述為：

定義 x 表示為一個 one-hot 向量，表示中間掩模 mt 中像素 (i, j) 的狀態(tài)。

設(shè)置 xi = [1, 0] 和 xi = [0, 1] 分別表示精細(xì)狀態(tài)和粗糙狀態(tài)。

因此，前向過程可以表示為公式：

where

為超參數(shù)，而 對應(yīng)了上述的狀態(tài)轉(zhuǎn)移概率。狀態(tài)轉(zhuǎn)移矩陣 的邊緣分布可以表示為：

整體上還是遵循 DDPM，但是引入 二元隨機(jī)變量 x 。

Reverse diffusion process

將粗糙 mask mr 逐漸修正為精細(xì) mask mo。

由于此時精細(xì) mask mo 和狀態(tài)轉(zhuǎn)移概率未知，對照著 SegRefiner 框架圖來看，訓(xùn)練一個神經(jīng)網(wǎng)絡(luò) 來預(yù)測精細(xì) ，表示為：

其中 I 是相應(yīng)的圖像。

和 分別表示預(yù)測的精細(xì) mask 和其置信度分?jǐn)?shù)。相應(yīng)的， 可視作 中每個像素處于“精細(xì)狀態(tài)”的概率。

反向狀態(tài)轉(zhuǎn)移概率

根據(jù)前向過程的設(shè)定和貝葉斯定理，延續(xù) DDPM 的方法，我們可以由前向過程的后驗概率和預(yù)測的 得到反向過程的概率分布，表示為：

where

其中 為反向過程的狀態(tài)轉(zhuǎn)移概率。

迭代修正過程

給定粗糙 mask mr 以及相應(yīng)的圖像 I，首先將所有像素初始化為粗糙狀態(tài) xi = [1, 0]。通過不斷迭代地狀態(tài)轉(zhuǎn)移，逐漸修正 m_T 中的預(yù)測值。

推理過程

給定一個粗糙的掩模和其對應(yīng)的圖像，我們首先初始化所有像素為粗糙狀態(tài)。我們在以下步驟之間迭代：

前向過程：以獲取 和 。

計算反向狀態(tài)轉(zhuǎn)移矩陣： 并得到 。

計算精細(xì)化的掩模：基于 ， 和 計算精細(xì)化的掩模 。

這個過程（1）-（3）迭代進(jìn)行，直到獲得精細(xì)的掩模。

實驗

分別訓(xùn)練了 LR-SegRefiner 和 HR-SegRefiner，數(shù)據(jù)集和具體的 settings 在上 。

定性對比其他方法，覺得對比 U-Net 和 ISNet 的效果的確很明顯。

總結(jié)

SegRefiner 是首個基于擴(kuò)散的圖像分割細(xì)化方法，采用了離散擴(kuò)散過程。SegRefiner 執(zhí)行模型不可知的分割細(xì)化，并在各種分割任務(wù)的細(xì)化中取得了強(qiáng)有力的實證結(jié)果。雖然它在準(zhǔn)確度上取得了顯著的提升，但其局限性在于擴(kuò)散過程由于多步迭代策略而導(dǎo)致推理速度變慢。