[논문 리뷰] Improved Unet model for brain tumor image segmentation based on ASPP-coordinate attention mechanism

Zixuan Wang, Yanlin Chen외 2명|arXiv (Cornell University)|2024. 09. 13.

Brain Tumor Detection and ClassificationNeuroscience인용 수 3

한 줄 요약

이 논문은 Atrous Spatial Pyramid Pooling (ASPP) 모듈과 좌표 attention 메커니즘을 통합하여 뇌 종양 세그멘테이션을 위한 향상된 U-Net 모델을 제안한다. 이는 특징 표현과 경계 정확도를 크게 향상시킨다. 모델은 최고 mIoU 0.76를 기록하며 기준 U-Net을 능가하여 의료 영상 데이터셋에서 뛰어난 세그멘테이션 성능을 보여준다.

ABSTRACT

In this paper, we propose an improved Unet model for brain tumor image segmentation, which combines coordinate attention mechanism and ASPP module to improve the segmentation effect. After the data set is divided, we do the necessary preprocessing to the image and use the improved model to experiment. First, we trained and validated the traditional Unet model. By analyzing the loss curve of the training set and the validation set, we can see that the loss value continues to decline at the first epoch and becomes stable at the eighth epoch. This process shows that the model constantly optimizes its parameters to improve performance. At the same time, the change in the miou (mean Intersection over Union) index shows that the miou value exceeded 0.6 at the 15th epoch, remained above 0.6 thereafter, and reached above 0.7 at the 46th epoch. These results indicate that the basic Unet model is effective in brain tumor image segmentation. Next, we introduce an improved Unet algorithm based on coordinate attention mechanism and ASPP module for experiments. By observing the loss change curves of the training set and the verification set, it is found that the loss value reaches the lowest point at the sixth epoch and then remains relatively stable. At the same time, the miou indicator has stabilized above 0.7 since the 20th epoch and has reached a maximum of 0.76. These results show that the new mechanism introduced significantly improves the segmentation ability of the model. Finally, we apply the trained traditional Unet model and the improved Unet model based on the coordinate attention mechanism and ASPP module to the test set for brain tumor image segmentation prediction. Compared to the traditional Unet, the enhanced model offers superior segmentation and edge accuracy, providing a more reliable method for medical image analysis with the coordinate attention mechanism and ASPP module.

연구 동기 및 목표

• 딥 러닝을 활용하여 MRI 영상에서 뇌 종양 세그멘테이션의 정확도를 향상시키기.
• 표준 U-Net이 다중 척도의 맥락을 포착하고 장거리 의존성을 모델링하는 데 한계를 보이는 문제를 해결하기.
• attention 메커니즘과 다중 비율 컨볼루션을 통해 특징 표현과 윤곽 검출을 향상시키기.
• 공개된 뇌 종양 데이터셋을 기반으로 제안된 아키텍처의 효과성을 검증하기.
• 임상 수준의 의료 영상 세그멘테이션을 위한 강력하고 신뢰할 수 있는 방법 제공하기.

제안 방법

• 제안된 모델은 인코더 특징에서 다중 척도의 맥락적 특징을 추출하기 위해 Atrous Spatial Pyramid Pooling (ASPP) 모듈을 통합한다.
• 좌표 attention 메커니즘이 동시에 채널과 공간 좌표에 주목함으로써 공간적 의존성을 명시적으로 모델링한다.
• ASPP와 좌표 attention 모듈이 U-Net 인코더-디코더 아키텍처 내부에 융합되어 특징 학습을 향상시킨다.
• 모델 일반화를 향상시키기 위해 전처리 단계에서 데이터 증강 및 정규화가 적용된다.
• 모델은 검증 손실 기반 조기 정지 기법을 사용한 Adam 옵timizer와 이진 교차 엔트로피 손실을 사용하여 훈련된다.
• 세그멘테이션 성능 평가를 위해 테스트 세트에서 추론가 수행되며, mIoU와 시각적 품질이 사용된다.

실험 결과

연구 질문

• RQ1 ASPP와 좌표 attention의 통합이 표준 U-Net 대비 뇌 종양 세그멘테이션에서 mIoU를 향상시키는가?
• RQ2 제안된 attention 메커니즘이 종양 영역의 특징 표현과 경계 정렬을 향상시키는가?
• RQ3 기준 U-Net과 비교해 향상된 모델의 훈련 수렴성과 안정성은 어떻게 다른가?
• RQ4 모델은 종양 세그멘테이션에서 거짓 양성과 거짓 음성의 수를 어느 정도 줄이는가?
• RQ5 향상된 모델은 높은 경계 정확도를 유지하면서도 미리 보지 않은 테스트 데이터에 잘 일반화되는가?

주요 결과

• 향상된 U-Net 모델은 최고 mIoU 0.76를 기록하며 기준 U-Net을 뚜렷이 능가했다.
• 향상된 모델의 mIoU는 20번째 에포크부터 0.7 이상으로 안정화되어 강력한 수렴성을 보였다.
• 향상된 모델의 손실 곡선은 6번째 에포크에 최소값에 도달하고 안정된 상태를 유지하여 더 빠르고 안정적인 훈련을 의미했다.
• 기준 U-Net은 46번째 에포크가 되어서야 mIoU가 0.65를 초과했으며, 이는 더 느린 수렴을 의미했다.
• 시각적 점검 결과, 향상된 모델은 더 선명한 종양 경계를 생성하고 세그멘테이션 노이즈를 감소시켰다.
• 테스트 세트에서 표준 U-Net 대비 제안된 모델이 뛰어난 경계 정확도와 종합적인 세그멘테이션 품질을 보였다.