여러분은 Medical AI에 대해 들어본 적이 있으신가요? 익숙하신 분들도 있겠지만, 처음 접한 분들도 많으실 텐데요. Medical AI는 최근 인공지능 분야에서 떠오르고 있는 유망한 분야로, MRI와 같은 의료 이미지에서 병변을 검출해내는 딥러닝 모델을 만들거나, 뇌파와 같은 신호에서 이상을 탐지하는 딥러닝 알고리즘 등을 개발하는 것이 대표적인 예시입니다. 저희의 프로젝트에서는 대표적인 이미지 분할 (Image Segmentation) 모델인 U-Net을 변형하여 뇌 MRI 이미지에서 병변을 검출하는 모델을 구현해보고자 하였습니다.
프로젝트 소개
안녕하세요, 저희는 논문 구현을 목적으로 한 🧱 보강재 🧱 팀입니다!
저희가 선택한 논문은 *<BU-net: rain Tumor Segmentation Using Modified U-Net Architecture> (이하 BU-net)*입니다. 이 연구는 뇌 종양 분할을 위해 기존의 변화된 모델을 사용하거나 새로운 손실 함수를 사용하는 등 개선된 아키텍처를 만들고 성능을 알아보는 연구입니다.
저희는 논문에 나온 다양한 방안을 최대한 이해해보고, 논문의 실험을 유사하게 구현하고자 했습니다. 이를 통해 새로운 모델 구조에 도입된 방법이 실제로 효과가 있는지, 성능은 좋은지 알아보고자 노력하였습니다.
데이터셋
매년 MICCAI 등의 학회에서 주최하는 BraTS (Brain Tumor Segmentation) challenge의 데이터셋을 활용하였습니다. 해당 데이터셋은 홈페이지에 가입하여 주최 측의 승인을 얻으면 누구나 접근 가능한 뇌 MRI 이미지 데이터셋이며, 이번 프로젝트에서는 BraTS 2018 데이터셋을 활용하였습니다.
데이터셋을 다운로드 받은 이후에 압축을 풀게 되면, 각 환자의 고유 ID 폴더를 확인할 수 있고 그 폴더 하위에는 T1, T1ce, T2, FLAIR MRI 이미지 데이터, 그리고 학습을 위한 뇌 종양 라벨 데이터가 포함되어 있습니다.
💡 T1, T1ce, T2, FLAIR가 무엇인가요?
T1 (T1-Weighted Imaging): 해부학적 구조를 시각화하는 데 사용되는 MRI 이미지입니다. 지방은 밝게, 물이나 뇌척수액은 어둡게 나타나는 것이 특징입니다.
T1ce (T1-Weighted Imaging with Contrast Enhancement): 조영제를 주입한 이후 T1 영상을 촬영한 것으로 병변이 더 밝게 나타나는 특징이 있습니다.
T2 (T2-Weighted Imaging): 병변을 시각화하는데 사용하는 이미지입니다. 물과 뇌척수액이 밝게 나타난다는 점이 특징입니다.
FLAIR (Fluid-Attenuated Inversion Recovery): 뇌척수액과 같은 유체를 억제하므로 뇌의 병변을 보는 데 사용할 수 있습니다. 뇌척수액이 어둡게 나타나고 병변이 더 뚜렷하게 나타납니다.
데이터셋에 포함된 라벨은 다음과 같았습니다.
GD-enhancing tumor (ET — label 4)
peritumoral edema (ED — label 2)
necrotic and non-enhancing tumor core (NCR/NET — label 1)
healthy tissue (label 0)
전처리
BU-net 논문에 나온 전처리 방식을 그대로 이용하였습니다.
전처리 단계는 크게 N4ITK 알고리즘 적용 - 극단값 제거 - 정규화의 3단계로 구분할 수 있습니다.
N4ITK 알고리즘 : 해당 알고리즘은 MRI 이미지를 사용하는 대부분의 프로젝트에서 전처리 방식으로 채택하고 있는 이미지 정규화 알고리즘입니다. 2010년 <N4ITK: improved N3 bias correction>이라는 논문에서 처음 소개된 이 알고리즘은 이미지의 밝기가 고르게 나타나지 않는다는 문제점을 해결하도록 도와줍니다. Python에서는 SimpleITK이라는 패키지 내에 구현되어 있어 누구나 편하게 이용 가능합니다.
극단값 제거 : 이미지를 배열 형식으로 불러올 때 극단적인 값이 있으면 모델의 성능에 영향을 주기 때문에 이러한 값을 제거해주어야 합니다. 논문에서는 상위 1%와 하위 1%의 값을 제거했으며, 이번 프로젝트에서도 동일한 범위의 값을 제거하였습니다.
정규화 : 배열로 불러온 이미지들의 값을 정규화하여 모델의 강건함 (Robustness) 을 살리고자 하였습니다.
전처리가 끝난 이미지는 다음과 같은 모습을 보이게 됩니다. 전처리 전과 후 이미지의 모습이 확실히 달라진 것을 확인할 수 있죠?
처리 전
처리 후
데이터셋 클래스 정의
전처리가 끝난 데이터를 모델에 학습시키기 위해서는 데이터셋 클래스를 정의하고 데이터로더를 이용하여 배치 단위로 데이터를 불러와야 합니다. 이 MRI 데이터는 3차원으로 구성되어 있지만, 프로젝트에 사용할 모델인 U-Net 그리고 BU-net은 모두 2차원 데이터만을 입력으로 받기 때문에 데이터셋을 불러올 때 3차원 이미지를 2차원 데이터로 쪼개는 과정이 필요했습니다. 우리는 이 2차원 데이터를 ‘슬라이스’라고 부르기로 하였습니다.
모델에 모든 슬라이스를 불러와서 사용하면 좋지만, 이미 T1, T2, T1ce, Flair 총 4가지 모달리티를 모두 학습시켜야 하는 상황에서 모든 슬라이스를 사용하는 것은 현실적으로 리소스 부족 문제로 인하여 어려워졌습니다. 따라서 저희는 MRI 이미지의 중앙 슬라이스를 기준으로 +-3개의 슬라이스만을 사용하기로 결정하였습니다.
또한 논문에 명시된 대로 이미지의 크기를 240\times240 에서 256\times256 사이즈로 늘렸습니다.
모델
U-Net이란?
먼저 BU-net을 보기 전에 U-Net에 대해 설명하겠습니다.
U-net의 구조
U-Net은 Medical Image 분야에서 이미지 분할을 용도로 널리 사용되는 모델입니다. 형태가 그림과 같이 U자의 구조를 가지고 있기 때문에 U-Net으로 불립니다.
U-Net은 크게 수축경로와 확장경로로 이루어져 있습니다. 그림을 절반으로 쪼갰을 때 왼쪽 부분이 수축경로, 오른쪽이 확장경로입니다.
수축경로에서는 Convolution과 Maxpooling을 통해서 feature map의 차원 축소가 일어나고 채널 개수는 많아지게 됩니다. 이 과정에서 이미지에 대한 특성을 추출해낼 수 있습니다.
확장경로에서는 Convolution과 Transpose Convolution을 통해서 수축경로에서 줄어들었던 차원을 다시 복원하는 과정을 거칩니다. 이 과정에서 이미지의 정확도를 향상시킬 수 있습니다. 확장경로에서 정확도를 향상하는 방안 중 하나는 Skip Connection을 사용하는 것입니다. Skip connection이란, 수축경로에서 나온 feature map을 확장경로에 연결하여 병합하는 과정을 의미합니다.
윗 단계의 수축경로에서 차원이 큰 feature map이 나오게 되고, pooling을 거듭할수록 즉 아랫 단계의 수축경로로 이동할수록 차원이 작은 feature map이 나오게 되는데요? 여러 단계의 Skip connection을 거치게 되면 차원이 큰 feature map에서 추출되는 low-level feature과 차원이 작은 feature map에서 추출되는 high-level feature을 동시에 얻을 수 있습니다! 이는 전반적인 이미지 구조를 이해함(high-level)과 동시에 세부적인 정보(low-level)를 제공 받아 더욱 폭 넓게 이미지를 이해할 수 있습니다.
BU-net이란?
BU-net은 U-Net의 구조와 비교해보면 중간에 몇 개의 Block이 추가된 모습임을 알 수 있는데요. 더 자세히 구조를 알아보도록 하겠습니다.
RES Block (Residual Extended Skip)
그림을 보면 본래 U-Net의 Skip connection에 중간에 RES Block이 껴있는 것을 볼 수 있습니다. 총 4개의 Skip connection이 있으니까 4개의 RES Block이 추가되었겠죠?
그렇다면 RES Block이란 무엇일까요?
위 이미지는 RES block의 모식도입니다. 총 4개의 병렬적인 Convolution 연산과 1개의 Skip Connection이 구현되어 있으며 이 결과를 합하여 Convolution 연산을 세 번 거쳐 최종 출력을 계산합니다. 이러한 복잡한 구조를 사용하는 이유는 모델에서 하위 단계 특징으로부터 중간 단계 특징을 추출해 내 정보 손실을 조절하기 위해서입니다.
하나의 층의 출력을 바로 다음 층의 입력으로 사용하지 않고 몇 개의 층을 건너뛴 이후 입력에 사용하는 방식입니다. 기존의 전통적인 깊은 딥러닝 모델을 사용할 경우 성능이 저하되는 문제가 존재합니다. 이 문제가 발생하는 원인은 Vanishing Gradient Problem 때문인데요, 층이 많아지면 미세한 Gradient가 점점 0에 가까워져 학습이 제대로 되지 않는다는 뜻입니다. 이러한 문제를 해결하기 위한 방식으로 Skip Connection이 제안되었습니다. Skip Connection을 통해 Gradient가 사라지지 않고 원활하게 흐를 수 있도록 도울 수 있습니다.
Skip Connection을 통해 Gradient가 사라지지 않고 원활하게 흐를 수 있도록 도울 수 있습니다.
하나의 층의 출력을 바로 다음 층의 입력으로 사용하지 않고 몇 개의 층을 건너뛴 이후 입력에 사용하는 방식입니다. 기존의 전통적인 깊은 딥러닝 모델을 사용할 경우 성능이 저하되는 문제가 존재합니다. 이 문제가 발생하는 원인은 Vanishing Gradient Problem 때문인데요, 층이 많아지면 미세한 Gradient가 점점 0에 가까워져 학습이 제대로 되지 않는다는 뜻입니다. 이러한 문제를 해결하기 위한 방식으로 Skip Connection이 제안되었습니다.
이 RES Block은 U-Net 모델의 Skip Connection 부분에서 사용됩니다. 암의 크기가 다양하므로 여러 크기의 convolution을 병렬적으로 사용해 병합을 합니다.
WC Block (Wide Context)
Wide Context (WC) Block은 RES block과 비슷한 형태이지만, 2개의 병렬적인 Convolution 연산만 진행되며 따로 Skip Connection은 수행하지 않습니다. WC Block은 입력에서 맥락적인 정보(Contextual Information)를 추출할 수 있으며, 이를 통해 종양의 하위 클래스를 구분하는 데 도움이 됩니다.
WC Block은 BU-net 모델의 제일 하단에 들어가 있는 것을 볼 수 있습니다. 가장 차원이 작은 Feature Map에서 특성을 추출할 때 사용하게 됩니다. 1개의 WC Block이 추가되었습니다.
BU-net은 U-Net 구조에서 2가지의 RES Block과 WC Block을 추가한 형태입니다. U-Net의 강점인 이미지를 폭 넓게 이해할 수 있다는 점을 뇌 이미지 분할의 특성에 맞춰서 더욱 극대화했다고 볼 수 있습니다.
손실 함수 정의
위의 수식은 U-Net의 손실 함수입니다. 그러나 사용할 데이터셋이 불균형 클래스 데이터라는 문제점이 있습니다.
이를 개선하기 위해 손실 함수(loss function)의 변화를 주었습니다.
Dice Loss Coefficient (DLC)
예측된 분할 영역과 실제 분할 영역의 최대 중첩을 찾기 위한 목적 함수입니다. 주로 이미지 분할 작업에서 사용되며, 분할 성능을 향상시킵니다.
Weighted Cross-Entropy Loss (WCE)
클래스 분류를 위한 손실 함수입니다. 클래스가 불균형한 경우, 각 클래스에 가중치를 부여하여 손실을 계산합니다.
이 둘을 더한 것이 새로운 손실 함수가 됩니다.
모델 학습
본 프로젝트에서는 총 세 가지 모델을 학습하였습니다. 이는 각각 U-Net 모델, U-Net에 WC block을 결합한 모델, 마지막으로 이에 Simple RES Block을 합친 모델입니다.
Simple RES Block은 원활한 학습을 위해 사용한 모델인데요, 기존의 RES Block의 기능을 살리되, 모델의 parameter 수를 줄여 학습에 용이하도록 하였습니다.
SImple RES Block
이렇게 3개의 모델을 비교하면서 각각의 Block이 효과적으로 이용되고 있는지 확인하였습니다.
에포크마다 진행 상황을 보며 학습이 잘 진행되는지 확인하고 평가하였으며 테스트 데이터셋으로 가시화하여 모델의 성능을 정성적으로 추가 평가하였습니다.
에포크 수에 대한 언급이 없었기에 이 수는 저희가 설정하여 다음과 같은 조건으로 학습을 진행하였습니다. 나머지 다른 조건들은 논문과 동일하게 진행하였습니다.
학습 결과
정량 평가 결과
위 이미지는 각 모델을 30 에포크까지 학습하는 과정에서의 손실(loss)와 정확도(accuracy) 그래프를 시각화한 것입니다. 시간과 리소스 부족으로 인해 30 에포크까지밖에 학습을 진행하지 못했으나 대체적으로 학습을 거듭할수록 손실이 줄어들고 있음을 확인할 수 있었고, 정확도 역시 95% 내외로 나타나는 것으로 확인할 수 있었습니다. 이를 통해서 loss function이 유효하다는 것을 알 수 있습니다.
첫번째 그래프에서 알 수 있는 점은 loss의 감소 속도가 Simple BU-Net이 제일 빠르고 U-Net도 비슷하지만, U-Net + WC block은 속도가 느리다는 점입니다. 이 과정에 대해 유추를 해보면 WC block이 추가됨으로써 parameter 수가 늘어났기 때문에 자연스럽게 loss가 늦게 감소하는 것 같습니다. 그러나 parameter 수가 더 많은 Simple RES block을 4개나 추가한 Simple BU-Net은 왜 빨리 감소하는 것일까요? 이는 아마 RES block의 특성 때문입니다. Residual connection (잔차 연결)은 gradient 소실에 대한 문제를 완화시켜줘서 깊은 신경망에도 효과적인 학습이 가능하도록 합니다. 이러한 특성 때문에 빠르게 loss가 감소한 것 같습니다.
두번째 그래프를 보면 accuracy가 U-Net + WC block 모델이 낮게 나오는 것을 볼 수 있습니다. 그렇다면 WC block은 효과가 없는 것일까요? 그렇게 보기만은 어려울 것 같습니다. 현재 첫번째 그래프에서도 보이듯이 loss가 아직 다른 모델에 비해 덜 줄어들었습니다. 학습이 더 진행됨에 따라 모델의 accuracy가 올라갈 것으로 예상됩니다. 다만 다른 두 모델만큼의 accuracy가 나올지는 미지수입니다. 이는 추후에 더 많은 학습을 통해 확인을 해봐야할 것 같습니다.
정량 평가 결과만으로 더 많은 정보를 얻기는 쉽지 않았기 때문에 정성 평가를 진행하였습니다.
정성 평가 결과
모델이 정말 제대로 학습되어 제대로 된 결과물을 도출해내는지 확인하기 위해서 테스트 데이터셋에 대한 추론을 진행하여 모델 간 결과물과 Ground Truth 이미지를 비교해보았습니다.
위 그림은 3가지 모델에 대한 차이점을 가장 잘 드러난 그림 중 하나입니다.
확인할 수 있듯 가장 Ground Truth와 비슷한 것은 Simple BU-net (U-Net + RES block) 모델의 결과물이었습니다. 반면 U-Net은 파란색 영역에 대해서는 비교적 Ground Truth와 비슷하게 탐지를 진행하였지만, 노란색 영역이 지나치게 넓게 나타났음을 확인할 수 있었고, U-Net과 WC를 결합한 모델의 경우 Ground Truth와 사뭇 다른 영역을 추론해낸 것을 확인할 수 있었습니다.
여기서 그치는 것이 아니라, 색의 양을 넘어 각각의 색의 경계를 주의해서 보시면 오른쪽의 모델로 갈수록 더욱 detail이 묻어난다는 것을 알 수 있습니다. 경계의 특징이 label과 유사하게 나타났습니다. U-Net 모델과 대조해보면 U-Net+ WC 모델의 경우는 파란색과 빨간색 라벨의 경계가 더욱 선명하게 나타나고, BU-Net의 경우는 노란색 라벨의 경계가 크게 차이남을 확인해볼 수 있습니다. 이를 통해서 block이 정말 효과가 있다는 것을 간접적으로 알 수 있었습니다.
물론 모든 그림이 Simple BU-Net이 이렇게 우수하게 나온 것은 아닙니다. 위의 그래프에서 accuracy를 보았듯이 U-Net의 성능도 좋기 때문에 단순한 형태의 label의 경우는 셋을 비교하기 어렵습니다. 사람마다 평가도 상이할 것입니다. 하지만 Dice Loss와 같이 차지하는 양을 제외하고 보더라도, 모양의 detail이 주는 효과는 정량적인 평가는 현재 안되었지만, 분명히 유효하게 작용하고 있습니다.
마치며
U-Net이 세상에 등장한지도 벌써 몇 년이 지났지만, 여전히 이미지 분할 분야에서 뛰어난 성능을 보이는 것을 확인할 수 있었습니다. 또한 RES block, WC block 등이 U-Net의 성능을 향상시키는 데 충분히 도움을 줄 수 있다는 사실도 확인할 수 있었습니다.
하지만 프로젝트를 진행하며 현실적인 문제에 많이 직면하게 되었는데요. <BU-net> 논문과 동일한 구조의 모델코드를 작성하는 데 성공했지만, 학습 여건 상 GPU와 RAM의 부족으로 학습이 불가하였다는 아쉬움이 남습니다. 또한 3D MRI 이미지를 슬라이스로 쪼개는 과정에서 데이터 수가 크게 늘어나 학습이 불가하다는 문제를 발견하여, 일부 슬라이스만을 사용했다는 점도, 에포크 수를 충분히 늘리지 못한 점도 아쉽습니다.
또한, <BU-net> 논문에 사용된 구조와 결과적으로 다른 구조로 학습을 진행하였기 때문에 논문에서 언급한 만큼의 성능을 확인하기에는 어려웠던 점 또한 개선해야 할 부분이라고 생각합니다.
추가로 저희가 이용한 Simple RES block에 대한 정확한 효율을 연구하지 못하고 사용했기 때문에 이에 대한 검증을 추후에 더 해야하며, 다른 모델의 경량화 방법도 알아보면 좋은 개선사항이라 생각이 듭니다.
참고 문헌
Rehman, Mobeen & Seungbin, Cho & Kim, Jee & Chong, Kil. (2020). BU-Net: Brain Tumor Segmentation Using Modified U-Net Architecture. Electronics. 9. 2203. 10.3390/electronics9122203.
