基於飛槳Res-Unet網路實現肝臟腫瘤分割任務

飛槳PaddlePaddle發表於2020-09-11

專案背景

近年來,快速發展的深度學習技術已經滲透進了各行各業,醫療方面也不例外。這篇文章我主要介紹如何使用深度學習計算機視覺方法對CT掃描中的肝臟和肝臟腫瘤進行分割。

根據2018年的統計資料[1],肝臟腫瘤是全球第7常見的腫瘤,但致死病例總數卻在所有腫瘤類疾病中排名第二。早發現早治療能有效提升肝臟腫瘤疾病的治癒率,但人工在大量的肝臟CT影像中尋找體積很小的腫瘤工作量極大,也很容易漏檢。這個場景下,使用深度學習演算法自動進行快速、準確的肝臟及肝臟腫瘤分割篩查是一個很好的解決方案。

基於飛槳PaddlePaddle框架,我使用Res-Unet網路結構在 LiTS 資料集[2]上訓練了一個分割網路,最終在肝臟和肝腫瘤上分別達到了 0.92 和 0.77 的分割準確率。LiTS資料集是目前最大的開源肝臟分割資料集,其中包含130名患者的CT掃描和醫生對患者肝臟及腫瘤的分割標註,下圖是資料集中的一個示例:

基於飛槳Res-Unet網路實現肝臟腫瘤分割任務圖1 肝臟分割示例專案在AI Studio上公開,提供包含資料集在內的完整環境,fork後可以直接執行。

https://aistudio.baidu.com/aistudio/projectdetail/250994

此外還有更適合命令列執行的Github開源專案medSeg,經過效能最佳化,訓練及推理速度更快。

https://github.com/davidlinhl/medSeg

網路結構介紹

本文中主要針對專案使用的網路結構,資料預處理及增強,Loss,訓練和推理步驟進行描述。

首先簡單介紹專案中用到的網路結構Res-Unet。在醫學影像領域,Unet[3]結構因為其網路引數規模較小,實現簡單,邊界分割比較準確被廣泛應用。其結構如下圖所示:

基於飛槳Res-Unet網路實現肝臟腫瘤分割任務圖2 Unet網路結構其採用編碼器-解碼器結構,是一個 U 的形狀,因此作者取名Unet。網路首先對輸入圖片進行了左邊的4組卷積和下采樣操作來獲取影像的抽象特徵,之後透過右邊的對稱的4組反摺積和上取樣將影像放大回接近輸入影像的大小。Unet的一個重要創新是在相同深度的下采樣和上取樣操作之間加入了跳轉連線(圖中橫向灰色箭頭所示),有效地提升了網路的分割精度。具體的實現方法一般是將左側卷積block的輸出拼接到右側同一深度反摺積block的輸入上。

這樣反摺積block的輸入特徵圖大小不變,但是厚度變成了原來的兩倍。其中一半是綠色箭頭代表的下層反摺積block的輸入,給網路提供更抽象的高階影像特徵;另一半是灰色箭頭代表的左側卷積block的輸出,給網路提供更準確的位置資訊,提升邊緣分割精度。

我使用的Res-Unet網路在Unet結構的基礎上引入了殘差連線,如下圖所示。具體的做法是新增一條從兩次卷積的輸入到輸出的連線,並做一次卷積操作。這種殘差結構改善了網路的梯度流通,避免網路退化,並能加速網路收斂


基於飛槳Res-Unet網路實現肝臟腫瘤分割任務圖3 殘差連線具體的網路構建程式碼比較複雜,這裡不做詳細展示,可以訪問AI Studio專案或Github repo檢視。

https://github.com/davidlinhl/medSeg/blob/master/medseg/models/unet.py

資料處理及增強

上述的Res-Unet結構是一個2D的分割網路,因此我們首先將LiTS資料集中3D的CT掃描分成2D的切片。CT在拍攝和重建的過程中會引入一些噪聲,因此我們只保留-1024到1024範圍內的資料。經過這兩步處理,可以得到大概1萬張CT掃描切片及對應的分割標籤,隨機選擇一組進行視覺化結果如下:

基於飛槳Res-Unet網路實現肝臟腫瘤分割任務圖4 2D切片在訓練深度神經網路的過程中,我們通常需要在訓練集上訓練多個epoch以讓網路達到一個比較高的訓練準確率。但是這樣做又容易使網路過擬合訓練集,其表現為網路在訓練集上準確率很高但是測試時準確率偏低。針對這個問題有多方面的解決方案,資料增強是其中重要的一種。這個專案中我們採用的資料增強策略包括隨機水平、垂直翻轉、隨機旋轉、隨機尺度縮放、隨機位置裁剪和彈性形變。在專案中可以看到具體程式碼,圖5是對圖4中資料進行資料增強的結果:

基於飛槳Res-Unet網路實現肝臟腫瘤分割任務圖5 資料增強效果CT影像和分割標籤共同進行了左右翻轉,逆時針15度旋轉,0.8倍尺度縮放和彈性形變。雖然一些簡單的資料增強步驟過後影像看起來沒有很大區別,但是隻要影像有變化對演算法來說就是新的資料,結合Droupout、權重正則化等方法能較好地抑制網路過擬合,提升測試準確率

開始訓練前的最後一個步驟是定義損失函式。飛槳PaddlePaddle框架為開發者準備了許多Loss函式,透過幾行程式碼就可以方便地呼叫。這裡我們採用交叉熵和Dice Loss結合作為模型的Loss。Dice評價我們網路分割輸出和資料集中的實際分割結果有多大程度的重合,是我們最終的最佳化目標。但是Dice Loss在訓練過程中不是很穩定,不利於網路收斂,因此加入了交叉熵來穩定訓練。

def create_loss(predict, label, num_classes=2):   predict = fluid.layers.transpose(predict, perm=[0, 2, 3, 1])   predict = fluid.layers.reshape(predict, shape=[-1, num_classes])   predict = fluid.layers.softmax(predict)   label = fluid.layers.reshape(label, shape=[-1, 1])   label = fluid.layers.cast(label, "int64")   dice_loss = fluid.layers.dice_loss(predict, label)  # 計算dice loss   ce_loss = fluid.layers.cross_entropy(predict, label) # 計算交叉熵   return fluid.layers.reduce_mean(ce_loss + dice_loss) # 最後使用的loss是dice和交叉熵的和

模型訓練

萬事俱備,下面可以開始訓練了。首先使用靜態圖API進行組網

with fluid.program_guard(train_program, train_init):   # 定義網路輸入   image = fluid.layers.data(name="image", shape=[3, 512, 512], dtype="float32")   label = fluid.layers.data(name="label", shape=[1, 512, 512], dtype="int32")   # 定義給網路訓練提供資料的loader   train_loader = fluid.io.DataLoader.from_generator(       feed_list=[image, label],       capacity=cfg.TRAIN.BATCH_SIZE * 2,   )   # 建立網路   prediction = create_model(image, 2)   # 定義 Loss   avg_loss = loss.create_loss(prediction, label, 2)   # 定義正則項   decay = paddle.fluid.regularizer.L2Decay(cfg.TRAIN.REG_COEFF)   # 選擇最佳化器   if cfg.TRAIN.OPTIMIZER == "adam":       optimizer = fluid.optimizer.AdamOptimizer(learning_rate=0.003, regularization=decay)   optimizer.minimize(avg_loss)

之後定義讀取資料的reader

def data_reader(part_start=0, part_end=8):   data_names = os.listdir(preprocess_path)   data_part=data_names[len(data_names) * part_start // 10: len(data_names) * part_end // 10] # 取所有資料中80%做訓練資料   random.shuffle(data_part) # 打亂輸入順序   def reader():     for data_name in data_part:         data=np.load(os.path.join(preprocess_path, data_name) )         vol=data[0:3, :, :]         lab=data[3, :, :]         yield (vol, lab)   return reader

將資料增強操作整合進一個函式

def aug_mapper(data):    vol = data[0]    lab = data[1]    vol, lab = aug.flip(vol, lab, cfg.AUG.FLIP.RATIO)    vol, lab = aug.rotate(vol, lab, cfg.AUG.ROTATE.RANGE, cfg.AUG.ROTATE.RATIO, 0)    vol, lab = aug.zoom(vol, lab, cfg.AUG.ZOOM.RANGE, cfg.AUG.ZOOM.RATIO)    vol, lab = aug.crop(vol, lab, cfg.AUG.CROP.SIZE, 0)    return vol, lab

資料增強操作涉及旋轉和彈性形變,計算比較複雜,耗時長,如果只使用單執行緒進行資料讀取和增強會拖慢網路的訓練速度。但使用飛槳PaddlePaddle框架,只需兩行程式碼就可以將單執行緒reader變成多執行緒,大幅提升訓練效率。在AI Studio的測試環境中,8執行緒reader讓訓練速度提升了7倍以上。

train_reader = fluid.io.xmap_readers(aug_mapper, data_reader(0, 8), 8, cfg.TRAIN.BATCH_SIZE * 2) train_loader.set_sample_generator(train_reader, batch_size=cfg.TRAIN.BATCH_SIZE, places=places)

最後一步就是進行訓練,以下是訓練中進行前向和反向梯度傳遞的核心程式碼,其餘的輸出,驗證等操作可以視需要新增。

step = 0 for pass_id in range(cfg.TRAIN.EPOCHS):  for train_data in train_loader():    step += 1    avg_loss_value = exe.run(compiled_train_program, feed=train_data, fetch_list=[avg_loss])    print(step, avg_loss_value)

LiTS資料集比較大,我們選擇的Res-Unet也比較複雜,整個訓練過程大概需要20個epoch,6個小時左右的時間完成。

推理預測

訓練完成後儲存模型,我們就可以對新的資料進行分割了。進行分割前我們同樣需要將資料轉化為2D切片,並保留相同的強度範圍。經過網路前向處理後將資料從2D合併為原來的3D形態

segmentation = np.zeros(scan.shape) with fluid.scope_guard(inference_scope):   # 讀取預訓練權重   [inference_program, feed_target_names, fetch_targets] = fluid.io.load_inference_model(infer_param_path, infer_exe)    for slice_ind in tqdm(range(1, scan.shape[2]-1)):     # 2.5D的輸入,每次取出CT中3個相鄰的層作為模型輸入     scan_slice = scan[:, :, slice_ind - 1: slice_ind + 2]       # 新增batch_size維度     scan_slice = scan_slice[np.newaxis, :, :, :]      # 模型的輸入是 CWH 的, 通道在第一個維度,因此需要將陣列中的第一和第三個維度互換     scan_slice = scan_slice.swapaxes(1,3)      result = infer_exe.run(inference_program, feed={feed_target_names[0]: scan_slice }, fetch_list=fetch_targets)       result = result[0][0][1].reshape([scan.shape[0], scan.shape[1]])     # 儲存分割結果     segmentation[:, :, slice_ind] = result.swapaxes(0,1) # 預測機率超過 0.5 的部分認為是前景,否則認為是背景 segmentation[segmentation >= 0.5] = 1 segmentation[segmentation < 0.5 ] = 0

基於飛槳Res-Unet網路實現肝臟腫瘤分割任務圖6 分割結果深度學習演算法對一組CT掃描進行分割大概耗時15S,其效率明顯高於醫生閱片的效率。而且從分割結果中,我們可以計算獲得肝臟體積,腫瘤數量,腫瘤體積,肝臟腫瘤負擔等數量化的指標,更好地輔助醫生進行診斷。

專案內容到這裡就介紹完了,如果你對深度學習醫療應用感興趣,歡迎加入 AI Studio醫療興趣小組和更多大佬一起學習進步,QQ群號:810823161 

·Reference·

[1] https://pubmed.ncbi.nlm.nih.gov/30207593/

[2] https://aistudio.baidu.com/aistudio/datasetdetail/10273

[3] https://arxiv.org/abs/1505.04597

·飛槳官網地址·

https://www.paddlepaddle.org.cn/

·飛槳開源框架專案地址·

GitHub: 

https://github.com/PaddlePaddle/Paddle 

Gitee: 

https://gitee.com/paddlepaddle/Paddle 

相關文章