HCHUNGW的部落格

8.2 圖像分割案例——腦MRI膠質瘤分割

前言

本案例以腦部MRI腫瘤資料為例，介紹圖像分割（本節特指語義分割）的訓練、推理過程。其中，涉及的知識點有：

1. 基於csv的資料集維護管理，及其dataset編寫；
2. smp庫的介紹與使用：segmentation_models_pytorch庫，是語義分割的高級API庫，提供9種分割架構、數百個encoder的backbone及預訓練權重，以及分割的loss和衡量指標計算函數，是語義分割的好幫手，使用它可以快速實現各語義分割功能；
3. 對smp中9中網路架構對比實驗，網路架構分別是：'Unet', 'UnetPlusPlus', 'MAnet', 'Linknet', 'FPN', 'PSPNet', 'DeepLabV3', 'DeepLabV3Plus', 'PAN';
4. 語義分割iou計算時，image-wise與整體計算的差異，當純陰性片時，iou統計應當採用image-wise更合理。
5. 探究不同backbone對於語義分割的效果差異；
6. 探究不同loss對語義分割的效果差異；
7. 探究encoder採用較小學習率時，模型的精度變化。

本案例將從資料介紹、訓練代碼、smp庫使用、對比實驗和模型推理，五個模組進行講解。

資料模組

資料集來自Kaggle，包含110位腦膠質瘤患者的MRI資料，總共3929張圖片，其中有腫瘤區域的圖片為2556張，陰性圖片1373張。

下圖為每個患者圖片數量，以及陰、陽圖片的比例匯總，從圖中可知，一個MRI序列包含20-40張圖片，其中出現腫瘤的圖片有60%左右。（不過這裡需要注意，腫瘤區域的圖元點還還是遠小於陰性圖元點區域的）

下圖為資料集示意，第一行為MRI圖像，第二行為人工標注的腦膠質瘤mask。

資料集劃分

首先下載資料集，解壓得到kaggle_3m資料夾，然後設置資料根目錄data_dir = args.data_path，運行下面代碼，即可得到對應csv

python 01_parse_data.py --data-path /mnt/chapter-8/data/kaggle_3m

Copy

資料集為110個資料夾形式，這裡需要進行資料集劃分，本案例採用csv對資料集進行維護，這裡將會通過01_parse_data.py對資料集進行解析，獲得以下三個csv

• data_info.csv：包含資料夾id、圖片路徑、標籤路徑；
• data_train.csv：根據資料夾id分組劃分的訓練集，比例為80%，有88位元患者的3151張圖片；
• data_val.csv：根據資料夾id分組劃分的驗證集，比例為20%， 有22位元患者的778張圖片；

知識點：資料劃分需要按患者維度劃分，不能基於圖片維度隨機劃分，基於圖片維度隨機劃分會使得模型存在作弊行為，導致模型在真實應用場景下效果很差。

為什麼不能基於圖片維度劃分資料？因為這樣做的話，以為患者的40張圖片，有32張用於訓練，另外8張用於測試，這8張與那32張是非常接近的，因為是同一個人的連續影像。這樣劃分的資料集是帶有偏差的，理所當然的效果很好，模型不容易出現過擬合。後續的實驗也證明了這一點，基於圖片維度劃分的精度要高出10個百分點。

Dataset編寫

dataset編寫就相當容易了，因為資料的路徑資訊已經獲得，因此只需要注意資料讀取進來之後，如何轉換稱為標籤的格式即可。

這裡要注意，對於語義分割，若採用的是交叉熵損失函數，Dice損失函數，它們要求標籤是一個long類型資料，不需要手動轉為one-hot向量，因此對於本實驗，mask要變成一個[256,256]的矩陣，其中每個元素是0或者1。對應的實現代碼如下：

mask = cv_imread(self.df.iloc[idx, 2])

mask[mask == 255] = 1  # 轉換為0, 1 二分類標籤

mask.long()

Copy

訓練代碼

訓練代碼整體結構仍舊沿用第七章第四節中的訓練腳本實現。

在此處需要做的修改主要是，語義分割模型的創建、分割模型的Loss創建、分割模型指標評價，以下四行代碼分別是對應的實現

model = smp.Unet(encoder_name=args.encoder,  encoder_weights="imagenet",  in_channels=3, classes=1)

criterion = smp.losses.DiceLoss(mode='binary')

tp, fp, fn, tn = smp.metrics.get_stats(outputs.long(), labels, mode="binary")

iou_score = smp.metrics.iou_score(tp, fp, fn, tn, reduction="macro")

Copy

可以發現，這裡面無一例外都用到了smp庫，下面簡單介紹smp庫的優點，以及使用方法。

smp庫介紹

segmentation-models-pytorch是pytorch的語義分割工具庫，提供9個分割框架，數百個encoder，常用的loss，指標計算函數，非常方便開發者進行語義分割開發。

這是smp庫的github連結與官方文檔

github: https://github.com/qubvel/segmentation_models.pytorch

docs:https://smp.readthedocs.io/en/latest/

它安裝很方便，只需要pip即可， pip install segmentation-models-pytorch。

掌握pytorch基礎知識的話，smp庫只需要10分鐘即可掌握上手，更系統應用建議配合smp庫的兩個案例進行學習。

下面將從模型創建、loss創建、指標計算三個部分介紹smp使用。

模型創建

語義分割模型發展至今，主要還是採用encoder-decoder的形式，通常會採用主流的CNN作為encoder，decoder部分則進行隨機初始化去訓練。

而encoder與decoder之間如何資訊交互、以及decoder由哪些元件構成等等一系列問題，就引出了不同的語義分割架構。

在smp中，提供了9種常用的語義分割模型架構，分別是'Unet', 'UnetPlusPlus', 'MAnet', 'Linknet', 'FPN', 'PSPNet', 'DeepLabV3', 'DeepLabV3Plus', 'PAN'。

在語義分割中，除了架構、encoder，輸入和輸出的維度也非常重要，這關係到可接收的資料形式是什麼，以及可以預測的類別有多少個。

因此，一個語義分割模型的創建，需要確定架構、選擇encoder、再確定輸入通道數、輸出通道數。

下面介紹unet的創建

import segmentation_models_pytorch as smp

model = smp.Unet(

    encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7

    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization

    in_channels=1,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)

    classes=3,                      # model output channels (number of classes in your dataset)

)

Copy

對於初學者來說，那麼多模型或許是個頭痛的問題，後續也進行了對比實驗，採用相同的encoder（resnet-18），分別訓練9個語義分割架構，觀察精度變化。

從經驗來說，凡是有系列的模型都是應用廣泛、相對可靠的模型，例如net系列，deeplab系列，yolo系列等等。

如何用代碼來實現類屬性的調用，這裡是一個值得學習的程式碼片段，這裡主要通過getattr()方法獲取module的屬性，然後對其進行產生實體即可。

archs = ['Unet', 'UnetPlusPlus', 'MAnet', 'Linknet', 'FPN', 'PSPNet', 'DeepLabV3', 'DeepLabV3Plus', 'PAN']

for arch_str in archs:

    model_class = getattr(smp, arch_str)

    model = model_class(encoder_name=args.encoder,  encoder_weights="imagenet",  in_channels=3, classes=1)

Copy

更多關於分割模型的介紹，可以參見：https://smp.readthedocs.io/en/latest/

損失函數創建

smp提供了8個損失函數，分別是JaccardLoss、DiceLoss、TverskyLoss、FocalLoss、LovaszLoss、SoftBCEWithLogitsLoss、SoftCrossEntropyLoss、MCCLoss。具體差異參見官方文檔，這裡要講講損失函數創建時，需要設置的模式。

損失函數創建時需要設置mode，smp庫提供了3種mode，分別是二分類、多分類、多標籤分類。

二分類：'binary'， 用於一個類的分割（陰性不算一個類，例如本案例），對於二分類，標籤需要是(N, H, W)的形式，元素為0或1，,它不需要通道維度，而模型輸出是跟著pytorch走的，因此仍舊是4D張量，形狀是(N, 1, H, W).

多分類：'multiclass'，多分類是更為常見的場景，例如VOC、COCO資料集，這時標籤元素為0, 1, ..., C-1，類似交叉熵損失函數（可以把語義分割看成是逐圖元分割），模型的輸出自然是(N, C, H, W)了，因為一個圖元點，需要用C維的分類概率向量來做分類。

多標籤：'multilabel'，多標籤語義分割指一個圖元即是A類又是B類，因此它的標籤需要借助C個通道來標注，對應類別設置為1，其它設置為0。所以標籤形式為(N, C, H, W)，模型輸出仍舊是(N, C, H, W)，多標籤需要注意模型輸出時就不再做softmax，而是對每個神經元做sigmoid，以此判斷該類是否存在。

對於loss的選擇，一般是交叉熵損失函數、Dice這兩個系列，其它的可以自行選擇，它就像模型架構一樣，學術界有幾十上百個選擇，但在工程領域，仁者見仁智者見智，更多的語義分割損失函數可參考SegLoss

指標計算

語義分割可以理解為逐圖元的圖像分類，因此圖像分類的各指標也可用於衡量分割模型，但分割與分類不同的是，它注重空間結構資訊，關注集合與集合之間的關係，因此更常用的是IoU或Dice係數來評估模型。

IoU（Intersection over Union，交並比）是用來衡量兩個集合重疊的情況，公式計算為：交集/並集，而dice係數（Dice similarity coefficient，又名dsc）也用於評估兩個集合重疊情況，但是計算公式不一樣，而且根據文章闡述, "Dice傾向於衡量平均性能，而 IoU 傾向於衡量最壞的表現。"

具體計算時，通常先獲得tn, tp, fn, fp，然後計算指標。藍色部分為TP(True Positives)，紅色部分為FN(False Negatives)，黃色部分為(False Positives)

smp工具庫中也是這麼做的，首先根據smp.metrics.get_stats函數，獲得tp, fp, fn, tn。隨後通過各指標計算函數獲得相應指標，

可計算的指標有fbeta_score、f1_score、iou_score、accuracy、precision、recal1、sensitivity、specificity、balanced_accuracy、positive_predictive_value、negative_predictive_value、false_negative_rate、false_positive_rate、false_discovery_rate、false_omission_rate、positive_likelihood_ratio、negative_likelihood_ratio。

來看一段使用示例：

import segmentation_models_pytorch as smp

# lets assume we have multilabel prediction for 3 classes

output = torch.rand([10, 3, 256, 256])

target = torch.rand([10, 3, 256, 256]).round().long()

# first compute statistics for true positives, false positives, false negative and

# true negative "pixels"

tp, fp, fn, tn = smp.metrics.get_stats(output, target, mode='multilabel', threshold=0.5)

# then compute metrics with required reduction (see metric docs)

iou_score = smp.metrics.iou_score(tp, fp, fn, tn, reduction="micro")

f1_score = smp.metrics.f1_score(tp, fp, fn, tn, reduction="micro")

f2_score = smp.metrics.fbeta_score(tp, fp, fn, tn, beta=2, reduction="micro")

accuracy = smp.metrics.accuracy(tp, fp, fn, tn, reduction="macro")

recall = smp.metrics.recall(tp, fp, fn, tn, reduction="micro-imagewise")

Copy

在使用時，需要注意的有2個地方，

1. 計算tp, fp, fn, tn時，模型輸出要轉換為類別標籤，而不是概率向量。
2. 對batch資料統計時，是否需要考慮樣本不平衡問題，進行加權平均，是否需要基於圖像維度進行計算後再平均。

對於問題1，get_stats函數提供了二分類、多標籤時的處理方法，只需要在model下設置'binary' 或 'multiclass'之後，設置threshold即可。從此可看出需要手動進行sigmoid()；

對於問題2，較為複雜一些，smp提供了6個模式，這些在sklearn中有的，分別是，'micro' ， 'macro' ， 'weighted' ， 'micro-imagewise' ， 'macro-imagewise' ， 'weighted-imagewise'。

'micro’ 用於計算總體的指標，不對每個類別進行計算，

'macro'計算每個類別的指標，然後求和平均，不加權。

’weighted’ 計算每個類別的指標，然後根據每類樣本數，進行加權求平均

可參考（https://blog.csdn.net/qq_27668313/article/details/125570210）

對於x-imagewise，表示根據圖片維度進行計算，然後將指標求取平均。

因此，可知道，首碼micro、macro、weighted是決定如何對類別進行求取平均，尾碼-imagewise表示如何對圖片之間進行求平均。

所以，對於binary來說，'micro' = 'macro' = 'weighted' ，並且 'micro-imagewise' = 'macro-imagewise' = 'weighted-imagewise'。

在這裡重點講一下，本案例採用的是macro來統計，因此iou比較低，如果是手寫iou統計，一般會基於圖片維度計算，然後再平均，也就是macro-imagewise。

本案例中，由於大量陰性圖片的存在，所以若不採用-imagewise的話，陰性片雖然預測正確，但實際上是tn很大，而tn缺不計入iou計算中。若採用imagewise，陰性預測正確時，iou為1，從而可以大幅度提高iou值。

下面可以通過代碼觀察，對於陰性片，模型預測完全正確，tp, fp, fn值都是0的情況下，iou也是會被計算為1的。

tp, fp, fn, tn = smp.metrics.get_stats(c, b, mode="binary")

tp

Out[12]: tensor([[0]], device='cuda:0')

fp

Out[13]: tensor([[0]], device='cuda:0')

fn

Out[14]: tensor([[0]], device='cuda:0')

tn

Out[15]: tensor([[65536]], device='cuda:0')

smp.metrics.iou_score(tp, fp, fn, tn, reduction="macro")

Out[16]: tensor(1., device='cuda:0')

Copy

對比實驗

此部分實驗沒有進過嚴格微調，只用於橫向比對，主要觀察不同架構、不同backbone、不同學習率策略之間的差異，可以大體上觀察出一定的趨勢，供大家參考，以便後續選擇模型。

在這裡主要做了4次實驗，分別是：

實驗一：不同backbone的實驗，猜想為越複雜的backbone，精度越高，這裡採用uent-resnet18/34/50來觀察。

實驗二：不同網路架構之間的實驗，猜想是有系列的模型，魯棒性更高，適用性更廣，如unet、deeplab系列。這裡backbone均採用resnet-18，訓練了九個模型。

實驗三：encoder採用小10倍學習率，讓decoder學習率高一些，猜想是要比相同學習率的效果更好，這裡就需要與實驗二中的九個模型精度對比。

實驗四：根據圖片隨機劃分與病人維度劃分的差異，猜想是根據圖片隨機劃分，精度要比病人高的，畢竟用到了極其類似的圖片進行訓練。

實驗完整代碼在github

實驗一：不同backbone的差異

python 02_train_seg.py --batch-size 16 --workers 8 --lr 0.01 --epochs 100 --useplateau --model unet --encoder resnet18

python 02_train_seg.py --batch-size 16 --workers 8 --lr 0.01 --epochs 100 --useplateau --model unet --encoder resnet34

python 02_train_seg.py --batch-size 16 --workers 8 --lr 0.01 --epochs 100 --useplateau --model unet --encoder resnet50

Copy

結論：可證實猜想基本正確，越複雜的backbone，精度越高。但resnet34的精度反而比resnet18要差，這個需要仔細研究，在此不做討論。

實驗二：對比不同模型架構差異

python 03_train_architecture.py --batch-size 16 --workers 4 --lr 0.01 --epochs 100 --useplateau --encoder resnet18

Copy

這裡可以看到unet++和deeplabv3+精度較高，過它們的訓練速度也是最慢的，側面反映出它的模型複雜，理所應當獲得最高精度。

實驗三：encoder採用更小學習率差異

# encoder採用小10倍學習率

python 03_train_architecture.py --batch-size 16 --workers 4 --lr 0.01 --epochs 100 --useplateau --encoder resnet18 --lowlr

Copy

要想獲得好的精度，往往需要各種trick來微調，對於語義分割模型中encoder部分，可以採用較小學習率，因為encoder提取的是共性的語義特徵，對於decoder才需要特有的特徵，因此可以對它們兩個進行差異化設置學習率。為此，對encoder學習率乘以0.1，觀察模型精度變化。

從實驗結果來看，簡單粗暴的調整大都未獲得精度提升，反而都存在3-4個點的掉點，deeplabv3, MAnet， PAN缺有提升，由此可見訓練的trick還是難以適應各個場景的。

綜合實驗二、實驗三，可以觀察到unet系列和deeplab系列都是比較穩定的，這也是它們一直被工業界認可、使用至今的原因，因此推薦首選Unet系列或deepalabv3+進行任務的初步驗證。

實驗四：根據圖片隨機劃分資料集

此實驗需要在01_parse_data.py 代碼中的data_split()函數下修改groupby的物件，需要改為

grouped = dff.groupby('image_path')  # bad method

Copy

並且將csv檔案名做相應修改。

很明顯，基於imgs維度劃分存在很明顯的性能提升，約6個點，但是這個提升是假性的，會迷惑工程師，誤以為模型很好。

這是實際業務場景中常碰到的問題，一定要注意業務資料的維度劃分問題。

模型推理

模型推理與圖像分類類似，沒有什麼特殊的地方，在這裡想重點講一下模型輸出的資料如何轉換為要用的類別mask。

由於是二分類，並且輸出是一通道的矩陣，因此會採用sigmoid將它變為分類概率的形式，然後再通過閾值（一般設為0.5）轉換為0/1的mask矩陣。如下代碼所示：

outputs = model(img_tensor_batch)

outputs_prob = (outputs.sigmoid() > 0.5).float()

outputs_prob = outputs_prob.squeeze().cpu().numpy().astype('uint8')

Copy

接著通過opencv尋找輪廓函數可以得到邊界，最後進行視覺化。

最後通過imageio實現gif圖的生成，可參見05-gen-gif.py。

即可得到下圖

小結

通過本案例，可以掌握：

1. 語義分割模型的訓練與推理流程
2. smp工具庫中的模型創建、損失函數創建、評價指標計算使用
3. 評價指標中，micro、macro、weighted， x-imagewise的差別
4. dice與iou評價指標的定義與差異
5. 9個語義分割架構實驗比對，為後續選擇模型提供參考
6. 醫學資料處理劃分維度，需要基於患者維度，不可基於圖片維度

基於本案例需要進一步學習的內容：

1. 採用多類分割任務進行實驗，瞭解softmax與sigmoid的處理差異；
2. 進一步瞭解unet系列、deeplab系列適合的場景
3. 語義分割後處理：影像處理的形態學操作