基於U-Net的眼底圖像血管分割實例

【英文說明】https://github.com/orobix/retina-unet#retina-blood-vessel-segmentation-with-a-convolution-neural-network-u-net

【更新】針對Python3版本對此部分代碼做了優化，已上傳到我的GitHub：點擊打開鏈接

【注意事項】

1.運行run_training.py或run_testing.py如果出錯，可以嘗試將src/retinaNN_training.py或src/retinaNN_predict.py拷貝到工程根目錄下運行。（代碼目前已更新至根目錄下）

2.針對模型Unet的可視化代碼在training的111、112行中，如果不需要模型可視化可以註釋掉。如果可視化過程出錯，可以參考CSDN博客進行修改。

3.如果出現提示test文件夾下的不存在的問題，請先在工程目錄下新建test文件夾。

4.大體流程：數據處理→訓練→預測生成

 

1  介紹

爲了能夠更好的對眼部血管等進行檢測、分類等操作，我們首先要做的就是對眼底圖像中的血管進行分割，保證最大限度的分割出眼部的血管。從而方便後續對血管部分的操作。

DRIVE數據集下載：百度網盤 （密碼：4l7v）

這部分代碼選用的數據集是DRIVE數據集，包括訓練集和測試集兩部分。眼底圖像數據如圖1所示。

圖1 DRIVE數據集的訓練集眼底圖像

DRIVE數據集的優點是：不僅有已經手工分好的的血管圖像（在manual文件夾下，如圖2所示），而且還包含有眼部輪廓的圖像（在mask文件夾下，如圖3所示）。

圖2 DRIVE數據集的訓練集手工標註血管圖像

圖3 DRIVE數據集的訓練集眼部輪廓圖像

DRIVE數據集的缺點是：顯而易見，從上面的圖片中可以看出，訓練集只有20幅圖片，可見數據量實在是少之又少。。。

所以，爲了得到更好的分割效果，我們需要對這20幅圖像進行預處理從而增大其數據量

 

2  依賴的庫

【2020.03.24注】因版本不一可能會導致某些不可預期的錯誤，現將包的版本固定（我的機器可以運行）

- numpy == 1.16.4
- Keras == 2.2.5
- Tensorflow == 1.13.1
- Pillow == 5.0.0
- opencv-python == 4.1.1.26 
- h5py == 2.7.1
- configparser == 3.5.0
- scikit-learn == 0.19.1

 

3  數據讀取與保存

數據集中訓練集和測試集各只有20幅眼底圖像（tif格式）。首先要做的第一步就是對生成數據文件，方便後續的處理。所以這裏我們需要對數據集中的眼底圖像、人工標註的血管圖像、眼部輪廓生成數據文件。這裏使用的是hdf5文件。有關hdf5文件的介紹，請參考CSDN博客（HDF5快速上手全攻略）。

數據處理重要代碼部分（prepare_datasets_DRIVE.py）：

def get_datasets(imgs_dir,groundTruth_dir,borderMasks_dir,train_test="null"):
    imgs = np.empty((Nimgs,height,width,channels))
    groundTruth = np.empty((Nimgs,height,width))
    border_masks = np.empty((Nimgs,height,width))
    for path, subdirs, files in os.walk(imgs_dir): #list all files, directories in the path
        for i in range(len(files)):
            #original
            print("original image: " +files[i])
            img = Image.open(imgs_dir+files[i])
            imgs[i] = np.asarray(img)
            #corresponding ground truth
            groundTruth_name = files[i][0:2] + "_manual1.gif"
            print("ground truth name: " + groundTruth_name)
            g_truth = Image.open(groundTruth_dir + groundTruth_name)
            groundTruth[i] = np.asarray(g_truth)
            #corresponding border masks
            border_masks_name = ""
            if train_test=="train":
                border_masks_name = files[i][0:2] + "_training_mask.gif"
            elif train_test=="test":
                border_masks_name = files[i][0:2] + "_test_mask.gif"
            else:
                print("specify if train or test!!")
                exit()
            print("border masks name: " + border_masks_name)
            b_mask = Image.open(borderMasks_dir + border_masks_name)
            border_masks[i] = np.asarray(b_mask)

    print("imgs max: " +str(np.max(imgs)))
    print("imgs min: " +str(np.min(imgs)))
    assert(np.max(groundTruth)==255 and np.max(border_masks)==255)
    assert(np.min(groundTruth)==0 and np.min(border_masks)==0)
    print("ground truth and border masks are correctly withih pixel value range 0-255 (black-white)")
    #reshaping for my standard tensors
    imgs = np.transpose(imgs,(0,3,1,2))
    assert(imgs.shape == (Nimgs,channels,height,width))
    groundTruth = np.reshape(groundTruth,(Nimgs,1,height,width))
    border_masks = np.reshape(border_masks,(Nimgs,1,height,width))
    assert(groundTruth.shape == (Nimgs,1,height,width))
    assert(border_masks.shape == (Nimgs,1,height,width))
    return imgs, groundTruth, border_masks

 

4  訓練

4.1  數據預處理

訓練過程，我們首先對眼底圖像數據進行數據預處理。調用lib/pre_processing.py下的my_PreProc()完成數據預處理相關工作。

預處理包括：灰度變換、標準化、對比度受限的自適應直方圖均衡化（CLAHE）以及伽馬變換。有關對比度受限的自適應直方圖均衡化可以參考CSDN博客，有關伽馬變換可以參考CSDN博客。

下面是對比度受限的自適應直方圖均衡化代碼：

# CLAHE (Contrast Limited Adaptive Histogram Equalization)
#adaptive histogram equalization is used. In this, image is divided into small blocks called "tiles" (tileSize is 8x8 by default in OpenCV). Then each of these blocks are histogram equalized as usual. So in a small area, histogram would confine to a small region (unless there is noise). If noise is there, it will be amplified. To avoid this, contrast limiting is applied. If any histogram bin is above the specified contrast limit (by default 40 in OpenCV), those pixels are clipped and distributed uniformly to other bins before applying histogram equalization. After equalization, to remove artifacts in tile borders, bilinear interpolation is applied
def clahe_equalized(imgs):
    assert (len(imgs.shape)==4)  #4D arrays
    assert (imgs.shape[1]==1)  #check the channel is 1
    #create a CLAHE object (Arguments are optional).
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    imgs_equalized = np.empty(imgs.shape)
    for i in range(imgs.shape[0]):
        imgs_equalized[i,0] = clahe.apply(np.array(imgs[i,0], dtype = np.uint8))
    return imgs_equalized

下面是伽馬變換的代碼：

def adjust_gamma(imgs, gamma=1.0):
    assert (len(imgs.shape)==4)  #4D arrays
    assert (imgs.shape[1]==1)  #check the channel is 1
    # build a lookup table mapping the pixel values [0, 255] to
    # their adjusted gamma values
    invGamma = 1.0 / gamma
    table = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
    # apply gamma correction using the lookup table
    new_imgs = np.empty(imgs.shape)
    for i in range(imgs.shape[0]):
        new_imgs[i,0] = cv2.LUT(np.array(imgs[i,0], dtype = np.uint8), table)
    return new_imgs

4.2  數據擴增

深度學習需要大量的數據來擬合模型參數，針對只有20張眼底圖像的DRIVE數據集，我們採用隨機切片的方式來對數據進行擴增。調用lib/extract_patches.py下的extract_random()來對數據進行切片。

每個尺寸爲48*48的貼片是通過在整個圖像內隨機選擇其中心獲得的。此外，選擇部分可能完全在視野（FOV）之外的斑塊，通過這種方式，神經網絡可以學習如何區分FOV邊界與血管。

通過以下代碼完成隨機切片，擴充數據：

for i in range(full_imgs.shape[0]):  #loop over the full images
        k=0
        while k <patch_per_img:
            x_center = random.randint(0+int(patch_w/2),img_w-int(patch_w/2))
            # print "x_center " +str(x_center)
            y_center = random.randint(0+int(patch_h/2),img_h-int(patch_h/2))
            # print "y_center " +str(y_center)
            #check whether the patch is fully contained in the FOV
            if inside==True:
                if is_patch_inside_FOV(x_center,y_center,img_w,img_h,patch_h)==False:
                    continue
            patch = full_imgs[i,:,y_center-int(patch_h/2):y_center+int(patch_h/2),x_center-int(patch_w/2):x_center+int(patch_w/2)]
            patch_mask = full_masks[i,:,y_center-int(patch_h/2):y_center+int(patch_h/2),x_center-int(patch_w/2):x_center+int(patch_w/2)]
            patches[iter_tot]=patch
            patches_masks[iter_tot]=patch_mask
            iter_tot +=1   #total
            k+=1  #per full_img

通過在20個DRIVE訓練圖像中的每一箇中隨機提取9500個patches來獲得一組190000個patches。儘管貼片重疊，即不同的貼片可以包含原始圖像的相同部分，但是不執行進一步的數據增強。然後對數據進行分佈，前90％的數據集用於訓練（171000個patches），而最後10％用於驗證（19000個patches）。

4.3  搭建網絡模型

神經網絡架構源自U-Net架構（參見論文）。損失函數是交叉熵，隨機梯度下降用於優化。每個卷積層之後的激活函數是整流器線性單元（ReLU），並且在兩個連續卷積層之間使用0.2的dropout。

這一部分用keras便可輕鬆完成，U-Net的結構代碼如下：

#Define the neural network
def get_unet(n_ch,patch_height,patch_width):
    inputs = Input(shape=(n_ch,patch_height,patch_width))
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same',data_format='channels_first')(inputs)
    conv1 = Dropout(0.2)(conv1)
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same',data_format='channels_first')(conv1)
    pool1 = MaxPooling2D((2, 2))(conv1)
    #
    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same',data_format='channels_first')(pool1)
    conv2 = Dropout(0.2)(conv2)
    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same',data_format='channels_first')(conv2)
    pool2 = MaxPooling2D((2, 2))(conv2)
    #
    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same',data_format='channels_first')(pool2)
    conv3 = Dropout(0.2)(conv3)
    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same',data_format='channels_first')(conv3)

    up1 = UpSampling2D(size=(2, 2))(conv3)
    up1 = concatenate([conv2,up1],axis=1)
    conv4 = Conv2D(64, (3, 3), activation='relu', padding='same',data_format='channels_first')(up1)
    conv4 = Dropout(0.2)(conv4)
    conv4 = Conv2D(64, (3, 3), activation='relu', padding='same',data_format='channels_first')(conv4)
    #
    up2 = UpSampling2D(size=(2, 2))(conv4)
    up2 = concatenate([conv1,up2], axis=1)
    conv5 = Conv2D(32, (3, 3), activation='relu', padding='same',data_format='channels_first')(up2)
    conv5 = Dropout(0.2)(conv5)
    conv5 = Conv2D(32, (3, 3), activation='relu', padding='same',data_format='channels_first')(conv5)
    #
    conv6 = Conv2D(2, (1, 1), activation='relu',padding='same',data_format='channels_first')(conv5)
    conv6 = core.Reshape((2,patch_height*patch_width))(conv6)
    conv6 = core.Permute((2,1))(conv6)
    
    conv7 = core.Activation('softmax')(conv6)

    model = Model(inputs=inputs, outputs=conv7)

    # sgd = SGD(lr=0.01, decay=1e-6, momentum=0.3, nesterov=False)
    model.compile(optimizer='sgd', loss='categorical_crossentropy',metrics=['accuracy'])

    return model

4.4  執行訓練

通過model.fit執行訓練，對在訓練過程中隨時保存模型就可以了。

訓練過程的我電腦信息如下，佔用內存還是挺高的。Windows10，內存一共24GB（一塊8GB，一塊16GB）。GPU使用的是GTX 1060（閹割版，顯存3GB）。

 

5  預測生成

5.1 準備數據

與訓練過程一樣，準備預測的數據。

5.2 讀取保存好的模型權重

在訓練過程中，會判斷當前模型權重是否最好，最好則會進行保存。預測時，讀取保存的權重。

model.load_weights(path_experiment+name_experiment + '_'+best_last+'_weights.h5')

5.3 預測

通過model.predict執行預測。

5.4 預測結果對比

通過代碼，將預測圖像與原圖像拼接起來，進行可視化對比。

5.5 計算評價結果

此部分預測的結果評價標準由準確率、召回率、AUC/ROC曲線進行評價。相關內容學習請見CSDN博客。

最後在test文件夾下，會有預測之後的結果圖，以及AUC/ROC曲線、準確率/召回率曲線等。