아티클 정리 - Medical images segmentation with Keras: U-net architecture

안녕하세요. 이번 포스팅에서는 제가 구독하는 사이트 중에 하나인 Medium이라는 사이트의 기사를 번역하는 일을 해보도록 하겠습니다. 오늘 번역할 기사의 링크는 아래와 같습니다.링크

Medical images segmentation using Keras

이 글을 써주신 Soriba D.님에게 먼저 감사의 인사를 하도록 하겠습니다. 

 

Segmentation은 의학 영상 처리에서 기초적이고 중요한 분야 중에 하나입니다. 이러한 Segmentation 작업은 주로 자동화나 반자동화된 작업으로써 2D, 3D 영상에서 어떤 객체를 찾는 것이 목표입니다. 이때, 어떤 객체라고 하면 진단하려는 위치마다 달라지겠죠. 예를 들어서 간을 중심으로 진단하고 싶다면 간을 segmentation할 수도 있고, 폐를 중심으로 진단하고 싶다면 폐를 segmentation할 수도 있습니다. 

 

이번 미디움 포스트에서는 간을 segmentation합니다. 이 포스트에서 데이터를 제공해주고 있는 데 확장자가 생소합니다. NifTi라는 확장자입니다. 이 확장자의 풀 네임은 Neuroimaging Informatics Technology Initiative라고 합니다. 위키피디아의 설명을 들어보면 MRI 데이터를 저장하기 위해 사용되는 오픈 파일 포맷이라고 합니다!! 평소에 의학 데이터를 자주 다루지 않기 때문에 이러한 문제가 발생하였습니다. 이를 통해서 알 수 있는 것은 특정 분야에 사용되는 특정 확장자가 존재할 수 있다는 것입니다. 의학 분야에서는 NifTi였습니다. 이 확장자는 .nii, .nii.gz, .hdr로 쓰여진다고 합니다. 또한 파이썬에서는 nibabel이라는 라이브러리를 이용해서 이러한 종류의 데이터셋을 다룰 수 있다고 하니, 역시 파이썬은 만능의 언어인거 같습니다.(그렇다고 해서 다른 언어의 공부를 게을리해서는 안되겠죠?)

https://en.wikipedia.org/wiki/Neuroimaging_Informatics_Technology_Initiative

Neuroimaging Informatics Technology Initiative - Wikipedia

이 데이터셋은 IRCAD라는 곳이 출처라고 합니다. 이곳은 프랑스의 의학 연구소 센터라고 하는 군요. 이 데이터셋은 3D로 구성된 20개의 데이터가 있습니다. segmentation을 위해서는 기본적으로 mask 데이터도 있어야합니다. 사실 20개의 데이터로는 딥 러닝 네트워크를 학습하기에는 턱없이 모자랍니다.

 

저희가 학습시킬 모델을 제가 논문 함께 읽기[4].U-net에서 설명한 U-net을 사용한다고  하였습니다. 혹시 모르시는 분들은 어렵지 않으니 논문을 한번 읽어보시는 것을 추천드립니다. 또한 segmentation의 성능은 주로 dice coefficient로 측정하게 됩니다. 간단하게 설명하면 Ground Truth와 Prediction가 얼마나 겹치는 지를 0 ~ 1 사이의 실수로  수식으로 표현한 것입니다. 1에 가까우면 성능이 좋은 것이고 0에 가까우면 성능이 좋지 않은 것입니다. 수식은 아래와 같습니다.

해당 포스트의 코드를 실행하기 위해서는 몇 가지 과정을 거쳐야합니다. 먼저, file structure부터 아래와 같이 맞추도록 하겠습니다. 저는 코랩에서 진행할 예정이기 때문에 제 구글 드라이브를 마운트하여 사용하면 됩니다.

 

제일 먼저 시작하는 것은 데이터 전처리입니다. 여러분들은 위의 포스트에 들어가셔서 데이터를 다운받으셔야 아래의 단계들을 수행할 수 있습니다. 

 그러면 아래와 같이 3D 영상을 2D 영상으로 쪼갠 뒤 .npy 파일로 변환된 것을 볼 수 있습니다. 여기서 주의하셔야할 점은 각 데이터끼리 번호가 맞지 않으면 오류가 발생하기 때문에 마스크와 기존 영상의 번호를 맞추기 위해서 저는 코드를 살짝 수정하였습니다. sorted 내장함술를 이용하면 쉽게 바꿀 수 있습니다!! 

 

이제 실제로 학습을 진행하도록 하겠습니다. 먼저 필요한 라이브러리를 임포트합니다.

from __future__ import print_function

import os
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
from skimage.segmentation import mark_boundaries
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, concatenate, Conv2D, MaxPooling2D, Conv2DTranspose
from tensorflow.keras.optimizers import Adam, SGD
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras import backend as K
from skimage.exposure import rescale_intensity
from tensorflow.keras.callbacks import History
from skimage import io
from data import load_train_data, load_test_data

그리고 이전에서 설명드렸다싶이 이 포스트에서 사용되는 성능 평가지표는 dice coef입니다. 하지만 이것은 keras에 구현되어 있지 않기 때문에 아래와 같이 커스텀 손실 함수로 직접 구현해야합니다. 

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)

#The functions return our metric and loss

이제 U-net을 구현해보도록 하겠습니다. 블럭 단위로 설명드리면 압축하는 부분은 32 -> 64 -> 128 -> 256 -> 512(bottel neck)으로 기존의 논문과는 살짝 다릅니다. 하지만 그 구조는 유사합니다. 이제 expanding 영역은 그 반대로 512(bottel neck) -> 256 -> 128 -> 64 -> 32로 특징 맵의 채널의 개수는 점점 줄어들고 가로, 세로 크기는 점점 늘어나게 됩니다. 최종 출력 계층은 semantic segmentation을 구현해야하기 때문에 sigmoid 출력 단위를 이용해서 1개의 유닛만 있으면 됩니다.

def get_unet():
    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-3), loss=dice_coef_loss, metrics=[dice_coef])

    return model

#The different layers in our neural network model (including convolutions, maxpooling and upsampling)

그 다음으로 약간의 전처리 과정으로 reshape 과정을 거쳐줍니다.

def preprocess(imgs):
    imgs_p = np.ndarray((imgs.shape[0], img_rows, img_cols), dtype=np.uint8)
    for i in range(imgs.shape[0]):
        imgs_p[i] = resize(imgs[i], (img_cols, img_rows), preserve_range=True)

    imgs_p = imgs_p[..., np.newaxis]
    return imgs_p

#We adapt here our dataset samples dimension so that we can feed it to our network

드디어 학습하는 부분입니다. 제 컴퓨터에는 GPU를 사용하고 있기 때문에 코드를 수정하여 학습 시 GPU를 사용할 수 있도록 바꿧습니다.

def train_and_predict():
    strategy = tf.distribute.MirroredStrategy()

    with strategy.scope():
        print('-'*30)
        print('Loading and preprocessing train data...')
        print('-'*30)
        imgs_train, imgs_mask_train = load_train_data()

        imgs_train = preprocess(imgs_train)
        imgs_mask_train = preprocess(imgs_mask_train)

        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 /= 255
    #     imgs_train /= std
        #Normalization of the train set
        print(imgs_train.shape)
        plt.imshow(imgs_train[1000].reshape(256, 256), cmap='gray')
        plt.show()
        plt.imshow(imgs_mask_train[1000].reshape(256, 256), cmap='gray')
        plt.show()
        imgs_mask_train = imgs_mask_train.astype('float32')

        print('-'*30)
        print('Creating and compiling model...')
        print('-'*30)
        model = get_unet()
        model_checkpoint = ModelCheckpoint('weights.h5', monitor='val_loss', save_best_only=True)
        #Saving the weights and the loss of the best predictions we obtained

        print('-'*30)
        print('Fitting model...')
        print('-'*30)
        history=model.fit(imgs_train, imgs_mask_train, batch_size=32, epochs=20, verbose=1, shuffle=True,
                  validation_split=0.2,
                  callbacks=[model_checkpoint])

        print('-'*30)
        print('Loading and preprocessing test data...')
        print('-'*30)
        imgs_test, imgs_maskt = load_test_data()
        imgs_test = preprocess(imgs_test)

        imgs_test = imgs_test.astype('float32')
        imgs_test -= mean
        imgs_test /= std
        #Normalization of the test set

        print('-'*30)
        print('Loading saved weights...')
        print('-'*30)
        model.load_weights('weights.h5')

        print('-'*30)
        print('Predicting masks on test data...')
        print('-'*30)
        imgs_mask_test = model.predict(imgs_test, verbose=1)
        np.save('imgs_mask_test.npy', imgs_mask_test)
        print('-' * 30)
        print('Saving predicted masks to files...')
        print('-' * 30)
        pred_dir = 'preds'
        if not os.path.exists(pred_dir):
            os.mkdir(pred_dir)

        for k in range(len(imgs_mask_test)):
            a=rescale_intensity(imgs_test[k][:,:,0],out_range=(-1,1))
            b=(imgs_mask_test[k][:,:,0]).astype('uint8')
            io.imsave(os.path.join(pred_dir, str(k) + '_pred.png'),mark_boundaries(a,b))
        #Saving our predictions in the directory 'preds'
        plt.plot(history.history['dice_coef'])
        plt.plot(history.history['val_dice_coef'])
        plt.title('Model dice coeff')
        plt.ylabel('Dice coeff')
        plt.xlabel('Epoch')
        plt.legend(['Train', 'Test'], loc='upper left')
        plt.show()
        #plotting our dice coeff results in function of the number of epochs

train_and_predict()

이제 위의 코드를 실행하면 바로 학습을 진행합니다. GPU가 없으신 분들은 COLAB(393~511)에서 진행하시는 것을 추천드립니다. 1에폭 당 약 4분 정도 걸리니 계산하면 약 80~90분 정도의 시간이 걸립니다. 그 동안 다른 논문을 읽고 있으면 어느 순간 끝나있겠죠. 

------------------------------
Creating and compiling model...
------------------------------
------------------------------
Fitting model...
------------------------------
Epoch 1/20
37/37 [==============================] - 233s 6s/step - dice_coef: 0.2343 - loss: 0.7665 - val_dice_coef: 0.3887 - val_loss: 0.5696
Epoch 2/20
37/37 [==============================] - 233s 6s/step - dice_coef: 0.6275 - loss: 0.3724 - val_dice_coef: 0.5277 - val_loss: 0.4157
Epoch 3/20
37/37 [==============================] - 232s 6s/step - dice_coef: 0.7630 - loss: 0.2364 - val_dice_coef: 0.5638 - val_loss: 0.3758
Epoch 4/20
37/37 [==============================] - 233s 6s/step - dice_coef: 0.7757 - loss: 0.2263 - val_dice_coef: 0.6475 - val_loss: 0.2831
Epoch 5/20
37/37 [==============================] - 234s 6s/step - dice_coef: 0.8342 - loss: 0.1636 - val_dice_coef: 0.6678 - val_loss: 0.2607
Epoch 6/20
37/37 [==============================] - 233s 6s/step - dice_coef: 0.8481 - loss: 0.1517 - val_dice_coef: 0.6733 - val_loss: 0.2546
Epoch 7/20
37/37 [==============================] - 233s 6s/step - dice_coef: 0.8689 - loss: 0.1290 - val_dice_coef: 0.6340 - val_loss: 0.2981
Epoch 8/20
37/37 [==============================] - 234s 6s/step - dice_coef: 0.8756 - loss: 0.1259 - val_dice_coef: 0.6086 - val_loss: 0.3262
Epoch 9/20
37/37 [==============================] - 237s 6s/step - dice_coef: 0.9026 - loss: 0.0981 - val_dice_coef: 0.6994 - val_loss: 0.2256
Epoch 10/20
37/37 [==============================] - 234s 6s/step - dice_coef: 0.9062 - loss: 0.0935 - val_dice_coef: 0.6384 - val_loss: 0.2933
Epoch 11/20
37/37 [==============================] - 233s 6s/step - dice_coef: 0.9127 - loss: 0.0885 - val_dice_coef: 0.6620 - val_loss: 0.2670
Epoch 12/20
37/37 [==============================] - 237s 6s/step - dice_coef: 0.9301 - loss: 0.0699 - val_dice_coef: 0.6110 - val_loss: 0.3237
Epoch 13/20
37/37 [==============================] - 234s 6s/step - dice_coef: 0.9268 - loss: 0.0736 - val_dice_coef: 0.6234 - val_loss: 0.3116
Epoch 14/20
37/37 [==============================] - 236s 6s/step - dice_coef: 0.9359 - loss: 0.0637 - val_dice_coef: 0.7227 - val_loss: 0.2000
Epoch 15/20
37/37 [==============================] - 235s 6s/step - dice_coef: 0.9428 - loss: 0.0576 - val_dice_coef: 0.6986 - val_loss: 0.2266
Epoch 16/20
37/37 [==============================] - 234s 6s/step - dice_coef: 0.9485 - loss: 0.0493 - val_dice_coef: 0.7225 - val_loss: 0.2000
Epoch 17/20
37/37 [==============================] - 234s 6s/step - dice_coef: 0.9405 - loss: 0.0599 - val_dice_coef: 0.7335 - val_loss: 0.1879
Epoch 18/20
37/37 [==============================] - 233s 6s/step - dice_coef: 0.9557 - loss: 0.0446 - val_dice_coef: 0.7413 - val_loss: 0.1793
Epoch 19/20
37/37 [==============================] - 238s 6s/step - dice_coef: 0.9620 - loss: 0.0379 - val_dice_coef: 0.7288 - val_loss: 0.1932
Epoch 20/20
37/37 [==============================] - 233s 6s/step - dice_coef: 0.9631 - loss: 0.0359 - val_dice_coef: 0.7539 - val_loss: 0.1655

최종 결과는 훈련 데이터는 거의 100%, 시험 데이터는 약 80% 정도로 과적합이 일어난 것으로 보입니다. 만약 이를 발전시켜서 과적합을 방지하려면 몇 가지 규제 기법을 적용하거나 데이터를 추가하는 방법을 생각해볼 수 있습니다. 가장 좋은 방법은 data augmentation이겠네요!! 이 아티클에서는 시험 데이터를 예측한 결과의 경계선을 함께 보여줄 수도 있습니다. 

 

저기 보이는 노란색 경계선이 예측한 결과입니다. 하지만, 실제 간과 많은 차이가 있음을 볼 수 있습니다. 따라서 아직은 성능이 좋지 않은 것을 볼 수 있습니다.