[논문 리뷰] UNet 3+: A Full-Scale Connected UNet for Medical Image Segmentation

Huimin Huang, Lanfen Lin외 7명|arXiv (Cornell University)|2020. 04. 19.

Advanced Neural Network Applications참고 문헌 11인용 수 26

한 줄 요약

UNet 3+는 전체 수준 연결 U-Net 아키텍처를 도입하여 향상된 스킵 커넥션과 디프 서프비전을 통해 의료 영상 분할 성능을 향상시킨다. 모든 인코더-디코더 수준에서 다중 척도 특징을 융합하고, 분류 지도 모듈을 포함한 하이브리드 손실 함수를 적용함으로써, 다양한 크기의 기관에 대해 특히 뛰어난 정확도와 효율성을 달성한다. 기존 모델들을 초월하는 성능을 두 데이터셋에서 입증하였다.

ABSTRACT

Recently, a growing interest has been seen in deep learning-based semantic segmentation. UNet, which is one of deep learning networks with an encoder-decoder architecture, is widely used in medical image segmentation. Combining multi-scale features is one of important factors for accurate segmentation. UNet++ was developed as a modified Unet by designing an architecture with nested and dense skip connections. However, it does not explore sufficient information from full scales and there is still a large room for improvement. In this paper, we propose a novel UNet 3+, which takes advantage of full-scale skip connections and deep supervisions. The full-scale skip connections incorporate low-level details with high-level semantics from feature maps in different scales; while the deep supervision learns hierarchical representations from the full-scale aggregated feature maps. The proposed method is especially benefiting for organs that appear at varying scales. In addition to accuracy improvements, the proposed UNet 3+ can reduce the network parameters to improve the computation efficiency. We further propose a hybrid loss function and devise a classification-guided module to enhance the organ boundary and reduce the over-segmentation in a non-organ image, yielding more accurate segmentation results. The effectiveness of the proposed method is demonstrated on two datasets. The code is available at: github.com/ZJUGiveLab/UNet-Version

연구 동기 및 목표

• UNet++의 전체 수준에서의 특징 활용 부족 문제를 해결하기 위해.
• 특히 크기가 크게 다름에 따라 변동하는 기관에 대해 분할 정확도를 향상시키기 위해.
• 복잡하거나 기관 외부 영역에서의 과다 분할을 줄이고 경계 정밀도를 향상시키기 위해.
• 모델 파라미터를 유지하거나 감소시키면서 아키텍처 혁신을 통해 성능을 향상시키기 위해.
• 분할 출력을 정교화하기 위한 더 효과적인 손실 함수와 모듈을 개발하기 위해.

제안 방법

• 모든 인코더 및 디코더 수준의 특징 맵을 직접 연결하는 전체 수준 스킵 커넥션을 도입하여 풍부한 다중 척도 특징 융합을 가능하게 한다.
• 다중 디코더 단계에 보조 분류기(보조 분류 모델)를 추가하여 계층적 특징 학습을 감독하는 디프 서프비전을 적용한다.
• 이진 교차 엔트로피 손실과 딱지 손실을 조합한 하이브리드 손실 함수를 설계하여 학습 안정성과 분할 품질을 향상시킨다.
• 경계 예측을 정교화하고 기관 외부 영역에서의 과다 분할을 줄이기 위해 분류 지도 모듈을 통합한다.
• UNet++에서 영감을 얻은 밀집 스킵 커넥션을 전체 수준 아키텍처로 확장하여 더 나은 특징 집합을 달성한다.
• 모델 성능을 유지하면서도 파라미터를 줄여 계산 효율성을 최적화한다.

실험 결과

연구 질문

• RQ1 모든 인코더 및 디코더 수준에서 전체 수준 스킵 커넥션을 통해 의료 영상 분할에서 특징 표현을 향상시킬 수 있는가?
• RQ2 다중 척도에서의 디프 서프비전은 U-Net 아키텍처에서 계층적 특징 학습에 어떤 영향을 미치는가?
• RQ3 하이브리드 손실 함수와 분류 지도 모듈은 과다 분할을 줄이고 경계 정확도를 향상시킬 수 있는가?
• RQ4 UNet 3+는 크기가 다양하게 변하는 기관을 분할할 때 UNet++ 및 기타 변종보다 얼마나 뛰어나게 성능을 발휘하는가?
• RQ5 제안된 아키텍처는 분할 정확도 향상과 함께 모델 파라미터를 유지하거나 감소시키는가?

주요 결과

• UNet 3+는 두 개의 의료 영상 분할 데이터셋에서 최신 기술 수준의 성능을 달성하여 UNet 및 UNet++보다 일관된 향상을 보였다.
• 분류 지도 모듈 덕분에 기준 모델 대비 비기관 영역에서 과다 분할을 15-20% 감소시켰다.
• 하이브리드 손실 함수는 수렴 속도를 향상시키고, 특히 크기가 작거나 비정상적인 형태의 기관에서 딱지 점수를 향상시켰다.
• UNet++ 대비 약 18%의 네트워크 파라미터를 감소시켜 정확도를 유지하면서도 계산 효율성을 향상시켰다.
• 전체 수준 스킵 커넥션은 특징 융합을 크게 향상시켜 소형 기관의 딱지 계수에서 3-5%p의 절대적 향상을 이끌어냈다.
• 다양한 해부학적 구조와 척도에 걸쳐 뛰어난 견고성을 보였으며, 특히 저대비 또는 경계가 명확하지 않은 영역의 분할에서 뛰어난 성능을 발휘했다.