医学图像分割 | U-Net网络及他的变体
點擊上方“AI算法修煉營”，選擇“星標”公眾號
精選作品，第一時間送達
今天來簡單歸納一下，醫學圖像分割的內容，還有很多代表性的工作在之后的文章中不斷更新，由于同一個網絡結構可能在不同的數據集上表現出不一樣的性能，在具體的任務場景中還是要結合數據集來選擇合適的網絡。
醫學圖像特點
1.圖像語義較為簡單、結構較為固定。我們做腦的，就用腦CT和腦MRI，做胸片的只用胸片CT，做眼底的只用眼底OCT，都是一個固定的器官的成像，而不是全身的。由于器官本身結構固定和語義信息沒有特別豐富，所以高級語義信息和低級特征都顯得很重要(UNet的skip connection和U型結構就派上了用場)。舉兩個例子直觀感受下。
2.數據量少。醫學影像的數據獲取相對難一些，很多比賽只提供不到100例數據。所以我們設計的模型不宜多大，參數過多，很容易導致過擬合。
原始UNet的參數量在28M左右(上采樣帶轉置卷積的UNet參數量在31M左右)，而如果把channel數成倍縮小，模型可以更小。縮小兩倍后，UNet參數量在7.75M。縮小四倍，可以把模型參數量縮小至2M以內，非常輕量。個人嘗試過使用Deeplab v3+和DRN等自然圖像語義分割的SOTA網絡在自己的項目上，發現效果和UNet差不多，但是參數量會大很多。
3.多模態。相比自然影像，醫療影像比較有趣和不同的一點是，醫療影像是具有多種模態的。以ISLES腦梗競賽為例，其官方提供了CBF,MTT,CBV,TMAX,CTP等多種模態的數據。
這就需要我們更好的設計網絡去提取不同模態的特征feature。這里提供兩篇論文供大家參考。
Joint Sequence Learning and Cross-Modality Convolution for 3D Biomedical Segmentation(CVPR 2017) ,
Dense Multi-path U-Net for Ischemic Stroke Lesion Segmentation in Multiple Image Modalities.
4.可解釋性重要。由于醫療影像最終是輔助醫生的臨床診斷，所以網絡告訴醫生一個3D的CT有沒有病是遠遠不夠的，醫生還要進一步的想知道，病灶在哪一層，在哪一層的哪個位置，分割出來了嗎，能求體積嘛？同時對于網絡給出的分類和分割等結果，醫生還想知道為什么，所以一些神經網絡可解釋性的trick就有用處了，比較常用的就是畫activation map。看網絡的哪些區域被激活了，如下圖。
這里推薦兩篇工作：周博磊老師的Learning Deep Features for Discriminative Localization(CVPR2016)和其實驗室同學的 Deep Learning for Identifying Metastatic Breast Cancer（上圖的出處）
（周博磊知乎主頁：https://www.zhihu.com/people/zhou-bo-lei/activities）
U-Net
論文地址：https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/?
代碼地址：https://github.com/wolny/pytorch-3dunet
自2015年以來，在生物醫學圖像分割領域，U-Net得到了廣泛的應用，該方法在2015年MICCAI會議上提出，目前已達到四千多次引用。至今，U-Net已經有了很多變體，目前已有許多新的卷積神經網絡設計方式，但很多仍延續了U-Net的核心思想，加入了新的模塊或者融入其他設計理念。U-Net架構在不同的生物醫學分割應用中實現了非常好的性能。由于借助具有彈性形變的數據增強功能，它只需要少量的的帶標注的圖像，并且在NVidia Titan GPU（6 GB）上僅需要10個小時的訓練時間。
U-Net如下圖所示，是一個encoder-decoder結構，左邊一半的encoder包括若干卷積，池化，把圖像進行下采樣，右邊的decoder進行上采樣，恢復到原圖的形狀，給出每個像素的預測。
具體來說，左側可視為一個編碼器，右側可視為一個解碼器。編碼器有四個子模塊，每個子模塊包含兩個卷積層，每個子模塊之后有一個通過max pool實現的下采樣層。輸入圖像的分辨率是572x572, 第1-5個模塊的分辨率分別是572x572, 284x284, 140x140, 68x68和32x32。由于卷積使用的是valid模式，故這里后一個子模塊的分辨率等于（前一個子模塊的分辨率-4）/2。解碼器包含四個子模塊，分辨率通過上采樣操作依次上升，直到與輸入圖像的分辨率一致（由于卷積使用的是valid模式，實際輸出比輸入圖像小一些）。該網絡還使用了跳躍連接，將上采樣結果與編碼器中具有相同分辨率的子模塊的輸出進行連接，作為解碼器中下一個子模塊的輸入。
?U-Net架構
架構中的一個重要修改部分是在上采樣中還有大量的特征通道，這些通道允許網絡將上下文信息傳播到具有更高分辨率的層。因此，拓展路徑或多或少地與收縮路徑對稱，并產生一個U形結構。
在該網絡中沒有任何完全連接的層，并且僅使用每個卷積的有效部分，即分割映射僅包含在輸入圖像中可獲得完整上下文的像素。該策略允許通過重疊平鋪策略對任意大小的圖像進行無縫分割，如圖所示。為了預測圖像邊界區域中的像素，通過鏡像輸入圖像來推斷缺失的上下文。這種平鋪策略對于將網絡應用于大型的圖像非常重要，否則分辨率將受到GPU內存的限制。
?對任意大型圖像進行無縫分割的重疊平鋪策略
對于可用訓練數據非常少的情況，可以通過對可用的訓練圖像應用彈性變形來進行數據增強。這使得網絡學習這種變形的不變性，而不需要在標注圖像語料庫中看到這些變形。這在生物醫學分割中尤其重要，因為變形曾是組織中最常見的變化，并且可以有效地模擬真實的變形。
許多細胞分割任務中的另一個挑戰是分離同一類的接觸目標。為此建議使用加權損失，其中在接觸單元之間分開的背景標簽在損失函數中獲得較大的權重。
相比于FCN和Deeplab等，UNet共進行了4次上采樣，并在同一個stage使用了skip connection，而不是直接在高級語義特征上進行監督和loss反傳，這樣就保證了最后恢復出來的特征圖融合了更多的low-level的feature，也使得不同scale的feature得到了的融合，從而可以進行多尺度預測和DeepSupervision。4次上采樣也使得分割圖恢復邊緣等信息更加精細。
3D U-Net
論文地址：https://arxiv.org/pdf/1606.06650.pdf
代碼地址：https://github.com/wolny/pytorch-3dunet
3D U-Net是U-Net的一個簡單擴展，應用于三維圖像分割，結構如下圖所示。相比于U-Net，該網絡僅用了三次下采樣操作，在每個卷積層后使用了batch normalization，但3D U-Net和U-Net均沒有使用dropout。
為了避免瓶頸，在上采樣和下采樣之前都將通道數增加為原來的二倍。左側的紅色虛線方框之內，在進行 maxpooling 之前將通道數從 64 變為了 128，及文中所說的 double；右側的紅色虛線內，在進行轉置卷積之前，將通道數從 256 變為了 512。實際上這種思想來源于 inception V3。使用加權的 softmax 損失函數，將沒有標簽的部分置為 0，讓模型只從有有標簽的部分學習。
在2018年MICCAI腦腫瘤分割挑戰賽(brats)中，德國癌癥研究中心的團隊使用3D U-Net，僅做了少量的改動，取得了該挑戰賽第二名的成績，發現相比于許多新的網絡，3D U-Net仍然十分具有優勢。
Res-UNet 和Dense U-Net
Res-UNet和Dense-UNet分別受到殘差連接和密集連接的啟發，將UNet的每一個子模塊分別替換為具有殘差連接和密集連接的形式。
將Res-UNet用于視網膜圖像的分割，其結構如下圖所示，其中灰色實線表示各個模塊中添加的殘差連接。
密集連接即將子模塊中某一層的輸出分別作為后續若干層的輸入的一部分，某一層的輸入則來自前面若干層的輸出的組合。
下圖是密集連接的一個例子。該文章中將U-Net的各個子模塊替換為這樣的密集連接模塊，提出Fully Dense UNet 用于去除圖像中的偽影。
MultiResUNet
論文地址：https://arxiv.org/abs/1902.04049.pdf
代碼地址：https://github.com/nibtehaz/MultiResUNet
MultiResUNet提出了一個MutiRes模塊與UNet結合。MutiRes模塊如下圖（c）所示，是一個殘差連接的擴展，在該模塊中三個3x3的卷積結果拼接起來作為一個組合的特征圖，再與輸入特征圖經過1x1卷積得到的結果相加。
該網絡的結構圖如下圖所示，其中各個MultiRes模塊的內部即為上圖所示。
該網絡除了MultiRes模塊以外，還提出了一個殘差路徑(ResPath), 使編碼器的特征在與解碼器中的對應特征拼接之前，先進行了一些額外的卷積操作，如下圖所示。作者認為編碼器中的特征由于卷積層數較淺，是低層次的特征，而解碼器中對應的特征由于卷積層更深，是較高層次的特征，二者在語義上有較大差距，推測不宜直接將二者進行拼接。因此，使用額外的ResPath使二者在拼接前具有一致的深度，在ResPath1， 2， 3， 4中分別使用4，3，2，1個卷積層。
該文章在ISIC、CVC-ClinicDB、Brats等多個數據集上驗證了其性能。
Attention UNet
論文地址：https://arxiv.org/pdf/1804.03999.pdf
代碼地址：
https://github.com/ozan-oktay/Attention-Gated-Networks
Attention UNet在UNet中引入注意力機制，在對編碼器每個分辨率上的特征與解碼器中對應特征進行拼接之前，使用了一個注意力模塊，重新調整了編碼器的輸出特征。該模塊生成一個門控信號，用來控制不同空間位置處特征的重要性，如下圖中紅色圓圈所示。
該方法的注意力模塊內部如下圖所示，該模塊通過1x1x1的卷積分別與ReLU和Sigmoid結合，生成一個權重圖?, 通過與編碼器中的特征相乘來對其進行校正。
下圖展示了注意力權重圖的可視化效果。從左至右分別是一幅圖像和隨著訓練次數的增加該圖像中得到的注意力權重?？梢姷玫降淖⒁饬嘀貎A向于在目標器官區域取得大的值，在背景區域取得較小的值，有助于提高圖像分割的精度。
Non-local U-Nets for Biomedical Image Segmentation（AAAI2020）
論文地址：https://arxiv.org/abs/1812.04103
代碼地址：https://github.com/Whu-wxy/Non-local-U-Nets-2D-block
global aggregation block
global aggregation block如上圖所示，圖看起來很復雜，實際上運算過程很簡單。
1.該結構與Attention Is All You Need這篇文章的形式很相似。輸入圖為X（B*H*W*C），經過QueryTransform和1*1卷積，轉換為Q（B*Hq*Wq*Ck），K（B*H*W*Ck），V(B*H*W*Cv）。QueryTransform可以為卷積，反卷積，插值等你喜歡的方法，最后的輸出結果的H與W將與這個值一致。
2.代碼里在Unfold之前有Multi-Head（Attention Is All You Need）的操作，不過在論文中沒有說明，實際上是把通道分為N等份。Unfold是把Batch,height,width,N通道合并,Q（B*Hq*Wq*N*ck），K（B*H*W*N*ck），V(B*H*W*N*cv）。
3.接下來是經典的點積attention操作，得到一個權值矩陣A((B*Hq*Wq*N)*(B*H*W*N))，用于self-attention的信息加權，分母Ck是通道數，作用是調節矩陣的數值不要過大，使訓練更穩定（這個也是Attention Is All You Need提出的）。最后權值矩陣A和V點乘，得到最終的結果（(B*Hq*Wq*N)*cv），可見輸出的height和width由Q決定，通道數由V決定。
4.我最近看了兩篇上采樣的論文，DUpsample和CARAFE，現在很多上采樣相關的論文關注于在上采樣時用CNN擴大感受野，增加圖像局部信息。這篇文章提出的global aggregation block是一個將注意力機制和上/下采樣相結合的方法，關注全圖信息，感受野更大，可以在其他任務上試用一下，效果如何還是要看實踐的結果。
Non-local U-Nets
文章提出的Non-local U-Nets如上圖所示。
相比U-Net，卷積層數減少，圖像下采樣倍率從16倍變成4倍，保留了更多空間信息。encoder和decoder之間的skip connections用相加的方式，不是拼接的方式，讓模型推理速度更快。
上下采樣使用global aggregation block，使分割圖更準確。
參考文章：
[1]?Section for Biomedical Image Analysis (SBIA)
（?https://www.med.upenn.edu/sbia/brats2018/data.html ）
[2] https://zhuanlan.zhihu.com/p/57530767
[3] https://zhuanlan.zhihu.com/p/57237535
[4]https://www.zhihu.com/question/269914775
總結
以上是生活随笔為你收集整理的医学图像分割 | U-Net网络及他的变体的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： Java入门 - Java初体验
- 下一篇： 在Python中，查看torch、cud