基于Diffusion Probabilistic Model的醫(yī)學圖像分割

前言

在前面的一篇文章ICLR 2023：基于 diffusion adversarial representation learning 的血管分割中，我們已經(jīng)介紹過了 diffusion model 在醫(yī)學圖像分割上的一個應用，推薦對 diffusion model 不了解的同學優(yōu)先閱讀，其中講了一些基本概念。上一篇文章是將 diffusion 應用到自監(jiān)督學習中，而 MedSegDiff 是一個有監(jiān)督的框架，現(xiàn)在已更新到 V2 版本， V2 版本區(qū)別于 V1 使用了 Transformer，且適用于多分類。MedSegDiff-V1 已被接收在 MIDL 2023。

MedSegDiff

MedSegDiff 在原版 DPM 的基礎上引入了動態(tài)條件編碼，增強 DPM 在醫(yī)學圖像分割上的分步注意力能力。特征頻率解析器（FF-Parser）可以消除分割過程中損壞的給定掩碼中的高頻噪聲。DPM 是一種生成模型，由兩個階段組成，正向擴散階段和反向擴散階段。在正向過程中，通過一系列步驟 T，將高斯噪聲逐漸添加到分割標簽 x0 中。在反向過程中，訓練神經(jīng)網(wǎng)絡通過反向噪聲過程來恢復原始數(shù)據(jù)：

其中 theta 是反向過程參數(shù)。從高斯噪聲開始，p(xT) 表示原始圖像，反向過程將潛在變量分布 p(xT) 轉換為數(shù)據(jù)分布 p(x0)。反向過程逐步恢復噪聲圖像，以獲得最終的清晰分割。該模型使用 U-Net 作為學習網(wǎng)絡，步長估計函數(shù)由原始圖像先驗條件確定：

其中 EI 是條件特征嵌入，即原始圖像嵌入，Ex 是當前步驟的分割映射特征嵌入。這兩個組件被添加并發(fā)送到 U-Net 的解碼器進行重建。步長索引 t 與新增的嵌入和解碼器功能集成在一起，使用共享的 look-up table 進行嵌入，這在 DDPM 的論文中有介紹。總而言之，MedSegDiff 模型基于 DPM，使用 U-Net 進行學習。步長估計函數(shù)由原始圖像先驗得到，步長索引與新增的嵌入和解碼器功能集成在一起。使 MedSegDiff 在三項具有不同圖像模式的醫(yī)學分割任務中表現(xiàn)不錯。先看下 MedSegDiff 整體流程圖：

下面我們將逐一介紹動態(tài)條件編碼和 FF-Parser。

動態(tài)條件編碼

對于 MRI 或超聲之類的低對比度圖像，很難將感興趣的對象與背景分開。所以使用動態(tài)條件編碼方法來解決這個問題?？梢宰⒁獾剑紙D像包含準確的目標分割信息，不過很難與背景區(qū)分開，而當前步驟的 grand truth 包含增強的目標區(qū)域，但不準確。

為了整合這兩個信息來源，使用類似注意力的機制將條件特征圖的每個尺度與當前步驟的編碼特征融合。這種融合是首先對兩個特征圖分別應用層歸一化，然后將它們相乘以獲得 affine map 來實現(xiàn)的，再將 affine map 與條件編碼特征相乘以增強注意力區(qū)域。如 MedSegDiff 流程圖所示，此操作應用于中間兩個階段，其中每個階段都是在 Resnet34 之后實現(xiàn)的卷積階段。但是，集成當前條件編碼功能可能會產(chǎn)生額外的高頻噪聲。為了解決這個問題，使用 FF-Parser 來限制特征中的高頻分量。

FF-Parser

FF-Parser 的流程如下圖所示：

FF-Parser 是頻率濾波器的可學習版本，它可以全局調整特定頻率的分量，以限制高頻分量進行自適應集成。首先使用二維 FFT（快速傅立葉變換）沿空間維度對解碼器特征圖 m 進行變換，生成頻譜 M。然后，將參數(shù)化的注意力地圖 A 與 M 相乘以調整頻譜，得出 M'。最后，使用逆向 FFT 將 M' 反向回空間域，以獲得修改后的特征圖 m'。使用 FF-Parser 可以學習適用于傅里葉空間特征的權重圖，該權重圖可用于全局調整特定頻率的分量。這種技術不同于空間注意力，后者調整特定空間位置的組成部分。

實驗

下圖分別是腦部 MRI、眼底視盤和甲狀腺結節(jié)的超聲圖像分割結果的可視化，可以看出 MedSegDiff 在簡單解刨結構的二分類上效果還是不錯的。

對比其他 SOTA 方法的結果如下表：

MedSegDiff-V2

MedsegDiff-v2 通過將 Transformer 機制整合到原始的U-Net骨干中，增強了基于擴散的 MedSegDiff-v1。具體來說，MedsegDiff-v2 引入了新的 Spectrum-Space Transformer（SS-former），對噪聲和語義特征之間的相互作用進行建模。驗證了 Medsegdiff-v2 對具有不同模態(tài)圖像的五個分割數(shù)據(jù)集的十八個器官的有效性。

概述

如下圖所示，MedsegDiff-v2 結合了錨點條件和語義條件兩種不同的條件方式，以提高擴散模型的性能。錨點條件將錨分割特征（條件模型的解碼分割特征）集成到擴散模型的編碼特征中。即允許使用粗略但靜態(tài)的參照來初始化擴散模型，有助于減少擴散方差。

然后將語義條件強加于擴散模型的 embedding，理解為將條件模型的語義 embedding 集成到擴散模型的 embedding 中。這種條件集成由 SS-former 實現(xiàn)，它彌合了噪聲和語義嵌入之間的鴻溝，并利用 Transformer 的全局和動態(tài)特性抽象出更強的特征表達形式。

Medsegiff-v2 是使用 DPM 的標準噪聲預測損失 Lnoise 和錨損失 Lanchor 進行訓練的。Lanchor 是 Dice loss 和 CE loss 的組合。總損失函數(shù)表示為：

其中 t ≡ 0 (mod α) 通過超參數(shù) α 控制監(jiān)督條件模型的時間，β 是另一個用于加權交叉熵損失的經(jīng)驗超參數(shù)?？偠灾?，Medsegdiff-v2 顯著提高了 MedsegDiff 的性能。該方法采用了新的基于 Transformer 的條件 U-Net 框架和兩種不同的條件方式，以提高擴散模型的性能。

Anchor Condition with Gaussian Spatial Attention

與卷積層相比，Transformer 具有更強的表示性，但對輸入方差更敏感。為了克服這種負面影響，所以使用了錨條件運算，如上面概述中的介紹，該運算將條件模型的解碼分割特征（錨點）集成到擴散模型的編碼器特征中。此外，還使用了高斯空間注意力來表示條件模型中給定分割特征的不確定性（概率）。

在第一個公式中，表示在錨點特征上應用高斯核以進行平滑激活，因為錨點可能不完全準確，且高斯核的均值和方差是可以學習的。選擇平滑 ground truth 和原始圖之間的最大值以保留最相關的信息，從而生成平滑的錨特征。在第二個公式中，將平滑錨點特征集成到擴散模型中以獲得增強特征。首先應用 1x1 卷積將錨特征中的通道數(shù)減少到 1（經(jīng)常作用于解碼器的最后一層）。最后，在錨點特征上使用 sigmoid 激活函數(shù)，將其添加到擴散模型的每個通道中，類似于空間注意力的實現(xiàn)。

Semantic Condition with SS-Former

關于 SS-Former 的作用，我們只做簡單的總結。對比 MedSegDiff-v1，是一種將條件模型分割 embedding 集成到擴散模型 embedding 中的新架構，其使用頻譜空間注意力機制來解決擴散和分割 embedding 之間的域差距。此外，注意力機制在傅里葉空間中合并語義和噪聲信息，和 MedSegDiff-v1 是類似的。

實驗

Medsegiff-v2 是可以進行多分類任務的，只是目前 Github 還沒有維護多分類的代碼，不過 Medsegiff-v2 的實現(xiàn)代碼已經(jīng)提交了。下圖為在腹部 CT 圖像中的多器官分割上的表現(xiàn)。

對比其他 SOTA 方法的結果如下表：

總結

關于 diffusion model 在醫(yī)學圖像分割上的應用，其是否能適應多分類且復雜的解刨結構還需要進一步探索，訓練和推理時的效率也低于常規(guī)的有監(jiān)督神經(jīng)網(wǎng)絡，可以進一步優(yōu)化。

審核編輯：劉清