SegRefiner:通過擴散模型實現高精度圖像分割

背景介紹

盡管圖像分割在過去得到了廣泛研究和快速發展，但獲得細節上非常準確的分割 mask 始終十分具有挑戰性。因為達成高精度的分割既需要高級語義信息，也需要細粒度的紋理信息，這將導致較大的計算復雜性和內存使用。而對于分辨率達到2K甚至更高的圖像，這一挑戰尤為突出。

由于直接預測高質量分割 mask 具有挑戰性，因此一些研究開始集中于 refine 已有分割模型輸出的粗糙 mask。為了實現高精度的圖像分割，來自北京交大、南洋理工、字節跳動等的研究者們引入了一種基于擴散模型Diffusion去逐步提高mask質量的方法。

現有方法

Model-Specific

一類常見的 Refinement 方法是 Model-Specific 的，其通過在已有分割模型中引入一些新模塊，從而為預測 Mask 補充了更多額外信息，從而增強了已有模型對于細節的感知能力。這一類方法中代表性的工作有 PointRend，RefineMask，MaskTransfiner等。然而，這些方法是基于特定模型的改進，因此不能直接用于 refine 其他分割模型輸出的粗糙 mask。

Model-Agnostic

另一類 Refinement 方法是 Model-Agnostic 的，其只使用原始圖像和粗糙mask作為輸入信息，如 BPR，SegFix，CascadePSP，CRM 等。由于這類方法在 Refinement 過程中未使用已有模型的中間特征，因此不依賴于特定分割模型，可以用于不同分割模型的 Refinement。然而，盡管這類方法能夠有效地提升分割準確度，但由于粗糙 mask 中存在多種多樣的錯誤預測（如下圖所示），導致模型無法穩定地修正粗糙 mask 中的全部預測錯誤。

實現目標

相比于 Model-Specific 的方法，Model-Agnostic 的方法能夠直接應用于不同分割模型的 Refinement，從而有著更高的實用價值。更進一步地，由于不同分割任務（語義分割，實例分割等）的結果都可以被表示為一系列 binary mask，具有相同的表征形式，在同一個模型中統一實現不同分割任務的 Refinement 同樣是可能的。因此，我們希望實現能夠應用于不同分割模型和分割任務的通用 Refinement 模型。

如前所述，已有分割模型產生的錯誤預測是多種多樣的，而想要通過一個通用模型一次性地更正這些多樣性的錯誤十分困難。面對這一問題，在圖像生成任務中取得巨大成功的擴散概率模型給予了我們啟發：擴散概率模型的迭代策略使得模型可以在每一個時間步中僅僅消除一部分噪聲，并通過多步迭代來不斷接近真實圖像的分布。這大大降低了一次性擬合出目標數據分布的難度，從而賦予了擴散模型生成高質量圖像的能力。

直觀地，如果將擴散概率模型的策略遷移到 Refinement 任務中，可以使得模型在進行 Refinement 時每一步僅關注一些“最明顯的錯誤”，這將降低一次性修正所有錯誤預測的難度，并可以通過不斷迭代來逐漸接近精細分割結果，從而使得模型能夠應對更具挑戰性的實例并持續糾正錯誤，產生精確分割結果。

在這一想法下，我們提出了一個新的視角：將粗糙 mask 視作 ground truth 的帶噪版本，并通過一個去噪擴散過程來實現粗糙 mask 的 Refinement，從而將 Refinement 任務表示為一個以圖像為條件，目標為精細 mask 的數據生成過程。

算法方案

擴散概率模型是一種由前向和反向過程表示的生成模型，其中前向過程通過不斷加入高斯噪聲得到不同程度的帶噪圖像，并訓練模型預測噪聲；而反向過程則從純高斯噪聲開始逐步迭代去噪，最終采樣出圖像。而將擴散概率模型遷移到 Refinement 任務中，數據形式的不同帶來了以下兩個問題：

(1) 由于自然圖像往往被視作高維高斯變量，將圖像生成的過程建模為一系列高斯過程是十分自然的，因此已有的擴散概率模型大多基于高斯假設建立；而我們的目標數據是 binary mask，通過高斯過程擬合這樣一個離散變量的分布并不合理。

(2) 作為一種分割 Refinement 方法，我們的核心思想是將粗糙 mask 視為帶有噪聲的 ground truth，并通過消除這種噪聲來恢復高質量的分割結果。這意味著我們擴散過程的結尾應當收斂到確定性的粗糙 mask（而非純噪聲），這也與已有的擴散概率模型不同。

針對上述問題，我們建立了如下圖所示的基于“隨機狀態轉移”的離散擴散過程。其中，前向過程將 ground truth 轉換為“不同粗糙程度”的 mask，并用于訓練；而反向過程用于模型推理，SegRefiner 從給出的粗糙 mask 開始，通過逐步迭代修正粗糙 mask 中的錯誤預測區域。以下將詳細介紹前向和反向過程。

前向過程

前向過程的目標是將 ground truth 提供的精細 mask 逐步消融為粗糙的 mask，記前向過程每一步的變量為 ，則前向過程應當滿足：

(1) 為 ground truth

(2) 為粗糙 mask

(3) 介于和之間，隨 t 增大逐漸向粗糙 mask 演變 ? ?

基于這些限制條件，我們用隨機狀態轉移來表述前向過程：假設變量中的每一個像素都有兩種可能的狀態：精細和粗糙，處于精細狀態的像素值與保持一致，處于粗糙狀態的像素則取 的值（即使二者一致）。我們提出了一個“轉移采樣”模塊來進行這一過程，如上圖右側所示。在每一個時間步，其以當前 mask ，粗糙 mask 以及狀態轉移概率作為輸入。

在前向過程中，狀態轉移概率描述了當前 mask 中每個像素轉移到中的狀態的概率。根據狀態轉移概率進行采樣，可以得到后一個時間步 每個像素的狀態，從而確定其取值。這一模塊確定了一個“單向”過程，即只會發生“轉移到目標狀態”的情況。這一單向性質確保了前向過程會收斂到（盡管每一步都是完全隨機的），從而滿足了上述限制 條件(2)，(3)。

通過重參數技巧，我們引入了一個二元隨機變量來描述上述過程：我們將 表示為一個 one-hot 向量，用來表示中間掩模 中像素 的狀態，同時設置 和 來表示精細狀態和粗糙狀態。因此，前向過程可以被表示為：

其中 為超參數，而 對應了上述的狀態轉移概率，是狀態轉移矩陣。則前向過程的邊緣分布可以表示為：

其中 。從而我們可以直接獲取任何中間時間步 的 mask 并用于訓練，而無需逐步采樣 。

反向過程

反向擴散過程用于模型推理，目標是將粗糙 mask 逐漸修正為精細 mask 。由于此時精細 mask 和狀態轉移概率未知，類似 DDPM 的做法，我們訓練一個神經網絡 ，來預測精細 mask ，表示為：

其中 是相應的圖像。和 分別表示預測的精細 mask 和其置信度分數。這里表示了網絡對于預測準確與否的置信度，故同樣可以被視作中每個像素處于“精細狀態”的概率。為了獲得反向狀態轉移概率，根據前向過程的設定和貝葉斯定理，延續 DDPM 的做法，我們可以由前向過程的后驗概率和預測的得到反向過程的概率分布，為：? ?

其中 為反向過程的狀態轉移概率。給定粗糙 mask 以及相應的圖像，我們首先將所有像素初始化為粗糙狀態 ，然后通過不斷迭代地狀態轉移，逐漸修正 中預測值。下圖為一個推理過程的可視化展示。

模型結構

任意滿足 形式的網絡均可滿足我們的要求，這里我們延續了之前工作的做法，采用 U-Net 作為我們的去噪網絡，將其輸入通道數修改為4（圖像和 在通道維度上串聯），并輸出1通道的改進掩模。

算法評估

由于 Refinement 任務的核心是獲取細節精確的分割結果，在實驗中我們選取了三個代表性的高質量分割數據集，分別對應Semantic Segmentation，Instance Segmentation 和 Dichotomous Image Segmentation。

Semantic Segmentation

如表1所示，我們在 BIG 數據集上將提出的 SegRefiner 與四種已有方法：SegFix，CascadePSP，CRM 以及 MGMatting 進行了對比。其中前三個為語義分割的 Refinement 方法，而 MGMatting 使用圖像和 mask 進行 Matting 任務，也可以用于 Refinement 任務。結果表明，我們提出的 SegRefiner 在 refine 四個不同語義分割模型的粗糙 mask 時，都在 IoU 和 mBA 兩項指標上獲得了明顯提升，且超越了之前的方法。

? ?

Instance Segmentation

實例分割中，我們選擇了之前的工作廣泛使用的 COCO 數據集進行測試，并使用了 LVIS 數據集的標注。與原始 COCO 標注相比，LVIS 標注提供了更高質量和更詳細的結構，這使得 LVIS 標注更適合評估改進模型的性能。

首先，在表2中，我們將提出的SegRefiner與兩種 Model-Agnostic 的實例分割 Refinement 方法 BPR 和 SegFix 進行了比較。結果表明我們的 SegRefiner 在性能上明顯優于這兩種方法。

然后在表3中，我們將 SegRefiner 應用于其他7種實例分割模型。我們的方法在不同準確度水平的模型上都取得了顯著的增強效果。值得注意的是，當應用于三種 Model-Specific 的實例分割 Refinement 模型（包括PointRend、RefineMask 和 Mask TransFiner）時，SegRefiner 依然能穩定提升它們的性能，這說明 SegRefiner 具有更強大的細節感知能力。

Dichotomous Image Segmentation

Dichotomous Image Segmentation 是一個較新提出的任務，如下圖所示，其數據集包含大量具有復雜細節結構的對象，因此十分適合評估我們 SegRefiner 對細節的感知能力。

在本實驗中，我們將 SegRefiner 應用于6種分割模型，結果如表4所示?？梢钥吹?，我們的SegRefiner在 IoU 和 mBA 兩項指標上都明顯提升了每個分割模型的準確度。

可視化展示

審核編輯：劉清