基於 kears的全卷積網絡u-net端到端醫學圖像多類型圖像分割（二）

1. 端到端的圖像分割網絡

首先，回顧下網絡模型

def get_unet(pretrained_weights=None):
    inputs = Input((img_rows, img_cols, 1))
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same')(inputs)
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same')(pool1)
    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same')(pool2)
    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

    conv4 = Conv2D(256, (3, 3), activation='relu', padding='same')(pool3)
    conv4 = Conv2D(256, (3, 3), activation='relu', padding='same')(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)

    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same')(pool4)
    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same')(conv5)

    up6 = concatenate([Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(conv5), conv4], axis=3)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same')(up6)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same')(conv6)

    up7 = concatenate([Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(conv6), conv3], axis=3)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same')(up7)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv7)

    up8 = concatenate([Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(conv7), conv2], axis=3)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same')(up8)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv8)

    up9 = concatenate([Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(conv8), conv1], axis=3)
    conv9 = Conv2D(32, (3, 3), activation='relu', padding='same')(up9)
    conv9 = Conv2D(32, (3, 3), activation='relu', padding='same')(conv9)

    conv10 = Conv2D(1, (1, 1), activation='sigmoid')(conv9)

    model = Model(inputs=[inputs], outputs=[conv10])

    model.compile(optimizer=Adam(lr=1e-5), loss=dice_coef_loss, metrics=[dice_coef])

    if (pretrained_weights):
        model.load_weights(pretrained_weights)

    return model

• 比較重要的是定義loss function 我們使用dice loss,主要思路就是求圖像重合部分佔比。相關代碼如下：

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 + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)


def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)

2. 數據準備

由於是端到端的訓練，我們的數據其實並沒有經過太多處理，將mask和原圖分別放到np array中即可。以下代碼段展示training data的準備方式，test data 基本類似。

data_path = 'data/' #設置路徑
image_rows = 256
image_cols = 256
def create_train_data():
    train_data_path = os.path.join(data_path, 'train/Image') #訓練文件路徑
    train_data_Label_path = os.path.join(data_path, 'train/Label') #mask文件路徑
    images = os.listdir(train_data_path)
    total = len(images)
    imgs = np.ndarray((total, image_rows, image_cols), dtype=np.uint8)
    imgs_mask = np.ndarray((total, image_rows, image_cols), dtype=np.uint8)
    i = 0
    print('Creating training images...')
    for image_name in images:
        img = imread(os.path.join(train_data_path, image_name), as_grey=True)
        img_mask = imread(os.path.join(train_data_Label_path, image_name), as_grey=True)
        img = np.array([img])
        img_mask = np.array([img_mask])
        imgs[i] = img
        imgs_mask[i] = img_mask
        if i % 100 == 0:
            print('Done: {0}/{1} images'.format(i, total))
        i += 1
    print('Loading done.')
    np.save('imgs_train.npy', imgs)
    np.save('imgs_mask_train.npy', imgs_mask)
    print('Saving to .npy files done.')

3. 訓練

訓練前，要先對相關數據歸一化，對於訓練數據，我們先求均值，與標準差，原始值與均值做差，再除以均值。

    imgs_train = imgs_train.astype('float32')
    mean = np.mean(imgs_train)  # mean for data centering
    std = np.std(imgs_train)  # std for data normalization
    imgs_train -= mean
    imgs_train /= std

對於mask數據，直接除以通道最大值255即可。

    imgs_mask_train = imgs_mask_train.astype('float32')
    imgs_mask_train /= 255.  # scale masks to [0, 1]

設置相關參數，開始訓練：

    model.fit(imgs_train, imgs_mask_train, batch_size=16, nb_epoch=20, verbose=1, shuffle=True,
              validation_split=0.2,
              callbacks=[model_checkpoint])

4. 結果

經過大約200 epoch 的訓練，我們的結果如下圖（從左至右依次爲原始圖，手工標註圖，算法分割圖）：

最終的dice_coef 在0.7左右。

5.多分類

可以看到最終的結果是3分類的，但訓練過程中發現，網絡結果仍舊是二分類，因此，我們仍舊需要針對不同分類訓練多個二分類模型，在預測時，先做多個二分類，再對數據進行相關融合，得到最終結果。相關預測及合併代碼如下：

    imgs_mask_test = model.predict(imgs_test, verbose=1)
    imgs_mask_test_t = modelt.predict(imgs_test, verbose=1)
    pred_dir = 'preds_a'
    if not os.path.exists(pred_dir):
        os.mkdir(pred_dir)
    for image,image_t,image_id in zip(imgs_mask_test,imgs_mask_test_t ,imgs_id_test):
        image = (image[:, :, 0] * 128.).astype(np.uint8)
        image_t = (image_t[:, :, 0] * 255.).astype(np.uint8)
        image[image_t>200]=255
        imsave(os.path.join(pred_dir, str(image_id) + '_pred.png'), image)

6.其它相關

二分類源代碼及已上傳至github keras-u-net 多分類，由於大家數據不同，分類類別不同，暫時沒有更新相關代碼。歡迎交流。

• 注：由於相關訓練數據不便提供，所以相關數據文件沒有上傳。