怎麼用 CNN 語意切割，來偵測核磁共振影像裡的腦部腫瘤？

急診、神經內外科的醫師，經常需要判讀 MRI 的 Brain Tumor，包含位置、大小、種類… 等，深度學習也想湊上一腳看能不能幫上什麼忙 🤓

在電腦視覺的應用上，除了分類（Classification）以外，還有不同的任務：物件偵測（Object Detection）、語意切割（Semantic Segmentation）、物件切割（Instance Segmentation）。用下圖來當例子：
1️⃣ 左上圖做分類：輸出圖中羊、狗、貓、馬的多元分類機率
2️⃣ 左下圖做物件偵測：訓練時會有已經 Label 好的種類，模型則預測圖片裡哪裡有出現預設的哪些種類，會把預測到的物件框起來，比分類任務還要多做預測分類物件的座標位置與長寬大小
3️⃣ 右上圖做語意切割：相對簡單且常用的切割任務，目的要抓出每個類別的輪廓，針對原始圖片的每個像素做個別分類，再把相同語意類別塗上相同的顏色，只要屬於「羊」這個語意的 Pixel 就都塗上藍色，不管是不是同一隻羊都會連在一起
4️⃣ 右下圖做物件切割：比語意切割複雜一點，還要知道每個物體的範圍，圖中三隻羊雖然都是羊，但我們希望切割出三隻羊個別的輪廓。做法是先把圖片做物件偵測，得到三隻羊的範圍，再把範圍內的小圖做個別切割，就能得到每隻羊的輪廓

我們今天要聊的是語意切割（Semantic Segmentation），用的是 Fully Convolutional Network（FCN）的概念，把全連接層（Fully Connected）替換為卷積層（Convolutional）。

上次我們聊到肺炎 X 光的分類任務，就像下圖裡上面有張 256 x 256 Pixel 大小的彩色圖片，經過 CNN 之後接上全連接層，得到 1000 個分類的個別機率值，例子中機率最高的類別是「Tabby Cat」，我們就知道這張圖是哪一個類別，可以想像 Output 是 1 x 1 大小、總共有 1000 個類別。而前面聊到，切割就是針對每個 Pixel 進行所屬類別的 Dense Prediction，拔掉分類曾、膨脹成大小變成 16 x 16 且深度 1000，預測 1000 張大小 16 x 16 的圖，把「Tabby Cat」這張圖拿出來，數值比較大的用紅色、小的用深藍色來塗，這樣我們就知道圖的右邊有一隻 Tabby Cat 的範圍，來完成切割任務，也就是 Output 變成針對長寬 16 x 16 大小的每一個像素、分 1000 類的結果。

中間的神經網絡我們可以自行設計串接，重點是輸出要把不斷在 CNN 時變小的 Feature Maps 還原成和原圖相同，256 x 256 x 21（深度要多深，是看我們要做切割分類的數量，這裡的 21 只是這篇論文中用 VOC 資料庫的類別切割數量）的大小。

我們在用 CNN 抽取圖片特徵時，越深層的特徵越抽象，這在做分類任務沒什麼問題，因為分類不用考量到物體的位置在哪裡；但是切割任務除了抽象的語意資訊之外，我們還需要空間上的資訊，Convolution 會不斷縮小 Feature Maps，縮小的過程中會濃縮、遺失空間的資訊，FCN 就會用跟富含空間資訊的淺層特徵圖合併的方式，來去捕捉遺漏的空間資訊。用這張肚子 CT 的影像為例，假設原圖是 320 x 320 的大小，經過了 5 次卷積和 MaxPooling 後得到 10 x 10 的 Feature Map（Conv 7）。但因為輸出的圖片大小要和原圖一樣，如果我們要直接暴力放大回原圖的大小，就是直接放大 32 倍；或是我們先把 10 x 10 放大 2 倍之後，跟 Pool 4 同樣大小 20 x 20 合併在一起，再放大 16 倍；或是我們把 Conv 7 放大 4 倍、Pool 4 放大 2 倍之後，跟 Pool 3 同樣大小 40 x 40 合併在一起，再放大 8 倍，這是更細緻的做法，把淺層（空間資訊）和深層（語意資訊）合併混疊來完成切割任務。

放大有兩種方法：UpSampling / UnSampling 和 Transposed Convolution，前者就像反向 MaxPooling，後者 Transposed Convolution 比較複雜，它也是一種卷積，只是反向把 2 x 2 卷積成 4 x 4，有參數、有 Filter 可以學習的 Upsampling，Stride 設成 2 就可以放大成 2 倍。

介紹完 FCN，我們來看看其他切割網路，大致上分成兩種架構 Encoder（資訊萃取）- Decoder（解碼）
👉 前面是一般的 Convolution，不斷去縮小 Feature Maps、濃縮資訊
👉 後面是還原且輸出和原圖一樣大小的結果

縮小我們用 MaxPooling、放大則用 UnPooling 來做。

2015 年有個知名的網路 UNet，架構簡單、對稱且漂亮，寫起來也不困難，形狀看起來像是英文字母 U。Input 是 216 x 256 的黑白圖片，經過簡單的 3 x 3 Convolution 和 ReLU 後做 MaxPooling 變一半的大小，持續重複做到 Encoder Output 變成 27 x 32 x 256，前面縮小幾遍、後面就要放大幾遍，UNet 採用 Up-Conv 來放大、再和前面同等大小的 Feature Maps 混疊、放大、Convolution… 不斷重複，跟前面聊到的架構差不多，對稱地去疊加特徵圖。倒數第二層的長寬就和原圖一樣了，不過深度是 32，所以透過 1 x 1 的Convolution，把 Feature Maps 的深度轉成我們想要的數值，這篇論文最後是分成 2 類。

在切割任務中，當然也可以計算一整張圖裡每個像素分類完的結果、跟正確答案比對來看，不過目前衡量 Segmentation 模型表現的指標有兩種：IOU 和 Dice Coefficient。IOU 是交集的面積除上聯集的面積，越準則 IOU 越大。Dice 計算和 IOU 很像，分子是 2 倍的交集（True Positive）面積、分母是預測和答案各自的面積大小相加，分子分母個別比 IOU 多一個 True Positive。

聊完 Segmentation 的網路和評估，我們來看個運用在腦部腫瘤（Brain Tumor）的實例吧！資料集是從 The Cancer Imaging Archive (TCIA)、The Cancer Genome Atlas (TCGA) 裡 110 個病人的影像資料中取得，內含 2D 的 MR Image，圖片的深度是 3 不代表是彩色，而是用不同參數照射出來的結果；我們目標是預測出腦部腫瘤的位置，答案是 Label 在 FLAIR Sequence（在檔案名稱裡有 _mask 的 tif 檔），大家可以去下載資料集來玩玩看，因為原始資料集沒有分 Train 或 Test，可以自己手動把要 Train 和 Test 的檔案放到設定的資料夾裡。

先把會用到的套件 Import 進來、把檔案放到環境中能讀取的空間。

import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
from glob import glob
from tqdm.auto import tqdm
import imgaug.augmenters as iaa
import imgaug as ia
from tensorflow.keras import *
import tensorflow.keras.backend as K

隨機地挑選一張影像來讀取，原始圖片的 Shape 是 256 x 256 x 3，因為有 3 個 Channel，畫圖時 matplotlib 會以為是彩色的圖，但實際上我們切 MRI 時是用不同的參數照同樣的位置 3 次再把結果對位後疊加在一起，到時候 Model 也是一次要看 3 張圖，所以讓我們把它們分開來看。

# show channelwise image 
plt.figure(figsize=(20, 5))
for i in range(3):
    plt.subplot(1, 4, i+1)
    plt.imshow(img[:,:,i], cmap='gray')
plt.subplot(1, 4, 4)
plt.imshow(mask)
plt.show()

左邊 3 張就是分開來看的 MR 結果，最右邊則是 Ground Truth，背景是黑色 (0)、正確答案是白色 (255)。

接下來要做資料的處理，我們自行設計 Data Generator，Image Size = 256 x 256、Batch Size = 64，把資料夾丟進去就能得到圖片和結果的路徑，並嘗試用 imgaug 套件做水平翻轉、平移、放大縮小等常見的 Augmentation、要同時針對圖和答案做等量的變化喔！再設定一個 Epoch 有幾個 Batch、取出一批資料再丟到 Data Generator 的 Function，要注意的是我們是對整張圖裡每一個 Pixel 做二元分類，所以 y 的深度 = 1。前處理的方式是 Resize、再 Normalize 轉換成 0–1 的區間，記得 Mask 讀進來時也要一起做正規化，結果才會落在 0–1 之間。

IMG_SIZE = 256
BS = 64

class DataGenerator(utils.Sequence):
    def __init__(self, folder_path, batch_size, img_size, shuffle=True, aug=False):
        self.folder_path = folder_path
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.img_size = img_size
        self.aug = aug
        self.seq = iaa.Sequential([
            iaa.Fliplr(0.5), # 50% horizontal flip
            iaa.Affine(
                rotate=(-10, 10), # random rotate -45 ~ +45 degree
                shear=(-16,16), # random shear -16 ~ +16 degree
                scale={"x": (0.8, 1.2), "y": (0.8, 1.2)} # scale x, y: 80%~120%
            ),
        ])

        self.mask_paths = glob(os.path.join(folder_path, '*_mask.tif'))
        self.img_paths = [p.replace('_mask', '') for p in self.mask_paths]
        self.indexes = np.arange(len(self.mask_paths))
        self.on_epoch_end()

    def __len__(self):
        return int(np.ceil(len(self.mask_paths) / self.batch_size)) # batches per epoch

    def __getitem__(self, index):
        # Generate indexes of the batch
        idxs = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]
        # Find list of IDs
        batch_img_paths = [self.img_paths[i] for i in idxs]
        batch_mask_paths = [self.mask_paths[i] for i in idxs]

        # Generate data
        X, y = self.__data_generation(batch_img_paths, batch_mask_paths)
        return X, y

    def on_epoch_end(self):
        # Updates indexes after each epoch
        if self.shuffle:
            np.random.shuffle(self.indexes)

    def __data_generation(self, img_paths, mask_paths):
        # Generates data containing batch_size samples
        x = np.empty((len(img_paths), self.img_size, self.img_size, 3), dtype=np.float32)
        y = np.empty((len(img_paths), self.img_size, self.img_size, 1), dtype=np.float32)

        for i, (img_path, mask_path) in enumerate(zip(img_paths, mask_paths)):
            img = cv2.imread(img_path)
            mask = cv2.imread(mask_path)
            # img and mask preprocess
            img = self.preprocess(img)
            mask = self.preprocess(mask)
            x[i] = img
            y[i] = mask[:,:,:1]
        if self.aug:
            x, y = self.seq(images=x, heatmaps=y)
        return x, y
    def preprocess(self, img):
        data = cv2.resize(img, (self.img_size, self.img_size))
        data = data / 255. # normalize to 0~1
        return data

我們來隨機看看做完資料處理之後的結果 👇

下一步就是建模啦！我們採取簡化版的 UNet，大家如果比對前面的原始 UNet 會看到他們的 Convolution 有好幾層，為了方便示範，我們就偷懶只用一層就好惹 🤪 不過重點是在架構：先縮小再放大（這裡用 Conv2DTranspose，所以要設定參數），為了不丟失重要的空間資訊，把後面富含語意資訊的特徵圖跟前面富含空間資訊的特徵圖疊加在一起，最後要輸出成符合語意切割任務的深度；以我們的例子，是要偵測「有或沒有腫瘤」，所以是個二元分類的任務，所以 output_layer 的 Convolution 用 1 x 1 大小、深度 = 1、用 Sigmoid 作為 Activation Function。

# Enlarge feature maps by Conv2DTranspose
input_layer = layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
c1 = layers.Conv2D(filters=8, kernel_size=(3,3), activation='relu', padding='same')(input_layer)
l = layers.MaxPool2D(strides=(2,2))(c1)
c2 = layers.Conv2D(filters=16, kernel_size=(3,3), activation='relu', padding='same')(l)
l = layers.MaxPool2D(strides=(2,2))(c2)
c3 = layers.Conv2D(filters=32, kernel_size=(3,3), activation='relu', padding='same')(l)
l = layers.MaxPool2D(strides=(2,2))(c3)
c4 = layers.Conv2D(filters=32, kernel_size=(3,3), activation='relu', padding='same')(l)

l = layers.concatenate([layers.Conv2DTranspose(filters=32, kernel_size=3, strides=2, padding='same', activation='relu')(c4), 
                        c3], 
                       axis=-1)
l = layers.Conv2D(filters=32, kernel_size=(2,2), activation='relu', padding='same')(l)
l = layers.concatenate([layers.Conv2DTranspose(filters=32, kernel_size=3, strides=2, padding='same', activation='relu')(l), 
                        c2], 
                       axis=-1)
l = layers.Conv2D(filters=24, kernel_size=(2,2), activation='relu', padding='same')(l)
l = layers.concatenate([layers.Conv2DTranspose(filters=32, kernel_size=3, strides=2, padding='same', activation='relu')(l), 
                        c1], 
                       axis=-1)
l = layers.Conv2D(filters=16, kernel_size=(2,2), activation='relu', padding='same')(l)
l = layers.Conv2D(filters=64, kernel_size=(1,1), activation='relu')(l)
output_layer = layers.Conv2D(filters=1, kernel_size=(1,1), activation='sigmoid')(l)
                                                         
model = models.Model(input_layer, output_layer)

建模之後，我們就來訓練吧！我們採取 Dice Coefficient 來作為評估的 Metric，不過 Keras 沒有 Dice，所以我們要自己寫，2 倍交集除上預測和答案的面積總合，訓練時就會額外幫我們算出 Dice。Optimizer 用基本的 Adam、因為是二元分類所以用 binary_crossentropy 作為 Loss Function，設定 Earlystop 在 Overfitting 之前就先停住。

# Dice coefficient
def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + K.epsilon()) / (K.sum(y_true_f) + K.sum(y_pred_f) + K.epsilon())
    
model.compile(optimizer=optimizers.Adam(), loss='binary_crossentropy', 
              metrics=[dice_coef])

weight_saver = callbacks.ModelCheckpoint('seg.h5', monitor='val_loss', save_best_only=True)
earlystop = callbacks.EarlyStopping(monitor='val_loss', patience=3)

logs = model.fit(train_gen,
                 validation_data = test_gen,
                 epochs=100,
                 callbacks = [weight_saver, earlystop])

訓練完之後我們來驗證看看 👇

history = logs.history
plt.plot(history['loss'])
plt.plot(history['val_loss'])
plt.title('loss')
plt.show()
plt.plot(history['dice_coef'])
plt.plot(history['val_dice_coef'])
plt.title('Dice')
plt.show()

下一步是要看預測的結果，隨機地從 Test Data 裡取出一筆，把圖片和答案都讀進來，讓剛剛訓練好的 Model 去預測讀進來的圖片，一一地顯示出來。

# Sample 1 batch
batch_idx = np.random.randint(len(test_gen))
print(batch_idx)
data = test_gen[batch_idx]
imgs, mask = data # (bs, 256, 256, 3), (bs, 256, 256, 1)
mask_pred = model_final.predict(imgs)

# show inputs
img_idx = np.random.randint(len(imgs)) # sample 1 image from batch
plt.figure(figsize=(20, 5))
for i in range(3):
    plt.subplot(1,3,i+1)
    plt.imshow(imgs[img_idx, :,:, i], cmap='gray')
plt.show()

# show ground truth & model prediction
plt.figure(figsize=(20, 5))
plt.subplot(1, 3, 1)
plt.imshow(mask[img_idx, :, :, 0])
plt.subplot(1, 3, 2)
plt.imshow(mask_pred[img_idx, :, :, 0])
plt.show()
# plt.imshow(mask_pred[img_idx, :, :, 0], cmap='gray')

上面的三張圖是 Input，左下是正確答案，右下是我們模型預測的結果，越靠近黃色數值越高，深紫色的數值則是 0。

我們預測結果都是介於 0–1 之間的數值，設定特定的 Threshold，低於該閾值一刀切成 0、高於該閾值則是 1，這樣才能銳利化。隨著閾值越來越高，我們就會發現偵測範圍越來越小了。

落落長寫了這麼多，大家有沒有對於 AI 在醫療應用的有趣之處呢？身為醫師，我不相信目前的 AI 能媲美甚至是取代醫療專業，但是細看日常工作的結構，會發現有些冗餘、重複、單調的任務，在自動化之後對我們的效能和品質帶來的幫助，最終能夠在技術的幫忙下，為人類健康創造更大的價值。下次，我們再來看看其他的醫療 AI 運用吧 😍