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

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

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

相關工作

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

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

Forward diffusion process

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

請添加圖片描述

右側是提出的 transition sample 模塊，它基于輸入狀態轉換概率從當前掩模中隨機采樣像素，并改變它們的值以匹配目標掩模中的值（也就是完成我們在引言中提到的“每個像素執行單向隨機狀態轉換”）。看這張框架圖的時候，注意區分不同 pipeline 的顏色區分，這里吐槽一下，我第一看 transition sample 沒有清楚的看到三種輸入。

在訓練期間，轉換樣本模塊將 Ground Truth 轉換為 Coarse Mask，因此 Coarse Mask 是目標掩模。在推理過程中，目標掩模指的是預測的細致掩模，這個模塊根據預測的細致掩模和轉換概率在每個時間步中更新 Coarse Mask 中的值。

接下來仔細介紹下前向過程：

這種方法涉及將一個精細的 mask 逐漸轉變為一個更粗糙的 mask，包含以下關鍵要素：

起始狀態 (mo)：與 ground truth 的精細 mask 相對應。

目標狀態 (mr)：一個粗糙的 mask。

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

轉移采樣模塊

引入了“轉移采樣”模塊，該模塊負責根據當前 mask mt、粗 mask mr 以及狀態轉移概率來進行狀態的轉移。狀態的轉移是單向的，保證了最終會收斂到粗糙 mask mr。

重參數技巧

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

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

設置 xi = [1, 0] 和 xi = [0, 1] 分別表示精細狀態和粗糙狀態。

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

where

為超參數，而 對應了上述的狀態轉移概率。狀態轉移矩陣 的邊緣分布可以表示為：

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

Reverse diffusion process

將粗糙 mask mr 逐漸修正為精細 mask mo。

由于此時精細 mask mo 和狀態轉移概率未知，對照著 SegRefiner 框架圖來看，訓練一個神經網絡 來預測精細 ，表示為：

其中 I 是相應的圖像。

和 分別表示預測的精細 mask 和其置信度分數。相應的， 可視作 中每個像素處于“精細狀態”的概率。

反向狀態轉移概率

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

where

其中 為反向過程的狀態轉移概率。

迭代修正過程

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

推理過程

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

前向過程：以獲取 和 。

計算反向狀態轉移矩陣： 并得到 。

計算精細化的掩模：基于 ， 和 計算精細化的掩模 。

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

實驗

分別訓練了 LR-SegRefiner 和 HR-SegRefiner，數據集和具體的 settings 在上 。

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

總結

SegRefiner 是首個基于擴散的圖像分割細化方法，采用了離散擴散過程。SegRefiner 執行模型不可知的分割細化，并在各種分割任務的細化中取得了強有力的實證結果。雖然它在準確度上取得了顯著的提升，但其局限性在于擴散過程由于多步迭代策略而導致推理速度變慢。